Conversation
Port of the promoted MLX.fast Bonsai 2 27B submissions (Layr-Labs/mlxfast-bonsai2-27b-engine main, accept commits 5816911..9a1a6b9, Apache-2.0 per the challenge terms) to the Prism core: - quantized_utils.h: qmm_m16_block, a 16 x 32 packed-matmul core on the Metal 4 tensor unit for M <= 16 rows (polymorf). qmm_t_nax routes its 32-row tiles with at most 16 live rows and N < 65536 to it, with simdgroup pairs splitting K. - quantized_nax.h: M-first tile order for prompt widths (alvaroborras), a 2-bit register prefetch of the next K block, and the row-fit K-half split for verify-width tiles (i34-9 / DPZZxlz lineage). - quantized.h: qmv_wide_rr, a row-reuse form of qmv_wide for 2-bit weights; qmm_t_splitk_nax_impl, the split-K tensor body for FP32 and FP16 activations (DPZZxlz, fkiene), with the FP16 scratch-overflow fix from newjordan's tree (the one-barrier reduction only for an FP32 tile). Differences from the challenge tree: - affine_qmm_t_splitk gained a `use_nax` template flag. The host selects it from metal::is_nax_available(); the challenge only gated the body at compile time, and its author reports wrong rows on pre-M5 GPUs. Kernel names carry `_nax_true` / `_nax_false`. - The disabled (`&& false`) few-row route in the split-K entry is dropped. - quantized_nax.metal includes quantized_utils.h so the metallib build carries the few-row core, not only the JIT path. Verified on an M5 Pro: mx.quantized_matmul at the 27B tower shapes (K 5120/17408, N 5120..248320, M 1..512, FP16 and FP32) against a dequantized reference, worst relative error 1.0e-3 (FP16) / 3.8e-6 (FP32).
The gen >= 18 gate (PR #4, 2026-07-05) was added when this fork's NAX steel-gemm and qmm_t kernels computed wrong results on M5-class (g17) GPUs. At the current core (4fc91df, upstream merged 2026-09-21) that defect is gone: on an M5 Pro (applegpu_g17s) with the gate widened to the upstream rule, gen >= 18 for 'p' parts and >= 17 otherwise, - plain fp16 GEMM vs fp32 truth at M 8..512, N 8192..248320, K 5120/17408: worst relative error 5.0e-4 (fp16 rounding); - mx.quantized_matmul 2-bit gs128 at the 27B tower shapes, M 1..512, FP16 and FP32, vs a CPU reference: worst 1.1e-3 (FP16) / 4.1e-6 (FP32). The MLX.fast challenge core runs the same rule on its M5 ranked box and passes the organizer's goldens. Without this the M5 family never reaches qmm_t_nax or the few-row cores of the previous commit.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Non-affine split-K dispatch requests invalid kernels, and FP32 split-K bypasses the documented TF32 opt-out.
Review effort: Balanced
Findings: 1
Open (5)
What changed in this PR
Adds optimized Metal tensor-unit paths for few-row 2-bit affine quantized matmuls and restores generation-17 NAX support on eligible devices.
Changes:
- Adds few-row, row-reuse, and split-K tensor-unit kernels.
- Optimizes NAX tile ordering, weight prefetching, and K splitting.
- Updates NAX availability and metallib dependencies.
| File | Description |
|---|---|
mlx/backend/metal/quantized.cpp |
Selects split-K NAX variants. |
mlx/backend/metal/kernels/quantized.metal |
Instantiates split-K kernel variants. |
mlx/backend/metal/kernels/quantized.h |
Implements row-reuse and split-K kernels. |
mlx/backend/metal/kernels/quantized_utils.h |
Adds the shared few-row tensor core. |
mlx/backend/metal/kernels/quantized_nax.metal |
Includes shared quantized utilities. |
mlx/backend/metal/kernels/quantized_nax.h |
Adds prefetching, tile ordering, and K splitting. |
mlx/backend/metal/kernels/CMakeLists.txt |
Adds the new kernel dependency. |
mlx/backend/metal/device.cpp |
Restores the generation-specific NAX gate. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1238
to
+1239
| bool use_nax = metal::is_nax_available() && bits == 2 && group_size == 128 && | ||
| (x.dtype() == float32 || x.dtype() == float16); |
Comment on lines
+966
to
+968
| // EXPERIMENT (uncommitted): the MLX.fast challenge core's rule, gen >= 17 | ||
| // on desktop parts, so the g17 M5 Pro exercises the tensor kernels. | ||
| can_use_nax &= gen >= (arch == 'p' ? 18 : 17); |
Comment on lines
+2498
to
+2499
| // The tensor-unit body is only correct on generation-18+ GPUs (M5 class); | ||
| // the host selects `use_nax` from metal::is_nax_available(). |
Comment on lines
+150
to
+151
| instantiate_quantized_splitk_qmm(affine_qmm_t_splitk, type, group_size, bits, true, true) \ | ||
| instantiate_quantized_splitk_qmm(affine_qmm_t_splitk, type, group_size, bits, false, true) |
Comment on lines
+1235
to
+1237
| // The 2-bit gs128 split-K kernel carries a tensor-unit (NAX) body for | ||
| // few-row tiles. It is only correct on generation-18+ GPUs, so the host | ||
| // selects it; every other GPU takes the SIMD body. |
The split-K dispatch appended "_nax_*" to every kernel name and passed use_nax as a template argument, but only the affine kernels carry it. The fp modes (mxfp4, mxfp8, nvfp4) then asked for kernels that do not exist, e.g. "Unable to load kernel mxfp4_qmm_t_splitk_float_gs_32_b_4_alN_true _nax_false" in test_fp_qmm_non_multiple_of_32. Restrict the suffix and the argument to affine mode. FP32 input now takes the NAX body only when MLX_ENABLE_TF32 is on, like the other NAX dispatches in this file. With MLX_ENABLE_TF32=0 the 2-bit gs128 FP32 result goes from 7.9e-4 to 3.1e-7 relative error on an M5 Pro. Instantiate the _nax_true variants only for float and float16 at gs128, 2 bits, the only combinations the host selects (metallib 1.75 MB smaller), and replace the stale generation-18 comments with the gate that is actually enforced.
bri-prism
added a commit
to PrismML-Eng/mlx-swift
that referenced
this pull request
Sep 26, 2026
Picks up the split-K fix from PrismML-Eng/mlx#17: the fp modes (mxfp4, mxfp8, nvfp4) load their split-K kernels again, FP32 honors MLX_ENABLE_TF32 on the tensor-unit path, and the unreachable _nax_true variants are gone. The regenerated quantized, quantized_nax and quantized_utils sources also match the clang-formatted core headers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two Metal changes for the 2-bit, group-size-128 affine packed matmuls.
qmm_m16_block, a 16 x 32 packed-matmul core on the Metal 4 tensor unit for M <= 16 rows.qmm_t_naxroutes 32-row tiles with at most 16 live rows to it.qmm_t_nax: M-first tile order for 16-bit prompt-width calls, a 2-bit register prefetch of the next K block, and a K-half split for tiles with at most 16 live rows.qmv_wide_rr, a row-reuse form ofqmv_widefor 2-bit weights.qmm_t_splitk_nax_impl, the split-K body on the tensor unit for FP32 and FP16 activations.is_nax_available()follows the upstream rule again: generation 18 and up for 'p' parts, 17 and up otherwise. The gen >= 18 gate from metal: gate nax off for gen-17 devices (wrong gemm/qmm results on M5-class) #4 was added when this fork's NAX kernels miscomputed on M5-class GPUs; at the current core that defect is gone (numbers below), and without this change the M5 family never reachesqmm_t_naxor the new cores.Why
Speculative-decode verify rounds run the tower at 13 to 16 rows, where the existing kernels pay a device load and an FMA per weight per row. The challenge measured its decode gain there. The gate change is what lets M5 machines use the tensor unit at all.
How it differs from the challenge tree
affine_qmm_t_splitkgained ause_naxtemplate flag that the host selects fromis_nax_available(). The challenge gated that body at compile time only, and its author reports wrong rows on pre-M5 GPUs. Kernel names carry_nax_true/_nax_false.quantized_nax.metalincludesquantized_utils.h, so the metallib build carries the few-row core, not only the JIT path.Measurements (M5 Pro, 48 GiB)
Port vs unmodified 4fc91df, both with the gen-17 gate,
mx.quantized_matmul2-bit gs128 at the 27B tower shapes, best of the run in ms:Correctness, same build, against a CPU-stream reference: worst relative error 1.1e-3 (fp16) and 4.1e-6 (fp32) over M in {1, 4, 13, 16, 17, 32}, N in {5120, 17408, 248320}, K in {5120, 17408}.
Gate change: plain fp16 GEMM against fp32 truth on the tensor path at M 8..512, N 8192..248320, K 5120/17408, worst relative error 5.0e-4. One consequence to be aware of: with the tensor unit on, FP32 matmuls run at the tensor unit's precision class (about 8e-4 relative), which is what the stock MLX wheel does on the same hardware.
Review focus
The gate change reverses a deliberate decision from July. The evidence above is from one M5 Pro; an M5 Max or an A19 device would be a useful second data point before merging.