Repository navigation
Rebase the fork onto v0.32.4; qmv_quad ragged-tile scale/bias read (gemma4-qat#302) - #12
Conversation
…-aligned-K qmv_fast tail test (gemma4-qat#282) Patch 1 of the numen-tech/mlx rebase onto upstream main (06eb748), which carries upstream's final 1-bit affine quantization (ml-explore#3161, e0408d4). It replaces the fork's five pre-review ml-explore#3161 snapshot commits (2fb3d1a..20bf2eb) with only the deltas we still carry over upstream. - qdot<bits=1>: read the packed codes as uint32 words, 32 values per load, instead of upstream's byte-at-a-time loop. Same little-endian bit order and accumulation order, so bit-exact with the per-byte path. Provenance: mlx-swift bc96a51 "metal: uint32 wide load in 1-bit qmv qdot (large mobile decode win)", sourced in the fork by b7e6d4d (numen-tech/gemma4-qat#179). Every 1-bit qdot caller has values_per_thread == 32 and 4-byte aligned weights (qmv_fast, qmv, qmv_quad at D == 128); qdot_safe, qouter and dequantize keep upstream's byte path. - qmv_fast_impl packs_per_thread: upstream already took the fork's `bits <= 2 ? 1 : 2`; restore the fork's comment (#3). - qmv_fast_impl partial tail: keep upstream's lane predicate `aligned_end + simd_lid * vpt < K` and the partial_rows clamp from ml-explore#4516. The fork's `aligned_end + (simd_lid + 1) * vpt <= K` is equivalent whenever K % vpt == 0, which the host K gate guarantees. Upstream's comment there was stale (it said the 1-bit block is 2048; it is 1024 with 1 pack/thread); replace it with the fork's OOB reasoning, reworded for this predicate. - dispatch_qmv: document why the qmv_quad K == 64 route excludes 1-bit (a quad lane would hold 16 values, less than one uint32 pack). The fork's affine_sym fall-through comment lands with affine_sym (patch 2). Test: test_qmv_fast_tail_non_aligned_k (GPU vs CPU, 1/4-bit, K = 17/33 groups, N = 67/256). It passes before and after the qdot port. Note that on this base the qmv_fast tail is unreachable from the host (the 1-bit K gate is 2048 = 2 * block; other widths gate at block), so these shapes route to qmv. With an experimental (uncommitted) 1-bit K gate of 512 the tail runs, the ported qdot plus upstream's predicate matched the CPU on 1-bit K = 1536/2560 tail shapes, and a predicate letting one lane read past K failed them.
…-proof (gemma4-qat#282)
The comment claimed the shapes run qmv_fast's partial tail and the
partial_rows clamp. They do not on this base: the host K gate is 2048 for
1-bit and 512 for 4-bit, so every shape here routes to plain qmv (or
qmv_wide for M = 3 on gen 15+). Say so, and say when it becomes the tail
guard: once the routing patch widens the gate to 1-bit K % 512 and affine
4-bit K % 256 (gemma4-qat#192).
Add the shapes that the widened gate makes tail-reachable, keeping the
existing ones: 1-bit K in {1536, 2560} x gs {32, 64, 128} and 4-bit
K in {768, 1280} x gs {32, 64}, with N in {67, 256} and M in {1, 3}.
Same GPU-vs-CPU comparison and tolerances (atol = rtol = 1e-2). Skip
without Metal like the neighbouring tests; ruff-format.
…m main (gemma4-qat#179, ml-explore#282)
…m main (gemma4-qat#190, ml-explore#282)
…tream main (gemma4-qat#192, ml-explore#282) Port of numen-tech/mlx PR #8 (923aef6, 7d565ec) onto upstream main. qmv_fast_k_alignment takes the mode and returns half a kernel block for affine / affine_sym 1-bit and 4-bit (1-bit K % 512, 4-bit K % 256), the whole block otherwise. nvfp4 / mxfp4 share bits == 4 but fp_qmv_fast_impl has no tail, so they keep the whole block; every other width keeps upstream's routing. qmv keeps upstream's gate shape from ml-explore#4516 (fast = K-aligned, partial_rows = fast && N % 8 != 0); gather_qmv keeps N % 8 == 0. Tail predicate: kept upstream's lane predicate aligned_end + simd_lid * values_per_thread < K and the ml-explore#4516 partial_rows row clamp, not the fork's aligned_end + (simd_lid + 1) * values_per_thread <= K. The two agree whenever K % values_per_thread == 0, which the host gate guarantees (block / 2 = 16 * values_per_thread, and K % group_size already implies it). Evidence on M3 Max (applegpu_g15s): - GREEN: test_qmv_fast_tail_non_aligned_k and test_qmv_fast_half_block pass; a kernel-name probe (removed) showed the 1-bit K 1536/2560 and 4-bit K 768/1280 M = 1 subtests on affine_qmv_fast (pr_1 at N = 67). - RED: with the predicate loosened to `< K + values_per_thread` (one extra lane reads past K) both tests fail: 20/20 tail-reachable subtests of test_qmv_fast_tail_non_aligned_k and every tail subtest of test_qmv_fast_half_block (165 + 3 gather), errors up to ~78. No read past K shows with upstream's predicate, so the fork's stricter form is not needed; the OOB-reasoning comment stays. Tests: test_qmv_fast_half_block (fork's, with the generic baseline moved to K + gs zero-padded because N % 8 != 0 no longer leaves qmv_fast after ml-explore#4516; the fast call's x is a view with non-zero memory past K so the test is a tail guard; fp32 allows 4 ulps for accumulation order, measured <= 1.1), K = 768 in test_implied_bias_matmul, the "affine quantized_matmul qmv_fast tail block" doctest (renamed from the fork's "test affine qmv_fast tail block" so the *quantiz* filter runs it), updated bias-free doctest case comments. Also: the stale path="qmv_generic" subTest label is now "n_not_mult_8", and affine_sym_qmv[_fast] note their unused has_global_scale / results_per_simdgroup slots.
…w verify (gemma4-qat#191, ml-explore#282) Port of cddb7be (route affine qmv_wide by bit width and batch) and numen-tech/mlx PR #9 (071a92a..43e6613: run an affine qmv_wide verify of 6 or 7 rows as one tile) onto upstream main. - use_qmv_wide takes bits and M: affine 1-bit, and affine 2-bit at M < 3, stay on qmv; other affine widths take qmv_wide on gen 15+ as before; fp modes unchanged; affine_sym still has no qmv_wide kernel. - qmv_wide runs affine M <= 7 as one tile (vecs_per_tg = M); M >= 8 keeps upstream's ceil(M / 5) tiling. quantized.metal instantiates the affine and implied-bias qmv_wide kernels for vecs_per_tg 6 and 7. No conflict with ml-explore#4516, ml-explore#4629 or ml-explore#4641: none touches use_qmv_wide or the qmv_wide tiling, and ml-explore#4629's M1 batch limit only applies below gen 15, where affine never takes qmv_wide. Tests: test_qmv_wide gains 1-bit affine (full sweep and tiny shapes); test_qmv_wide_tile_matches_split (M = 6/7/8/11/12 bit-identical to the split calls, verbatim from the fork); the "affine quantized_matmul qmv_wide single tile of 6 and 7 vectors" doctest (renamed from the fork's "test affine qmv_wide single tile of 6 and 7 vectors" so the *quantiz* filter runs it). Routing comments in the tests updated for the 1/2-bit gate.
…#282) Process-wide Metal work counters (metal::counters() / metal::reset()): dispatches, commits, explicit stream synchronizations, and host waits on GPU work. Source: numen-tech/mlx PR #3, `git diff 08121fd 0f4257d` (5447675, aaefa50, 2e5d6cc, 1fee32b, 9cb5069). Conflicts with upstream main since v0.32.2: - fence.cpp: upstream dropped ~FenceImpl (fence is now an std::optional<allocator::Data>) and Fence::wait takes the target value as a parameter. The fork's `Device device` member and the host-wait count in the CPU spin-wait are kept, spinning on upstream's `value`. - tests/CMakeLists.txt: upstream renamed METAL_TEST_SOURCES to GPU_TEST_SOURCES; metal_counters_tests.cpp is appended to the new list. No new host-wait or dispatch site upstream: CommandEncoder:: dispatch_threadgroups/dispatch_threads and CommandEncoder::synchronize are still the only dispatch and waitUntilCompleted sites in mlx/backend/metal, and eval error recovery still flushes via gpu::synchronize(s, false). Tests: ./build/tests/tests -tc='*counters*' -> 4 cases, 26 assertions pass. test_quantized.py + test_fast.py: 102 passed, 2 skipped.
…gemma4-qat#189, ml-explore#282) fast::metal_kernel's closure keeps a per-kernel memo (shared by all copies of the closure, mutex-guarded) from each call variant (input dtypes and pass mode, output dtypes, template argument kind + name + value) to its kernel name, source and source hash. CustomKernel holds the source in a shared_ptr<const std::string> with a precomputed hash. The memo holds at most metal_kernel_max_cached_variants (64) entries and is cleared when full; entries are shared_ptrs, so a pending array whose entry was evicted still evaluates. Source: numen-tech/mlx PR #5, f5428f8 + b2c93bd. Merged with upstream ml-explore#4584 (908148e, `-` sanitized in make_template_hash): the memo calls make_template_hash unchanged when it builds a variant, so negative template values still give a valid kernel name, and the memo key encodes the value's bytes, so -1 and 1 stay separate variants. Both behaviors are kept; the only textual conflict was tests/CMakeLists.txt (upstream's METAL_TEST_SOURCES -> GPU_TEST_SOURCES rename). Tests: test_fast.py -k "custom_kernel or metal_kernel" -> 11 passed (incl. ml-explore#4584's test_custom_metal_kernel_negative_template_int and the fork's test_custom_kernel_many_template_values). ./build/tests/tests -tc='*metal kernel*,*counters*' -> 8 cases, 778 assertions pass. test_quantized.py + test_fast.py: 103 passed, 2 skipped.
…l-explore#282) Source: numen-tech/mlx 73db231 ("gate nax off for gen-17 devices", fork PR #4). M5-class gen-17 parts (applegpu_g17*) compute wrong results in the NAX steel-gemm and qmm_t kernels; the fallback bit-matches stock mlx. is_nax_available() now requires gen >= 18 on every device class (upstream: gen >= (arch == 'p' ? 18 : 17)). The unused `arch` local is dropped instead of the fork's `(void)arch`. The gate is one predicate. Audit of every NAX dispatch decision on upstream main 06eb748 + patches 1-7 (`git grep -n -iE nax mlx/backend/metal`, including kernels/*.h, jit/ and the JIT/no-JIT kernel getters), each site -> the predicate it asks: matmul.cpp:964 steel_matmul_axpby use_nax (thin ml-explore#4654 :995, NAX split-K :1050, regular fused :1087) -> is_nax_available() matmul.cpp:2964 gather_mm_rhs_nax -> is_nax_available() matmul.cpp:3043 segmented_mm use_nax (:3060, :3103) -> is_nax_available() quantized.cpp:1243 qmm -> qmm_nax (qmm_t_nax/qmm_n_nax, incl. the _ib and affine_sym variants from patches 2-3) -> is_nax_available() quantized.cpp:1522 gather_qmm -> gather_qmm_nax -> is_nax_available() quantized.cpp:1917 gather_qmm_rhs -> gather_qmm_rhs_nax -> is_nax_available() quantized.cpp:2468 global-scale gather_qqmm K alignment (ml-explore#4481; the matrix kernels then dispatch via :1917) -> is_nax_available() scaled_dot_product_attention.cpp:29 pre-M5 D512 non-NAX path (ml-explore#4518) -> !is_nax_available() scaled_dot_product_attention.cpp:412 sdpa_full_self_attention_nax -> is_nax_available() scaled_dot_product_attention.cpp:436 D72/D80 pad to NAX (ml-explore#4455) -> is_nax_available() scaled_dot_product_attention.cpp:964 D512 support check (ml-explore#4487) -> is_nax_available() scaled_dot_product_attention.cpp:1436 D512 fused eligibility (ml-explore#4487) -> is_nax_available() scaled_dot_product_attention.cpp:1445 D256 masked/causal fused NAX (ml-explore#3842, array mask ml-explore#4416) -> is_nax_available() scaled_dot_product_attention.cpp:1452 D256 non-NAX fused (ml-explore#4505) -> !is_nax_available() gated_delta_update.cpp:18, :24 chunk C = 16 (the only value that selects gated_delta_*_nax kernels, :133/:335/:369/:416) only when NAX (ml-explore#4020, ml-explore#4565) -> is_nax_available() No other path selects a NAX kernel: the get_*_nax_kernel getters (jit_kernels.cpp, nojit_kernels.cpp) are only called from the functions above, get_sdpa_vjp_nax_kernel has no caller, the MPP tensor_ops code in kernels/ lives only in *_nax.h/.metal, and nothing outside mlx/backend/metal mentions NAX. The remaining raw get_architecture_gen()/get_architecture() uses pick tiles or non-NAX kernels (qmv batch limits quantized.cpp:99-134, nvfp4 narrow qmv quantized.cpp:550, affine qmv_wide gen 15+ quantized.cpp:633, gemv stream matmul.cpp:1416, GEMM_TPARAM/devc tile tunings incl. ml-explore#4447). This host is an M3 Max (gen 15, no NAX), so the change is a no-op here; the gen-17 runtime check is the iPad M5 leg (gemma4-qat ml-explore#282 T9/T11). test_quantized/test_fast/test_blas/test_fast_sdpa: 162 passed, 5 skipped.
…#282) Source: numen-tech/mlx c731626 (applied unchanged). On macOS 15 / iOS 18+, set_compile_options set only mathMode, and MTLCompileOptions.mathFloatingPointFunctions (independent, default Fast) stayed Fast, so a Safe/Relaxed runtime compile (JIT kernels, compiled fusions, custom kernels) still emitted the approximate exp/log/sin/cos/pow. It is now Fast only for MathMode::Fast and Precise otherwise; the pre-macOS-15 branch keeps setFastMathEnabled. Checked against upstream ml-explore#4461 (365bd0f, "Use precise::exp in Sigmoid so compiled and eager sigmoid agree"): that fix only respells Sigmoid's exp as metal::precise::exp. It does not change compile options, and other runtime-compiled ops still spell unqualified metal:: transcendentals (binary_ops.h LogAddExp exp/log1p, Power pow, complex log/exp/sin/cos/ atan2, ...), so upstream covers one op of what this rule covers. Nothing is dropped. The two compose: an explicit metal::precise:: call stays precise under every mode, including Fast. Tests: test_ops.py -k "exp or log or sin or cos or tanh or erf or softmax or sigmoid or power or sqrt" -> 29 passed. test_compile.py (incl. ml-explore#4461's test_compile_sigmoid_matches_eager) + test_quantized.py + test_fast.py (incl. test_custom_metal_kernel_math_mode) -> 174 passed, 2 skipped.
…re#282) mx.fast.spec_decode_verify(draft_tokens [B, K], target_logits [B, K+1, V]) -> (n_accepted [B], committed [B, K+1]): greedy speculative verify. The target argmax runs as a regular op; the SpecDecodeVerify primitive does the per-row prefix match in a one-thread-per-row Metal kernel, with an op-composition fallback on CPU / no-GPU / CUDA. Source: 3e1d206 (v1 composition), 28efbe6 (fused Metal v2), 36030f6 (docs), 49840c9 (fuzz), e197191 (test moved into test_fast.py), taken as `git diff 3e1d206^ e197191` minus quantized.cpp / test_quantized.py (cddb7be in that range is patch 5). Applied unchanged, except: - add/add conflicts with upstream: kernels/CMakeLists.txt (after the sdpa_vjp / gated_delta build_kernel lines), no_gpu/primitives.cpp (after GatedDeltaUpdate[VJP]::use_fallback), test_fast.py (after ml-explore#4584's test_custom_metal_kernel_negative_template_int); both sides kept; - clang-format (fast.cpp, fast.h) and ruff 0.16.10 format (test_fast.py) of the new code, whitespace only. No caller in NumenKit or Python; carried as is (spec §2.10). Tests: test_fast.py -k spec -> 1 passed, 76 subtests. test_quantized.py + test_fast.py -> 104 passed, 2 skipped.
Comment-only (gemma4-qat ml-explore#282 ruling R12, Task 3 review minors): - test_qmv_fast_half_block: the "x is a view whose memory past K is non-zero" claim holds for B = 0 only. At B = 2 ensure_row_contiguous_matrix copies the (2, 1, K) slice, so past K batch 0 reads batch 1's x (still non-zero) and batch 1 reads past the buffer. - test_implied_bias_matmul: affine qmv_wide routing needs gen 15+.
…xplore#282) Codex round 1, P1 (pre-existing in the fork's 28efbe6). The fused kernel indexes dense rows (draft + row * K, target + row * (K + 1)), but eval_gpu bound its inputs as given, and the wrapper's astype(int32) keeps an int32 view a view. A transposed draft (logically [[1,2],[3,4]]) gave n_accepted [1, 0] on Metal vs [2, 2] on CPU; a broadcast [1,2] -> [2,2] draft made row 1 read past the stored elements; a column slice read the wrong elements. eval_gpu now copies a non-row-contiguous input (draft or target) into a dense temporary (contiguous_copy_gpu + add_temporary, as in gated_delta_update.cpp) before setting the kernel's pipeline. A row-contiguous input is returned untouched before the encoder is used: same path as before, no extra copy or dispatch. Targets come from argmax, which already returns dense rows from strided logits; the check covers them anyway. Test: test_spec_decode_verify_strided_inputs (transposed, broadcast, sliced and mismatching transposed drafts, plus a transposed logits view; Metal and CPU vs explicit expected results). RED (before): 4 failed -- every Metal draft-view subtest, e.g. case='transposed' path='default': Lists differ: [1, 0] != [2, 2] case='broadcast' path='default': Lists differ: [2, 0] != [2, 2] case='sliced' path='default': Lists differ: [2, 0] != [2, 2] case='mismatch-transposed' path='default': [1, 0] != [1, 2] (CPU and the transposed-logits subtests passed.) GREEN: test_fast.py -k spec -> 2 passed, 86 subtests passed.
…e#282) Codex round 1, P2 (pre-existing in the fork's 3e1d206). With K = 0 the CPU / no-GPU composition took min over an empty axis, which ops.cpp rejects, so draft [B, 0] with logits [B, 1, V] threw on stream=mx.cpu while the fused kernel returns a result. Per ruling R16, K = 0 is not rejected: the composition now returns what Metal returns, measured first: n_accepted = 0 for every row and committed [B, 1] = the target argmax token (Metal: B=1 -> [0], [[5]]; B=2 -> [0, 0], [[5], [3]]). Test: test_spec_decode_verify_empty_draft (B = 1 and 2; Metal and CPU vs the explicit result, shapes included). RED (before): 2 failed -- both CPU subtests: ValueError: [min] Cannot min reduce over axis 1 with size 0. (The Metal subtests passed.) GREEN: test_fast.py -k spec -> 3 passed, 90 subtests passed; CPU and Metal both give int32 [0] / [[5]].
…atching (ml-explore#282) Ruling R17 (fix round 1). With B = 0 (drafts [0, K], logits [0, K+1, V], K = 3 and K = 0) the op threw on every backend before reaching the fused kernel: `[argmax] Cannot argmax reduce zero size array.` (also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1, where nothing else is reported: the throw precedes any encode). The op now returns int32 zeros [0] and [0, K+1] right after the shape checks, with no argmax and no SpecDecodeVerify primitive, so nothing is dispatched. B > 0 inputs take exactly the previous path. Test: test_spec_decode_verify_empty_batch (K = 3 and 0; default device and CPU; shape and dtype of both outputs; default vs CPU equal). RED (before): 6 failed -- the 4 path subtests with the argmax ValueError, the 2 compare subtests with KeyError (no result to compare). GREEN: test_fast.py -k spec -> 4 passed, 96 subtests passed; both backends give n (0,) int32 and committed (0, K+1) int32, also clean under the Metal debug/shader-validation layers.
…kend (ml-explore#282) Found while running the R17 B = 0 check under the Metal validation layer (fix round 1; pre-existing in the fork's 28efbe6). With K = 0 and B > 0 the fused kernel was dispatched with an empty draft, which has no buffer, so index 0 was left unbound. Without validation the kernel never reads it (K = 0) and the result was right, but MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 aborts: validateComputeFunctionArguments:1043: failed assertion `Compute Function(spec_decode_verify): missing Buffer binding at index 0 for draft[0].' K = 0 now takes the op composition on every backend (n_accepted = 0, committed = the argmax bonus token, the K = 0 result from 52d3ef5), so the kernel is never bound to an empty draft. K > 0 is unchanged. Test: test_spec_decode_verify_empty_draft (from 52d3ef5), run under the validation layer. RED (before): pytest -k empty_draft under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 -> rc 134, the assertion above. GREEN: the whole -k spec suite under the validation layer -> 4 passed, 96 subtests (rc 0); without it -> 4 passed, 96 subtests.
Review feedback on #10: on Metal with B > 0 and K > 0, mx.export_function of a function calling mx.fast.spec_decode_verify threw "Unable to serialize primitive SpecDecodeVerify" because the PrimitiveFactory did not list it. Registered like the other fast primitives: state() already puts nullptr in the fallback slot, so the imported primitive has no fallback, as RMSNorm, RoPE and ScaledDotProductAttention do after import. The fallback is not needed there: eval_gpu never calls it, and the CPU, K = 0 and B = 0 paths are chosen when the op is built, so they export as plain op compositions. test_export_import_spec_decode_verify (full accept, reject at 0, batch of 2, K = 0, B = 0; imported results equal the eager call): RED 3 failed (the three B > 0, K > 0 rows), GREEN after the fix, also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1.
Review feedback on #10 (AGENTS.md code standards: 1-2 line comments). Comment-only. The block is reached only when the host's K gate (qmv_fast_k_alignment) is finer than block_size: affine 1/4-bit take K % (block_size / 2) == 0. Masking the out-of-row lanes instead (zeroed x_thread, scale and bias) was rejected: qdot would still run, and its weight reads land past the row end and, for the last output rows, past the end of the weight buffer. A K gate that did not keep in_vec_size a multiple of values_per_thread would need the stricter predicate aligned_end + (simd_lid + 1) * values_per_thread <= in_vec_size. The comment came in with 7c1362e (numen-tech/gemma4-qat#179, ml-explore#282).
…e#291) On an M3 Max (applegpu_g15s) the D512 sdpa_vector_2pass path from ml-explore#4459 decodes Gemma 4 E2B (8 query heads, 1 KV head) 3-4% slower than the unfused fallback beyond 1024 keys; forcing the fallback with MLX_SDPA_D512_MIN_KL recovers it. Upstream measured gains on M5 only. Keep the default on generation 17+; elsewhere use the vector kernel only when MLX_SDPA_D512_MIN_KL is set.
…tile (gemma4-qat#302) qmv_quad runs ceil(N / 64) tiles with no N % 64 restriction. On a ragged last tile, the quad rows past out_vec_size loaded their scale (and, for an explicit bias, their bias) before the row predicate, reading past the end of the scales/biases buffers. The qdot and the store were already guarded, so the output was never wrong. Move both loads inside the existing predicate; the pointer arithmetic stays outside (forming a pointer reads nothing). v0.32.4's qmv_quad_impl has the same unguarded structure, so this is not an upstream change taken over. test_qmv_quad_ragged_n covers N = 65 (one valid row in tile 2) and N = 72 (eight), at bits 1/2/4/8, gs 32/64, K 64/128 (1-bit K = 128 only), with an explicit and an implied (0-d factor) bias, against a CPU dequantize reference. affine_sym never reaches qmv_quad. Validation: with MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 MTL_SHADER_VALIDATION_GLOBAL_MEMORY=1, the unfixed kernel reported nothing, at those shapes nor at fp32 K = 128 gs = 32 N = 1025. A deliberately out-of-bounds custom kernel (1M elements past a 16-element input) was not reported either, so GPU validation does not see out-of-bounds device loads in this setup; the fix rests on the bounds argument above. The test passes before and after the fix (an allclose cannot see a discarded read).
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: 33ef381f3b
ℹ️ 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".
| // Fused greedy speculative-decoding verify. Inputs: draft_tokens int32 [B, K] | ||
| // and target_tokens int32 [B, K+1] (the target argmax). Outputs: n_accepted | ||
| // [B] int32 and committed [B, K+1] int32. The op-level fallback (a plain | ||
| // composition) runs on CPU and for autodiff; the Metal eval_gpu runs a trivial | ||
| // per-row prefix-match kernel. |
There was a problem hiding this comment.
Condense this five-line comment to the required one or two concise lines. It repeats the input/output types already expressed by the class API and adds implementation detail available in the methods, contrary to the repository’s explicit standards for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 4c4d61c: the class comment is one line now.
| bool is_equivalent(const Primitive& other) const override; | ||
| std::vector<Shape> output_shapes(const std::vector<array>& inputs) override; | ||
| auto state() const { | ||
| return std::make_tuple(nullptr); |
There was a problem hiding this comment.
When a Metal trace containing this primitive is exported and imported, this state reconstructs SpecDecodeVerify with an empty fallback_. Applying mx.vmap to that imported function—an explicitly supported imported-function transformation—then reaches Custom::vmap, invokes the empty function, and throws std::bad_function_call instead of vectorizing. Serialize enough state to rebuild the fallback or provide a transform override that does not depend on it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 4c4d61c. The composition now takes B and K from its input shapes, and SpecDecodeVerify builds its fallback from the stream alone, so an imported primitive has a working fallback. state() is gone; the serialized form is unchanged, because nullptr wrote no bytes. test_export_import_spec_decode_verify has a vmap case: it threw std::bad_function_call before the fix, and it passes now.
…#282) Review feedback on #12: an imported SpecDecodeVerify had an empty fallback (state() put nullptr in its slot), so mx.vmap over an imported function threw std::bad_function_call from Custom::vmap. The composition now takes B and K from its input shapes, and the primitive builds its fallback from the stream alone, so import needs no state. The serialized form does not change: nullptr wrote no bytes. Also condenses the class comment to one line. test_export_import_spec_decode_verify vmap case (vmap of the imported function equals vmap of the eager call): RED std::bad_function_call, GREEN after the fix. test_fast + test_quantized + test_vmap + test_export_import 167 passed, 2 skipped; C++ JIT and metallib 311/311.
…Verify fallback after import)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4c4d61c5c2
ℹ️ 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".
| auto j = broadcast_to( | ||
| reshape(arange(K, int32, s), Shape{1, K}, s), Shape{B, K}, s); |
There was a problem hiding this comment.
When spec_decode_verify is wrapped in mx.compile(..., shapeless=True) on CPU, these Shape{B, K} values and arange(K) are fixed using the first invocation. Reusing the compiled function with another draft length therefore builds comparisons and outputs using the old K, causing shape errors or incorrect results; construct this fallback without embedding trace-time dimensions, or specialize the compiled graph by shape.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1e1f993 (test gate in c134ee2). The K > 0 CPU case threw 'Slice cannot infer output shapes' on the first call rather than giving wrong results. A K = 0 first trace did give wrong results, on both CPU and Metal: it returned n_accepted = [0] for every later shape. SpecDecodeVerify now has a native eval_cpu, and the Metal kernel takes K = 0, so every B > 0 call builds the one primitive, whose output shapes come from its inputs. B = 0 keeps its early return, because argmax rejects a zero-size input. CUDA GPU streams still use the composition. test_spec_decode_verify_shapeless_compile traces K > 0 first and K = 0 first, over 5 shapes, on CPU and on Metal: 13 subtests failed before the fix, all pass now.
| // 1-bit: read 32 packed weights as one 32-bit word (one load instead of | ||
| // four byte loads). Little-endian bit order and accumulation order are | ||
| // those of the per-byte path, so the result is bit-exact with it. Needs | ||
| // values_per_thread % 32 == 0 and a 4-byte aligned w, which every 1-bit | ||
| // caller (qmv_fast, qmv, qmv_quad with D == 128) provides. |
There was a problem hiding this comment.
Condense this six-line implementation note to one or two lines; it repeats the load behavior and caller constraints in excessive detail, contrary to the repository's explicit requirement for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in cbbbe2f: the comment is two lines now.
…l-explore#282) Review feedback on #12: the CPU path and the K = 0 path built the op composition, whose shapes (Shape{B, K}, arange(K)) are fixed at trace time. Under mx.compile(shapeless=True) a K = 0 first trace then returned n_accepted = [0] for every later shape on CPU and Metal, and a K > 0 CPU trace threw "Slice cannot infer output shapes". SpecDecodeVerify now has a native eval_cpu and the Metal kernel takes K = 0 (target is bound at the empty draft's slot; the kernel reads no draft for K = 0), so every B > 0 call is the one primitive, whose output shapes come from its inputs. The composition stays as the transform fallback and for CUDA GPU streams. B = 0 keeps its early return (argmax rejects a zero-size input). test_spec_decode_verify_shapeless_compile (cpu and gpu, K > 0 and K = 0 traced first, 5 shapes each): RED 13 failed, GREEN after the fix, also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1. test_fast + test_quantized + test_vmap + test_export_import + test_compile 239 passed, 2 skipped; C++ JIT and metallib 311/311.
Review feedback on #12 (AGENTS.md: concise comments). Comment only; test_quantized 73 passed, 1 skipped.
…ith Metal (ml-explore#282) Codex on 1e1f993: mx.stream(mx.gpu) throws on a CPU-only build, and a CUDA GPU stream still takes the shape-dependent composition. CPU coverage stays unconditional. test_fast -k spec_decode 6 passed.
…e_verify primitive for every K and on CPU; qdot comment)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c134ee26d0
ℹ️ 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 (B == 0) { | ||
| return {zeros(Shape{0}, int32, s), zeros(Shape{0, K + 1}, int32, s)}; |
There was a problem hiding this comment.
When a shapeless-compiled function is first traced with B == 0, this trace-time branch records two fixed empty zeros results with no dependency on the inputs. Shapeless compilation reuses that graph for later inputs of the same rank and dtype, so a subsequent call with B > 0 still returns empty arrays instead of verifying its tokens. Route empty batches through a shape-dynamic primitive or otherwise avoid branching on the concrete trace-time batch size.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 560cd02. B = 0 now takes the same primitive path as every other B, so one shapeless graph serves every batch size. The argmax op rejects any zero-size input, so the target ids now come from an ArgReduce (ArgMax) primitive built directly, and V = 0 is rejected up front with a clear error.
The B > 0-first order had a second bug, which only the Metal validation layer showed. Replaying that trace at B = 0 dispatched a zero-size grid in both ArgReduce and SpecDecodeVerify. Both evals, CPU and Metal, now return early on an empty output.
CUDA GPU streams still use the composition and its B = 0 early return. That path bakes B and K in as constants, which predates this change; a comment now says so. test_spec_decode_verify_shapeless_compile now also traces B = 0 first and B > 0 first, on CPU and on Metal. Its 10 failing subtests now pass.
Results: Python 965 passed, 57 skipped; C++ 311/311 on Metal, 311/311 on Metal JIT, 274/274 on CPU-only; the new test is clean under MTL_DEBUG_LAYER=1.
ml-explore#282) The B = 0 early return recorded two constant empty results, so a shapeless compile first traced with B = 0 returned empty arrays for every later B. B = 0 now takes the same primitive path as any B: the target ids come from an ArgMax primitive built directly (the argmax op rejects any zero-size input), V = 0 is rejected up front, and SpecDecodeVerify's CPU/Metal evals and Metal ArgReduce return early on an empty output instead of dispatching a zero-size grid (a B > 0-first trace replayed at B = 0 tripped the Metal validation layer). CUDA GPU streams keep the composition and its B = 0 early return; that path is not shapeless-safe for a change of B or K.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 560cd028a5
ℹ️ 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".
| // Rows past out_vec_size exist only on a ragged last tile (N % 64 != 0): | ||
| // read no scale or bias for them (gemma4-qat#302; the store below was | ||
| // already guarded). |
There was a problem hiding this comment.
Condense this three-line note to one or two lines. It repeats the bounds predicate and guarded load behavior visible directly below and includes issue-tracking detail, contrary to the repository requirement for concise, nonredundant code comments.
AGENTS.md reference: AGENTS.md:L31-L32
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7b7e8d2. The comment is now two lines, without the issue reference.
| if (inputs.size() < 4) { | ||
| throw std::runtime_error( | ||
| "[QuantizedMatmul::eval_cpu] Bias-free affine is Metal-only for " | ||
| "now; pass biases (e.g. derive_biases) on CPU."); |
There was a problem hiding this comment.
When a 1/2-bit bias-free affine matmul is evaluated on CPU, this error directs users to derive_biases, but a repo-wide search finds no definition or binding with that name, so the suggested recovery cannot be followed through MLX. Recommend an actual supported operation for materializing the biases or omit the nonexistent helper name.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 9482135. derive_biases doesn't exist. The CPU kernels already accept a 0-d implied-bias factor, so the error now names it: -0.5 for 1-bit and -1 for 2-bit. Those values give the same bias that the Metal bias-free path derives in-kernel (sym_derived_bias). Checked on random 1- and 2-bit weights: on Metal, the bias-free and implied-factor results are bit-identical; on CPU, the implied factor and explicitly materialized biases are bit-identical. test_quantized.py and test_fast.py pass: 112 tests, 2 skipped.
…explore#282) derive_biases does not exist. The CPU kernels take a 0-d implied-bias factor, and -0.5 (1-bit) / -1 (2-bit) gives the same bias the Metal bias-free path derives (sym_derived_bias).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b7e8d2b1f
ℹ️ 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".
| // draft [B, K] int32 — drafted tokens | ||
| // target [B, K+1] int32 — target argmax tokens over [last] + draft | ||
| // Writes: | ||
| // n_accepted [B] — longest matching prefix length (0..K) | ||
| // committed [B, K+1] — draft prefix (< n) then the corrected token at n. |
There was a problem hiding this comment.
Condense this six-line header to one or two concise lines. The argument names and types directly below already express most of the input/output layout, so the current block is redundant and exceeds the repository's explicit comment-length standard.
AGENTS.md reference: AGENTS.md:L31-L32
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c52689c. The header is now two lines.
| if (fallback) { | ||
| return spec_decode_verify_composition(prim_inputs, s); |
There was a problem hiding this comment.
On a CUDA GPU, a function wrapped in mx.compile(..., shapeless=True) records this composition with the trace-time B and K; a later call with another batch size or draft length reuses fixed slice, arange, broadcast_to, and zeros shapes, causing shape errors or stale results. In particular, tracing first with B == 0 records the constant-empty return above, so every later nonempty batch remains empty; use a shape-dynamic primitive on CUDA or specialize the compiled graph by shape.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7474d2c, by refusing a shapeless trace on that path rather than making it shape-dynamic. Making the CUDA path shape-dynamic needs a SpecDecodeVerify CUDA kernel, and this fork has no CUDA build or test machine to verify one on. Until it has one, the CUDA fallback throws a clear invalid_argument under detail::in_dynamic_tracing(), so mx.compile(..., shapeless=True) on a CUDA GPU stream now fails loudly instead of returning stale or empty results. That covers the B = 0-first trace too. A plain mx.compile and eager calls are unchanged.
Metal and CPU never reach this path: they take the primitive for every B and K, and test_spec_decode_verify_shapeless_compile already covers them. The guard itself can't be exercised here, because the fallback is CUDA-only.
| // Block of qmv_fast_impl (kernels/quantized.h): pack_factor * | ||
| // packs_per_thread * SIMD_SIZE, with one pack per lane for bits <= 2 (must | ||
| // match the kernel). The affine kernel's partial last block only needs | ||
| // K % values_per_thread == 0, which K % group_size already implies; block / 2 | ||
| // is the cutoff measured for 1/4-bit, and the other widths stay whole-block |
There was a problem hiding this comment.
Condense this seven-line tuning note to one or two concise lines. It mixes the expression's derivation, kernel synchronization requirements, benchmark history, and behavior of a separate FP implementation, exceeding the repository's explicit standard for concise inline comments.
AGENTS.md reference: AGENTS.md:L31-L32
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c52689c. The note is now two lines: the block must match the kernel, and affine 1/4-bit takes half a block.
| if (!variant) { | ||
| std::string kernel_name = "custom_kernel_" + name; |
There was a problem hiding this comment.
When multiple threads request the same previously unseen variant concurrently, they can all miss the initial locked lookup and independently execute write_template, write_signature, source allocation, and hashing before the second lock discards every result but one. This defeats the new build-once memoization precisely for concurrent first use, potentially multiplying large source-generation work by the number of racing callers; coordinate construction per key rather than doing it outside the protected section.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 713bf19. A missing variant is now built while cache->mutex is held, so concurrent first callers build it once and the rest reuse it. The lock covers only source generation and hashing: the Metal library is compiled later, in CustomKernel::eval_gpu, so the extra serialization is small. Python suite: 965 passed, 57 skipped. C++: 311/311 on Metal, 311/311 on Metal JIT, 274/274 on CPU-only.
…ition (ml-explore#282) The CUDA GPU composition fixes B and K at trace time (slice, arange, full, zeros), so a shapeless compile replayed at another B or K gave shape errors or stale results (a B = 0-first trace kept returning empty arrays). It now throws under detail::in_dynamic_tracing(); Metal and CPU take the primitive for every B and K.
…e#282) Concurrent first callers of one variant each built its source and hash before the second lookup kept one. The build now runs under cache->mutex, so it happens once. It covers source generation only; the Metal library is compiled later in CustomKernel::eval_gpu.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 713bf19828
ℹ️ 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".
| bool fast::SpecDecodeVerify::use_fallback(Stream s) { | ||
| return s.device == Device::gpu; | ||
| } | ||
| NO_GPU_MULTI(fast::SpecDecodeVerify) |
There was a problem hiding this comment.
When a function is exported on Metal, SpecDecodeVerify is serialized as a primitive; importing that file in a CUDA build maps its GPU stream to CUDA and evaluates the primitive directly. This NO_GPU_MULTI implementation then throws without consulting use_fallback, even though a newly traced CUDA call supports the operation through its composition. Implement the composition in CUDA evaluation or export a backend-portable graph so Metal-produced .mlxfn files remain usable on CUDA.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not changing this one, because exported fast primitives are backend-specific upstream too. An import keeps each primitive and its device and never swaps in the composition. Upstream's GatedDeltaUpdate is the same case: use_fallback is true on CUDA and its CUDA eval_gpu is NO_GPU_MULTI, so a Metal-exported gated_delta_update throws on CUDA in the same way. Fixing it needs a SpecDecodeVerify CUDA kernel, and this fork has no CUDA build or test machine to verify one on. Until it does, the import fails with a clear error (fast::SpecDecodeVerify has no CUDA implementation.) rather than a wrong result. Exporting on CUDA records the composition and works there.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 713bf19828
ℹ️ 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".
| // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused | ||
| // primitive does only the tiny per-row prefix-match. Primitive inputs are | ||
| // int32 token arrays, so the kernel needs no large-V reduction. The ArgMax | ||
| // primitive is built directly because the argmax op rejects any zero-size | ||
| // input, B = 0 included; V > 0 is checked above. |
There was a problem hiding this comment.
Condense this five-line comment to one or two lines. It repeats details visible in the direct ArgReduce construction and mixes the optimization rationale with the empty-batch workaround, contrary to the repository requirement for concise, nonredundant comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in bb8fa5f. The note is now two lines: why ArgMax is built directly, and what the fused kernel does.
| // affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. | ||
| inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { | ||
| return mode != "affine" || d.get_architecture_gen() >= 15; | ||
| // affine_sym has no qmv_wide kernel. The 1-bit path has so little weight | ||
| // traffic that unpacking into registers dominates even when several input | ||
| // rows reuse the result. The 2-bit path breaks even after two rows and wins | ||
| // once three or more rows share a block. |
There was a problem hiding this comment.
Condense this five-line tuning note to one or two lines. The branches below already express the architecture, bit-width, and row-count gates, while the comment adds detailed performance rationale that exceeds the repository's concise-comment standard.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in bb8fa5f. The note is now two lines and keeps only the measured gates.
| // 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 four-line implementation note to one or two lines. The assignment below already shows the Fast versus Precise selection, so the current block repeats implementation detail and exceeds the repository's explicit concise-comment standard.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| // The generated source depends on the dtypes of the inputs and outputs | ||
| // and on how each input is passed (see `write_signature`). Include them | ||
| // in the kernel name so that a given name always maps to the same | ||
| // source. |
There was a problem hiding this comment.
Condense this four-line comment to one or two lines. The dtype and input-location suffix construction immediately below already shows what distinguishes generated sources, so the block repeats the implementation and exceeds the repository's concise-comment standard.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in bb8fa5f. The note is now two lines. The same commit condenses every other source comment in this PR that was longer than two lines.
1 participant
Rebases the fork onto ml-explore/mlx
v0.32.4and adds one kernel fix on top:qmv_quadread scales and biases for rows past N on a ragged last tile (gemma4-qat#302).Base
prism-v0.32.4-fixes=v0.32.4=36969c09(PyPImlx0.32.4). This branch is that tag plus 23 commits: the 22 fork patches fromprism-main-06eb7483-fixes(06eb7483..317df531), replayed in order with their original authorship and messages, then patch 23.The 22 replayed patches
All 22 cherry-picked cleanly.
git range-diff 06eb7483..317df531 36969c09..ee0b1bcfcshows 20=, 2!, no<or>.The two
!are patch 1 (metal: 1-bit uint32 wide-load qdot over upstream #3161; ...) and patch 20 (metal: condense the qmv partial-last-block comment). Both differ only in a context line. Upstream ml-explore#4675 rewrote thein_boundsline under the fork's comment tostatic_cast<int>(simd_lid). The patch content is unchanged.Overlap review
Each file that both
06eb7483..v0.32.4and the fork touch:mlx/backend/metal/device.cpp: ok. [BUG] Metal fence deadlock ml-explore/mlx#4689 addsend_encoding(wait_in_flight), which callsencoder_->waitForFenceon in-flight output fences. That is a GPU-side MTLFence wait: no host wait, no commit, and the pre-existingwaitForFencewas never counted either. The fork's counter sites are untouched and still sit on their events:dispatch_countindispatch_threadgroups/dispatch_threads,commit_countafterbuffer_->commit(), andsync_count+count_host_waitbeforewaitUntilCompletedinsynchronize.is_nax_available()keeps the gen >= 18 gate.mlx/backend/metal/device.h: ok. Upstream'send_encoding(bool wait_in_flight = false)and the fork'ssynchronize(bool explicit_sync = false)/count_host_wait()are independent declarations.mlx/backend/metal/fence.cpp: ok. [BUG] Metal fence deadlock ml-explore/mlx#4689's GPU path callsend_encoding()before thefence_waitspin dispatch andend_encoding(true)after it. The spin kernel goes throughdispatch_threads, so it is counted once, as before. Neither new call commits or host-waits. The fork's CPU-pathcount_host_wait()(f.device == gpu) is unchanged.mlx/backend/metal/quantized.cpp: ok.qvm_split_ktakex/outby value and collapse batch dims to [M,K]/[M,N] when B == 1. It sits before the fork's implied-biasb_strides, which are unaffected: those are zero strides sized froms_strides, and scales/w are not reshaped.qvm_split_kis reached only for!transpose && K >= 1024, which is disjoint from the fork'sdispatch_qmv/use_qmv_widerouting.transpose &&from thegather_qmm_rhsNAX branch. That branch asksmetal::is_nax_available(), the fork's gen-17-gated predicate.mlx/backend/metal/kernels/quantized.h: ok. Fix signedness warning in quantized Metal tail bounds check ml-explore/mlx#4675'sstatic_cast<int>(simd_lid)lands inside the fork's partial-last-blockin_boundspredicate, and the fork's comment still describes it.mlx/ops.cpp: ok. Fix missing stream in gather matmul indices cast ml-explore/mlx#4687 (astype(..., uint32, s)inindices_or_default) and Adding minmax method to MLX to follow Numpy/Jax API ml-explore/mlx#4674 (minmax) are disjoint from the fork'svalidate_mode/quantized_matmul/affine_dequantize/gather_qmmimplied-bias hunks. The fork'sgather_qmmstill reachesindices_or_default, now with the stream fix.mlx/ops.h: ok. Upstream'sminmaxdeclaration and the fork'squantize/dequantize/quantized_matmulsignature hunks are disjoint.python/src/ops.cpp: ok. Upstream'sarctan2/radians/degreesscalar bindings (python: Make arctan2/radians/degrees support scalar inputs ml-explore/mlx#4679) and theminmaxbinding are disjoint from the fork's binding hunks.python/tests/test_quantized.py: ok. Fix split-K test compatibility on M5, M6 ml-explore/mlx#4677 (M 16 -> 25 in the split-K test) and fix batched row count mismatch in metal qvm_split_k ml-explore/mlx#4673'stest_quantized_matmul_batched_split_kare disjoint from the fork's added tests. All of them pass.NAX site list
The fork's gate lives inside
metal::is_nax_available()(device.cpp:1004-1022,can_use_nax &= gen >= 18;), so every site that asks that function is gen-17-gated. Fromgit grep -n -iE "nax" mlx/backend/metal/*.cppatee0b1bcfc:matmul.cpp:964:steel_matmul_axpbyuse_nax(thin :995, NAX split-K :1050, regular fused :1087)matmul.cpp:2964:gather_mm_rhs_naxmatmul.cpp:3043:segmented_mmuse_nax(:3060, :3103, :3149)quantized.cpp:1252:qmm->qmm_nax(has_nax_kernel)quantized.cpp:1532:gather_qmm->gather_qmm_naxquantized.cpp:1927:gather_qmm_rhs->gather_qmm_rhs_nax, now also non-transposed after Enable non-transposed NAX gather_qmm_rhs ml-explore/mlx#4684. It is gated, so it needs no fix.quantized.cpp:2478: global-scalegather_qqmmK alignmentscaled_dot_product_attention.cpp:29(the pre-M5 D512 path, negated), :412, :436, :964, :1447, :1456, :1463 (negated)gated_delta_update.cpp:18, :24 (chunk C = 16 selects the_naxkernels at :135/:339/:375/:422)jit_kernels.cpp/nojit_kernels.cpp:get_*_nax_kernelgetters only, called from the functions aboveEvery one of these goes through
is_nax_available(). No site picks a NAX kernel with a raw architecture check, so no NAX fix commit was needed. The test host is an M3 Max (gen 15), so the NAX paths are checked statically only.Patch 23: qmv_quad reads no scale or bias for rows past N (gemma4-qat#302)
qmv_quad(K = 128, or K = 64 at bits >= 2; power-of-two bits; notaffine_sym) runs ceil(N / 64) tiles with no N % 64 restriction. On a ragged last tile, the quad rows pastout_vec_sizeloaded their scale, and their bias when it is explicit, before the row predicate, so they read past the end of the scales and biases buffers. Theqdotand the store were already guarded, so the output was never wrong. The patch moves both loads inside the existing predicate. The pointer arithmetic stays outside, since forming a pointer reads nothing.v0.32.4'sqmv_quad_implhas the same unguarded structure, so nothing here comes from upstream.The new test is
test_qmv_quad_ragged_n. It runs N = 65 (one valid row in tile 2) and N = 72 (eight) at bits 1/2/4/8, gs 32/64 and K 64/128 (1-bit at K = 128 only), with an explicit and an implied (0-d factor) bias, against a CPU dequantize reference. That is 56 subtests.affine_symnever reachesqmv_quad;test_qmv_affine_symcovers it on its own path.Validation outcome. Metal validation (
MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 MTL_SHADER_VALIDATION_GLOBAL_MEMORY=1) could not see the read:mx.fast.metal_kernel(it reads 1M elements past a 16-element input) was not reported either. "Metal API/GPU Validation Enabled" was printed each time.Tests (at the patch-23 head, M3 Max)
pytest python/tests/test_quantized.pypytest python/tests/test_quantized.py python/tests/test_fast.py./build/tests/tests(MLX_METAL_JIT=ON)./build-lib/tests/tests(metallib)At the bare
v0.32.4tag,test_quantized.py+test_fast.pygave 84 passed, 2 skipped. The extra tests are the fork's own.Review
Codex round (whole
v0.32.4..HEADdiff, source-only): no fork-introduced defects; the 22 replays match the old base; no adverse interaction with ml-explore#4689/ml-explore#4673/ml-explore#4675/ml-explore#4672/ml-explore#4687/ml-explore#4660/ml-explore#4004; gen 17 stays off every NAX path. One inherited P2 in upstream v0.32.4 (ml-explore#4684: the non-transposed NAXgather_qmm_rhschecks only tile rows before loading scales/biases/weights for N not a multiple of the tile, reachable on gen 18+ only) is tracked as a gemma4-qat follow-up, not fixed here