Skip to content

Add MLX's affine 4-bit quantizer and a check against shipped tensors - #709

Draft
philip-pentatonic wants to merge 1 commit into
ashhart:mainfrom
Pentatonic-Ltd:mlx-affine-quantizer
Draft

philip-pentatonic wants to merge 1 commit into
ashhart:mainfrom
Pentatonic-Ltd:mlx-affine-quantizer

Conversation

@philip-pentatonic

Copy link
Copy Markdown
Contributor

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 that mlx_lm used for the shipped MLX checkpoints, and tf-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 to edge / q0 so 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

  • Environment: TensorFold e7c523eb plus this commit, Zig 0.17.0, Linux x86_64 (AMD Threadripper PRO, 32 cores), host only, no GPU. Shipped checkpoint TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit at d9d758fb (every file's sha256 checked against the Hub); dense source nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 at a9904d24, only the named tensors fetched by HTTP range.
  • Exactness: not applicable. This is an offline tool; no lane round, cache, draft, sampling or serving path is touched.
  • Quantizer against the shipped checkpoint, 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.
  • Independent cross-check: a separate float32 implementation of the same published algorithm (not part of this pull request) reproduces these Nemotron tensors and two tensors of the shipped TensorFold/Qwen3.8-27B-MLX-4bit from Qwen/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.
  • Decode speed, prompt speed: not applicable, no serving change.
  • Tests run: 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 --check on 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.
  • Not run: loading or serving a model quantized this way; Metal and macOS builds (the build.zig test-step wiring is configure-checked for aarch64-macos only); Windows; group sizes 32 and 128 on real checkpoints (unit tests only); Qwen3.8-27B through the Zig check; CI.

Checklist

  • Every token still goes through the lane rounds. Not applicable, no serving path changes.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Not applicable.
  • No precision traded for speed. The quantizer is exact against the shipped checkpoint where the dense source and the conversion scheme match.
  • Prompt processing is no slower than the last release. Not applicable.
  • The description says which platforms I ran: Linux x86_64 host only; no Metal, no CUDA, no GPU.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. The vector test fails for a reciprocal-multiply variant; the real-tensor check takes two paths and needs no model in the repository.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

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`.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant