Repository navigation
metal: affine_sym host+kernels, Precise transcendentals and tests on the 0.32.2 line (gemma4-qat #86) - #2
Conversation
quantized_matmul(biases=None, mode='affine') now routes to affine_sym_qmv/ _fast kernel variants (kernel source lives in mlx-swift's generated JIT preamble): ops accepts 3-input affine (bits<=2), eval_gpu maps it to the affine_sym mode string on the 1/2-bit decode path, vjp/jvp guarded, CPU throws. Mirrors PrismML-Eng/llama.cpp-fork mlx commit c1dcd54.
ml-explore#179) Port the three hand-edits that so far existed only in numen-tech/mlx-swift's regenerated Source/Cmlx/mlx-generated/{quantized.cpp,metal/quantized.h} into their real source, mlx/backend/metal/kernels/quantized.h, so that mlx-swift's tools/update-mlx.sh regenerates the JIT preamble byte-identically from this fork and nobody has to hand-edit generated files again. The header is byte-for-byte the mlx-swift origin/prism (37bd6fb) copy of Source/Cmlx/mlx-generated/metal/quantized.h, which fda5e49 had already synced with the JIT source. Relative to prism-0.31.1-fixes (b2b8a3d) the delta is exactly three logical hunks: 1. qdot<...,bits=1>: read the 1-bit codes as uint32 words (32 values per load) instead of byte-at-a-time, using select() per bit. Provenance: mlx-swift bc96a51 "metal: uint32 wide load in 1-bit qmv qdot (large mobile decode win)" (Brian, 2026-06-07). 2. sym_derived_bias<U, bits, bias_free>(scale) (-0.5*scale for 1-bit, -scale for 2-bit; static_assert'ed to those two widths) and a `bool bias_free = false` template parameter threaded through qmv_fast_impl / qmv_impl so the biases buffer is never read on the symmetric path. Provenance: mlx-swift 8950e67 "metal: bias-free (symmetric) 1/2-bit decode kernels (affine_sym_qmv/_fast)" (Brian, 2026-06-09) and fda5e49 "review: sync metal/quantized.h with JIT source; guard sym_derived_bias; fix FMT_CONSTEVAL comment" (Brian, 2026-06-10). 3. The [[kernel]] entry points affine_sym_qmv_fast / affine_sym_qmv (13 buffers: the biases slot is dropped and the scale strides are reused for the dummy biases pointer), which the host side already dispatches to since d90771c (QuantizedMatmul::eval_gpu selects mode "affine_sym" for affine weights without biases; jit_kernels.cpp pulls metal::quantized() for every "affine*" mode). Provenance: mlx-swift 8950e67 / fda5e49 as above. quantized_nax.h and quantized_utils.h were compared against the mlx-swift generated copies and are already identical; no other file changes. Validation: explicit `xcrun -sdk macosx metal -c` instantiation of affine_sym_qmv_fast / affine_sym_qmv for bits 1 and 2, group sizes 32/64/128, float/float16_t/bfloat16_t, batched 0/1 (72 kernels) compiles with the kernels/CMakeLists.txt warning flags, and the 4-bit instantiation trips the sym_derived_bias static_assert as intended. Note that quantized.metal still does not instantiate affine_sym_* kernels, so the non-JIT metallib does not carry them; that is a pre-existing gap of the fork, not of this change, and mlx-swift consumes the JIT path only. No fork tests cover the affine_sym kernels (there is no tests/quantize_tests.cpp in this fork); adding them is tracked separately (gemma4-qat ml-explore#86). Refs numen-tech/gemma4-qat#179. 0.32.2 port note (numen-tech/gemma4-qat#86): the kernel text is byte-identical to 4aa785f (fork/179-affine-sym-kernels) except for the two adaptations that 0.32.2's host dispatch forces on the [[kernel]] entry points: * affine_sym_qmv_fast / affine_sym_qmv take the same six template parameters as 0.32.2's affine_qmv_fast / affine_qmv (<T, group_size, bits, batched, has_global_scale = false, results_per_simdgroup = 4>): qmv() now instantiates every qmv kernel with global_scale.has_value() and results_per_simdgroup appended, and both are unused on the affine paths exactly as in affine_qmv. * Buffer indices follow 0.32.2's fixed affine_qmv layout (x at buffer 3 ... s_strides at buffer 13) with the biases slot (buffer 2) left unbound, because qmv() now binds x at 3 whether or not a biases / global_scale buffer is present (0.31.1 bound sequentially, which is where the 13-buffer "biases slot skipped" layout of 4aa785f came from). The qdot 1-bit wide load, sym_derived_bias and the bias_free threading through qmv_fast_impl / qmv_impl applied without conflict.
MTLCompileOptions has two independent knobs: mathMode (Safe / Relaxed / Fast, the -ffast-math contract) and mathFloatingPointFunctions (Fast / Precise, which transcendental implementations the compiler emits). MLX's set_compile_options only ever set mathMode, so on macOS 15 / iOS 18 and later every kernel compiled with MathMode::Safe or Relaxed still picked up the default mathFloatingPointFunctions = Fast, i.e. the approximate exp/log/sin/cos/pow variants, exactly what a Safe/Relaxed request is supposed to opt out of (this is the fork's math-mode finding from numen-tech/gemma4-qat#86). Set mathFloatingPointFunctions next to mathMode, under the same availability guard: Fast when MathMode::Fast is requested, Precise otherwise. The pre-macOS-15 branch keeps setFastMathEnabled, whose semantics already cover both knobs. The vendored metal-cpp (the version fetched by CMake) has the setter (MTLLibrary.hpp:266).
Dispatch audit for the bias-free (affine_sym) 1/2-bit decode path on top of
0.32.2's restructured QuantizedMatmul::eval_gpu:
* eval_gpu selects mode "affine_sym" (or throws for bits > 2, a
non-transposed weight, or M >= the device's qmv batch limit) BEFORE the
M >= vector_limit branch, so neither qmm_splitk nor qmm can ever see it.
* dispatch_qmv skips qmv_quad for affine_sym (there is no bias-free quad
kernel; d90771c), and now also skips qmv_wide: use_qmv_wide() returned
true for every mode other than "affine", which under 0.32.2's
"M >= 2 -> qmv_wide" rule (332d459) would have sent an affine_sym
matvec with 2 <= M < vector_limit to a non-existent affine_sym_qmv_wide
kernel. affine_sym now always reaches qmv() and its
affine_sym_qmv / affine_sym_qmv_fast entry points.
Add test_qmv_affine_sym (1/2-bit, group sizes 32/64/128, K in
{64, 128, 512, 2048} to cover the qmv_quad bypass and both the _fast and
generic variants, M in {1, 2, 3, 8} to cover the qmv_wide bypass, fp32 vs
the biased kernels fed the derived biases and vs dequantize-then-matmul,
plus fp16/bf16 spot checks) and test_qmv_affine_sym_throws (bits > 2,
non-transposed, matrix regime, CPU stream).
…me throws
test_qmv_affine_sym compared the bias-free result against the biased
kernel bit-for-bit (1e-5 fp32, 1e-2 half). That only holds when both
calls hit the same kernel; the biased reference takes qmv_quad for
K in {64, 128} and, since 332d459, qmv_wide for 2-bit M >= 3 on
gen-15+ GPUs, whose accumulation order differs by a few ulps of the
output dtype (observed: 1 bf16 ulp at |y| ~ 140, 2-4 fp32 ulps at
|y| ~ 300). Compare in fp32 with a tolerance of 4 output ulps at the
largest output magnitude instead; the 1e-3 check against the
dequantize-then-matmul reference is unchanged. Also skip the
K < group_size combinations, which mx.quantize rejects.
test_qmv_affine_sym_throws hung on the CPU case: it built fresh inputs
inside evals that were expected to throw, and an array computed inside
an eval that throws never has its completion event signaled, so the
later CPU-stream call waited on that event forever
(Scheduler::wait_event -> IOSurfaceSharedEvent waitUntilSignaledValue).
Evaluate every input up front before the throwing evals.
python/tests/test_quantized.py: 41 passed, 3843 subtests passed.
…active-lane reads in the qmv_fast tail (numen-tech/gemma4-qat#179, Codex P1) Two P1 findings from the adversarial Codex review of #1: - quantized.metal had no affine_sym instantiations, so a non-JIT build (MLX_METAL_JIT=OFF, the metallib path) could not resolve the affine_sym_qmv[_fast]_* names the host asks for when mode == "affine_sym" (bias-free 1/2-bit quantized_matmul). Add instantiate_quantized_sym_groups for bits 1 and 2 only (sym_derived_bias static_asserts the other widths away and the host only requests bits <= 2), all three types, group sizes 128/64/32, batched 0/1, with the same host-name format the JIT path builds: affine_sym_qmv[_fast]_<type>_gs_<gs>_b_<bits>_batch_<0|1>. 72 new entries. - qmv_fast_impl's partial tail block (in_vec_size % block_size != 0: 1-bit with K = 512 mod 1024, since the host's K % 512 fast gate is coarser than the 1-bit block size of 1024; also bits=5) zeroed x_thread/scale/bias for lanes past the end of the row but still ran qdot, whose weight reads land past the end of the row and, for the last output rows, past the end of the weight buffer. Skip the row loop for those lanes instead: they contributed exactly +0, so results are unchanged, and every lane still joins the simd_sum. The aligned loop is untouched; gather_qmv_fast shares the impl. The host-side `&& bits >= 2` fast-path guard is deliberately NOT restored (the fork keeps the 1-bit fast path: mlx#3, ~+11% decode).
…emma4-qat#179, Codex P1) Pins the affine_sym decode kernels in both the metallib (MLX_METAL_JIT=OFF) and JIT builds: M=1, transpose=true, bits=1 with K in {512, 1536, 1024} and bits=2 with K=512 (N=8, plus N=64 multi-row-group cases, group sizes 128/64, and K=544 for the non-fast affine_sym_qmv), float16 and float32, against matmul(x, dequantize(w, scales, derived bias)) with the bias the kernel synthesizes (-0.5*scale for 1-bit, -scale for 2-bit). K=512/1536 end in the partial 1-bit qmv_fast tail block; K=1024 and the 2-bit cases cover the no-tail path. GPU-only: the CPU backend has no bias-free path. Verified: 54/54 assertions in both the MLX_METAL_JIT=OFF and =ON builds (Xcode 27.0). Note the metallib build needs -std=metal4.0 with Xcode 27's metal front end: a pre-existing steel/gemm/mma.h incompatibility, independent of this change and not addressed here.
…-bit kernel (ml-explore#86) Upstream 0.32.2 replaced the K % 512 qmv_fast gate with an exact one (ml-explore#3965, qmv_fast_k_alignment = pack_factor * (bits == 2 ? 1 : 2) * SIMD_SIZE, "must match the K step in qmv_fast_impl"). The fork's 1-bit kernel reads one pack per lane as well (packs_per_thread = bits <= 2 ? 1 : 2, mlx#3), so its block is 1024 values, but on prism-0.32.2-fixes the host formula still said 2048 for 1-bit: every 1-bit matmul with K = 1024 mod 2048 (Bonsai-27B's down_proj, K = 21504) silently took the generic qmv instead of qmv_fast, and K = 1024/1536 shapes never reached the fast kernel at all. Mirror the kernel's packs_per_thread in the host formula (bits <= 2), which also keeps the partial-block tail of qmv_fast_impl (cherry-picked from #1) unreachable on this line: the gate is exact for every width. Comment-only follow-ups for the 0.32.2 line: the tail-guard note in quantized.h and the routing comments of the bias-free decode doctest (the 0.31.1 line's K % 512 gate sends 1-bit K = 512 mod 1024 through the tail block; here those shapes run the generic affine_sym_qmv and K = 1024 the fast kernel; the cases themselves stay identical to #1's), and test_qmv_affine_sym adds K = 1024 so the 1-bit affine_sym_qmv_fast kernel is exercised at its block size. Verified: tests/tests -tc='test bias-free affine quantized_matmul decode' 54/54 in both the MLX_METAL_JIT=ON and =OFF (metallib) builds, Xcode 27.0; python/tests/test_quantized.py; 1-bit K=21504 N=5376 fp16 decode qmv microbench before/after in the PR.
Bump the Source/Cmlx/mlx submodule from 2c2d8a4e7 to the numen-tech/mlx#2 head 7072c9a2f, which carries the numen-tech/mlx#1 kernel work over to the 0.32.2 line (ab914a5b5: affine_sym qmv/qmv_fast metallib instantiations for bits 1/2 + the inactive-lane skip in the qmv_fast tail; ee91d6acb: the bias-free decode doctest) and makes the host fast gate exact for the fork's 1-pack-per-lane 1-bit kernel (7072c9a2f: qmv_fast_k_alignment = 1024 for 1-bit, 512 for 2-bit; the previous 2048 routed every K = 1024 mod 2048 1-bit matmul to the generic qmv). LC_ALL=C CC=$(xcrun -f clang) CXX=$(xcrun -f clang++) ./tools/update-mlx.sh touched only metal/quantized.h and quantized.cpp (the tail-block comment under qmv_fast_impl); the host-gate change lives in the submodule's mlx/backend/metal/quantized.cpp, which SwiftPM compiles directly. The "Autogenerated from commit" stamp update-mlx.sh writes into include-framework/Metal.hpp (the mlx-swift HEAD, not a metal-cpp commit) is restored to upstream 9019419's value, as prism-179-regenerate does; no host paths leaked into the generated sources.
…l-explore#86) The submodule must now sit at the numen-tech/mlx#2 head (fork/86-affine-sym-0.32.2; prism-0.32.2-fixes once merged), never at the pre-#2 prism-0.32.2-fixes tip e1971918, which has no affine_sym. The patch history no longer claims the 1-bit qmv_fast handles a K % 512 remainder in-kernel: on this line qmv_fast_k_alignment (numen-tech/mlx 7072c9a2f) keeps the host gate exact, so the tail block is unreachable. The pinned mlx-c is c74db53 (v0.6.0 + 7), not the v0.6.0 tag.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7072c9a2fb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (group_size != 32 && group_size != 64 && group_size != 128 && | ||
| group_size != 256 && group_size != 512) { |
There was a problem hiding this comment.
For GPU builds using the default MLX_METAL_JIT=OFF, accepting group sizes 256 and 512 makes quantize request affine_quantize_*_gs_256/512_* kernels, but mlx/backend/metal/kernels/quantized.metal instantiates affine kernels only for groups 32, 64, and 128. Calls such as mx.quantize(..., group_size=256, stream=mx.gpu) therefore pass validation and fail later while loading the kernel; either add the required metallib instantiations or continue rejecting these sizes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: affine_quantize rejects group sizes 256 and 512 again, so only sizes with metallib kernels pass validation.
| if (qmode == QuantizationMode::Affine) { | ||
| inputs = { | ||
| astype(x, dtype), w, astype(scales, dtype), astype(*biases, dtype)}; | ||
| inputs = {astype(x, dtype), w, astype(scales, dtype)}; | ||
| if (biases) { | ||
| inputs.push_back(astype(*biases, dtype)); |
There was a problem hiding this comment.
When biases is absent, this creates a three-input QuantizedMatmul, but requesting a VJP with respect to scales still evaluates zeros_like(primals[3]) in QuantizedMatmul::vjp, causing an out-of-bounds vector access. That branch also needs the symmetric derivative (q - 0.5 for 1-bit or q - 1 for 2-bit), not the existing independent-bias derivative, so it should implement that derivative or reject scale differentiation before indexing the missing input.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: the scale VJP no longer indexes primals[3]; for the bias-free case it uses the symmetric derivative (q - 0.5 for 1-bit, q - 1 for 2-bit). Covered by test_qmv_affine_sym_grad.
| (mode_ == QuantizationMode::Affine && primals.size() > 3) | ||
| ? std::optional<array>(primals[3]) | ||
| : std::nullopt, | ||
| !transpose_, |
There was a problem hiding this comment.
For every supported bias-free call with transpose=true, the VJP with respect to x forwards biases=nullopt while flipping transpose to false. QuantizedMatmul::eval_gpu explicitly rejects bias-free non-transposed operations, so a normal mx.grad of this new path always fails during evaluation; derive and pass the implicit biases for the reverse multiplication, or reject autodiff explicitly if it is intentionally unsupported.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: the VJP wrt x derives the implicit bias (-scale/2 or -scale) and passes it to the reverse product, so it runs on the biased kernels. Covered by test_qmv_affine_sym_grad.
| // Bias-free (symmetric) affine quantized_matmul: the 1/2-bit decode path where | ||
| // no biases buffer is bound and the kernel derives the bias from the scale | ||
| // (-scale/2 for 1-bit, -scale for 2-bit; sym_derived_bias in | ||
| // backend/metal/kernels/quantized.h). Metal only: the CPU backend has no | ||
| // bias-free path. Pins two things (numen-tech/gemma4-qat#179, Codex P1): |
There was a problem hiding this comment.
This 17-line block documents old branch behavior, external issue history, implementation details, and each test case instead of giving the test's brief purpose. Reduce it to a concise description and let the case names or short inline comments explain routing; the repository explicitly limits comments to roughly one or two lines and prohibits redundant or excessive commentary.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: the comment is down to two lines.
| if (qmode == QuantizationMode::Affine && !biases && bits > 2) { | ||
| throw std::invalid_argument( | ||
| "[quantized_matmul] Bias-free affine supports bits in {1, 2} only."); |
There was a problem hiding this comment.
The new bias-free branch claims to allow only {1, 2}, but this one-sided check also accepts zero and negative bits. In particular, bits=0 reaches validate_quantized_input, which computes w.shape(-1) * 32 / bits and can terminate with an integer divide-by-zero instead of reporting an invalid argument; check explicitly for bits != 1 && bits != 2 before extracting dimensions.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: the check is now bits != 1 && bits != 2, and test_qmv_affine_sym_throws covers bits 0 and -1.
| if (mode == "affine" && !biases) { | ||
| // Bias-free (symmetric) affine: 1/2-bit formats where the bias is a fixed | ||
| // function of the scale and is derived in-kernel. | ||
| qmode = QuantizationMode::Affine; | ||
| dtype = scales.dtype(); |
There was a problem hiding this comment.
This device-independent API branch also creates a three-input affine primitive on CUDA. For a 2-bit call, the CUDA qmv and naive qmm support predicates allow the missing optional bias, but their affine uint2b_t kernels select has_bias=true at compile time and dereference the null bias pointer, which can cause an illegal GPU memory access; add a CUDA-side rejection like the CPU backend or implement the derived-bias behavior there.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 84b4217: the CUDA QuantizedMatmul::eval_gpu rejects bias-free affine before any kernel is selected, matching the CPU backend.
- quantized_matmul: reject bits outside {1, 2} for bias-free affine
- QuantizedMatmul::vjp: derive the implicit bias for the reverse product
and use the symmetric scale derivative (q - 0.5 / q - 1)
- affine_quantize: reject group sizes 256 and 512 again (no kernels)
- cuda: reject bias-free affine like the CPU backend
- tests: condense the doctest comment; add bias-free vjp/jvp checks
… Codex P2) test_qmv_affine_sym tolerated any "Bias-free affine" RuntimeError for M > 1 and skipped the numerics, so an erroneous rejection of, say, M = 8 passed. Every shape it runs (K <= 2048, N <= 256, M <= 8) sits under get_qmv_batch_limit on every GPU generation (minimum 13), so the eval must succeed: drop the escape. test_qmv_affine_sym_throws now pins the rejection at M = 33 (the largest limit on any generation, rejected as M >= limit everywhere) as well as M = 512, and matches the message.
… Codex fold) Bump the Source/Cmlx/mlx submodule from 7072c9a2f to the numen-tech/mlx#2 head ff1deee55: 84b42176f (bias-free affine: reject bits outside {1, 2}, derived-bias vjp/scale derivative, affine_quantize back to group sizes 32/64/128 only, CUDA rejects bias-free affine like the CPU backend) and ff1deee55 (Python test: no RuntimeError escape in test_qmv_affine_sym, explicit M >= limit rejection cases). tools/update-mlx.sh regenerated Source/Cmlx/mlx-generated/ with no content change (the bump touches ops.cpp / primitives.cpp / cuda host dispatch / tests only); Metal.hpp's autogenerated-from line follows this checkout's HEAD as on every regen.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ff1deee551
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.assertEqual(y_q.shape, y_hat.shape) | ||
| self.assertLess((y_q - y_hat).abs().max(), 1e-3) | ||
|
|
||
| def test_qmv_affine_sym(self): |
There was a problem hiding this comment.
These three new tests run unconditionally, but the implementation explicitly throws for bias-free affine matmul on both CPU and CUDA. Consequently, test_qmv_affine_sym fails as soon as its result is evaluated in CPU-only or CUDA test runs, and the later throw/gradient tests have the same problem. Guard all three with mx.metal.is_available() so the cross-platform Python suite continues to pass.
Useful? React with 👍 / 👎.
| if (!is_available(Device::gpu)) { | ||
| return; | ||
| } |
There was a problem hiding this comment.
Checking generic GPU availability is insufficient on a CUDA build: Device::gpu is available there, so this Metal-specific test proceeds, while the newly added CUDA QuantizedMatmul::eval_gpu rejection throws for every bias-free affine call. Gate on Metal availability rather than generic GPU availability to avoid breaking the CUDA C++ test suite.
Useful? React with 👍 / 👎.
| // Partial last block (in_vec_size % block_size != 0). Reachable only from | ||
| // a host fast gate coarser than block_size (the 0.31.1 line's K % 512 gate | ||
| // with 1-bit K = 512 mod 1024); qmv_fast_k_alignment (quantized.cpp) keeps | ||
| // the gate exact here. Lanes past the end of the row contribute +0 (their | ||
| // x_thread, sum and scale would all be zero), so they skip the row loop |
There was a problem hiding this comment.
Reduce this long explanation to one or two lines describing the bounds guard. The current block records old-release routing, implementation history, detailed masking behavior, an external issue, and a prior review priority, contrary to the repository requirement for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L32
Useful? React with 👍 / 👎.
…l comment (PR #2, Codex P1) - python: skip the three affine_sym tests unless mx.metal.is_available() - ops_tests: check metal::is_available() instead of Device::gpu - quantized.h: two-line comment on the partial last block
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b4750de841
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // with packs_per_thread = bits <= 2 ? 1 : 2: upstream reads one pack per lane | ||
| // only for 2-bit, the fork's 1-bit kernel does too (mlx#3), so its block is | ||
| // 1024 values, not upstream's 2048. An exact gate keeps the partial-block tail | ||
| // of qmv_fast_impl unreachable; a stricter one (2048 for 1-bit) silently | ||
| // routed every K = 1024 mod 2048 1-bit matmul to the generic qmv. |
There was a problem hiding this comment.
Reduce this block to one or two lines stating that the calculation must match qmv_fast_impl. The current comment records upstream and fork behavior, an issue reference, old alignment values, and the prior failure mode, which is excessive branch history rather than concise guidance.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| // static_asserts the other widths away, and the host (quantized.cpp, | ||
| // mode = "affine_sym") only requests them for bits <= 2. Same name format as | ||
| // the affine set, so the non-JIT metallib serves the same names the JIT path | ||
| // builds: affine_sym_qmv[_fast]_<type>_gs_<gs>_b_<bits>_batch_<0|1>. | ||
| // (numen-tech/gemma4-qat#179, Codex P1) |
There was a problem hiding this comment.
Condense this block to a brief statement that these macros instantiate the 1-bit and 2-bit symmetric kernels. The host/JIT naming explanation and the issue/reviewer reference are implementation and review history that violate the repository requirement for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
… comments (PR #2, Codex P1)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aa25bdde5a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // 1-bit (scale=2d, values ±d): bias = -scale/2. 2-bit ternary (scale=d, | ||
| // values {-d,0,+d}): bias = -scale. Removes the bias stream (~10% of the | ||
| // weight-stream bytes at 1-bit/g128); same ALU (sum_x is computed anyway). | ||
| // bias_free is threaded through so the assert fires only when the symmetric | ||
| // path is actually selected — the ternary at the call sites instantiates this |
There was a problem hiding this comment.
Condense this block to one or two lines stating how the bias is derived. The performance estimate, ALU discussion, and template-instantiation history are excessive for this small helper and conflict with the repository’s requirement for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| // Wider load: read 4 packed bytes as one 32-bit word (1 load instead of 4 | ||
| // byte loads) and extract the same 32 one-bit weights. Little-endian bit | ||
| // order + accumulation order are preserved, so this is bit-exact vs the | ||
| // per-byte path while quartering the weight-load instruction count. |
There was a problem hiding this comment.
Reduce this block to a brief statement that the 1-bit path uses a 32-bit packed load. The instruction-count comparison, endianness explanation, and bit-exactness rationale make this explicit loop carry four lines of redundant implementation commentary, contrary to the repository’s concise-comment rule.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| // a few fp32 ulps of the result; 1e-4 relative is ~100x that. fp16: the | ||
| // kernel accumulates in fp32 and rounds the result once to fp16 (2^-11 | ||
| // relative), and x/scales are exact in both; 4e-3 relative is ~8x that. | ||
| // A wrong bias, a missed tail or a misread weight word moves outputs by | ||
| // O(scale * K^0.5) ~ 1 or more, far outside either bound. |
There was a problem hiding this comment.
Replace this seven-line numerical derivation with a one- or two-line statement that the tolerance covers accumulation-order and output-rounding differences. The ULP multipliers, fp16 rounding analysis, and expected failure magnitude are excessive inline test commentary under the repository’s documented standard.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| // MTLCompileOptions.mathFloatingPointFunctions is independent of mathMode | ||
| // and defaults to Fast, so a Safe/Relaxed request would still compile the | ||
| // fast (approximate) transcendental variants. Keep them Precise unless | ||
| // fast math was explicitly asked for. |
There was a problem hiding this comment.
Condense this block to one or two lines stating that non-fast modes require precise floating-point functions. The discussion of the independent API properties, default behavior, and resulting approximate implementation is longer than needed beside the explicit setter and violates the repository’s concise-comment standard.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
…d compile-option comments (PR #2, Codex P1)
1 participant
Part of numen-tech/gemma4-qat#86 (epic ml-explore#178, Phase 0): move the prism line from MLX 0.31.1 to 0.32.2.
Merge order: #1 (
fork/179-affine-sym-kernels->prism-0.31.1-fixes) first, then this PR intoprism-0.32.2-fixes.prism-0.32.2-fixes(e197191) is the new base branch for the 0.32.2 line: v0.32.2 (1f8e74e) + the fork deltas rebased fromprism(ed00ee2, the 0.32.1 line). After this merges, numen-tech/mlx-swiftprism-86-mlx-0.32.2re-points itsSource/Cmlx/mlxsubmodule at the merge result.What is on this branch (8 commits on top of
prism-0.32.2-fixes)quantized_matmul(biases=None, mode="affine")accepted for bits <= 2,eval_gpumaps it to modeaffine_symon the transposed small-M decode path, vjp/jvp guarded, CPU throwsprism-0.31.1-fixes/fork/179-affine-sym-kernels). d90771c is not onprism, so theprism->prism-0.32.2-fixesrebase never carried itmlx/backend/metal/kernels/quantized.h(ml-explore#179)set_compile_options: setmathFloatingPointFunctions= Precise unlessMathMode::Fastuse_qmv_wide()returns false foraffine_sym;test_qmv_affine_sym+test_qmv_affine_sym_throwsquantized.metal:instantiate_quantized_sym_{batched,types,groups}addaffine_sym_qmv_fast/affine_sym_qmv(bits 1 and 2, gs 32/64/128, float/half/bfloat, batched 0/1) to the metallib;qmv_fast_implpartial tail: lanes past the row end skip the row loop instead of masking it (no out-of-row weight reads)qmv_fast/qmv, neverqmv_wide/qmv_quadtests/ops_tests.cpp:test bias-free affine quantized_matmul decodedoctest (9 shapes x fp16/fp32, 54 assertions)qmv_fast_k_alignment(bits)=pack_factor * (bits <= 2 ? 1 : 2) * 32: the host fast gate follows the fork's 1-pack-per-lane 1-bit kernel (block 1024), where upstream 0.32.2 (ml-explore#3965) assumes 2 packs (2048); comments and test shapes updated for the exact gateBase branch:
prism-0.32.2-fixesvsprismrange-diff (git range-diff v0.32.1..prism v0.32.2..prism-0.32.2-fixes, 246 lines)All 12 fork commits survive: 8
=(identical), 4!(context-only rewrites: 9b9e0f0 "Add 1-bit affine quantization support" against 0.32.2'sdequantize<W,U,N,bits>float-cast intermediates, 54f772b, 4a8b667 fusedspec_decode_verify, dfe0540 docs, 332d459 "Route affine qmv_wide by bit-width and batch size"), and 065585d "Guard fast-path Metal kernel dispatch for 1-bit quantization" re-applied as 9a28ad1 on 0.32.2's restructureddispatch_qmv(range-diff pairs it as drop + add because the surrounding code moved). No fork behaviour dropped. The gen>=18 NAX gate (4fc445f) sits onis_nax_available()at its new location.Kernel interdiff vs 4aa785f (
fork/179-affine-sym-kernels)Kernel text is byte-identical except for what 0.32.2's
qmv()host forces on the two[[kernel]]entry points:affine_sym_qmv_fast/affine_sym_qmvtake the same six template parameters as 0.32.2'saffine_qmv_fast/affine_qmv(<T, group_size, bits, batched, has_global_scale = false, results_per_simdgroup = 4>), both unused on the affine paths exactly as inaffine_qmv.affine_qmvlayout (xat buffer 3 ...s_stridesat buffer 13) with the biases slot (buffer 2) left unbound, becauseqmv()now bindsxat 3 whether or not a biases / global_scale buffer is present (0.31.1 bound sequentially, hence 4aa785f's 13-buffer "biases slot skipped" layout).qdot1-bit uint32 wide load,sym_derived_bias, and thebias_freethreading throughqmv_fast_impl/qmv_implapplied without conflict.Host fast gate vs the fork's 1-bit block (7072c9a)
Upstream 0.32.2 (ml-explore#3965) made the
qmv_fastgate exact instead ofK % 512:qmv_fast_k_alignment(bits) = get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32, i.e. the K step ofqmv_fast_impl(pack_factor * packs_per_thread * SIMD_SIZE) with upstream'spacks_per_thread = bits == 2 ? 1 : 2. The fork's kernel reads one pack per lane for 1-bit as well (packs_per_thread = bits <= 2 ? 1 : 2, mlx#3), so its 1-bit block is 1024 values, and the upstream formula (2048 for 1-bit) silently routed every 1-bit matmul withK = 1024 mod 2048-- e.g. Bonsai-27B'sdown_proj, K = 21504 -- to the genericaffine[_sym]_qmv. The gate now usesbits <= 2. Either way the partial-tail block inqmv_fast_impl(kept from c431c50) is unreachable on this line: the gate is exact, so theK % 512remainder case of the 0.31.1 line (#1) does not arise; the doctest case comments say which entry point each shape reaches here.Decode microbench (M=1 fp16 gs=128, 1-bit,
mx.quantized_matmulwith and without biases, 300 iters, min of 3 runs,$SP/86-qmv1bit-bench.py), before = 2c2d8a4, after = 7072c9a. The after runs were taken while three unrelated xcodebuild suites were running on the machine (single runs varied by up to 35% on the shapes whose routing does not change), so only the minimum of three is quoted:Metallib (non-JIT) build
MLX_METAL_JIT=OFFconfigures and builds clean with Xcode 27.0's Metal toolchain on this line (185 MBmlx.metallib, no-std=metal4.0); the 0.31.1-line failure insteel/attn/mma.h/steel_attention.hdoes not reproduce on 0.32.2. The doctest above passes in both the JIT and the metallib build, which is what ab914a5's instantiations are for.Dispatch audit (0.32.2
QuantizedMatmul::eval_gpu)eval_gpuselects modeaffine_sym(or throws for bits > 2, non-transposed weights, or M >= the device's qmv batch limit) before theM >= vector_limitbranch, so neitherqmm_splitknorqmmcan see it.dispatch_qmvskipsqmv_quadforaffine_sym(no bias-free quad kernel; from d90771c) and, new here, skipsqmv_wide:use_qmv_wide()returned true for every mode other thanaffine, which under 0.32.2'sM >= 2 -> qmv_widerule would have sent a 2 <= M < vector_limitaffine_symmatvec to a non-existentaffine_sym_qmv_wide.affine_symnow always reachesqmv()and itsaffine_sym_qmv/affine_sym_qmv_fastentry points.Math mode (c731626)
MTLCompileOptionshas two independent knobs:mathMode(Safe/Relaxed/Fast) andmathFloatingPointFunctions(Fast/Precise, which transcendental implementations get emitted). MLX only setsmathMode, so on macOS 15+/iOS 18+ every JIT kernel compiled withMathMode::Safe/Relaxed(the default) still gets the approximatefast::exp/log/pow. The precompiled metallib is built with-fno-fast-math, so the Python wheel's built-in ops are precise while every JIT kernel (which is everything mlx-swift runs) is not. Measured withmx.fast.metal_kernel(a metal_kernel ulp probe), max ulp vs float64:metal::expmetal::logfast::exp)fast::log)precise::exp)precise::log)(
sinis 2.1 ulp in both namespaces on this GPU.) Filed upstream as ml-explore#4553; the patch keepsMathMode::Fastfast and makes Safe/Relaxed actually precise.Tests (Apple Silicon, JIT build of this branch)
python/tests/test_quantized.py(at 7072c9a): 41 passed, 3891 subtests passed (includestest_qmv_affine_sym: 1/2-bit x group 32/64/128 x K {64,128,512,1024,2048} x M {1,2,3,8} x N {256,67} fp32, plus fp16/bf16 spot checks; andtest_qmv_affine_sym_throws).tests/tests -tc='*quantiz*'(doctest): 2 test cases, 66 assertions passed in the JIT build and in the metallib build (test bias-free affine quantized_matmul decode: 54/54 in both).python/tests/test_fast.py -k "spec or sdpa or rope or rms or layer_norm": 19 passed, 76 subtests (covers the fork'sspec_decode_verify).python/tests/test_ops.pytranscendental subset (exp/log/sin/cos/tanh/erf/softmax/sigmoid/power/sqrt): 29 passed.Two test-design notes worth knowing (329b29e):
qmv_quad) and 2-bit M >= 3 on gen-15+ (qmv_wide, 332d459), so sym-vs-biased is compared in fp32 with a tolerance of 4 output-dtype ulps at the largest output magnitude (observed: 1 bf16 ulp at |y| ~ 140, 2-4 fp32 ulps at |y| ~ 300); the 1e-3 check against dequantize-then-matmul is unchanged.mx.evalthat throws never gets its completion event signaled; a later consumer on another stream (the CPU-stream throw case) then blocks forever inScheduler::wait_event. The throws test evaluates all inputs up front. This is pre-existing MLX behaviour, not introduced here.Consumers
prism-86-mlx-0.32.2) pinsSource/Cmlx/mlxat this branch head 7072c9a and is regenerated from it. After this PR merges, that pin moves to the merge commit andtools/update-mlx.shis re-run (expected no-op).phase0/86-mlx-0.32.2consumes it through mlx-swift.Codex review round (gemma4-qat ml-explore#86 stack, 2026-09-23)
Findings and disposition (full report in the gemma4-qat ml-explore#86 PR):
af01cc55d(this PR's first commit) wideningaffine_quantizeto accept group sizes 256/512 andquantized_matmulto acceptbiases=Noneon every backend; the packer, metallib instantiations and CUDAqmv.cuare byte-identical to upstream v0.32.2. Resolved by84b42176f: group sizes back to 32/64/128, bias-free affine rejects bits outside {1, 2}, CUDAeval_gputhrows "Bias-free affine is Metal-only for now." like the CPU backend, derived-biasvjp+test_qmv_affine_sym_grad.test_qmv_affine_symtolerated a "Bias-free affine" throw for any M > 1 (every tested shape sits under every architecture's vector limit, so the escape only hid failures). Resolved byff1deee55: the escape is gone andtest_qmv_affine_sym_throwsasserts the rejection at M = 33 and M = 512.Verified against the in-tree Python build on an M3 Max (applegpu_g15s):
python -m unittest test_quantized→ 42 tests OK, including the threeaffine_symtests. Not verified: CUDA compile, JIT-build variant, non-gen15 GPUs (reasoned from theget_qmv_batch_limittable).