Repository navigation
Add MLX's affine 4-bit quantizer and a check against shipped tensors - #709
Draft
philip-pentatonic wants to merge 1 commit into
Draft
philip-pentatonic wants to merge 1 commit into
philip-pentatonic wants to merge 1 commit into
Conversation
quantizeMlx quantizes a dense tensor to MLX affine 4-bit with groups of 32, 64 or 128 the way MLX does: float32 arithmetic with a true division, the bias at the larger-magnitude edge of a group, a signed scale chosen so that zero is a grid point, round half away from zero, codes taken from the unrounded scale and bias, and bf16 rounding only for storage. When the dense source and the conversion scheme match how a checkpoint was made, it reproduces that checkpoint's words, scales and biases exactly. A non-finite value, an overflowing range, a zero adjusted scale and a scale or bias that rounds to a bf16 infinity are refused, naming the group. `tf-merge-check --quantizer DENSE SHIPPED MODULE...` quantizes each named dense tensor and compares it with the shipped one, reporting the share of equal scales, biases, codes and whole groups. The dense tensor must be an n x k f32, bf16 or f16 matrix of the shipped projection's shape and payload size, every required input must be a regular file, and a failed allocation unwinds cleanly. Tests pin the quantizer to nine groups produced by an independent float32 implementation of the published algorithm: two where true division and reciprocal multiplication differ, one that needs the unrounded scale and bias, a q0 tie at -8.5, the larger edge at the minimum and at the maximum, a negative flat group and a tiny group with q0 = 0; they also check a local reciprocal variant disagrees with them. Other tests cover group sizes 32, 64 and 128, the refusals above, and the check on synthetic dense and shipped folders. Build: `zig build test-merge`, `zig build tf-merge-check`.
This was referenced Oct 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 1 of 2 for #666. The merge command in the second pull request builds on this one.
What this changes
Adds
affine4_host.quantizeMlx, a Zig implementation of the 4-bit affine quantizer thatmlx_lmused for the shipped MLX checkpoints, andtf-merge-check --quantizer DENSE SHIPPED MODULE..., which quantizes a dense bf16 tensor and compares the words, scales and biases with the shipped checkpoint's. The scheme: float32 arithmetic on the bf16 input, true division by 15, the bias at the group's larger-magnitude edge with a signed scale, the scale reset toedge / q0so zero is an exact grid point, codes from the unrounded scale and bias, round half away from zero, bf16 only for storage. Nothing in the engine calls it yet. One change: the quantizer and its check.Receipt
e7c523ebplus this commit, Zig 0.17.0, Linux x86_64 (AMD Threadripper PRO, 32 cores), host only, no GPU. Shipped checkpointTensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bitatd9d758fb(every file's sha256 checked against the Hub); dense sourcenvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16ata9904d24, only the named tensors fetched by HTTP range.tf-merge-check --quantizer DENSE SHIPPED backbone.layers.{0,2,30}.mixer.{in_proj,out_proj}: scales, biases, codes and whole groups 100.0000% identical on all six tensors (10304x2688 and 2688x4096, group 64), exit 0,quantizer check: byte-identical on every tensor.TensorFold/Qwen3.8-27B-MLX-4bitfromQwen/Qwen3.8-27B(public main, 11 October 2026) bit for bit. The Qwen tensors were not run through the Zig check. A reciprocal multiply instead of the division by 15 misses about 0.5% of groups, which the vector test asserts.zig build test-merge: 6/6 steps succeeded, 33/33 tests passed.zig build test: 39/39 steps succeeded, 463/463 tests passed.zig fmt --checkon every touched path: clean.python3 tools/zig/lean_check.py: 534 problems before and after, none added. Nine pinned vectors from the independent implementation (synthetic groups: division against reciprocal, a q0 tie, the edge at the minimum and at the maximum, a negative flat group, a tiny q0 = 0 group); non-finite input, an overflowing range, a zero adjusted scale and a bf16 overflow are refused.build.zigtest-step wiring is configure-checked foraarch64-macosonly); Windows; group sizes 32 and 128 on real checkpoints (unit tests only); Qwen3.8-27B through the Zig check; CI.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Not applicable.