wh-iron kernel-op coverage audit
Esta página aún no está disponible en tu idioma.
Snapshot of the kernels shipped by wh-iron-std as of dev c017c94. Comparison columns track parity with ml-explore/mlx@main (commit 2414e5df) and ekryski/mlx@alpha (commit 4919270e).
Summary
Section titled “Summary”- Total kernels (
iron build): 651 — all compiled unconditionally; the 7 NAX kernels are runtime-gated to Apple10+ (M4 family and newer). See § NAX kernels for what NAX is, which M-series chips activate it, and how it interacts with CI. (The jump from 374 is the block-scaled precision matrix — see § Quantization precision coverage.) - 90 / 91 kernel-op rows ported — 90 ✓, 0 partial, 1 intentionally out of scope (
fence; see § Fence ops). - Every floating-point kernel exposes f32 / f16 / bf16. Completing bf16 coverage also migrated every cooperative-tensor (NAX) kernel from hand-built
Op::InlineMslIR to the#[kernel]DSL via thecoop_tile_*intrinsics +coop_stage(T)(bf16 →halfstaging because Apple’smatmul2dmishandlesbfloatcooperative tensors). - int4 and int8 quantized perf paths are at parity. The int8 dense GEMM (
qmv/qmm/qmm_mma/qmm_mpp/qmm_nax) and int8 MoE BGEMM (mma/bm{8,16,64}_mpp) paths were built out, plus int4 polish (rms_norm_qgemv_fast,batched_qkv_qgemv_fast,iron_dequant_gemv_int4_fast,qvm_int4_fast). - Full block-scaled precision matrix across every weight-bearing kernel. The
ek/precision-supportwork added spec-conformant nvfp4 / mxfp4 / mxfp8 (e4m3+e5m2) / nvfp8, legacy float-scale fp4 / fp8 (e4m3+e5m2), the complete symmetric integer matrix —int2/3/4/5/6/8(FP32 group) +mxint2/3/4/5/6/8(E8M0 block) — and FP16-scale twins of every FP32-scaled format (*_f16): 30 formats to every quantized family (matmul/MoE/attention/embedding/conv, on the reduction, simdgroup-MMA, MPP, and NAX paths). Each(family × format × dtype)is 1:1#[test_kernel]-verified against thequant::formatCPU oracle and benched. The full integer family is present in all of them (incl. the fast tensor-engine paths). See § Quantization precision coverage. - Attention coverage spans every production head_dim. Added
steel_attention_nax_d{64,128,256},iron_sdpa_vector_d{64,96,192,256},sdpa_vector_2pass_d{64,96,256}, andflash_quantized_sdpa_d{96,512}(GPT-NeoX d=96, Gemma 4 global d=512). - Vision / STT / TTS front-end has MMA-tiled perf paths. Shipped
iron_conv2d_mma/iron_conv3d_mma/patch_embed_mma(implicit-im2col + 4-SG 2×2 simdgroup-matrix MMA) and Bluestein non-pow2 FFT (iron_fft_bluestein_*, covers Whisper n_fft=400/480). - Tail items: added
iron_sort_segmented,iron_scan_{prod,max,min}(+ exclusive),sdpa_decode_batched_q8, and 12flash_quantized_sdpa_{bool,float}_mask_*variants. - VLM vision-tower attention added:
iron_sdpa_bidirectional_d{32,64,72}— multi-query bidirectional SDPA (no causal mask) for SigLIP / CLIP / FastViT / PaliGemma encoders. d=72 uses a ragged 3-elements-per-lane layout (24 active lanes × 3 = 72, 8 lanes idle) for PaliGemma SigLIP-So400m.
NAX kernels
Section titled “NAX kernels”NAX (“Neural Acceleration”) is the cooperative-tensor matmul path exposed by Apple’s MetalPerformancePrimitives.framework — the mpp::tensor_ops::matmul2d<desc, execution_simdgroup> intrinsic. NAX kernels invoke it directly to get tensor-core-class throughput on the Apple GPU’s MMA units; the non-NAX equivalents fall back to simdgroup_matmul (8×8 frag MMA) or scalar code.
Hardware support
Section titled “Hardware support”NAX requires macOS 26+ (Metal 4) and Apple GPU family ≥ 10:
- M4 family (M4, M4 Pro, M4 Max, M4 Ultra, iPad M4) — Apple10 ✓
- M5 family — Apple11 ✓
- M1 / M2 / M3 — Apple7/8/9, no NAX; correctness tests use a
skip_unless_apple10runtime gate so the suite still passes on pre-M4 hardware
The runtime gate lives in crates/wh-iron-runtime/src/context.rs::Context::chip_family() — it reports the highest supported MTLGPUFamily value the device claims (returning None when no Metal device is available or on the virtualised GPU GitHub macOS runners expose).
Build-time gating
Section titled “Build-time gating”None. NAX kernels compile unconditionally and register their inventory::submit! BenchSpecs alongside every other kernel. iron build reports the full 374-kernel count on every host that can compile wh-iron-std. The decision to dispatch them is made at runtime through Context::chip_family() (see § CI coverage for the macOS Paravirtual-GPU caveat) and skip_unless_apple10 guards in the GPU-correctness tests.
The previous wh-iron-std/nax Cargo feature was removed — there’s no longer a way to opt out at build time. Dispatching a NAX kernel on pre-M4 hardware will fail at pipeline-creation time when the device rejects the mpp::tensor_ops::matmul2d symbol; callers should consult chip_family() before selecting the NAX path.
NAX kernels
Section titled “NAX kernels”The 7 NAX kernels:
| Kernel | File | Role |
|---|---|---|
iron_qmm_nax |
kernels/gemm/quantized_nax.rs |
int4 quantized matmul prefill |
iron_qmm_nax_int8 |
kernels/gemm/quantized_nax_int8.rs |
int8 quantized matmul prefill |
iron_fp_qmm_nax |
kernels/gemm/fp_quantized_nax.rs |
fp4 (E2M1) quantized matmul prefill |
iron_steel_gemm_fused_nax |
kernels/gemm/steel/steel_gemm_fused_nax.rs |
plain fused GEMM |
iron_steel_gemm_gather_nax |
kernels/gemm/steel/steel_gemm_gather_nax.rs |
MoE gather GEMM |
iron_steel_gemm_splitk_nax + _accum_nax |
kernels/gemm/steel/steel_gemm_splitk_nax.rs |
split-K GEMM (pass1 + pass2) |
iron_sdpa_prefill_nax |
mlx/steel/attn/steel_attention_nax.rs |
FlashAttention-2 prefill |
The quantized_mpp family (iron_qmm_mma_mpp, iron_qmm_mma_mpp_int8, the four MoE *_mpp variants) uses the same MPP cooperative-tensor primitive and is similarly runtime-gated via skip_unless_apple10. The distinction between *_mpp and *_nax: quantized_mpp and its MoE siblings have working MXU-fallback paths on M1–M3 via Apple’s matmul2d itself (slower than NAX hardware but functionally correct), whereas the *_nax kernels were authored specifically to exercise the M4+ tensor-core descriptor and have no fallback.
CI coverage
Section titled “CI coverage”GitHub’s macOS runners expose an Apple Paravirtual GPU that doesn’t claim Apple10+, so the NAX kernels and their tests are skipped at runtime via skip_unless_apple10. The Tile workflow’s iron build step still compiles them — if MetalPerformancePrimitives headers are unavailable on the runner’s Xcode, the build will surface that breakage immediately rather than silently dropping coverage.
Local verification of NAX kernels is the developer’s responsibility on M4+ hardware. The make test target runs the full suite; tests behind skip_unless_apple10 execute on real Apple10+ chips and auto-skip elsewhere.
Op coverage table
Section titled “Op coverage table”| Op | MLX (upstream) | MLX (ekryski@alpha) | wh-iron | Notes |
|---|---|---|---|---|
| arange | ✓ | ✓ | ✓ | kernels/ops/arange.rs → iron_arange. Generic T. |
| arg_reduce (argmax/argmin → u32 index) | ✓ | ✓ | ✓ | kernels/ops/arg_reduce.rs → iron_argmax<T> + iron_argmin<T>. Both generic over T (values widened to f32 for comparison); winning index emitted as u32; ties take the smallest index. (The former iron_argmax was a byte-for-byte duplicate of iron_argmax and was dropped.) |
| binary (elementwise add/sub/mul/div/min/max) | ✓ | ✓ | ✓ | kernels/ops/binary.rs → 6 kernels. Generic T. |
| binary_two (fused two-output elementwise) | ✓ | ✓ | ✓ | kernels/ops/binary_two.rs → iron_binary_two<T>. |
| copy (contiguous) | ✓ | ✓ | ✓ | kernels/ops/copy.rs → iron_copy<T>. |
| copy (strided / general) | ✓ | ✓ | ✓ | kernels/ops/strided.rs → iron_strided_copy (2-D padded) + iron_strided_copy_nd (arbitrary-rank). Each output element unravels its flat index against a runtime shape array and gathers src[Σ coord_d · strides[d]]. Covers padded copies, transposes, broadcasts (stride 0), and dilated slices in one kernel. |
| ternary (select) | ✓ | ✓ | ✓ | kernels/ops/ternary.rs → iron_select<T>. |
| unary (exp/log/sqrt/rsqrt/abs/silu/etc.) | ✓ | ✓ | ✓ | kernels/ops/unary.rs → 7+ kernels including iron_silu. Plus iron_scalar_fma_chain8 (fused 8-way scalar FMA) and iron_add_rms_norm (residual-add + RMSNorm fusion, Reduction mode). |
swiglu (silu(gate)·up fused MLP activation) |
✗ | ✗ | ✓ | kernels/ops/gated_activation.rs → iron_swiglu<T> (+ iron_clamped_swiglu<T>, per-layer activation-clip limit). Standard modern-transformer MLP activation (Llama 4, Qwen3 dense + MoE, Gemma, Mistral). |
| random (key hash → u32) | ✓ | ✓ | ✓ | kernels/ops/random.rs → iron_random_hash. |
| reduce (sum/prod/max/min — all + row + col) | ✓ | ✓ | ✓ | kernels/ops/reduce.rs covers all_reduce*, row_reduce*, col_reduce*, and seg_reduce* (Grid3D one-thread-per-segment, contiguous fixed-length runs) — all four ops for each shape. |
| sort | ✓ | ✓ | ✓ | kernels/sampling/sort.rs → iron_sort<T> (single-block bitonic) + iron_merge<T> (multi-block merge) + iron_sort_segmented<T> (per-row bitonic for [batch, n] matrices, n ≤ 1024, one TG per row). |
| scan (prefix sum/prod/max/min) | ✓ | ✓ | ✓ | kernels/ops/scan.rs → iron_scan<T> + iron_scan_exclusive<T> (sum), iron_scan_prod<T> / iron_scan_max<T> / iron_scan_min<T> + exclusive variants. Sum pair uses hardware simd_scan_exclusive; the prod/max/min pairs use a tgs[lsize] threadgroup buffer for sequential cross-thread prefix reads. |
| softmax | ✓ | ✓ | ✓ | kernels/sampling/softmax.rs → iron_softmax<T> (looped + single-row collapsed). |
| logsumexp | ✓ | ✓ | ✓ | kernels/ops/logsumexp.rs → iron_logsumexp<T>. |
| layer_norm | ✓ | ✓ | ✓ | kernels/norm/layer_norm.rs → iron_layer_norm<T>. |
| rms_norm | ✓ | ✓ | ✓ | kernels/norm/rms_norm.rs → iron_rms_norm<T> + iron_rms_norm_small<T> (2-elem/thread, small-head_dim per-head q_norm/k_norm) + iron_rms_norm_wide<T> (strided wide-row variant for head_dim > 4096, e.g. Gemma 4 31B hidden=5376). |
| rope (standard) | ✓ | ✓ | ✓ | kernels/rope/base.rs → iron_rope. |
| rope (frequency-band scaled) | ✗ | ✗ | ✓ | kernels/rope/rope_banded.rs → iron_rope_banded<T>. Per-row positions tensor + row grid axis (decode = T=1, prefill = all tokens in one dispatch); generic dtype, optional frequency-band scaling (Llama-3 / Qwen). |
| rope (YaRN context extension) | ✗ | ✗ | ✓ | kernels/rope/rope_yarn.rs → iron_rope_yarn<T> (decode, scalar position constexpr) + iron_rope_yarn_many<T> (prefill, per-row positions tensor + row grid axis + row_stride), one variants(...) definition. Per-dimension extrapolation/interpolation ramp between low and high plus the attn_factor mscale; factor == 1 is plain RoPE. Nemotron-Labs-Diffusion, GPT-OSS. |
| sdpa_vector (prefill / generic) | ✓ | ✓ | ✓ | mlx/scaled_dot_product_attention.rs → iron_sdpa<T>. Scalar SDPA for short sequences. |
| sdpa_vector (GQA decode, single pass) | ✓ | ✓ | ✓ | mlx/sdpa_vector.rs → iron_sdpa_vector<T> (d=128) + iron_sdpa_vector_d{64,96,192,256} (every production head_dim). Each scales the per-lane element count (2/3/6/8 elements). TPG=1024 throughout. |
| sdpa_vector_2pass | ✓ | ✓ | ✓ | iron/sdpa_decode_2pass.rs → pass1/pass2 pairs for d ∈ {64, 96, 128, 256}. d=256 uses 4-buffer TG reuse to stay within the 32 KB cap. |
sdpa_decode (Iron production decode, decoupled kv_stride) |
✗ | ✗ | ✓ | iron/sdpa_decode.rs → iron_sdpa_decode<T> + sdpa_decode_d{64,256,512}.rs. Iron-only; kv_stride ≠ n_kv (pre-allocated max-seq cache). Covers head_dim ∈ {64, 128, 256, 512}, sliding-window + sink-token (sink_end / window_start), and the GPT-OSS-20B learned-sink path (has_sink / sink_logit constexprs fold the per-head learned-sink logit into the cross-simdgroup softmax denominator on-GPU). Ring-aware for windowed caches: trailing ring_size / kv_begin constexprs walk absolute positions [kv_begin, kv_begin + n_kv) through slot a mod ring_size (ring_size == 0 = linear, the pre-ring behaviour); the same contract is on iron_sdpa_decode_sink_buf, iron_sdpa_decode_d512_sink, iron_aura_flash_pass1{,_causal}, and (ring_size only, absolute walk) iron_sdpa_prefill_*_sink / iron_aura_flash_sdpa. |
| sdpa_decode_batched (speculative-decode batched-Q) | ✗ | ✗ | ✓ | iron/sdpa_decode_batched.rs → sdpa_decode_batched_q{2,4,8}<T> + sdpa_decode_batched_prefill.rs. K query positions share one KV walk per dispatch, amortising KV memory bandwidth K×. q8 dispatches at TPG=256 due to register pressure. Iron-only. |
| sdpa_bidirectional (VLM vision-tower SDPA, no causal mask) | ✗ | ✗ | ✓ | iron/sdpa_bidirectional.rs → iron_sdpa_bidirectional_d{32,64,72}<T>. Multi-query bidirectional SDPA — each query attends [0, base_kv + n_query) with no causal gating. Covers SigLIP-base/large + CLIP-L (d=64), FastViT-HD (d=32), and PaliGemma SigLIP-So400m (d=72). d=72 uses a ragged 3-elements-per-lane layout: lanes 0..23 own the 72 valid indices (24 × 3), lanes 24..31 bounds-mask their q·k contribution to 0 and skip their output stores (25 % lane-occupancy loss, acceptable vs the cost of a wholly different parallel decomposition). TPG=1024, one threadgroup per (query, q_head). Online softmax in fp32. Iron-only. |
| steel_attention (Flash, prefill) | ✓ | ✓ | ✓ | mlx/steel/attn/steel_attention.rs → iron_sdpa_prefill<T>. Scalar-flash prefill (BQ=4, online softmax, causal), generic T, head_dim=128. |
| steel_attention_mma (Flash prefill, simdgroup-MMA) | ✓ | ✓ | ✓ | mlx/steel/attn/steel_attention_mma.rs → iron_sdpa_prefill_mma<T>. Real simdgroup-matrix MMA path; head_dim=128. A pre-M3 bf16-tuned sibling (steel_attention_mma_bf16.rs) is selected by sdpa_prefill_mma_for(). |
| steel_attention_nax | ✓ | ✓ | ✓ | mlx/steel/attn/steel_attention_nax.rs → iron_sdpa_prefill_nax<T> (d=32 base) + iron_sdpa_prefill_nax_d{64,128,256}. Flash-attention prefill via Apple mpp::tensor_ops::matmul2d. The wide variants loop the QK contraction over head_dim/32 consecutive 32-wide D-chunks inside the outer K-block loop (first chunk uses overwrite descriptor, subsequent chunks accumulate); PV stores each chunk to a scratch Opv tile then accumulates into the full-width O buffer. Causal masking + GQA. Runtime-gated to Apple10+. |
| steel_gemm_fused | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_fused.rs → iron_steel_gemm_{32x32x16_1x2,32x64x16_1x2,32x32x16_2x2,64x64x16_2x2}<T> (4 block shapes via instantiate_gemm_shapes_helper). |
| steel_gemm_fused_nax | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_fused_nax.rs → iron_steel_gemm_fused_nax<T>. Plain fused GEMM C = A·B via NAX cooperative-tensor matmul2d. Runtime-gated to Apple10+. |
| steel_gemm_gather | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_gather.rs → iron_steel_gemm_gather_{64x64x16_2x2,32x32x16_2x2}<T>. Row-major C = A_gathered·B_gathered (MLX gather_mm, the dense matmul of a MoE FFN). |
| steel_gemm_gather_nax | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_gather_nax.rs → iron_steel_gemm_gather_nax<T>. Gather GEMM via NAX matmul2d. Runtime-gated to Apple10+. |
| steel_gemm_masked | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_masked.rs → iron_steel_gemm_masked_{64x64x16_2x2,32x32x16_2x2}<T>. Block-masked C = A·B (output-block mask zeros whole BM×BN blocks; operand-block mask scales each K-block contribution). |
| steel_gemm_segmented | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_segmented.rs → iron_steel_gemm_segmented_{64x64x16_2x2,32x32x16_2x2}<T>. Ragged-K batched matmul (MLX segmented_mm); each segment sums over its own [k_start, k_end) range. |
| steel_gemm_splitk + accum | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_splitk.rs → pass 1 iron_steel_gemm_splitk_{64x64x16_2x2,32x32x16_2x2}<T> + pass 2 iron_steel_gemm_splitk_accum<T> / iron_steel_gemm_splitk_accum_axpby<T>. Partials stay fp32 for cross-split precision on f16/bf16 inputs. |
| steel_gemm_splitk_nax | ✓ | ✓ | ✓ | kernels/gemm/steel/steel_gemm_splitk_nax.rs → pass 1 iron_steel_gemm_splitk_nax<T> + pass 2 iron_steel_gemm_splitk_accum_nax<T>. Split-K via NAX matmul2d; partials fp32. Runtime-gated to Apple10+. |
| steel_conv 2D (implicit-GEMM) | ✓ | ✓ | ✓ | iron/conv2d.rs → iron_conv2d_patch14 / iron_conv2d_patch16 / iron_conv2d_generic. Direct conv (implicit im2col, one thread per output). MMA-tiled perf path: iron/conv2d_mma.rs → iron_conv2d_mma<T> — implicit-im2col + 4-SG 2×2 simdgroup-matrix MMA, 32×32 output tile (stride=1/dilation=1/pad=0, out_ch and n_pixels divisible by 32). |
| steel_conv 3D | ✓ | ✓ | ✓ | iron/conv3d.rs → iron_conv3d_generic + iron_conv3d_grouped (depthwise + dilation). 5D NCDHW / OIDHW. MMA-tiled perf path: iron/conv3d_mma.rs → iron_conv3d_mma<T> — same MMA scaffold as 2D, decomposed over (kd, kh, kw, ic). |
| steel_conv_general (strides/dilation/groups) | ✓ | ✓ | ✓ | iron/conv2d.rs → iron_conv2d_grouped<T>. Fully general 2D conv: strides, dilation (atrous), padding, grouped channels. |
| conv (winograd + naive_unfold + depthwise) | ✓ | ✓ | ✓ | iron/conv2d.rs / iron/conv3d.rs cover naive_unfold + depthwise (via _generic / _grouped for both 2D and 3D). Winograd fast-conv: iron/winograd_conv.rs → iron_winograd_conv2d_3x3<T> (F(2×2, 3×3) minimal-filtering, one thread per 2×2 output tile, requires even output dims) + iron_winograd_filter_transform_3x3 + iron_winograd_conv2d_3x3_split (pre-transformed filters, removes O(tiles) redundant transform). |
| gemv | ✓ | ✓ | ✓ | kernels/gemm/gemv.rs → iron_gemv<T>. |
| gemv_masked | ✓ | ✓ | ✓ | kernels/gemm/gemv_masked.rs → iron_gemv_masked<T>. |
| quantized (affine_quantize / affine_dequantize) | ✓ | ✓ | ✓ | kernels/gemm/quantized.rs — quantize + dequantize for all widths: int2/int4/int8 (pack-aligned) + int3/int5/int6 (byte-stream). 12 kernels (iron_affine_{quantize,dequantize}_int{2,3,4,5,6,8}). int3/5/6 quantize uses bit-stream OR (lane 0 ORs codes into u32 words) to handle straddling — no atomics. |
| quantized (affine_qmv / qvm / qmm — matvec / matmul) | ✓ | ✓ | ✓ | kernels/gemm/quantized.rs — int4 perf: iron_qmv (8-row-per-TG decode, mirrors MLX qmv_fast) + iron_qmm / _bm2 / _bm4 (M-batched prefill) + iron_qmm_mma / _m16 (simdgroup-matrix MMA prefill) + iron_qmm_mma_mpp (MPP) + iron_qmm_nax (NAX). int8 perf: iron_qmv_int8_fast, iron_qmm_int8_fast / _bm2 / _bm4, iron_qmm_mma_int8 / _m16_int8, iron_qmm_mma_mpp_int8, iron_qmm_nax_int8 — pack-aligned (4 bytes/u32, byte-shift extract), closes the ~6–8× int8-vs-int4 perf gap. Odd-bitwidth MMA: iron_qmm_mma_b{3,5,6} — straddle-aware two-word bit-stream dequant in the 4-SG MMA body. All bit-widths × all dtypes: iron_{qmv,qvm,qmm}_b{3,4,5,6,8} (correctness-first scalar family). qvm perf: iron_qvm_int4_fast — 8-col-per-TG, MLX qvm_fast shape. |
| quantized (gather_qmv / gather_qmm — gather variants) | ✓ | ✓ | ✓ | kernels/moe/gather_qmm.rs → iron_moe_gather_qmm_int4 (int4 affine grouped-gather) + iron_moe_gather_qmm_b{3,5,6,8} (all bit-widths, scalar). int4 perf: iron_moe_gather_qmm_mma_int4{,_bm16} + _m8 (decode) + _m{16,32} (short-prefill, hand-unrolled acc0..accN cells — the DSL doesn’t lower runtime-indexed mutable arrays), MPP scale-ups bm{8,16,64}_mpp (kernels/moe/mpp{,_bm8,_bm64}.rs). int8 perf: pack-aligned iron_moe_gather_qmm_mma_int8 (1-SG MMA decode) + _bm16_mpp + _bm8_mpp (direct-input cooperative tensors, M=8 forbids coop-tensor) + _bm64_mpp (4-SG 2×2 long-context prefill). All MPP kernels stage bf16 through half cooperative tensors via coop_stage(T). Bare-tensor kernels/ops/gather.rs exists but is non-quantized. Expert-indexed dequant GEMV: iron_dequant_gemv_int4_expert_indexed — per-output-row expert selection for the gate/up FFN dispatch shape. |
| moe (router top-k + permute + unpermute orchestration) | ✗ | ✓ | ✓ | kernels/moe/moe_router_topk.rs (iron_moe_router_topk<T>, and its BIAS = 1 row iron_moe_router_topk_softmax_bias<T> — the same top-k ranked on softmax(logits) + expert_bias, for LFM2-MoE) + kernels/moe/moe_permute.rs (iron_moe_permute<T>, iron_moe_unpermute<T>). MoE expert-routing orchestration. The grouped quantized BGEMM that fuses per-expert FFN matmuls is counted under the quantized (gather_*) row. |
| moe routed-expert decode fuse (NVFP4) | ✗ | ✗ | ✓ | kernels/moe/moe_routed_qmv_nvfp4.rs → iron_moe_routed_swiglu_qmv_nvfp4<T> / iron_moe_routed_relu2_qmv_nvfp4<T> (one ACT variant axis: SwiGLU reads gate+up, ReLU² strips the gate buffers via #[optional]) and kernels/moe/moe_routed_down_reduce_nvfp4.rs → iron_moe_routed_down_reduce_nvfp4<T>. The whole decode-step routed-expert block in TWO dispatches instead of ~5·top_k: threadgroup-per-(slot, tile) gate/up QMV with the activation fused into the epilogue, then a top_k-simdgroup down QMV whose simdgroup 0 folds the router weights in T precision. Ported from the Butter MSL literals laguna_moe_routed_{swiglu_qmv,down_reduce}_nvfp4_bf16 / nemotron_moe_routed_{relu2_qmv,down_reduce}_nvfp4_bf16, bit-exact with them; in_dim / out_dim / top_k / slot_stride / global are runtime uniforms where the literals baked Laguna XS and NemotronH-Lightning geometry. Decodes through kernels/primitives.rs. Iron-only. |
| dequant_gather (quantized embedding-table gather) | ✗ | ✗ | ✓ | iron/dequant_gather.rs. int{3,4,5,6,8} all bit-widths. Iron-only. |
| dequant_gemv (quantized GEMV, Iron flavour) | ~ | ~ | ✓ | kernels/gemm/dequant_gemv.rs → iron_dequant_gemv_int{2,3,4,5,6,8}<T> (one-row-per-TG) + iron_dequant_gemv_int4_fast<T> (8-row-per-TG, mirrors MLX qmv_fast). The non-fast int4 kernel stays because Iron’s GPU-router opts into its indirect Swift wrapper. |
| fp_quantized (fp4/fp8 quant + dequant) | ✓ | ✓ | ✓ | kernels/gemm/fp_quantized.rs → iron_fp4_quant_dequant (fp4 E2M1) + iron_fp8_e4m3_quant_dequant / iron_fp8_e5m2_quant_dequant (fp8). Pure arithmetic transform (per-group max-scale + mantissa rounding via floor(log2)/exp2/round); exact for fp8 normals/subnormals, saturating (no NaN/Inf). |
| fp_quantized_mma | ✗ | ✗ | ✓ | kernels/gemm/fp_quantized_mma.rs → iron_fp4_qmm_mma<T> + iron_fp8_e4m3_qmm_mma<T>. Simdgroup-matrix BM=BN=BK=32 MMA — same 4-SG 2×2 scaffold as iron_qmm_mma_b{3,5,6} but with fp4 codebook lookup / fp8 E4M3 biased-exp decode. Not NAX-gated — runs on any M1+. Fills the M>1 perf slot between the scalar round-trip kernels and the NAX-gated fp_quantized_nax. fp4 decode goes through the e2m1_decode intrinsic; fp8 through e4m3_decode. The spec-conformant block-scaled MMA family (all 30 formats) lives in kernels/gemm/block_scaled_mma.rs (its float-scale fp4 kernel is iron_fp4_float_qmm_mma). |
| fp_quantized_nax | ✓ | ✓ | ✓ | kernels/gemm/fp_quantized_nax.rs → iron_fp_qmm_nax<T>. fp4 (E2M1) quantized matmul via NAX matmul2d. Same dequant-into-TG-memory + one cooperative matmul2d per simdgroup per K-block, with fp4 codebook lookup ({0,0.5,1,1.5,2,3,4,6} + sign bit, scale-only). 8 fp4 codes per u32 pack; GROUP_SIZE = 32. Runtime-gated to Apple10+. |
| quantized_nax | ✓ | ✓ | ✓ | kernels/gemm/quantized_nax.rs → iron_qmm_nax<T> (int4) + iron_qmm_nax_int8 (int8, in kernels/gemm/quantized_nax_int8.rs). MPP counterpart of iron_qmm_mma: same int4-dequant-into-TG-memory algorithm, one cooperative matmul2d per simdgroup per K-block; int8 variant uses byte-shift extract (2 packs/lane). Runtime-gated to Apple10+. |
| fft (radix + readwrite + non-pow2) | ✓ | ✓ | ✓ | kernels/kv_cache/fft.rs → iron_fft_n{32,64,128,256,512,1024}<T> (iterative radix-2 Cooley–Tukey, forward + inverse via inv constexpr; complex via parallel real/imag planes). Non-pow2 Bluestein: iron_fft_bluestein_preprocess<T> + iron_fft_bluestein_chirp_filter + iron_fft_bluestein_cmul<T> + iron_fft_bluestein_postprocess<T> — chirp-Z transform wrapping the existing pow2 FFT for arbitrary N in O(N log N); covers Whisper n_fft=400 / 480 with M=1024 padding. Prime-length (Rader) remains a follow-up. |
| hadamard (hadamard_n + hadamard_m) | ✓ | ✓ | ✓ | kernels/ops/hadamard.rs → iron_hadamard_n{64,128,256,512,1024}<T> (FWHT, log2(N) butterfly passes). kernels/ops/hadamard_m.rs → iron_hadamard_m{12,20,28}<T> (non-pow2 M factor, Sloane-table bitmask accumulate). Generic over T. |
| fence | ✓ | ✓ | — | Intentionally out of scope — a GPU-side sync primitive, not a compute kernel. See § Fence ops. |
| gather (bare-tensor embedding lookup) | ✓ | ✓ | ✓ | kernels/ops/gather.rs → iron_gather<T>. Iron’s embedding-table gather. |
| indexing (scatter, scatter_axis, gather_axis, gather_front, masked_scatter) | ✓ | ✓ | ✓ | kernels/ops/gather_axis.rs + kernels/ops/scatter_axis.rs → iron_gather_axis / iron_scatter_axis; kernels/ops/indexing.rs → iron_gather_front, iron_scatter, iron_masked_scatter. All one-thread-per-output Grid3D with bounds guards. |
| aura_encode (codebook quantize, fused) | ✗ | ✓ | ✓ | iron/aura_encode.rs. Bit-widths 2/3/4/6/8. |
| aura_dequant_rotated (bulk dequant to rotated codec space) | ✗ | ✓ | ✓ | iron/aura_dequant_rotated.rs. bits ∈ {2,3,4,6,8}. |
| aura_score (compressed-domain Q·K) | ✗ | ✓ | ✓ | iron/aura_score.rs. bits ∈ {2,3,4,6,8}. Generic over T. |
| aura_value (compressed-domain value aggregation) | ✗ | ✓ | ✓ | iron/aura_value.rs. Sparsity-threshold guard mirrors MLX upstream. Generic over T. |
| aura_flash_pass1 (compressed-domain flash pass 1) | ✗ | ✓ | ✓ | iron/aura_flash_pass1.rs → non-causal aura_flash_pass1_kb{2,3,4,6,8}_vb{2,3,4,6,8}_d{64,128,256} (75 instantiations) + causal aura_flash_pass1_causal_kb{2,3,4,6,8}_vb{2,3,4,6,8}_d{64,128,256} (75) — every (key, value) width the codec ships. Codebooks stack-cached at ≤ 4 bits, buffer-read at 6 / 8. Generic over T. |
| aura_flash_pass2 (cross-block online-softmax merge) | ✗ | ✓ | ✓ | iron/aura_flash_pass2.rs. fp32 accumulators → T final. Generic over T. |
| aura_flash_sdpa (fused single-pass SDPA, sinks variant) | ✗ | ✓ | ✓ | iron/aura_flash_sdpa.rs → aura_flash_sdpa_kb{2,3,4,6,8}_vb{2,3,4,6,8}_d{64,128,256}<T> (75 instantiations). Single-pass online-softmax over compressed K/V with attention sinks + sliding-window causal mask. Codebooks stack-cached at ≤ 4 bits, buffer-read at 6 / 8. |
| flash_quantized_sdpa (single-pass quantized SDPA, affine cache) | ✗ | ✓ | ✓ | iron/flash_quantized_sdpa.rs → base flash_quantized_sdpa_b{4,8}_d{64,96,128,256,512}<T> (10 kernels) + flash_quantized_sdpa_{bool,float}_mask_b{4,8}_d{64,128,256}<T> (12 mask-variant kernels). d=96 = GPT-NeoX (group_size=32 since 96 isn’t a multiple of 64); d=512 = Gemma 4 global attention (dispatches at 256 threads/TG because 16 elems/lane pushes maxTotalThreadsPerThreadgroup below 1024). Bool mask = Tensor<u32> segment-skip, combined with the causal gate; float mask = Tensor<T> per-token logit bias (ALiBi / T5-relative). Bool/float at d={96,512} are follow-ups. |
| gated_delta (GatedDeltaNet recurrence) | ✗ | ✓ | ✓ | kernels/ssm/gated_delta.rs → iron_gated_delta_step<T> (decode) + iron_gated_delta_chunk<T> (chunked-prefill). GDN linear-attention for Qwen3.5 / 3.6 hybrid models. MMA-tiled iron_gated_delta_wy_chunk and fused prep+recurrence iron_gated_delta_prep_step (kernels/ssm/gated_delta_prep.rs) are landed — the latter cuts 3 host commit+wait pairs per GDN layer down to 1. |
| gated_delta_replay (tape capture + state replay) | ✗ | ✓ | ✓ | kernels/ssm/gated_delta_replay.rs → gated_delta_step_record<T> + state_replay<T>. Speculative-decode rollback on GDN. |
| ssm_step (Mamba 2 SSD single-token decode) | ✗ | ✓ | ✓ | kernels/ssm/scan.rs → iron_ssm_step<T>, iron_ssm_step<T> (scalar A). 2D-A_log variant iron_ssm_step_a2d<T> (Jamba): per-(channel, state) A_log, moves Mamba 1 selective scan onto the GPU (previously host-side). Multi-token forms of the same decode contract: kernels/ssm/scan_inplace.rs → iron_ssm_scan_inplace<T> and iron_ssm_scan_a2d_inplace<T> — t_total steps of the identical recurrence over the in-place f32 [H, Ds, Dh] state, bit-for-bit equal to t_total sequential iron_ssm_step / iron_ssm_step_a2d dispatches. The _a2d form is Jamba’s chunked prefill (raw 2-D A_log, A = -exp(A_log) in-kernel). |
| conv1d_causal_step (depthwise SSM conv stream) | ✗ | partial | ✓ | kernels/convolution/conv1d_causal.rs → iron_conv1d_causal_step<T> (streaming decode) + iron_conv1d_causal_prefill (batched, SiLU-fused). The Mamba/SSM short-conv; a plain causal conv1d, so it lives with the convolution family (fused sibling: conv1d_causal_step_silu_cast_many). fp32 state recurrence. Multi-token form of the same contract: kernels/convolution/conv1d_causal_many.rs → iron_conv1d_causal_many<T> — t_len tokens in one dispatch, generic kernel_size (LFM2 uses 3, Mamba 2 / Qwen 3.5 use 4), no activation and no cast, bit-for-bit equal to t_len sequential iron_conv1d_causal_step dispatches with the state threaded through. |
| ssm_replay (sequential tape capture + replay) | ✗ | ✓ | ✓ | kernels/ssm/ssm_replay.rs → ssm_step_record<T> (SSD forward + dA/dBx tape) + ssm_replay<T> (re-fold first k entries). |
| fused_gate_activation (silu/gelu × up gate) | ✗ | ✓ | ✓ | kernels/ops/gated_activation.rs → iron_fused_gate_gelu (gelu-tanh); the silu variant ships in the same file as iron_swiglu. The GPT-OSS activation (min(g, limit), clamp(u, ±limit), g·sigmoid(α·g)·(u+1), limit and α as runtime uniforms) is kernels/ops/swiglu_limit.rs → iron_swiglu_limit_alpha<T>, alongside the DSv4 iron_swiglu_limit<T>; the literal-baked iron_fused_gate_clipped_swiglu was retired in its favour. |
| rms_norm_residual (RMSNorm + residual add fused) | ✗ | ✓ | ✓ | kernels/norm/rms_norm_residual.rs → iron_rms_norm_residual<T>. Reduction-mode, N = TPG*4. ~90 saved dispatches/token on Gemma4-30. |
| rms_norm_rope (RMSNorm + RoPE fused) | ✗ | ✓ | ✓ | kernels/norm/rms_norm_rope.rs → iron_rms_norm_rope<T>. Paired-layout RoPE; Q/K post-projection norm+rope in one dispatch. |
| rms_norm_qgemv (RMSNorm + quantized GEMV fused) | ✗ | ✓ | ✓ | kernels/norm/rms_norm_qgemv.rs → iron_rms_norm_qgemv<T> (int4, one-row-per-TG correctness shape) + iron_rms_norm_qgemv_fast<T> (int4, 8-row-per-TG perf path) + iron_rms_norm_qgemv_int8_fast<T> (int8, 8-row-per-TG). |
| batched_qkv_qgemv (Q/K/V 4-bit qGEMV → 1 dispatch) | ✗ | ✓ | ✓ | kernels/gemm/batched_qkv_qgemv.rs → iron_batched_qkv_qgemv<T> (one-row-per-TG) + iron_batched_qkv_qgemv_fast<T> (8-row-per-TG, GQA-guarded). program_id::<2>() selects Q/K/V, output concatenated [Q|K|V]. |
| kv_cache_update (raw bf16/fp16 single-token append) | ✗ | ✗ | ✓ | kernels/kv_cache/cache.rs → iron_kv_cache_update<T>. Iron-only; raw cache append. |
| kv_cache (affine-quant int4/int8/fp8 quantize + bulk dequant) | ~ | ~ | ✓ | kernels/kv_cache/cache.rs — iron_quantize_kv + iron_bulk_dequant_kv for int4/int8. fp8: iron_quantize_kv_fp8_{e4m3,e5m2} + iron_bulk_dequant_kv_fp8_{e4m3,e5m2}. Per-group amax → scale quantize, byte-shift extract + biased-exp decode. E4M3: mantissa_bits=3, e_bias=-6, max=448; E5M2: mantissa_bits=2, e_bias=-14, max=57344. Closes the host-side fp8 KV round-trip. |
| sampling (softmax + categorical inverse-CDF) | ✗ | ✗ | ✓ | kernels/sampling/categorical_sample.rs → row-aware iron_softmax_categorical_sample. Companion to iron_argmax for T > 0 decode. |
| logits processors (temperature, repetition penalty, top-k / top-p / min-p masks) | ✗ | ✗ | ✓ | kernels/sampling/logits_{processors,topk,top_p,min_p}.rs — in-place decode-form sampler stages composed before iron_softmax_categorical_sample. |
| sdpa_decode + learned attention sink (GPT-OSS-20B) | ✗ | ~ | ✓ | iron/sdpa_decode.rs has_sink / sink_logit constexprs. GPT-OSS-20B’s per-head learned attention-sink logit folds into the cross-simdgroup softmax denominator on-GPU as a virtual key — removing the host-side post-hoc rescale that previously cost a CPU sync per attention layer. |
| gated_rmsnorm (fp32-in gated RMSNorm → activation dtype) | ✗ | ✗ | ✓ | kernels/norm/gated_rmsnorm.rs → iron_gated_rmsnorm<T>. Fused Qwen3.5 / 3.6 GDN post-step out = w·rmsNorm(y)·silu(z); y arrives fp32 (the gated_delta recurrence output). Closes the per-GDN-layer host-side CPU sync (~75 % of Qwen3.5/3.6 layers). |
| conv2d (vision patch conv — im2col + tiled GEMM) | ✓ | ✓ | ✓ | iron/conv2d.rs → iron_conv2d_patch14 / iron_conv2d_patch16 + iron_conv2d_generic. NCHW input, OIHW weight; direct conv (implicit im2col, one thread per output). VLM front-end. |
| patch_embed (fused image unfold + linear projection) | ✗ | ✗ | ✓ | kernels/gemm/patch_embed.rs → iron_patch_embed<T>. Fused image-unfold + linear projection — gathers each patch’s pixels and dots them with one weight row, no intermediate unfolded buffer. MMA-tiled perf path: kernels/gemm/patch_embed_mma.rs → iron_patch_embed_mma<T> — implicit-patch-unfold + 4-SG 2×2 simdgroup-matrix MMA (hidden and num_patches divisible by 32); targets ViT-L/H shapes. |
| rope_2d (2D positional RoPE for vision tokens) | ✓ | ✓ | ✓ | kernels/rope/rope_2d.rs → iron_rope_2d<T>. 2D RoPE over a (row, col) token grid; head_dim split into row half + column half, each running rotate-half RoPE. VLM front-end. |
| mel_spectrogram (STFT + log-Mel filterbank) | ✓ | ✓ | ✓ | kernels/audio/mel_spectrogram.rs → iron_mel_spectrogram<T> (single-dispatch direct-DFT) + radix-FFT path iron_mel_stft_window<T> → iron_fft_n{n_fft}<T> → iron_mel_filterbank<T> (three kernels, O(N log N)). Generic over T. STT front-end. All four kernels are bounds-guarded (idx < n_out) for threadgroup-rounded dispatch. Correctness of the two direct-DFT kernels (iron_mel_spectrogram, iron_mel_spectrogram_magnitude) is gated at f32: their in-thread DFT hits spectrum cancellation nulls where the GPU’s approximate sin/cos diverge from libm by orders of magnitude relative to the (near-zero) true power, which low-precision input rounding moves onto the null — flaky O(6–16) log error on a correct kernel. The kernels stay generic over T; f16/bf16 are covered by iron_mel_filterbank (post-FFT, no in-thread cancellation). |
| audio_conv1d (wide-stride 1D conv — STT patch embed) | ✓ | ✓ | ✓ | iron/audio_conv1d.rs → audio_conv1d<T>. Dense wide-stride multi-channel 1D conv (NCL); distinct from depthwise conv1d_causal_step. STT front-end. |
| vocoder / iSTFT (TTS waveform synthesis) | ✓ | ✓ | ✓ | kernels/audio/vocoder.rs → iron_vocoder_istft<T>. Inverse-STFT overlap-add — one thread per output sample gathers every covering frame, inverse-DFTs with Hermitian symmetry, COLA-normalises. TTS waveform synthesis. |
Quantization precision coverage
Section titled “Quantization precision coverage”The op-coverage table above records which ops exist; this section records which precisions each weight-bearing op supports.
A quantized weight stores a small code per element plus a scale that restores magnitude. Formats vary along three orthogonal axes — bit-width is not one of them:
- Element — a signed integer (
int2…int8) or a micro-float codebook (E2M1 / E4M3 / E5M2). The element is what the code decodes to. - Scale — how the per-block scale is stored: a raw FP32 per large group, a
compact
E8M0power-of-two per small block (OCP MX), anE4M3micro-scale × a tensor-wide FP32 global (NVFP4), or FP16 (the memory-halving*_f16twins). - Zero-point — symmetric (no zero-point) or asymmetric (scale + bias). Every Track-1 format below is symmetric; the asymmetric integer track (Track 2) carries a zero-point for MLX-checkpoint interop.
The two “tracks” are just the practical split of that space: Track 1 = symmetric block-/group-scaled (the spec float formats + the integer family); Track 2 = asymmetric affine integers.
Track 1 — symmetric block-scaled / float-scale (quant::format::QFormat)
Section titled “Track 1 — symmetric block-scaled / float-scale (quant::format::QFormat)”A weight [N, K] is quantized in contiguous K-blocks: per-element codes +
one block scale, no zero-point. 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 for int2/3/5/6).
The in-kernel E2M1 and E4M3 decodes are bit-pattern constructions, not
value ladders: both codebooks are exactly a slice of the IEEE binary16 bit
space, so kernels::primitives::iron_decode_e2m1 / iron_decode_e4m3 build a
half pattern with a handful of shifts, reinterpret it (Op::Bitcast to F16),
and restore the exponent offset with an exact power-of-two multiply. Bit
equality with the ladder forms they replaced is pinned exhaustively — all 16
E2M1 codes, all 256 E4M3 bytes — in
tests/nvfp4_fast_decode_gpu_correctness.rs.
The Track-1 formats:
| format | element | block | block scale | global |
|---|---|---|---|---|
nvfp4 |
E2M1 | 16 | E4M3 (1 B) | FP32 |
mxfp4 |
E2M1 | 32 | E8M0 pow-2 (1 B) | — |
mxfp8_e4m3 |
E4M3 | 32 | E8M0 (1 B) | — |
mxfp8_e5m2 |
E5M2 | 32 | E8M0 (1 B) | — |
nvfp8 |
E4M3 | 16 | FP32 (4 B) | — |
fp4 (legacy) |
E2M1 | 32 | FP32 per-group | — |
fp8_e4m3 (legacy) |
E4M3 | 32 | FP32 per-group | — |
fp8_e5m2 (legacy) |
E5M2 | 32 | FP32 per-group | — |
int2 / int3 / int4 / int5 / int6 / int8 (symmetric) |
int N | 64 | FP32 per-group | — |
mxint2 / mxint3 / mxint4 / mxint5 / mxint6 / mxint8 |
int N | 32 | E8M0 pow-2 (1 B) | — |
*_f16 (twins of the FP32-scaled formats above) |
as twin | as twin | FP16 (2 B) | — |
30 formats. The integer members are symmetric int-N: a plain FP32 group scale
(int*, group 64) or an OCP-MX-style E8M0 power-of-two block scale (mxint*,
block 32 — mxint8 is OCP-ratified MXINT8, the rest follow the same
construction; these 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. Every FP32-scaled format additionally has an FP16-scale twin
(nvfp8_f16, fp4_f16, fp8_e4m3_f16, fp8_e5m2_f16, int2-6_f16, int8_f16)
— same element + block, scale stored as a 2-byte IEEE half (the layout real
checkpoints use; the host encoder does correct round-to-nearest-even with
subnormal support, since wide-range elements like E5M2 push scales into f16’s
subnormal range).
Track 2 — asymmetric affine int (scale + bias)
Section titled “Track 2 — asymmetric affine int (scale + bias)”The asymmetric integer track: int2 / int3 / int4 / int5 / int6 / int8,
per-group (64) scale + bias (zero-point), in kernels/gemm/quantized.rs,
kernels/gemm/dequant_gemv.rs, iron/dequant_gather.rs, kernels/kv_cache/cache.rs, and the
int4+int8 MoE / MMA / MPP / NAX perf kernels. The defining difference from the
Track-1 integers (int* / mxint*) is the zero-point — Track 2 is the only
track that can represent a lopsided range.
This track is current, not legacy — it predates Track 1 (it’s where 4/8-bit support started, for KV-cache + MLX model quant) and the zero-point keeps it irreplaceable:
- it is the on-disk interop format for MLX-quantized checkpoints (
mlx_lm.convert -qemits asymmetric affine codes +scalesandbiases;w = scale·q + bias). The Track-1 integers are symmetric (no bias), so even though block-scaledint4/mxint4now exist, they cannot represent an asymmetric MLX checkpoint — affine int4 (with its zero-point) is what’s required to load every MLX 4-bit model, by design, not for lack of a 4-bit block-scaled format; - it is the right scheme for per-decode-step KV-cache quant (cheap min/max → scale+bias; block-scaled is a static-weight format whose per-step encode would need GPU encode intrinsics).
Track 1 is a parallel symmetric family (the spec float formats + the full symmetric integer
matrix); it does not replace Track 2. (The float-scale fp4/fp8 within Track 1 — raw f32
group scale — are legacy, superseded by spec mxfp4/nvfp4/mxfp8/nvfp8, kept as labeled
comparison variants.)
Block-scaled coverage — every family supports all 30 Track-1 formats
Section titled “Block-scaled coverage — every family supports all 30 Track-1 formats”Each (family × format) ships a 1:1 #[test_kernel] (GPU-verified vs
quant::format::dequant) + a #[bench] with .flops() so the latency / GFLOP/s /
roofline columns rank precisions side by side. fp8_e4m3 reuses each family’s
nvfp8 kernel (identical 8-bit-E4M3 + FP32-scale shape); the rest decode in their
own kernel. The integer formats (int2-6, mxint2-8) are generated by a
parameterized (bit-width × scale-kind) decode macro per family — the same
straddle-aware bit-stream extract + float sign-extend everywhere — so they reuse
each family’s proven dispatch geometry verbatim (no new freeze surface).
| family | path | file(s) | also on affine int track |
|---|---|---|---|
| dequant (standalone) | elementwise | mlx/block_scaled_dequant.rs |
int2–8 |
| qgemv (GEMV decode) | reduction | kernels/gemm/block_scaled_matmul.rs |
int2–8, int4/int8-fast |
| qmm (GEMM prefill) | reduction | kernels/gemm/block_scaled_qmm.rs |
int2–8 |
| qmm — simdgroup-MMA | simdgroup-matrix | kernels/gemm/block_scaled_mma.rs |
int4, int8 |
| qmm — MPP (tensor engine) | MPP matmul2d |
kernels/gemm/block_scaled_qmm_mpp.rs |
int4, int8 |
| qmm — NAX | NAX matmul2d |
kernels/gemm/block_scaled_qmm_nax.rs |
int4, int8 |
| MoE gather-qmm | reduction | kernels/moe/block_scaled.rs |
int3–8 |
| MoE gather — MPP (bm8/16/64) | MPP | kernels/moe/mpp{,_bm8,_bm64}_block_scaled.rs |
int4, int8 |
| expert-indexed GEMV | reduction | kernels/moe/dequant_gemv_expert_indexed_block_scaled.rs |
int4 |
| fused RMSNorm + GEMV | reduction | kernels/norm/rms_norm_block_scaled_qgemv.rs |
int4, int8-fast |
| fused gated-RMSNorm + GEMV | reduction | kernels/norm/gated_rms_norm_block_scaled_qgemv.rs |
int4 |
| batched-Q/K/V qgemv + qmm | reduction | kernels/gemm/batched_qkv_block_scaled_{qgemv,qmm}.rs |
int4, int8-fast |
| batched-4 qgemv + qmm | reduction | kernels/gemm/batched_4_block_scaled_{qgemv,qmm}.rs |
int4 |
| embedding gather | elementwise | iron/dequant_gather_block_scaled.rs |
int3–8 |
| flash SDPA (block-scaled KV) | flash | iron/flash_block_scaled_sdpa.rs (d64/96/128/256/512, all 30 formats¹) |
affine int4/int8 KV, same dims |
| patch embed (linear projection) | reduction | kernels/gemm/patch_embed_block_scaled.rs |
— |
| patch embed (simdgroup-MMA) | simdgroup-matrix | kernels/gemm/patch_embed_mma_block_scaled.rs |
— |
| conv2d / conv3d (direct) | reduction | iron/{conv2d,conv3d}_block_scaled.rs |
— |
| conv2d / conv3d (im2col-MMA) | simdgroup-matrix | iron/{conv2d,conv3d}_mma_block_scaled.rs |
— |
| depthwise conv2d | reduction | iron/depthwise_conv2d_block_scaled.rs |
— |
| audio conv1d (STT front-end) | reduction | iron/audio_conv1d_block_scaled.rs |
— |
| fishspeech conv1d (TTS front-end) | reduction | iron/fishspeech_conv1d_block_scaled.rs |
— |
¹ Flash KV covers every production head dim (d64/96/128/256/512), each × all 30 formats —
no holes. int8’s group size (64) doesn’t divide d96, so that case tiles with a ragged
trailing block: n_blocks = ceil(dim/block_size) (a 64-block + a 32-block), with the host
packer and kernel rounding up identically so codes + scales stay self-consistent. The
geometry is one simdgroup per query (grid [32, n_query, 1]), identical across dims (only
the per-lane dim count changes).
The full integer matrix everywhere
Section titled “The full integer matrix everywhere”The complete symmetric integer family — int2/3/4/5/6/8 (FP32 group scale) and
mxint2/3/4/5/6/8 (E8M0 block scale) — is 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 where the mxint* E8M0 layout maps to tensor-core
block-scaling on future NVIDIA / AMD targets. The core matmul / MoE / RMSNorm-GEMV /
batched-QKV / KV-cache / attention families additionally carry the pre-existing
asymmetric affine integers (scale + bias) for MLX-checkpoint interop. No weight-bearing
family lacks an integer path.
Model-format decode — one codec, no per-oracle drift
Section titled “Model-format decode — one codec, no per-oracle drift”The Track-1 QFormat matrix above is wh-iron’s own block-scaled layout. Alongside it,
several kernels consume external model formats with their own on-disk byte layouts but
the same element/scale arithmetic. These all now decode through the shared
quant::codec primitives (e2m1_decode,
e4m3_decode, e8m0_decode, int8_decode, f16_scale_decode) and the
quant::gguf host packer/oracle, the single
source of truth the kernel, the host quantizer, and the CPU correctness oracle all read,
so an oracle can no longer drift from its kernel (the bug class fixed twice, independently,
once per duplicated copy of the layout map):
| Format | Provenance | Element × scale | Decode source |
|---|---|---|---|
q8_0 |
GGUF / llama.cpp | int8 × f16 block scale | gguf::{pack,dequant}_q8_0 → codec |
q2_k |
GGUF k-quant | 2-bit × two-level super-block (d·scale − dmin·min) | gguf::{pack,dequant}_q2_k, gguf::q2_k_qpos → codec |
| DSv4 fp8-block | DeepSeek-V3 safetensors | e4m3 × per-(128×128) f32 | codec::e4m3_decode (NaN sentinel kept explicit) |
| DSv4 mxfp4 | DeepSeek-V3 safetensors | = OCP mxfp4 (e2m1 × e8m0) |
codec::{e2m1,e8m0}_decode |
The MoE Q2_K correctness oracles (moe_gather_down_q2k, moe_bgemm_q2k_mpp) import the one
gguf::q2_k_qpos index map rather than each carrying a private copy.
iq2_xxs(GGUF i-quant) is a codebook format — the “element” is a 256×8 signed-octet grid lookup plus a 7-bit sign-parity expansion, not anelement × scaledecode — so it sits outside the codec matrix. The kernel is a WIP scaffold (grid table unlanded; ABI/shape smoke-test only) and its MoE siblings (moe_*_iq2xxs) share that status.gemm_q8/gemv_q8consume theq8_0format but are bench-only (no correctness oracle yet) — a test-coverage gap, not a decode-drift risk.
Gaps / deliberate exclusions
Section titled “Gaps / deliberate exclusions”- Winograd conv — the filter pre-transform (
GgGᵀ) amplifies quantization error; quantized Winograd is non-standard and counterproductive, so it stays f16/bf16/f32. - Activation-only ops (RoPE, SSM / GatedDeltaNet recurrence, standalone norms, dense
SDPA / GEMM, elementwise / reduction / softmax / sort / scan / FFT / gather-axis) carry
no persistent weight tensor → activation-precision (
<T>) only. - KV-cache write quantizes per decode step → affine int4/int8/fp8 (cheap min/max encode) is the right scheme; block-scaled is a static-weight format whose per-step encode would need GPU encode intrinsics.
Notes on counting decisions
Section titled “Notes on counting decisions”A few rows mix multiple .metal files into one op or split one file into multiple ops:
sdpa_vector*is counted as two ops:sdpa_vector(single pass) +sdpa_vector_2pass(two-pass pair). Upstreamsdpa_vector.hdefinessdpa_vector,sdpa_vector_2pass_1,sdpa_vector_2pass_2.- AURA stack — each codec stage (
encode,dequant_rotated,score,value,flash_pass1,flash_pass2) is a separate row;aura_flash_sdpa(sinks-fused single-pass) is also its own row. steel/— each kernel file becomes one op row; per-block-shape instantiations are not counted separately.steel_attention(scalar) andsteel_attention_mma(simdgroup-MMA) are two rows because they are separately compiled kernels with different lowering strategies.quantized.metal— split into four rows by semantic operation (quant/dequant, qmv/qvm/qmm matmul, gather-qmv/qmm, fp4/fp8). The Apple10+ variants (quantized_nax,fp_quantized_nax) are separate rows because they live in separate modules with runtime-only dispatch gating.fp_quantized_mmais its own row (runs on M1+, no Apple10 gating).indexing/is one row covering scatter / scatter_axis / gather_axis / gather_front / masked_scatter. Baregatheris its own row (Iron-specific).moeis the routing (kernels/moe/router_topk.rs) + permute/unpermute (kernels/moe/permute.rs). The grouped quantized BGEMM lives under thequantized (gather_*)row.logits processorsis one row for the Iron sampler-stage kernels (temperature,repetition_penalty,topk/top_p/min_pmasks).- Cells marked
~indicate a partial port (typically one bit-width, one dtype, or one block shape where upstream has many) — see the notes column for the specific gap.
Out-of-tree micro-optimization proposals
Section titled “Out-of-tree micro-optimization proposals”Some hot-path patterns require codegen-layer support to land cleanly and are documented as proposals rather than landed kernels. See specs/PROPOSED_OPTIMIZATIONS.md for full rationale and implementation sketches:
simd_broadcastfor scale/bias — int4/int8 GEMV kernels where 4 (int4) / 16 (int8) consecutive lanes share a group scale/bias. Hardware already coalesces same-address loads from one simdgroup, so the optimization is opportunistic (no measured profile signal yet).fast::math intrinsics —iron_mel_spectrogram,iron_softmax,iron_logsumexp,iron_vocoder_istftuse IEEE-precise built-ins. Switching tofast::exp/fast::log/fast::sin/fast::coswould give ~1.5–2× speedup at 1–3 ULP. Needs newUnaryOpKindIR variants + precision validation against existing test tolerances.- K-loop software pipelining — overlap next K-block load with current MMA in MMA-tiled K-loop kernels. ~15–25 % throughput win on M3+. Needs a new
Op::PrefetchAsyncIR op + aprefetch.rscodegen pass.
Already in place: float4 / half4 vectorized X loads via the existing VectorizePass (crates/wh-iron-codegen/src/passes/vectorize.rs). fp32 accumulators are correctness-required across all production shapes; the f16/bf16-accumulator proposal was rejected.
Fence ops — intentionally out of scope
Section titled “Fence ops — intentionally out of scope”MLX’s fence.metal (mlx/backend/metal/kernels/fence.metal, ~52 lines) is not a compute kernel — it’s a GPU-side synchronisation primitive. Deliberately not ported to wh-iron; the fence audit row is marked — rather than ✗.
What the fence ops are
Section titled “What the fence ops are”Three kernels: input_coherent (force input-buffer visibility), fence_update (bump a counter in a shared buffer), and fence_wait — a compute kernel that spin-loops reading that counter until it changes. Together they order work across command buffers / streams without a CPU round-trip.
How MLX uses them
Section titled “How MLX uses them”mlx/backend/metal/fence.cpp’s FenceImpl has two paths:
- Default:
device->newSharedEvent()— a standardMTLSharedEvent. The wait executes in the GPU command processor, not a shader core. use_fastpath (thefence.metalspin-wait kernels): gated behindGPUFamilyMetal3+ macOS 15 + an opt-in env var (metal_fast_synch). Off by default.
So MLX itself treats the GPU spin-wait fence as an opt-in latency micro-optimization for its multi-stream async_eval workloads — not a primitive every pipeline needs.
Why Iron doesn’t need it
Section titled “Why Iron doesn’t need it”- Iron’s current pipeline is single-stream autoregressive decode. Within a forward pass, Metal’s automatic hazard tracking orders kernels in a command buffer for free; across command buffers on one queue, submission order suffices.
- CPU/GPU pipelining (build command buffer N+1 while the GPU runs N) is
commit+ completion handlers, not a fence. - For genuine cross-queue / cross-stream GPU sync,
MTLEvent/MTLSharedEvent(encoder-level —encodeWaitForEvent/encodeSignalEvent) are the correct, power-efficient primitive, and they belong inwh-iron-runtime’s dispatch layer, not as a#[kernel]. - A
fence_waitspin-wait is a deliberate near-infinite GPU loop: it burns a shader core + power, and a counter that never updates (a bug, a wrong dispatch) is a permanent GPU pin → hard reboot.
When this could change
Section titled “When this could change”If Iron later runs multiple concurrent GPU streams — e.g. speculative decoding (draft/target overlap), prefill/decode overlap, or ANE+GPU concurrency — it will need cross-stream ordering. The right implementation is MTLEvent-based encoder-level sync added to wh-iron-runtime (MLX’s own default), not a spin-wait #[kernel]. Only if profiling later shows that MTLEvent’s command-processor latency is a measured bottleneck for an ultra-fine-grained sync pattern would the opt-in spin-wait become worth revisiting — and even then it’s a runtime concern, not a wh-iron kernel.
