feat(quant): complete op × quant-format matrix — 30 Track-1 formats across every weight-bearing kernel - #250
Merged
Conversation
ekryski
changed the title
Foundation for the spec-conformant nvfp4/mxfp4/mxfp8/nvfp8 formats
(BENCH_METRICS_SPEC.md Appendix B). Pure-Rust encode/decode mirroring the
pinned MLX fp4.h/fp8.h bit math (E5M2 per OCP = the high byte of an IEEE
half), to be the single source of truth shared by the host weight packer,
the CPU correctness oracle, and the math the Metal kernels emit — so a
kernel that disagrees with the oracle is a real bug, not an oracle mismatch.
- E2M1 fp4 codebook {0,.5,1,1.5,2,3,4,6} + midpoint-threshold encode
- E4M3 fp8 (PyTorch-style, ±448 saturation, no inf)
- E5M2 fp8 (IEEE-half high byte, wide exponent range)
- E8M0 pow-2 block scale (2^(bits-127))
- round-trip + saturation + codebook unit tests
QFormat {Nvfp4, Mxfp4, Mxfp8E4, Mxfp8E5, Nvfp8} with pack() (f32 weights →
per-block codes + scales) and dequant() (the CPU correctness oracle).
Built on the codec primitives so the oracle and the eventual Metal kernels
share one spec-exact reference.
- Per-block amax scaling along K; two-level global FP32 for nvfp4 so its
E4M3 micro-scales fit ±448. Scale stored as E8M0 (mxfp4/mxfp8), E4M3
(nvfp4), or raw FP32 (nvfp8).
- pack quantizes against the *effective* (post-encode) scale, so codes +
stored scale are self-consistent (no oracle drift).
- e5m2_encode now saturates finite inputs (a block scale that rounds down
can't turn a weight into inf and poison a matmul).
- Tests: pack→dequant cosine ≥ per-format floor (block quant is judged by
direction, not per-element worst case), byte-layout sizes, zero blocks.
First spec-conformant block-scaled dequant kernel: mt_mxfp4_dequant<T> (E2M1 elements, block 32, E8M0 pow-2 scale). Reads packed nibble codes (8/u32) + a u8 E8M0 scale per block, reconstructs out[i] = e2m1(code) * 2^(scale_bits - 127). - Freeze-safe Grid3D elementwise dispatch (one thread/element, `if i<n` guard) — not a reduction kernel, so no n_simd==0 hazard. - The decode mirrors quant::codec, and the #[test_kernel] expects quant::format::dequant — kernel vs spec-exact host oracle, no drift. Passes on GPU (f32/f16/bf16) via the kernel_tests_harness. - Establishes the pattern the remaining formats (nvfp4 E4M3-scale, mxfp8 E4M3/E5M2, nvfp8 FP32-scale) follow.
…tests Completes the standalone block-scaled dequant kernels (all 5 formats now GPU-verified against the spec-exact host oracle): - mt_nvfp4_dequant — E2M1, block 16, E4M3 micro-scale × global FP32 - mt_mxfp8_e4m3_dequant / mt_mxfp8_e5m2_dequant — 8-bit element, block 32, E8M0 pow-2 scale - mt_nvfp8_dequant — E4M3, block 16, per-block FP32 scale E4M3/E5M2 elements decode arithmetically ((1+mant/2^mb)·2^(exp-bias) with a subnormal branch), which equals codec.rs's half-reinterpret decode — so the kernels match quant::format::dequant exactly. All freeze-safe Grid3D elementwise. The shared #[test_kernel] setup binds codes as u32 (4-bit) or uchar (8-bit), scales as u8 (E8M0/E4M3) or f32 (nvfp8), and the nvfp4 global as a constexpr. 15 GPU checks (5 formats × f32/f16/bf16) pass.
First block-scaled *matmul* kernel: mt_mxfp4_qgemv<T> computes output[row] = Σ_k dequant(weight[row,k]) · input[k] for E2M1+E8M0 weights. Reuses the proven pack-strided Reduction geometry from ffai/dequant_gemv verbatim (one TG/row, TPG≥32 multiple-of-32, reduce_sum, tid==0 store) — so no new dispatch shape and no new reduction freeze-hazard surface; only the per-element decode (E2M1 codebook × E8M0 scale) replaces the int-affine one. #[test_kernel] oracle = dequant-then-dot via quant::format, with the input round-tripped through the dtype so GPU and oracle agree. Passes on GPU (f32/f16/bf16). Establishes the Phase B pattern the other formats + qmm/qvm variants follow.
…PU tests Completes the block-scaled dequantizing-GEMV set — all 5 formats now do output[row] = Σ_k dequant(weight[row,k]) · input[k] and pass on GPU (f32/f16/bf16) against the dequant-then-dot oracle: - mt_nvfp4_qgemv — E2M1 / block 16 / E4M3 micro-scale × global (pack-strided) - mt_mxfp8_e4m3_qgemv / mt_mxfp8_e5m2_qgemv — 8-bit element / block 32 / E8M0 (element-strided, one byte per code) - mt_nvfp8_qgemv — E4M3 / block 16 / per-block FP32 scale All reuse the proven Reduction geometry (TG/row, reduce_sum). The 4-bit formats pack-stride over u32 words; the 8-bit formats element-stride over uchar codes. Shared #[test_kernel] setup binds weight u32/uchar and scales u8/f32 per format, nvfp4's global as a constexpr.
… GPU tests Adds the multi-row matmul counterpart of the GEMVs: output[m,n] = Σ_k dequant(weight[n,k]) · x[m,k] for mxfp4 / nvfp4 / mxfp8(e4m3,e5m2) / nvfp8. Each (m,n) output element is one threadgroup reducing over K, flattened to a 1-D grid (tg → mr=tg/out_dim, n=tg−mr·out_dim) so it depends only on program_id::<0>() + the proven reduction geometry — no new dispatch shape. Completes the core matmul family: dequant + qgemv + qmm × 5 formats = 15 GPU-verified kernels (45 checks across f32/f16/bf16) against the quant::format dequant-then-matmul oracle.
Per-token expert-routed matmul: output[m,n] = Σ_k dequant(weight[expert_ids[m], n, k]) · x[m,k] for mxfp4 / nvfp4 / mxfp8(e4m3,e5m2) / nvfp8. The qmm kernel with one indirection — the expert stack is one [E·out_dim, in_dim] packed tensor, so the weight row is expert_ids[m]·out_dim + n (keeps nvfp4's single global valid across experts). Same proven reduction geometry. #[test_kernel] oracle = gather-dequant-matmul via quant::format; passes on GPU (f32/f16/bf16). 20 block-scaled kernels now verified: dequant + qgemv + qmm + gather_qmm × 5.
The M≥32 ALU-throughput path, adapted from mlx/quantized.rs::mt_qmm_mma (int4 affine MMA). Dispatch geometry, threadgroup-memory layout, 8×8 frag mapping, and the 4-k-inner × 4-frag MMA loop are copied verbatim — only the per-pack W dequant staging changes to E2M1 codebook × E8M0 pow-2 scale, and the scales param becomes u8 E8M0 (no biases). Reusing the proven geometry keeps it off the reduction freeze-hazard surface; correctness is the only risk and is gated by the #[test_kernel] oracle. Establishes the block-scaled MMA template (32×32 tile, 128 tpg, KernelMode Reduction). Passes on GPU (f32/f16/bf16) vs quant::format qmm oracle.
Add first-class `e2m1_decode` / `e4m3_decode` / `e5m2_decode` unary intrinsics so block-scaled dequant is terse in kernel source and the compiler emits the decode math — instead of a ~24-line `select` ladder inlined at every call site. A macro can't factor this cleanly (inner-macro-in-#[kernel] hazard + macro hygiene), so making it a DSL op is the right fix. - UnaryOpKind::DecodeE2m1/E4m3/E5m2 (op.rs) lower to emitted `mt_decode_*` preamble helpers (preamble.rs), gated by per-op KernelFeatures flags (features.rs) — same mechanism as erf/expm1/erfinv. - `#[bench]`/`#[kernel]` body parser maps the names (macros body.rs). - type_check: decode ops always yield f32 regardless of the integer operand (yields_f32), so `e2m1_decode(nib)` is f32. - The emitted helpers mirror quant::codec bit-for-bit, so kernel output stays identical to the host oracle. - Adopt across all 21 block-scaled kernels (dequant/qgemv/qmm/MoE/MMA × formats): every decode ladder collapses to a one-liner. `e4m3_decode` serves both fp8 elements and nvfp4 micro-scales. All 21 GPU correctness tests still pass.
Completes the 4-bit MMA pair. Same verified geometry as mt_mxfp4_qmm_mma (programmatically duplicated to avoid transcription error); the only delta is the block scale — E4M3 micro-scale × a global FP32 (e4m3_decode intrinsic) instead of the E8M0 pow-2 scale — plus the `global` constexpr. Passes on GPU (f32/f16/bf16) vs the quant::format qmm oracle.
…GPU tests Completes the block-scaled MMA path for all 5 formats. The 8-bit variants reuse the verified mxfp4 MMA geometry but swap the W-dequant staging to element-strided (1 byte/code → 8 contiguous bytes per lane, mirroring the X-load) instead of 4-bit pack-striding; weight is Tensor<u8>, decode via e4m3_decode/e5m2_decode, scale E8M0 (mxfp8) or per-block FP32 (nvfp8). Geometry (frags, MMA loop, writeback) is byte-identical to the proven kernel, so no new freeze surface. All pass on GPU (f32/f16/bf16) vs the quant::format qmm oracle. Matmul family now complete: dequant + qgemv + qmm + MoE gather_qmm + MMA × 5 formats = 25 GPU-verified block-scaled kernels.
Register #[bench]es for the 5 block-scaled qgemv kernels at the canonical N=K=4096 decode shape, with .flops(2·N·K) so tile bench's GFLOP/s + roofline columns rank the precisions side by side — the spec's "which precision is fastest" goal. Throughput is data-independent, so packed weight/scale buffers use random bytes (matching the int dequant_gemv benches). Confirmed on M1 Max: 4-bit (mxfp4/nvfp4) ~100-240 GFLOP/s beats 8-bit (mxfp8/nvfp8) ~83-111 (¼ vs ½ the weight bytes); all read latency-bound at low %BW — the scalar one-row-per-TG geometry is under-utilized, exactly the kind of insight the bench-metrics work surfaces. 15/15 correct.
Register #[bench]es for the batched-decode GEMM (qmm, m=32) and the prefill simdgroup-matrix GEMM (qmm_mma, m=n=k=4096) across all 5 formats, with .flops(2·M·N·K) so tile bench's GFLOP/s + roofline rank the precisions on the compute-bound path (where the MMA throughput + the M5 NA story matter). Random packed buffers (throughput is data-independent). All 15 block-scaled benches (qgemv + qmm + qmm_mma × 5) codegen + pass no-undefined-symbol checks.
nvfp4 / mxfp4 / mxfp8 (e4m3,e5m2) / nvfp8 are now spec-conformant block-scaled (quant::codec/format) and wired across the weight-bearing matmul families (dequant + qgemv + qmm + MoE gather + MMA = 25 GPU-verified kernels), decoded via first-class DSL intrinsics. Notes the follow-ups (more bench shapes, the fused perf-kernels that compose these primitives, block-scaled-KV attention, and the ekryski/mlx@alpha oracle audit).
…oundation) Generalize the quant framework so one set of templates + decode intrinsics covers all formats across the 3 scaling schemes, en route to full op×precision coverage: - QFormat += Fp4, Fp8E4m3, Fp8E5m2 (legacy per-group FP32 float-scale, reusing the e2m1/e4m3/e5m2 element decoders) + Int8 (symmetric affine, scale-only, group 64). All use the existing F32 per-group scale path — no new scale kinds. - codec: int8_encode/decode (symmetric, ±127); round-trip + saturation tests. - DSL: int8_decode intrinsic (op.rs/body.rs/features.rs/preamble.rs → emitted mt_decode_int8 sign-extend helper), mirroring the fp decode intrinsics. - format.rs pack/dequant + tests cover all 9 formats (cosine round-trip per format; byte-layout sizes; zero-block). Unblocks the int8/fp4/fp8 kernel variants across every weight-bearing family.
…tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of the fused RMSNorm + dequantizing GEMV (ffai/rms_norm_block_scaled_qgemv.rs), filling the first block-scaled gap in the fused families. Composition: the freeze-safe one-row-per-TG Reduction geometry of the int4 `ffai_rms_norm_qgemv` (phase 1 Σx² → `mt_rms_inv_scalar`) + the pack-/element-strided block-scaled weight decode of `mlx/block_scaled_matmul.rs` (phase 2), accumulating over the normalized activation `x·norm_weight·inv_rms`. Block-scaled formats carry no bias. 5 kernels, 15 GPU checks (×f32/f16/bf16) vs a `quant::format::dequant` RMSNorm oracle, all green; 15 decode-shape benches with `.flops(2·N·K)`. Reuses the DSL decode intrinsics — no new dispatch geometry.
…+ GPU tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of the fused gated-RMSNorm + dequantizing GEMV (ffai/gated_rms_norm_block_scaled_qgemv.rs). Reuses the proven 2-simdgroup per-row gated-RMSNorm staging (silu gate × per-row inv_rms × norm_weight → tg_inner) verbatim from the int4 fast variant — it is weight-format-independent — and swaps phase 2 for the simple one-output-row-per-TG block-scaled reduction GEMV over tg_inner. 5 kernels, 15 GPU checks (×f32/f16/bf16) vs a gated-RMSNorm + quant::format::dequant oracle, all green; 15 decode-shape benches.
…tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of the fused batched Q/K/V dequantizing GEMV (ffai/batched_qkv_block_scaled_qgemv.rs): one dispatch projects [Wq·x | Wk·x | Wv·x] over a shared activation. Mirrors the int4 scalar `ffai_batched_qkv_qgemv` geometry verbatim — grid [max(out_*),1,3], z-dim selects the matrix, one row per TG, Reduction — and swaps the int affine for the block-scaled decode (no bias). Establishes the batched z-dispatch exemplar for the remaining batched families. 5 kernels, 15 GPU checks (×f32/f16/bf16) vs a quant::format::dequant oracle, all green; 15 decode-shape benches.
…formats + GPU tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of two more fused families, extending the batched z-dispatch exemplar: - batched_4_block_scaled_qgemv.rs — 4-way projection (w_a..w_d → 4 separate output buffers), grid [max_out,1,4], one row/TG. - batched_qkv_block_scaled_qmm.rs — multi-token (M>1) Q/K/V GEMM, flattened per-output-element reduction over K with grid [max_out, M, 3]; x is [M,in_dim], outputs [M,out_*]. Both reuse their int4 families' proven Reduction geometry verbatim, swap the int affine for the block-scaled decode (no bias), and pin correctness with quant::format::dequant oracles. 10 kernels, 30 GPU checks (×f32/f16/bf16), all green; 30 decode-shape benches.
Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of the multi-token batched-4 dequantizing GEMM (ffai/batched_4_block_scaled_qmm.rs): 4-way projection (w_a..w_d → 4 separate [M,out_*] outputs) over an [M,in_dim] activation, grid [max_out, M, 4], per-output-element reduction over K. Completes the block-scaled batched matmul families (batched_4 + batched_qkv, qgemv + qmm). 5 kernels, 15 GPU checks vs quant::format::dequant oracle, all green; 15 decode-shape benches.
…ts + GPU tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 variants of the quantized embedding-table gather (ffai/dequant_gather_block_scaled.rs): per output element, gather the row, decode the block-scaled element × block scale. Grid3D one-thread-per-element (pure dequant, no reduction). 5 kernels, 15 GPU checks (×f32/f16/bf16) vs a gather-composed quant::format::dequant oracle (non-monotonic repeating gather order), all green; 15 decode-shape benches.
… + GPU tests Adds mxfp4 / nvfp4 / mxfp8 (e4m3,e5m2) / nvfp8 block-scaled KV-cache variants of the single-pass online-softmax attention (ffai/flash_block_scaled_sdpa.rs, head-dim 128). Reuses the affine flash_quantized_sdpa geometry verbatim — one simdgroup per query, lane-parallel over dim, grid [1,B·nQ,1] tpg [32,1,1] — and swaps the inline K/V decode from affine (q·scale+bias) to block-scaled (element_decode·block_scale, no bias). Completes the block-scaled (mxfp4/nvfp4/mxfp8/nvfp8) coverage across every weight-bearing family: dequant, qgemv, qmm, qmm-MMA, MoE gather, fused rms_norm / gated_rms_norm, batched-4 + batched-qkv (qgemv+qmm), attention KV, and embedding gather. 5 kernels, 21 GPU checks (base ×5 + sinks/window on mxfp4, ×f32/f16/bf16) vs a dequant-composed softmax-attention oracle, all green; 15 decode-shape benches.
…xis start) Extends the block-scaled qgemv family with the per-group-FP32-scale formats: fp4 (E2M1), fp8_e5m2 (E5M2), int8 (symmetric, int8_decode) — each a new decode-here kernel — plus fp8_e4m3, which reuses mt_nvfp8_qgemv (identical 8-bit-E4M3 + f32-scale shape, only block_size differs). Establishes the int8/fp4/fp8 second-axis pattern: f32 per-group scale, QFormat-parametrized test/bench, oracle via quant::format::dequant. 4 formats added (3 new kernels + 1 reuse), 12 GPU checks (×f32/f16/bf16) all green; 12 benches. block_scaled_qgemv now covers all 9 formats.
…U tests Extends the standalone dequant, dequantizing GEMM (qmm), and MoE gather-GEMM families with the per-group-FP32-scale formats fp4 (E2M1), fp8_e5m2 (E5M2), int8 (symmetric int8_decode) — new decode-here kernels — plus fp8_e4m3 reusing each family's nvfp8 kernel (same 8-bit-E4M3 + f32-scale shape). Adds bench modules to block_scaled_dequant + block_scaled_moe (were test-only). 3 families × 4 formats, GPU-verified vs quant::format::dequant oracles, all green. These core matmul families now cover all 9 formats.
…families Extends the fused RMSNorm/gated-RMSNorm GEMV, batched-4 + batched-Q/K/V (qgemv + qmm), flash SDPA (KV), and embedding-gather families with the per-group-FP32-scale formats fp4 (E2M1), fp8_e5m2 (E5M2), int8 (symmetric int8_decode) — new decode-here kernels — plus fp8_e4m3 reusing each family's nvfp8 kernel (identical 8-bit-E4M3 + f32-scale shape). Multi-matrix families apply the f32-scale decode across all Q/K/V (3-way) and a/b/c/d (4-way) dispatch blocks; flash changes both K and V decode sites. 8 families × 4 formats, all GPU-verified vs quant::format::dequant oracles, clippy/fmt/typos clean. With the core matmul families (prior commit), every weight-bearing family now covers all 9 formats except qmm-MMA (next).
…ix complete Extends the simdgroup-matrix qmm (qmm-MMA) family with fp4 / fp8_e5m2 / int8 (per-group FP32 scale; int8 via int8_decode) + fp8_e4m3 reusing mt_nvfp8_qmm_mma. Each clone reuses its source kernel's MMA geometry byte-for-byte, changing only the scale param type, the scale-decode line, and (int8) the element decode — freeze-safe. This closes the full op×format matrix: every weight-bearing family (dequant, qgemv, qmm, qmm-MMA, MoE gather, fused rms_norm / gated_rms_norm, batched-4 + batched-Q/K/V qgemv+qmm, flash SDPA KV, embedding gather) now supports all 9 quant formats (int8, fp4, fp8 e4m3/e5m2, mxfp4, nvfp4, mxfp8 e4m3/e5m2, nvfp8) in fp16/bf16/fp32 activation. All GPU-verified vs quant::format::dequant oracles.
Comment on lines
+1100
to
+1106
| 16, | ||
| 4, | ||
| 4, | ||
| 4, | ||
| 256, | ||
| 2, | ||
| dt, |
Owner
There was a problem hiding this comment.
These params are the same for all kernels, create a struct defining the default inputs and use it here or hardcode values into setup. This happens in many kernels.
Comment on lines
+809
to
+844
| if matrix == 1u32 { | ||
| if row < out_b { | ||
| for it in range(0u32, iters, 1u32) { | ||
| let c = it * lsize + tid; | ||
| if c < in_dim { | ||
| let elem = e5m2_decode(load(w_b[row_off + c]).cast::<u32>()); | ||
| let scale = load(scales_b[row_block_off + c / block_size]); | ||
| acc = acc + (elem * scale) * load(x[x_row_off + c]).cast::<f32>(); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| if matrix == 2u32 { | ||
| if row < out_c { | ||
| for it in range(0u32, iters, 1u32) { | ||
| let c = it * lsize + tid; | ||
| if c < in_dim { | ||
| let elem = e5m2_decode(load(w_c[row_off + c]).cast::<u32>()); | ||
| let scale = load(scales_c[row_block_off + c / block_size]); | ||
| acc = acc + (elem * scale) * load(x[x_row_off + c]).cast::<f32>(); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| if matrix == 3u32 { | ||
| if row < out_d { | ||
| for it in range(0u32, iters, 1u32) { | ||
| let c = it * lsize + tid; | ||
| if c < in_dim { | ||
| let elem = e5m2_decode(load(w_d[row_off + c]).cast::<u32>()); | ||
| let scale = load(scales_d[row_block_off + c / block_size]); | ||
| acc = acc + (elem * scale) * load(x[x_row_off + c]).cast::<f32>(); | ||
| } | ||
| } | ||
| } | ||
| } |
Owner
There was a problem hiding this comment.
This is bad syntax or a DSL issue if not currently possible, this should be expressible via a single loop rather than 3 independent if statements each containing their own duplicate logic.
Collaborator
Author
|
Current status for complete coverage. Close. Working on closing the gap. 
|
…, axes) Adds the host-side foundation for the full integer block-scaled matrix: - codec: generalize symmetric int8 to intN encode/decode (N in 2..6,8) + intn_max; sign-extend from bit N-1, clamp to +/-(2^(N-1)-1). - format: 11 new QFormat variants - symmetric int2/3/4/5/6 (FP32, group 64) and MXINT2/3/4/5/6/8 (E8M0, block 32). Add the element()/is_integer()/ symmetric() axis accessors (the element x scale x zero-point model) and make ScaleKind/scale_kind public. - packing: unify all sub-8-bit codes onto a tight LSB-first bit-stream (a 4-bit stream is byte-identical to the old nibble layout, so existing 4-bit kernels are unaffected) with a guard word for straddling 3/5/6-bit reads; 8-bit stays one byte/code. pack/dequant route through write_code/read_code. All 20 formats round-trip in the CPU oracle; new intn codec + bit-stream layout + cosine-fidelity unit tests pass. No kernels yet (Phase 0b).
Proves the parameterized integer block-scaled kernel pattern end-to-end: - int_dequant_f32! / int_dequant_e8m0! macros emit the straddle-aware two-word bit-stream extraction (mirroring the affine dequant_gemv) + float sign-extend (subtract 2^N when the top bit is set), keyed by bit width; FP32 vs E8M0 scale read is the only delta. mxint8 uses the byte path + int8_decode. - 11 new dequant kernels x f32/f16/bf16, each 1:1 #[test_kernel] vs the quant::format::dequant oracle + #[bench]. All GPU-verified. - dequant_setup/dequant_bench now type code/scale buffers via element_bits()/ scale_kind() instead of hardcoded format lists. - fix: bitstream_words only reserves a guard word for straddling widths (3/5/6); 2/4-bit pack to exact size so the 4-bit MoE/expert per-expert code concatenation stays byte-aligned (regression caught by the harness).
Roll symmetric int2-6 (FP32 group) + MXINT2-8 (E8M0 block) across three reduction families, reusing each family's proven geometry and the GPU-verified bit-stream-extract + float-sign-extend decode from block_scaled_dequant. 33 new kernels, each 1:1 #[test_kernel] vs the quant::format oracle + #[bench]. Setup helpers now type buffers via element_bits()/scale_kind(). Full GPU harness green.
… (wave 2) Roll symmetric int2-6 + MXINT2-8 across rms_norm_qgemv, gated_rms_norm_qgemv, patch_embed, batched_qkv_qgemv, batched_4_qgemv. 55 new kernels, each 1:1 #[test_kernel] vs oracle + #[bench]. Geometry unchanged; helpers axis-driven. Full GPU harness green.
Roll symmetric int2-6 + MXINT2-8 across batched_qkv_qmm, batched_4_qmm, conv2d/conv3d (direct), depthwise_conv2d, audio_conv1d, fishspeech_conv1d. 77 new kernels, each 1:1 #[test_kernel] vs oracle + #[bench]. Conv families decode via the flat global bit-stream offset (robust for non-32-multiple contraction); geometry unchanged; helpers axis-driven. Full GPU harness green.
Roll symmetric int2-6 + MXINT2-8 across the simdgroup-MMA, MPP cooperative-tensor, and NAX cooperative qMM families. W decoded into threadgroup staging via the bit-stream extract + float sign-extend, then the existing cooperative matmul. 33 new kernels, each 1:1 #[test_kernel] + #[bench]. Coop tile/extent geometry byte-for-byte unchanged; helpers axis-driven. Full GPU harness green.
…e 5) Roll symmetric int2-6 + MXINT2-8 across the simdgroup-matrix MMA patch-embed and conv2d/conv3d families, mirroring the block_scaled_mma staging-decode template. 33 new kernels, each 1:1 #[test_kernel] + #[bench]; MMA tile geometry byte-for-byte unchanged; helpers axis-driven. Full GPU harness green.
…exed (wave 6) Roll symmetric int2-6 + MXINT2-8 across flash block-scaled SDPA (all 5 head dims, K+V decode), MoE gather-qmm, and expert-indexed GEMV. Concat families pack the stacked expert weight once (byte-identical for 4/8-bit; correct guard handling for sub-byte 3/5/6). ~77 new kernels, each 1:1 #[test_kernel] + #[bench]. Geometry unchanged; helpers axis-driven. Full GPU harness green.
Roll symmetric int2-6 + MXINT2-8 across the three MoE gather-qMM MPP cooperative families. Combines the coop-staging integer decode with a single stacked expert pack() (replacing per-expert concat, which would misalign sub-byte guard words). 33 new kernels, each 1:1 #[test_kernel] + #[bench]. Coop geometry byte-for-byte unchanged. Full GPU harness green. Completes Phase 1: the full integer matrix (int2-6 + MXINT2-8) across all 24 weight-bearing block-scaled families.
…nteger matrix Rewrite the precision-coverage docs around the three orthogonal axes (element, scale, zero-point). Track-1 is now 20 symmetric formats (8 float + int2-6/int8 FP32-group + mxint2-8 E8M0-block); document the parameterized (bit-width x scale-kind) decode. Correct the 'no block-scaled int4' line: block-scaled int4/ mxint4 now exist (symmetric) but MLX 4-bit checkpoints still need affine int4 for its zero-point - a deliberate design boundary, not a missing format. FP16 group scale noted as the remaining scale-axis option.
…ase 2a/b) Adds ScaleKind::F16 + 10 fp16-scale QFormats (nvfp8/fp4/fp8_e4m3/fp8_e5m2 + int2-6/int8 twins of the FP32-scaled formats) - the memory-halving (2 B vs 4 B) scale layout real checkpoints use. codec f16_scale_encode does correct round-to-nearest-even WITH subnormal support (E5M2's tiny amax/57344 block scales land in f16's subnormal range; flushing them to zero wiped whole blocks). format pack/dequant handle the 2-byte F16 scale via ScaleKind::bytes(). 9 new dequant kernels (Tensor<f16> scale, .cast::<f32>()) + 10 #[test_kernel]s + 10 #[bench]es, all GPU-verified - the Apple-GPU half load matches the host decode exactly. 30 quant unit tests + full dequant harness green.
Add the 10 fp16-scale twins (nvfp8_f16/fp4_f16/fp8_e4m3_f16/fp8_e5m2_f16 + int2-6/int8_f16) to the core matmul reduction families by cloning each family's FP32-scaled kernels with the scale tensor as Tensor<f16> + .cast::<f32>(); geometry unchanged. 20 new kernels, each 1:1 #[test_kernel] + #[bench]. Setup/bench scales_dt now matches on ScaleKind::F16. Full GPU harness green.
…ched qgemv (Phase 2 wave 2) Add the 10 fp16-scale twins to embedding-gather, rms_norm/gated_rms_norm qgemv, patch_embed, and batched_qkv/batched_4 qgemv by cloning each family's FP32-scaled kernels with Tensor<f16> scale + .cast::<f32>(). 6 families, geometry unchanged, each 1:1 #[test_kernel] + #[bench]. Full GPU harness green.
…e 2 wave 3) Add the 10 fp16-scale twins to batched_qkv/batched_4 qmm, conv2d/conv3d (direct), depthwise_conv2d, audio_conv1d, fishspeech_conv1d by cloning each family's FP32-scaled kernels with Tensor<f16> scale + .cast::<f32>(). 7 families, geometry unchanged, each 1:1 #[test_kernel] + #[bench]. Full GPU harness green.
…ave 4) Add the 10 fp16-scale twins to the simdgroup-MMA + MPP/NAX cooperative qmm, patch-embed-MMA, and conv2d/3d-MMA families by cloning the FP32-scaled kernels with Tensor<f16> scale + .cast::<f32>(). 6 families, coop/MMA tile geometry byte-for-byte unchanged, each 1:1 #[test_kernel] + #[bench]. Full GPU harness green.
…-MPP (Phase 2 wave 5) Add the 10 fp16-scale twins to the MoE gather-qmm, expert-indexed GEMV, and the three MoE-MPP cooperative families by cloning the FP32-scaled kernels with Tensor<f16> scale + .cast::<f32>() (single stacked pack already in place from the integer wave). 5 families, coop geometry byte-for-byte unchanged, each 1:1 #[test_kernel] + #[bench]. Completes Phase 2: the fp16 group-scale axis across all 24 weight-bearing families. Full GPU harness green.
Update the precision-coverage docs: the 10 fp16-scale twins are landed across all 24 families, so Track-1 is 30 formats (8 float + 11 integer + 10 fp16 twins, with fp8_e4m3 sharing nvfp8's kernel per twin). Document the fp16 scale axis (2-byte IEEE half, RNE with subnormal support) in the format table + 3-axis description.
…xis (Phase 2 wave 6) Add the 10 fp16-scale twins to flash block-scaled SDPA across all 5 head dims (d64/96/128/256/512), cloning the FP32-scaled flash macros with Tensor<f16> K/V scale + .cast::<f32>() on both the score and value-accumulation scale reads. Geometry (one simdgroup per query) unchanged. Closes the one family the fp16 pass had missed — fp16 now covers all 24 weight-bearing families. Full GPU harness green.
ekryski
changed the title
Collaborator
Author
|
Full coverage is now complete. |
Collaborator
Author
|
Merging. Will address code length in a follow up PR. |
ekryski
added a commit
that referenced
this pull request
Jun 3, 2026
…262) ## What this is **Docs-only.** Four design specs for taking MetalTile's `#[kernel]` DSL / IR to new code-generation + runtime backends beyond Metal/MSL, plus a small repo reorg that gives long-form specs their own top-level home. All four specs are scoped to **kernel / compute lowering** — model loading, graph execution, and checkpoint readers stay out of scope (MetalTile is an optimized-kernel generator, not an inference engine). No code changes; no kernels added. ## Repo reorg: `specs/` Long-form design/spec docs now live in a top-level **`specs/`** folder, separate from the user-facing onboarding guides in `docs/`: - **Moved `docs/` → `specs/`:** `ARCHITECTURE`, `TOOLCHAIN_DESIGN`, `BENCH_METRICS_SPEC`, `KERNEL_AUDIT`, `KERNEL_ORGANIZATION`, `PROPOSED_OPTIMIZATIONS`, and the four new backend specs below. - **`docs/` keeps** the onboarding guides (`getting-started`, `developing`, `testing`, `publishing`, `cli`) + `README.md` as the master index. `docs/README.md` links across to `../specs/` and gains a **Backend ports** section. - Cross-references updated: `ARCHITECTURE.md`'s links to the staying guides became `../docs/…`; every `docs/<SPEC>.md` mention in `crates/**` doc-comments now points at `specs/<SPEC>.md`. `specs/` sits at the same depth as `docs/`, so spec↔spec and `../crates/…` links stay valid. ## The specs (`specs/`) | Spec | Role | Hardware block-scaling payoff | |---|---|---| | `CUDA_BACKEND_SPEC.md` | NVIDIA, peak perf — IR → CUDA C++ → NVRTC/PTX | Blackwell `tcgen05` scaled-MMA (MXFP/MXINT) | | `AMD_BACKEND_SPEC.md` | AMD, peak perf — a *delta* on the CUDA spec via HIP/ROCm | CDNA4 / MI350 OCP-MXFP microscaling | | `VULKAN_BACKEND_SPEC.md` | portable / fallback — IR → SPIR-V → Vulkan (Intel, Qualcomm, ARM, + Apple via MoltenVK) | none portable (software dequant) | | `ANE_BACKEND_SPEC.md` | Apple Neural Engine graph-offload — Core ML/MIL route + the direct private route | Core ML blockwise quant (conceptual) | ## Shared design All four build on **one backend seam** (a `CodegenBackend` + `Device` trait pair): the IR (`metaltile-core`), the `#[kernel]` macro, and the **entire 30-format `quant::{codec,format}` layer are reused unchanged** — only a codegen emitter and a runtime device are new per target. A recurring theme: the PR #250 block-scaled formats transfer. The `mx*` / `mxint*` **E8M0 / block-32** layout maps directly onto NVIDIA **Blackwell** *and* AMD **CDNA4** hardware microscaling, and aligns conceptually with Core ML's blockwise quant — so the precision work already spans Apple GPU (today) + the proposed targets. Each spec is honest about its hazards: - **CUDA** — the `mpp::`/`InlineMsl` cooperative kernels need CUTLASS reimplementation; MMA tiles need retuning. Evaluates NVlabs **`cuda-oxide`** (Rust→PTX + host runtime + Blackwell intrinsics) as a tooling option. - **AMD** — the **wavefront 32 vs 64** split (RDNA vs CDNA) is the main porting hazard; MFMA (CDNA) vs WMMA (RDNA3+); thinner Rust/ROCm ecosystem. - **Vulkan** — **subgroup size is variable *and* optional** (portable shared-memory reduction baseline + `VK_EXT_subgroup_size_control` fast path); MMA via the optional, runtime-shape-queried `VK_KHR_cooperative_matrix` with a fallback ladder; no portable hardware microscaling. - **ANE** — the crux: the ANE is **not custom-kernel-programmable**; it's a *graph-offload* target through Apple's compiler, not a code generator. Covers Route A (Core ML / MIL, recommended) and Route B (the direct private stack — `Espresso`/`ANECompiler`/`.hwx`/`ANEServices`/`H11ANE`/`_ANEClient`…, research-only). Folds in field evidence from *The ANE Book*, `hollance/neural-engine`, `skyfallsin/apple-neural-engine-field-guide`, `thebasedcapital/ane-infer`, and `john-rocky/CoreML-LLM`. Each spec carries a phased rollout, a risks section, and references.
0xClandestine
pushed a commit
that referenced
this pull request
Jun 3, 2026
…cross every weight-bearing kernel (#250) ## What this is PR 2 of the precision series: make fp4/fp8 **spec-conformant block-scaled** and bring the **full op × quant-format matrix to completion** — *every* weight-bearing kernel family supports **all 30 Track-1 quant formats** in fp16/bf16/fp32 activation, so the bench columns (GFLOP/s, %-peak, roofline) reveal the fastest precision per op. **No holes.** > Originally stacked on #247 (`ek/bench-metrics`); that has since merged into `dev`, so this branch is now precision-only on top of `dev`. ## Quantization is three orthogonal axes A quantized weight stores a small **code** per element + a **scale** that restores magnitude. Formats vary along three axes — **bit-width is not one of them**: 1. **Element** — a signed integer (`int2…int8`) *or* a micro-float codebook (E2M1 / E4M3 / E5M2). 2. **Scale** — how the per-block scale is stored: raw **FP32** per large group, a compact **E8M0** power-of-two per small block (OCP MX), an **E4M3** micro-scale × a tensor-wide FP32 global (NVFP4), or **FP16** (the memory-halving twin). 3. **Zero-point** — **symmetric** (no zero-point) or **asymmetric** (scale + bias). Track 1 (this PR) is the **symmetric** block-/group-scaled family; Track 2 is the pre-existing **asymmetric** affine integers (zero-point for MLX-checkpoint interop). The two are complementary, not a hierarchy. ## The 30 Track-1 formats | group | formats | element | block | scale | |---|---|---|---|---| | spec float | `nvfp4` | E2M1 | 16 | E4M3 + global FP32 | | | `mxfp4` / `mxfp8_e4m3` / `mxfp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | E8M0 pow-2 | | | `nvfp8` | E4M3 | 16 | FP32 | | legacy float | `fp4` / `fp8_e4m3` / `fp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | FP32 per-group | | **symmetric int** | `int2 / int3 / int4 / int5 / int6 / int8` | int N | 64 | FP32 per-group | | **MXINT** | `mxint2 / mxint3 / mxint4 / mxint5 / mxint6 / mxint8` | int N | 32 | E8M0 pow-2 | | **fp16-scale twins** | `*_f16` of every FP32-scaled format above | as twin | as twin | FP16 (2 B) | **30 formats** = 8 float + 6 symmetric int (FP32) + 6 MXINT (E8M0) + 10 fp16-scale twins. `mxint8` is OCP-ratified MXINT8; `mxint*` map to tensor-core block-scaling units for future NVIDIA / AMD targets. Sub-byte codes (2/3/4/5/6-bit) tight-bit-pack LSB-first into u32 words (a 4-bit stream is byte-identical to the classic nibble layout); 8-bit is one byte/code. All formats share the `quant::codec` bit primitives — one source of truth for the host packer, the CPU correctness oracle, and the in-kernel decode (the `e2m1_decode` / `e4m3_decode` / `e5m2_decode` / `int8_decode` intrinsics + the straddle-aware sub-byte bit-stream extraction). The integer + fp16 formats are emitted by a parameterized **(bit-width × scale-kind)** decode macro per family, so every variant reuses its family's proven dispatch geometry byte-for-byte (no new freeze surface). ## Coverage matrix — COMPLETE Every family below supports **all 30 formats**. Each `(family × format × dtype)` ships a `#[kernel]`, a 1:1 `#[test_kernel]` (GPU-verified vs `quant::format::dequant`), and a `#[bench]` with `.flops()`. | family | path | file(s) | |---|---|---| | dequant (standalone) | elementwise | `mlx/block_scaled_dequant.rs` | | qgemv (GEMV decode) | reduction | `mlx/block_scaled_matmul.rs` | | qmm (GEMM prefill) | reduction | `mlx/block_scaled_qmm.rs` | | qmm — simdgroup-MMA | simdgroup-matrix | `mlx/block_scaled_mma.rs` | | qmm — MPP | tensor engine | `mlx/block_scaled_qmm_mpp.rs` | | qmm — NAX | tensor engine | `mlx/block_scaled_qmm_nax.rs` | | MoE gather-qmm | reduction | `mlx/block_scaled_moe.rs` | | MoE gather — MPP (bm8/16/64) | tensor engine | `ffai/moe_mpp{,_bm8,_bm64}_block_scaled.rs` | | expert-indexed GEMV | reduction | `ffai/dequant_gemv_expert_indexed_block_scaled.rs` | | fused RMSNorm + GEMV | reduction | `ffai/rms_norm_block_scaled_qgemv.rs` | | fused gated-RMSNorm + GEMV | reduction | `ffai/gated_rms_norm_block_scaled_qgemv.rs` | | batched-Q/K/V qgemv + qmm | reduction | `ffai/batched_qkv_block_scaled_{qgemv,qmm}.rs` | | batched-4 qgemv + qmm | reduction | `ffai/batched_4_block_scaled_{qgemv,qmm}.rs` | | embedding gather | elementwise | `ffai/dequant_gather_block_scaled.rs` | | flash SDPA (block-scaled KV) | flash | `ffai/flash_block_scaled_sdpa.rs` — **d64/96/128/256/512** | | patch embed (linear projection) | reduction | `ffai/patch_embed_block_scaled.rs` | | patch embed (simdgroup-MMA) | simdgroup-matrix | `ffai/patch_embed_mma_block_scaled.rs` | | conv2d / conv3d (direct) | reduction | `ffai/{conv2d,conv3d}_block_scaled.rs` | | conv2d / conv3d (im2col simdgroup-MMA) | simdgroup-matrix | `ffai/{conv2d,conv3d}_mma_block_scaled.rs` | | depthwise conv2d | reduction | `ffai/depthwise_conv2d_block_scaled.rs` | | audio conv1d (STT) / fishspeech conv1d (TTS) | reduction | `ffai/{audio_conv1d,fishspeech_conv1d}_block_scaled.rs` | (`fp8_e4m3` / `fp8_e4m3_f16` reuse each family's `nvfp8` / `nvfp8_f16` kernel — identical 8-bit-E4M3 + scale shape.) **Flash KV covers every production head dim** (d64/96/128/256/512 × all 30 formats) — **no holes**. int8's group size (64) doesn't divide d96, so that case tiles with a **ragged trailing block** (a 64-block + a 32-block; `n_blocks = ceil(dim/block_size)`), the host packer and kernel rounding up identically so codes + scales stay self-consistent. ### The full integer matrix is a first-class citizen `int2/3/4/5/6/8` (FP32 group) + `mxint2/3/4/5/6/8` (E8M0 block) are present in **every** family above — including the fast tensor-engine paths (**simdgroup-MMA, MPP, NAX, MoE-MPP**), where integer throughput is highest on Apple GPUs / the ANE and the `mxint*` E8M0 layout maps to tensor-core block-scaling for future hardware. The core matmul / MoE / fused-norm / batched / KV-cache / attention families *also* carry the pre-existing **asymmetric affine** integers (scale + bias) for MLX-checkpoint interop. ### Symmetric block-scaled int4 now exists — affine int4 is still required, by design Block-scaled `int4` / `mxint4` are now present (symmetric). MLX 4-bit checkpoints are **asymmetric** affine (`w = scale·q + bias`), which a symmetric format cannot represent — so **affine int4 (Track 2, with its zero-point) remains required to load every MLX 4-bit model**. That's a deliberate design boundary (the zero-point), not a missing format. ("Legacy" applies only to the float-scale `fp4`/`fp8` within Track 1, superseded by spec mxfp4/nvfp4/mxfp8/nvfp8 but kept as labeled comparison variants.) ## fp16-scale axis Every FP32-scaled format has an **FP16-scale twin** (`*_f16`): same element + block, but the per-block scale is stored as a 2-byte IEEE half — the layout real checkpoints use, halving scale-buffer traffic. The host encoder does correct round-to-nearest-even **with full subnormal support**: wide-range elements (E5M2's `element_max` is 57344) push block scales into f16's subnormal range, and flushing them to zero would wipe whole blocks. The Apple-GPU `half` load matches the host decode exactly (subnormals included), so the CPU oracle holds. ## fp4 simdgroup-MMA f32 fix (from earlier in the series) Making the `#[test_kernel]` harness non-vacuous (it had been iterating an empty set — `metaltile::harness::registry::all_tests()` → `metaltile_std::all_tests()`) exposed two intertwined fp4-MMA defects, both fixed: 1. **Undefined shift.** `fp_quantized_mma::mt_fp4_qmm_mma` hand-rolled the E2M1 magnitude as `(mantissa + 2) << (exp - 1)`, undefined at `exp == 0` → miscompiled on f32 (output tile left unwritten). Replaced with the `e2m1_decode` intrinsic. 2. **Kernel name collision.** The block-scaled fp4 MMA kernel was *also* named `mt_fp4_qmm_mma`, colliding in the MSL/pipeline cache → order-dependent wrong pipelines (the real source of the ~0.46 f32 anomaly, and of an apparent `ffai_sdpa_multi_d256_causal` failure that was a *victim* of the contamination). Renamed to `mt_fp4_float_qmm_mma`. ## Foundation + verification - **`quant::codec`** — pure-Rust element + scale codecs (E2M1/E4M3/E5M2, E8M0, FP32, FP16, symmetric intN), the single source of truth shared by packer, oracle, and kernels. - **`quant::format`** — the 30-format `QFormat` with `element()` / `scale_kind()` / `symmetric()` axis accessors; host pack/dequant with ragged-tail blocks + unified sub-byte bit-stream packing. - Every `(family × format × dtype)` is GPU-verified 1:1 against the oracle. Full workspace gate (`cargo test --workspace` + `clippy`) green; geometry audited byte-for-byte unchanged per family (no new freeze surface). The MPP/NAX/MoE-MPP cooperative kernels are runtime-gated to Apple10+ and auto-skip on the CI Paravirtual GPU.
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
2 participants
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 this is
PR 2 of the precision series: make fp4/fp8 spec-conformant block-scaled and bring the full op × quant-format matrix to completion — every weight-bearing kernel family supports all 30 Track-1 quant formats in fp16/bf16/fp32 activation, so the bench columns (GFLOP/s, %-peak, roofline) reveal the fastest precision per op. No holes.
Quantization is three orthogonal axes
A quantized weight stores a small code per element + a scale that restores magnitude. Formats vary along three axes — bit-width is not one of them:
int2…int8) or a micro-float codebook (E2M1 / E4M3 / E5M2).Track 1 (this PR) is the symmetric block-/group-scaled family; Track 2 is the pre-existing asymmetric affine integers (zero-point for MLX-checkpoint interop). The two are complementary, not a hierarchy.
The 30 Track-1 formats
nvfp4mxfp4/mxfp8_e4m3/mxfp8_e5m2nvfp8fp4/fp8_e4m3/fp8_e5m2int2 / int3 / int4 / int5 / int6 / int8mxint2 / mxint3 / mxint4 / mxint5 / mxint6 / mxint8*_f16of every FP32-scaled format above30 formats = 8 float + 6 symmetric int (FP32) + 6 MXINT (E8M0) + 10 fp16-scale twins.
mxint8is OCP-ratified MXINT8;mxint*map to tensor-core block-scaling units for future NVIDIA / AMD targets. Sub-byte codes (2/3/4/5/6-bit) tight-bit-pack LSB-first into u32 words (a 4-bit stream is byte-identical to the classic nibble layout); 8-bit is one byte/code.All formats share the
quant::codecbit primitives — one source of truth for the host packer, the CPU correctness oracle, and the in-kernel decode (thee2m1_decode/e4m3_decode/e5m2_decode/int8_decodeintrinsics + the straddle-aware sub-byte bit-stream extraction). The integer + fp16 formats are emitted by a parameterized (bit-width × scale-kind) decode macro per family, so every variant reuses its family's proven dispatch geometry byte-for-byte (no new freeze surface).Coverage matrix — COMPLETE
Every family below supports all 30 formats. Each
(family × format × dtype)ships a#[kernel], a 1:1#[test_kernel](GPU-verified vsquant::format::dequant), and a#[bench]with.flops().mlx/block_scaled_dequant.rsmlx/block_scaled_matmul.rsmlx/block_scaled_qmm.rsmlx/block_scaled_mma.rsmlx/block_scaled_qmm_mpp.rsmlx/block_scaled_qmm_nax.rsmlx/block_scaled_moe.rsffai/moe_mpp{,_bm8,_bm64}_block_scaled.rsffai/dequant_gemv_expert_indexed_block_scaled.rsffai/rms_norm_block_scaled_qgemv.rsffai/gated_rms_norm_block_scaled_qgemv.rsffai/batched_qkv_block_scaled_{qgemv,qmm}.rsffai/batched_4_block_scaled_{qgemv,qmm}.rsffai/dequant_gather_block_scaled.rsffai/flash_block_scaled_sdpa.rs— d64/96/128/256/512ffai/patch_embed_block_scaled.rsffai/patch_embed_mma_block_scaled.rsffai/{conv2d,conv3d}_block_scaled.rsffai/{conv2d,conv3d}_mma_block_scaled.rsffai/depthwise_conv2d_block_scaled.rsffai/{audio_conv1d,fishspeech_conv1d}_block_scaled.rs(
fp8_e4m3/fp8_e4m3_f16reuse each family'snvfp8/nvfp8_f16kernel — identical 8-bit-E4M3 + scale shape.)Flash KV covers every production head dim (d64/96/128/256/512 × all 30 formats) — no holes. int8's group size (64) doesn't divide d96, so that case tiles with a ragged trailing block (a 64-block + a 32-block;
n_blocks = ceil(dim/block_size)), the host packer and kernel rounding up identically so codes + scales stay self-consistent.The full integer matrix is a first-class citizen
int2/3/4/5/6/8(FP32 group) +mxint2/3/4/5/6/8(E8M0 block) are present in every family above — including the fast tensor-engine paths (simdgroup-MMA, MPP, NAX, MoE-MPP), where integer throughput is highest on Apple GPUs / the ANE and themxint*E8M0 layout maps to tensor-core block-scaling for future hardware. The core matmul / MoE / fused-norm / batched / KV-cache / attention families also carry the pre-existing asymmetric affine integers (scale + bias) for MLX-checkpoint interop.Symmetric block-scaled int4 now exists — affine int4 is still required, by design
Block-scaled
int4/mxint4are now present (symmetric). MLX 4-bit checkpoints are asymmetric affine (w = scale·q + bias), which a symmetric format cannot represent — so affine int4 (Track 2, with its zero-point) remains required to load every MLX 4-bit model. That's a deliberate design boundary (the zero-point), not a missing format. ("Legacy" applies only to the float-scalefp4/fp8within Track 1, superseded by spec mxfp4/nvfp4/mxfp8/nvfp8 but kept as labeled comparison variants.)fp16-scale axis
Every FP32-scaled format has an FP16-scale twin (
*_f16): same element + block, but the per-block scale is stored as a 2-byte IEEE half — the layout real checkpoints use, halving scale-buffer traffic. The host encoder does correct round-to-nearest-even with full subnormal support: wide-range elements (E5M2'selement_maxis 57344) push block scales into f16's subnormal range, and flushing them to zero would wipe whole blocks. The Apple-GPUhalfload matches the host decode exactly (subnormals included), so the CPU oracle holds.fp4 simdgroup-MMA f32 fix (from earlier in the series)
Making the
#[test_kernel]harness non-vacuous (it had been iterating an empty set —metaltile::harness::registry::all_tests()→metaltile_std::all_tests()) exposed two intertwined fp4-MMA defects, both fixed:fp_quantized_mma::mt_fp4_qmm_mmahand-rolled the E2M1 magnitude as(mantissa + 2) << (exp - 1), undefined atexp == 0→ miscompiled on f32 (output tile left unwritten). Replaced with thee2m1_decodeintrinsic.mt_fp4_qmm_mma, colliding in the MSL/pipeline cache → order-dependent wrong pipelines (the real source of the ~0.46 f32 anomaly, and of an apparentffai_sdpa_multi_d256_causalfailure that was a victim of the contamination). Renamed tomt_fp4_float_qmm_mma.Foundation + verification
quant::codec— pure-Rust element + scale codecs (E2M1/E4M3/E5M2, E8M0, FP32, FP16, symmetric intN), the single source of truth shared by packer, oracle, and kernels.quant::format— the 30-formatQFormatwithelement()/scale_kind()/symmetric()axis accessors; host pack/dequant with ragged-tail blocks + unified sub-byte bit-stream packing.(family × format × dtype)is GPU-verified 1:1 against the oracle. Full workspace gate (cargo test --workspace+clippy) green; geometry audited byte-for-byte unchanged per family (no new freeze surface). The MPP/NAX/MoE-MPP cooperative kernels are runtime-gated to Apple10+ and auto-skip on the CI Paravirtual GPU.