mlx: opt-in dense GEMM for nvfp4 prefill (dequantise to bf16 above a row threshold) - #287
glennneuber wants to merge 21 commits into
Conversation
|
Runner measurements are in, and one of them is a blocker. Please do not merge this as-is. Flag off vs on, same binary (
Criterion 4 fails. The description's "one transient copy of one layer, ~230 MB" is wrong for a lazily evaluated graph: MLX builds a chunk's whole forward pass before evaluating, so the dequantised copies of many layers are live at once (~460 MB per layer for the 31b MLP pair, 60 layers). The long-text rows hide it because their peak is set by the attention transient at a different moment. Admission (#276) prices weights + KV + a per-architecture headroom. Several GiB of unpriced prefill transient on a ~40 GiB budget is exactly the failure mode the prefill-transient task documents, so this must not be enabled anywhere, even opt-in, until the copies are bounded (force evaluation per layer, or dequantise into one reused buffer). The code comment and the task doc now say so. Criterion 2 is inconclusive, not passed: think-off T1 with the flag on differs from the flag-off cells in 7 of 35 quality cells, three of them large and all on Criterion 3 is met. Criterion 1 is met for the ratios at ≥ 2048 rows; the crossover is now partly answered by the 1134-row image cell, which gains 2.13×, so the 1024 threshold is not obviously wrong. |
|
The control is in and criterion 2 fails too. This should not be merged as an enable-able flag; it needs the bounded-buffer rework first.
Each configuration is bit-reproducible on this arm; the two differ far outside that spread. So the 7 of 35 T1 cell differences are all the flag, not noise — exactly what a changed accumulation order should do: deterministic, and different. Two of three models score higher with it on and one lower, so there is no evidence of a systematic quality change, but "output-preserving" is false, and criterion 2 as written cannot be met by any threshold. This path has to be judged as a numerics change with the vision suite as its gate, not as a free optimisation. Scorecard: 1 ☑ (ratios; crossover partly answered by the 1134-row image cell at 2.13×) · 2 ✗ changes output deterministically · 3 ☑ 1.16–2.13× image, 1.17–1.37× long text · 4 ✗ up to +6.8 GiB peak · 5 ☑ no change with the flag unset. Prior art worth folding in before the rework. llama.cpp already does this and bounds it correctly: Neither ollama nor llama.cpp has solved the FP4 GEMM question: two llama.cpp issues (improve NVFP4 GEMM speed on sm120; generalise MMQ to floating-point data) went stale unsolved, and the live work is an open PR quantising activations into NVFP4 (+40 % PP dense) whose author is explicit that accuracy still needs input scales and a Hadamard transform. "NVGEMM" itself is a PyTorch Inductor backend name, not a concept in either project. |
|
Held, not merged (2026-09-10), while #285, #286, #288 and #289 went in. Two of the five acceptance criteria fail and both are properties of the approach, not of the threshold: the path changes model output deterministically (each flag state is bit-reproducible, the two states differ far outside that spread) and it raises per-request peak by up to 6.8 GiB on gemma4:31b, which admission does not price. Merging an inert-but-enable-able flag with those properties invites someone to turn it on. Upstream v0.34.0 does not touch The rework that would make it mergeable is known: dequantise into a reused, pool-allocated buffer the way |
…row threshold) On the RTX PRO 6000 MLX's mixed-input nvfp4 kernel is slower than plain bf16 cuBLASLt once a matmul carries a prefill chunk's worth of rows: 55-73 TFLOP/s against 100-160 at 4096 rows on the MLP projections. Dequantising the layer to bf16 per call and using the dense GEMM recovers 1.4-2.1x at >= 2048 rows with the error identical to today's kernel (same weights, bf16 math), at the cost of one transient copy of the layer. QuantizedLinear.Forward takes that path when OLLAMA_MLX_PREFILL_DEQUANT_ROWS is set, the mode is nvfp4, the device is CUDA and the activation carries at least that many rows; off by default, decode never qualifies. Policy and env parsing unit-tested; the branch body tested against the mixed-input kernel on a GPU (2080 Ti, in the image). The task doc carries the bench tables (Blackwell beside the campaign, 2080 Ti idle), the sizing against the production request shape (prefill is 4 % of request time there, so the opt-in is worth 1-2 % on gemma4 and nothing on qwen; it pays on long prompts and batched prefill), the acceptance criteria (quiet-GPU crossover, T1 parity, prefill tok/s, peak memory) and the FP4 GEMM as layer 2 behind an accuracy gate. Bench records committed under preflight/runs. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…4.6x at >= 2048 rows, repeatable Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… GiB peak Flag off vs on on the same binary: image prefill gains 1.16-2.13x and long text 1.17-1.37x, larger than the microbenchmark suggested. The memory result is the one that matters: the runner's per-request peak rises up to 6.8 GiB on gemma4:31b. The PR's first claim of one transient layer copy was wrong -- MLX evaluates a chunk's forward pass as one graph, so the dequantised copies of many layers are live at once. Admission does not price that, so criterion 4 fails and the flag must not be enabled until the copies are bounded. Criterion 3 met; criterion 2 inconclusive pending the name_bbox repeat control. Probe records committed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…iterion 2 fails document_single x3 per flag state, one container each: every run within a state is identical (spread 0), and the two states differ far outside it — 12b 0.714 off / 0.622 on, qwen3.8 0.542 / 0.697, qwen3.6 0.504 / 0.613. So the T1 differences are the flag, not noise, and the path is a numerics change rather than an output-preserving optimisation. Two models score higher with it on and one lower, so no systematic quality change, but criterion 2 as written cannot be met. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The prefill loop evaluates a whole chunk as one graph and sweeps once at the end, and a live Go handle retains its buffer through that eval, so every layer's dequantised copy stayed alive for the chunk -- the +6.8 GiB measured on gemma4:31b. mlx.Release frees a handle the caller exclusively owns; the graph keeps its own reference, so the buffer returns to MLX's allocator once that layer's matmul has run and the next layer reuses it. That is what llama.cpp gets from the pool-allocated scratch buffer in ggml_cuda_mul_mat_cublas_impl, expressed in MLX's model. Release is idempotent and panics on a pinned array. Tests cover both, and that a released input does not disturb the node that consumes it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Same probe on the reworked binary: peak with the flag on is 37.7 GiB against 37.5 before, still +7.0 over the flag-off control, which itself reproduced the earlier run to 0.1 GiB on all four shapes. So the buffers are not retained by the Go handle, and the mechanism this task stated is wrong. mlx.Release stays (correct and tested) but is not the fix. The per-layer Eval experiment that distinguishes graph-lifetime retention from 'these are not the buffers' is running. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
a75cdeb to
6c75037
Compare
…ort does not work Attempt 2 at bounding the dense path's transient: evaluate each layer's matmul inside denseGEMM so its dequantised copy cannot outlive it. On gemma4:31b's first one-image request the peak went to 64.7 GiB and the request died, against 37.7 without it; the flag-off control in the same run was unchanged. So the lazy whole-graph evaluation was helping, and the transient is neither the Go handles nor something eager evaluation can bound. llama.cpp's bound comes from ggml owning a reused pool buffer and evaluating node by node; MLX's Go API has no out-parameter matmul, no donation control and no scratch arena, so the fix does not port. What is left: price the transient in admission per architecture, or take a fused dequantise-matmul upstream to MLX. Both recorded, neither is this PR. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…own kernel Implements the streaming kernel the 'make it fast upstream' option assumes -- 4-bit weights unpacked in registers, bf16 activations, bf16 tensor cores -- and benchmarks it against bf16 cuBLAS on the same dequantised weights, autotuned over 24 configurations per shape. It reaches 23-33% of cuBLAS, below MLX's own 40-60%. So it does not measure the ceiling; it refutes the cheap hypothesis that MLX's kernel is badly written. Correctness checked against the dequantised reference (3.2e-3, bf16 rounding). Mixed-input is capped at bf16-dense speed by construction, since the 4-bit tensor cores need both operands in 4-bit. The 4-bit path already exists for sm120 in CUTLASS and is worth more, but its blocker is accuracy. The kernel work is not where the leverage is. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…an borrow Checked in the MLX source: on Metal the quantised GEMM is the dense steel template with a QuantizedBlockLoader swapped in for the weight operand, so it inherits every dense tuning; on CUDA it is a standalone CuTe kernel while dense is cuBLAS, so it inherits nothing and competes against decades-tuned kernels. That, not the format, is why nvfp4 keeps pace with q4_K_M on Metal and lags on CUDA -- and the Triton probe supports it, since a competent from-scratch kernel landed below MLX's own. Borrowable: build the CUDA quantised path on a CUTLASS collective with a dequantising loader (CUTLASS's Hopper mixed-dtype example is that shape) rather than hand-rolling, and parameterise one loader by mode instead of a kernel zoo behind capability gates. Caveat: Blackwell's higher FLOP-per-byte makes parity-with-dense a harder target than on Metal. The Metal half of Glenn's premise is not verifiable on this host and is labelled as such. Also records that the Triton probe needs a quiet-GPU rerun. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…or the Triton probe Production MLX is built and JIT-compiles with CUDA 13.0 and targets sm_120a natively; its Ampere-era MMA and cp.async are source-level choices no toolkit changes. The Triton prototype was compiled by Triton's bundled ptxas-blackwell 12.8.61, the first sm_120 release, so the hypothesis is plausible there and testable: the host has a CUDA 13.0 ptxas and Triton reads TRITON_PTXAS-BLACKWELL_PATH (hyphenated, so via env). The A/B is staged for the quiet-GPU rerun. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… host's ptxas vLLM pins weight-only NVFP4 to Marlin on RTX 5090/PRO 6000 because native FP4 GEMM needs FP4 activations, and measured ~31% prefill loss when a W4A4 checkpoint was misrouted to a dequantise-to-16-bit kernel. On our exact card, CUTLASS ollama#3096 needed compute_120f from CUDA 13.0 to get native NVFP4 MoE correct, and Marlin still beat it. MLX ollama#4339 reports gather_qmm ~24x slower than dense on consumer Blackwell, which our two MoE models use. Local ptxas 12.8 and 13.0 both reject wgmma and tcgen05 on sm_120a, so no toolkit unlocks Hopper or datacenter instructions here; sm_120f exists only from 13.0, the one version-gated item. The card allows ~100 KiB shared memory per SM, which bounds tile size. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…nel route viable CUTLASS builder composition is impossible on sm120, but the tuned TMA mainloop underneath is generic over the MMA atom; a ~580-line derivative with register-side e2m1->bf16 dequant (branch feat/sm120-mixed-input-gemm) is numerically exact and reaches 217-238 TF/s at M=2048, 0.65-0.84x cuBLAS bf16 in-run. Probes show the loader swap is free and the whole gap is conversion ALU on the MMA warps; next step is producer-warp dequant. Corrects this doc's earlier 'kernel work is not the leverage', which was drawn from a naive Triton prototype. States the MLX comparison as a within-harness ratio (~1.3-1.8x), not the agent's cross-harness ~3x, and names the veto. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… verified Moving the e2m1->bf16 conversion onto idle producer warps makes the MMA warps run the dense loop (1.97 non-MMA instructions per HMMA vs 10.73 for v3) and beats v3 by 8-14% on every row; 0.71-0.92x cuBLAS bf16. In-kernel cycle accounting shows the conversion mostly overlaps MMA (0-7% cost) and the main loss was pipeline latency, fixed by a third TMA stage that only fits with a shared-memory-free epilogue. Reads MLX's native weight layout with no repack. Numerically exact. Corrects 'mixed input tops out at bf16-dense speed': the cap is tensor-core throughput, and with less data per tile this design ran above the dense kernel when the conversion was compiled out. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…rness Triton 3.3's do_bench defaults to the mean, which skewed the first table under contention and made it incomparable with the sm120 kernels, which were timed with the median. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…put change Any kernel that accumulates in a different order from MLX's qmm changes model output deterministically, as the document_single control showed for the dequant path. The sm120 kernels avoid the full-size weight copy but not that; they need the same vision-suite judgement. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
On an idle GPU0 the producer-warp kernel sits at 0.86-0.91x cuBLAS on all eight rows (spread 0-1%) and ~1.8x MLX's qmm; the harnesses now agree on cuBLAS, so earlier gaps were contention and Triton's default mean. Stall attribution reframes the next step: the kernel is not issue-bound (29% of issue slots) and the conversion is off the critical path, so trimming it buys little; the measurable losses against dense are the register-to-global epilogue (STG, lg_throttle) and ring synchronisation. UIADD3 stalls are compiler padding, present in dense too. Triton script: TRITON_FORCE_PTX_VERSION pins ptx_version, working around Triton 3.3.1's missing CUDA 13 branch; the rerun uses identical PTX 87. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…quiet window The corrected assembler A/B fed byte-identical PTX 8.7 to ptxas 12.8.61 and 13.0. 13.0 produced different machine code that runs at 0.97-1.00 of 12.8's speed, at the same 0.24-0.26x of cuBLAS, so the toolkit is not what limits the prototype. The quiet rerun is audited against production's request log: two bursts landed on this GPU during the chain and one during the A/B. Steps 1 and 3 and the Nsight profiles fell in the gaps. The MLX bench step's caveat and phase A's standing are recorded where their numbers appear. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
….5 % on every row) The third agent's kernel (feat/sm120-mixed-input-gemm 556c211) runs the TMA epilogue inside the ring entry the MMA warps read last and lets them skip the TMA-barrier poll. In three quiet runs it beats the previous best by 1.6-6.5 % on all eight rows, the runs agreeing within 0.4 %, at 0.87-0.95x cuBLAS and 0.94-1.01x the dense sm120 kernel, exact on every shape. The tables are rendered from the run logs by a script. Also recorded: what was tried and dropped, the Nsight before/after, the unverified memory-model step, and the recommended next step (MLX integration above a row threshold, gated by the vision suite). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…caveat resolved In a clean window (production unloaded, no requests, no other GPU process) the new kernel runs at 285-324 TF/s, 2.1-6.8 % ahead of the previous best, which reproduces its 16:45 figures within 1.3 %. Against MLX's qmm from a clean rerun of the #286 bench it is 1.6-1.95x by ratio to each harness's cuBLAS. The rerun also answers the 16:47 caveat: every row within 10 %, the quoted rows within 5 %, no mark from the burst. Corrections: the gate M=2048 spread is the first round after an idle GPU, not one kernel's slow mode; the agent's per-run production counts came from a window 10 hours early (a timestamp without Z, which docker reads as local time). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…writes compute-sanitizer racecheck finds 0 hazards in the previous best, the epilogue-only build and the new kernel on two shapes. Two positive controls bound what that means: racecheck reports all 30 races of a build whose MMA warps skip both waits, but every one is a ring store against an MMA-warp load and none involves the TMA loads, and it misses a real race (tail barrier removed, 4 % of elements wrong at production shapes). So the ring hand-off is verified and the skipped TMA wait is not; it still rests on the memory-model argument. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Blocked on a design decision, not on a mechanical conflictTried merging current The two branches apply the nvfp4 global scale in different places. This branch ( out = mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
...
out = quantizedLinearOutputScale(out, ql.GlobalScale) // raw m
out := mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode, ql.GlobalScale)
// ops_extra.go: Mul(out, DivScalar(scale, Nvfp4MaxProduct))
// quant.go: MulScalar(flat, Nvfp4MaxProduct)So resolving Two measurements bearing on it, both from #312/#310:
To be explicit about what is not wrong here: this branch does apply the global scale on the dense-GEMM path too ( Not merging this on my own, because the resolution picks a global-scale representation for the whole fork and that was explicitly reserved. Everything else is green — all 11 checks pass on this branch. Once the representation is settled, the merge is small: take the chosen |
ADR 0037 (proposed): a measured model is identified by its manifest digest, not its tag. gemma4:31b-nvfp4 was re-published with a bf16 vision tower — 194 layers different, config blob unchanged, ollama show silent — and the old artifact cannot be fetched back, so the store is the archive. Records cite digests, benchmark pulls go to a separate store, and a cross-host comparison runs the audit on both hosts first. ADR 0038 (proposed): nvfp4 global scales are stored as the checkpoint multiplier and converted to MLX's m x 2688 form only at the call sites that hand the scale to MLX. The stored-MLX-form round trip f32(f32(m x 2688)/2688) misses m by one ulp for 17 of 31b's 191 vision scales, which is the whole of the CUDA-side encoder move in #312; the three wrappers that apply the scale themselves — QuantizedMatmul, Dequantize, GatherQMM off Metal — all pay it. #287's prefill dequantisation runs through Dequantize, so it inherits the fix; it stays held on its other grounds. SPEC vision-harness-reuse H15: a reported model is named by its manifest digest, and a slice is named by what it contains. The survey's claim that a 200-item OCRBench slice is comparable to a model card is corrected in place: the set is ordered by task, so rows 0-200 are four of ten categories. The ladder doc now states which checkpoint produced the mlx-cuda 0.860 (this store's 4-bit-tower copy, pulled 2026-08-17, digest 637cc0ff1570), what makes the Metal 0.875 comparable to it, and the one command that settles it. The MLX bf16 arm is recorded as refused by admission — 75.0 GiB asked against 62.7 GiB available with the 16 GiB reserve — not forced. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Parked 2026-09-19 at the maintainer's request. The blocker hasn't changed: resolving Converting to draft so it doesn't read as review-ready. All 11 checks were green on this branch when it went quiet; the merge itself is small once the representation is decided (see the earlier comment). |
|
Correction to the park note above: the global-scale ADR landed as ADR 0039, not 0038. When #323 merged, the manifest-digest ADR took 0038 and The park itself is unchanged: this PR waits on that representation decision. |
This branch has not been deployed
1 participant
The cheaper answer to the FP4 GEMM question, as an opt-in, with the sizing.
Change.
QuantizedLinear.Forwarddequantises the nvfp4 layer to the activation dtype and runs the dense cuBLASLt GEMM whenOLLAMA_MLX_PREFILL_DEQUANT_ROWSis set, the mode isnvfp4, the device is CUDA and the activation carries at least that many rows. Otherwise, and by default, today's mixed-input kernel. Global-scale and bias handling shared. One transient copy of the layer per call (~230 MB for a 31b MLP projection), nothing resident. Decode never qualifies.Why. Measured with the bench from #286 on the RTX PRO 6000 (beside the campaign, so only the M ≥ 2048 rows carry signal) and on the 2080 Ti (idle): at 4096 rows today's kernel reaches 55–73 TFLOP/s on the MLP projections against 100–160 for bf16; the dequant-then-GEMM path is 1.4–2.1× faster than today's kernel at ≥ 2048 rows with the error identical in every cell. The native FP4 GEMM (
qqmm) is 3.6–4.7× at the cost of activation quantisation; that is layer 2, behind an accuracy gate. Turing gets nothing from either.Sizing (in the task doc). Production requests (245 in the log tail): prompt p50 3531 tokens, generation p50 2136, prefill 4 % of request time (p90 5 %). On the MLX path the same shape gives prefill shares of 5 % on gemma4:31b down to 0.5 % on qwen3.8, so the opt-in is worth about 1–2 % per request on gemma4 and nothing measurable on qwen. It pays where prefill dominates: long text prompts on qwen3.8 (41.6k tokens in 31 s, matmul-heavy), and batched prefill. Off by default until the acceptance criteria are met.
Verification. Env parser and policy table unit-tested; the branch body tested against the mixed-input kernel on the same nvfp4 weights, run on a GPU inside the image (2080 Ti): PASS.
go test,go vet,golangci-lintclean. Not yet run in the runner: criteria 1–5 in the task doc (quiet-GPU crossover, T1 parity on the five nvfp4 models, prefill tok/s on 1-image / 3-image / 41.6k-token prompts, runner peak memory, no change with the flag unset).Bench records committed under
vision-suite/preflight/runs/like the calibration data.🤖 Generated with Claude Code