Aller au contenu

wh-iron kernel-op coverage audit

Ce contenu n’est pas encore disponible dans votre langue.

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).

  • 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::InlineMsl IR to the #[kernel] DSL via the coop_tile_* intrinsics + coop_stage(T) (bf16 → half staging because Apple’s matmul2d mishandles bfloat cooperative 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-support work added spec-conformant nvfp4 / mxfp4 / mxfp8 (e4m3+e5m2) / nvfp8, legacy float-scale fp4 / fp8 (e4m3+e5m2), the complete symmetric integer matrixint2/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 the quant::format CPU 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}, and flash_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 12 flash_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 (“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.

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_apple10 runtime 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).

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.

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.

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 MLX (upstream) MLX (ekryski@alpha) wh-iron Notes
arange kernels/ops/arange.rsiron_arange. Generic T.
arg_reduce (argmax/argmin → u32 index) kernels/ops/arg_reduce.rsiron_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.rsiron_binary_two<T>.
copy (contiguous) kernels/ops/copy.rsiron_copy<T>.
copy (strided / general) kernels/ops/strided.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_softmax<T> (looped + single-row collapsed).
logsumexp kernels/ops/logsumexp.rsiron_logsumexp<T>.
layer_norm kernels/norm/layer_norm.rsiron_layer_norm<T>.
rms_norm kernels/norm/rms_norm.rsiron_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.rsiron_rope.
rope (frequency-band scaled) kernels/rope/rope_banded.rsiron_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.rsiron_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.rsiron_sdpa<T>. Scalar SDPA for short sequences.
sdpa_vector (GQA decode, single pass) mlx/sdpa_vector.rsiron_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.rsiron_sdpa_decode<T> + sdpa_decode_d{64,256,512}.rs. Iron-only; kv_striden_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.rssdpa_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_steel_gemm_gather_nax<T>. Gather GEMM via NAX matmul2d. Runtime-gated to Apple10+.
steel_gemm_masked kernels/gemm/steel/steel_gemm_masked.rsiron_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.rsiron_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.rsiron_conv2d_patch14 / iron_conv2d_patch16 / iron_conv2d_generic. Direct conv (implicit im2col, one thread per output). MMA-tiled perf path: iron/conv2d_mma.rsiron_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.rsiron_conv3d_generic + iron_conv3d_grouped (depthwise + dilation). 5D NCDHW / OIDHW. MMA-tiled perf path: iron/conv3d_mma.rsiron_conv3d_mma<T> — same MMA scaffold as 2D, decomposed over (kd, kh, kw, ic).
steel_conv_general (strides/dilation/groups) iron/conv2d.rsiron_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.rsiron_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.rsiron_gemv<T>.
gemv_masked kernels/gemm/gemv_masked.rsiron_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.rsint4 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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_hadamard_n{64,128,256,512,1024}<T> (FWHT, log2(N) butterfly passes). kernels/ops/hadamard_m.rsiron_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.rsiron_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.rsiron_gather_axis / iron_scatter_axis; kernels/ops/indexing.rsiron_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.rsaura_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.rsiron_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.rsgated_delta_step_record<T> + state_replay<T>. Speculative-decode rollback on GDN.
ssm_step (Mamba 2 SSD single-token decode) kernels/ssm/scan.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsssm_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_kv_cache_update<T>. Iron-only; raw cache append.
kv_cache (affine-quant int4/int8/fp8 quantize + bulk dequant) ~ ~ kernels/kv_cache/cache.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsiron_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.rsaudio_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.rsiron_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.

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:

  1. Element — a signed integer (int2…int8) or a micro-float codebook (E2M1 / E4M3 / E5M2). The element is what the code decodes to.
  2. Scale — how the per-block scale is stored: a 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 *_f16 twins).
  3. Zero-pointsymmetric (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 -q emits asymmetric affine codes + scales and biases; w = scale·q + bias). The Track-1 integers are symmetric (no bias), so even though block-scaled int4/mxint4 now 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 complete symmetric integer familyint2/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_0codec
q2_k GGUF k-quant 2-bit × two-level super-block (d·scale − dmin·min) gguf::{pack,dequant}_q2_k, gguf::q2_k_qposcodec
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 an element × scale decode — 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_q8 consume the q8_0 format but are bench-only (no correctness oracle yet) — a test-coverage gap, not a decode-drift risk.
  • 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.

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). Upstream sdpa_vector.h defines sdpa_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) and steel_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_mma is its own row (runs on M1+, no Apple10 gating).
  • indexing/ is one row covering scatter / scatter_axis / gather_axis / gather_front / masked_scatter. Bare gather is its own row (Iron-specific).
  • moe is the routing (kernels/moe/router_topk.rs) + permute/unpermute (kernels/moe/permute.rs). The grouped quantized BGEMM lives under the quantized (gather_*) row.
  • logits processors is one row for the Iron sampler-stage kernels (temperature, repetition_penalty, topk / top_p / min_p masks).
  • 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.

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_broadcast for 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 intrinsicsiron_mel_spectrogram, iron_softmax, iron_logsumexp, iron_vocoder_istft use IEEE-precise built-ins. Switching to fast::exp/fast::log/fast::sin/fast::cos would give ~1.5–2× speedup at 1–3 ULP. Needs new UnaryOpKind IR 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::PrefetchAsync IR op + a prefetch.rs codegen 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.

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 .

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.

mlx/backend/metal/fence.cpp’s FenceImpl has two paths:

  • Default: device->newSharedEvent() — a standard MTLSharedEvent. The wait executes in the GPU command processor, not a shader core.
  • use_fast path (the fence.metal spin-wait kernels): gated behind GPUFamilyMetal3 + 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.

  • 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 in wh-iron-runtime’s dispatch layer, not as a #[kernel].
  • A fence_wait spin-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.

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.