feat(bench): wall-clock latency, GFLOP/s, roofline %-of-peak + bottleneck verdict - #247
Merged
Merged
Conversation
This was referenced Jun 1, 2026
ekryski
force-pushed
the
ek/bench-metrics
branch
3 times, most recently
from
June 1, 2026 22:37
c24c5d8 to
2a14b7c
Compare
Collaborator
Author
|
Rebased off of #249 now that it has landed. |
ekryski
commented
Jun 1, 2026
| cmd::diff as diff_cmd, | ||
| git, | ||
| matches_filter, | ||
| suite_printer::{ProfileRow, SuitePrinter}, |
Collaborator
Author
There was a problem hiding this comment.
Superseded by the new DerivedMetrics.
ekryski
commented
Jun 2, 2026
ekryski
force-pushed
the
ek/bench-metrics
branch
4 times, most recently
from
June 2, 2026 20:17
96dbd39 to
a173deb
Compare
9 tasks
Collaborator
|
[Corrected — my original explanation below was wrong; keeping it visible for the record.] These failures are not M5-hardware-gated, and not a regression in this PR. Two distinct things were mixed in the list:
So: dtype-invariance ≠ "hardware absent" as I first claimed — for (1) it's a runtime-compiler/OS-version issue, for (2) it was my test harness. Apologies for the earlier misdiagnosis. Original (incorrect) commentOn the 57 Test failures here — these are hardware-gated, not a regression in this PR (they're red on Root cause: the Test job runs on Why they hard-fail rather than skip: the integration Possible fix (separate, framework-level): add a |
Collaborator
Author
|
I merged PR #256, which should resolve it. |
Collaborator
Author
|
Rebasing off of #257 that just landed. |
Rebased onto the subprocess refactor in #249 (squashed). Adds per-bench wall-clock latency (us), compute throughput (GFLOP/s), roofline percent-of-peak bandwidth and compute, arithmetic intensity, and a combined occupancy/register/roofline bottleneck verdict, rendered under tile bench -v / -vv. Port (bench + runner internals moved in #249): - profile.rs and device_specs.rs relocated metaltile-std/src -> metaltile/src/runner/. The BenchResult builder lives in the metaltile crate now and metaltile-std depends on metaltile, so the analysis must sit crate-side to avoid a dependency cycle. Both are pure CPU analysis. - Metrics computed in cmd/bench.rs derive_metrics, attached to each OpResult via DerivedMetrics (replacing the old per-kernel ProfileRow map + set_profile_map); SuitePrinter renders them from result.metrics(). - 54 kernels keep their flops #[bench] annotations; the macro flops key plus BenchSetup::flops carry through. Protocol made subprocess-ready (per the schema agreed with the v4 plan): - protocol::ProfileInfo extended to carry the full metric set (gflops, pct_peak_bw, pct_peak_flops, arith_intensity) and its occ_pct/regs/bottleneck fields made Option for clean serialization on spec-less devices (CI's virtualized GPU). All new fields serde-default. - From<&DerivedMetrics> for ProfileInfo in bench_types — single source of truth so the future __tile_runner emit path produces identical numbers to the in-process render. - SuitePrinter::print_bench_result (Phase-2 wire path) renders a compact GFLOP/s + bottleneck suffix from the parsed profile. Verified across the workspace lib + CLI test suites; tile bench -v renders every column.
End-to-end map of how a #[kernel] becomes a compiled shader and how the bench runner, test runner, and kernel profiling work. Covers the crate graph, command dispatch via the TileCommand trait, the sequential bench loop with its timing-isolation guarantee (metrics are computed after the timed region, never inside it), the rayon-parallel test oracle pass, and the current in-process execution vs the scaffolded (not yet active) __tile_runner subprocess model. Linked from the docs index.
…table Replace the device_specs roofline ceilings with the published per-chip figures (GPU FP32 vector, matrix compute block, memory bandwidth) for M1–M5: - Correct stale FP32 vector peaks (M3 Max 17.5→12.8, M4 Max 21.1→17.2, M5 Max 24→16.6, M3 base 4.1→3.5, M3 Pro 6.2→7.1, M4 Pro 10.4→8.6, M2 Max 15.8→13.6, M3 Ultra 31.8→28.4, M1 Pro/Max 5.2/10.4→5.3/10.6). - M5 GPU Neural Accelerator FP16 matmul ceilings now match the table (base 16.8 / Pro 33.2 / Max 66.4 / Ultra 132.8 TFLOPS) — add M5 Pro (was unconfirmed → SIMD-only) and the projected M5 Ultra (Ultra = 2× Max). - Memory bandwidth takes the upper bound of any published range. Add `ane_tops`: the separate Apple Neural Engine "Matrix Compute Block" (M1–M4 Standalone/Dual ANE Block, 11–38 TOPS), recorded for the upcoming ANE kernels/benches. Distinguishes the standalone ANE from the M5+ GPU-embedded Neural Accelerator (`na_f16_tflops`). No M4 Ultra (never shipped). Per-family comments corrected; tests updated. Allowlist `ane` in typos config.
`tile bench` compiled + dispatched the MLX reference kernel for every kernel with a counterpart, on every run (~2× the work). The metaltile kernels have superseded the MLX references and per-kernel correctness now lives in `tile test` (the #[test_kernel] CPU-oracle harness), so the MLX A/B is now opt-in rather than always-on: - Add `tile bench --mlx` (alias `--reference`); default off. - Thread `compare_mlx` through `run_kernel_bench` and gate `run_reference_bench` on it. Default output benches only the metaltile kernels (the Ref / MT% columns stay blank); `--mlx` restores the side-by-side speed + output-equivalence comparison. Help/docs updated.
Reflect the new bench default in docs/cli.md: `tile bench` benches only the MetalTile kernels; the MLX A/B (speed + output-equivalence) is opt-in via `--mlx` (alias `--reference`). Updates the section title, intro, usage line, flags table (new `--mlx` row), and the metrics description (`Ref` / `MT%` only populated with `--mlx`). Also fixes a stale `device_specs.rs` path.
Align the bench table with the rest of the CLI's styling (cf. `tile device`): - column headers BrightWhite-bold → BrightBlack-bold, so the muted headers recede and the BrightWhite data values stand out (the label convention used everywhere else post-#255). - GFLOP/s value Cyan → BrightWhite — Cyan is reserved for the op/title row.
ekryski
added a commit
that referenced
this pull request
Jun 3, 2026
…cross every weight-bearing kernel (#250) ## What this is PR 2 of the precision series: make fp4/fp8 **spec-conformant block-scaled** and bring the **full op × quant-format matrix to completion** — *every* weight-bearing kernel family supports **all 30 Track-1 quant formats** in fp16/bf16/fp32 activation, so the bench columns (GFLOP/s, %-peak, roofline) reveal the fastest precision per op. **No holes.** > Originally stacked on #247 (`ek/bench-metrics`); that has since merged into `dev`, so this branch is now precision-only on top of `dev`. ## Quantization is three orthogonal axes A quantized weight stores a small **code** per element + a **scale** that restores magnitude. Formats vary along three axes — **bit-width is not one of them**: 1. **Element** — a signed integer (`int2…int8`) *or* a micro-float codebook (E2M1 / E4M3 / E5M2). 2. **Scale** — how the per-block scale is stored: raw **FP32** per large group, a compact **E8M0** power-of-two per small block (OCP MX), an **E4M3** micro-scale × a tensor-wide FP32 global (NVFP4), or **FP16** (the memory-halving twin). 3. **Zero-point** — **symmetric** (no zero-point) or **asymmetric** (scale + bias). Track 1 (this PR) is the **symmetric** block-/group-scaled family; Track 2 is the pre-existing **asymmetric** affine integers (zero-point for MLX-checkpoint interop). The two are complementary, not a hierarchy. ## The 30 Track-1 formats | group | formats | element | block | scale | |---|---|---|---|---| | spec float | `nvfp4` | E2M1 | 16 | E4M3 + global FP32 | | | `mxfp4` / `mxfp8_e4m3` / `mxfp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | E8M0 pow-2 | | | `nvfp8` | E4M3 | 16 | FP32 | | legacy float | `fp4` / `fp8_e4m3` / `fp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | FP32 per-group | | **symmetric int** | `int2 / int3 / int4 / int5 / int6 / int8` | int N | 64 | FP32 per-group | | **MXINT** | `mxint2 / mxint3 / mxint4 / mxint5 / mxint6 / mxint8` | int N | 32 | E8M0 pow-2 | | **fp16-scale twins** | `*_f16` of every FP32-scaled format above | as twin | as twin | FP16 (2 B) | **30 formats** = 8 float + 6 symmetric int (FP32) + 6 MXINT (E8M0) + 10 fp16-scale twins. `mxint8` is OCP-ratified MXINT8; `mxint*` map to tensor-core block-scaling units for future NVIDIA / AMD targets. Sub-byte codes (2/3/4/5/6-bit) tight-bit-pack LSB-first into u32 words (a 4-bit stream is byte-identical to the classic nibble layout); 8-bit is one byte/code. All formats share the `quant::codec` bit primitives — one source of truth for the host packer, the CPU correctness oracle, and the in-kernel decode (the `e2m1_decode` / `e4m3_decode` / `e5m2_decode` / `int8_decode` intrinsics + the straddle-aware sub-byte bit-stream extraction). The integer + fp16 formats are emitted by a parameterized **(bit-width × scale-kind)** decode macro per family, so every variant reuses its family's proven dispatch geometry byte-for-byte (no new freeze surface). ## Coverage matrix — COMPLETE Every family below supports **all 30 formats**. Each `(family × format × dtype)` ships a `#[kernel]`, a 1:1 `#[test_kernel]` (GPU-verified vs `quant::format::dequant`), and a `#[bench]` with `.flops()`. | family | path | file(s) | |---|---|---| | dequant (standalone) | elementwise | `mlx/block_scaled_dequant.rs` | | qgemv (GEMV decode) | reduction | `mlx/block_scaled_matmul.rs` | | qmm (GEMM prefill) | reduction | `mlx/block_scaled_qmm.rs` | | qmm — simdgroup-MMA | simdgroup-matrix | `mlx/block_scaled_mma.rs` | | qmm — MPP | tensor engine | `mlx/block_scaled_qmm_mpp.rs` | | qmm — NAX | tensor engine | `mlx/block_scaled_qmm_nax.rs` | | MoE gather-qmm | reduction | `mlx/block_scaled_moe.rs` | | MoE gather — MPP (bm8/16/64) | tensor engine | `ffai/moe_mpp{,_bm8,_bm64}_block_scaled.rs` | | expert-indexed GEMV | reduction | `ffai/dequant_gemv_expert_indexed_block_scaled.rs` | | fused RMSNorm + GEMV | reduction | `ffai/rms_norm_block_scaled_qgemv.rs` | | fused gated-RMSNorm + GEMV | reduction | `ffai/gated_rms_norm_block_scaled_qgemv.rs` | | batched-Q/K/V qgemv + qmm | reduction | `ffai/batched_qkv_block_scaled_{qgemv,qmm}.rs` | | batched-4 qgemv + qmm | reduction | `ffai/batched_4_block_scaled_{qgemv,qmm}.rs` | | embedding gather | elementwise | `ffai/dequant_gather_block_scaled.rs` | | flash SDPA (block-scaled KV) | flash | `ffai/flash_block_scaled_sdpa.rs` — **d64/96/128/256/512** | | patch embed (linear projection) | reduction | `ffai/patch_embed_block_scaled.rs` | | patch embed (simdgroup-MMA) | simdgroup-matrix | `ffai/patch_embed_mma_block_scaled.rs` | | conv2d / conv3d (direct) | reduction | `ffai/{conv2d,conv3d}_block_scaled.rs` | | conv2d / conv3d (im2col simdgroup-MMA) | simdgroup-matrix | `ffai/{conv2d,conv3d}_mma_block_scaled.rs` | | depthwise conv2d | reduction | `ffai/depthwise_conv2d_block_scaled.rs` | | audio conv1d (STT) / fishspeech conv1d (TTS) | reduction | `ffai/{audio_conv1d,fishspeech_conv1d}_block_scaled.rs` | (`fp8_e4m3` / `fp8_e4m3_f16` reuse each family's `nvfp8` / `nvfp8_f16` kernel — identical 8-bit-E4M3 + scale shape.) **Flash KV covers every production head dim** (d64/96/128/256/512 × all 30 formats) — **no holes**. int8's group size (64) doesn't divide d96, so that case tiles with a **ragged trailing block** (a 64-block + a 32-block; `n_blocks = ceil(dim/block_size)`), the host packer and kernel rounding up identically so codes + scales stay self-consistent. ### The full integer matrix is a first-class citizen `int2/3/4/5/6/8` (FP32 group) + `mxint2/3/4/5/6/8` (E8M0 block) are present in **every** family above — including the fast tensor-engine paths (**simdgroup-MMA, MPP, NAX, MoE-MPP**), where integer throughput is highest on Apple GPUs / the ANE and the `mxint*` E8M0 layout maps to tensor-core block-scaling for future hardware. The core matmul / MoE / fused-norm / batched / KV-cache / attention families *also* carry the pre-existing **asymmetric affine** integers (scale + bias) for MLX-checkpoint interop. ### Symmetric block-scaled int4 now exists — affine int4 is still required, by design Block-scaled `int4` / `mxint4` are now present (symmetric). MLX 4-bit checkpoints are **asymmetric** affine (`w = scale·q + bias`), which a symmetric format cannot represent — so **affine int4 (Track 2, with its zero-point) remains required to load every MLX 4-bit model**. That's a deliberate design boundary (the zero-point), not a missing format. ("Legacy" applies only to the float-scale `fp4`/`fp8` within Track 1, superseded by spec mxfp4/nvfp4/mxfp8/nvfp8 but kept as labeled comparison variants.) ## fp16-scale axis Every FP32-scaled format has an **FP16-scale twin** (`*_f16`): same element + block, but the per-block scale is stored as a 2-byte IEEE half — the layout real checkpoints use, halving scale-buffer traffic. The host encoder does correct round-to-nearest-even **with full subnormal support**: wide-range elements (E5M2's `element_max` is 57344) push block scales into f16's subnormal range, and flushing them to zero would wipe whole blocks. The Apple-GPU `half` load matches the host decode exactly (subnormals included), so the CPU oracle holds. ## fp4 simdgroup-MMA f32 fix (from earlier in the series) Making the `#[test_kernel]` harness non-vacuous (it had been iterating an empty set — `metaltile::harness::registry::all_tests()` → `metaltile_std::all_tests()`) exposed two intertwined fp4-MMA defects, both fixed: 1. **Undefined shift.** `fp_quantized_mma::mt_fp4_qmm_mma` hand-rolled the E2M1 magnitude as `(mantissa + 2) << (exp - 1)`, undefined at `exp == 0` → miscompiled on f32 (output tile left unwritten). Replaced with the `e2m1_decode` intrinsic. 2. **Kernel name collision.** The block-scaled fp4 MMA kernel was *also* named `mt_fp4_qmm_mma`, colliding in the MSL/pipeline cache → order-dependent wrong pipelines (the real source of the ~0.46 f32 anomaly, and of an apparent `ffai_sdpa_multi_d256_causal` failure that was a *victim* of the contamination). Renamed to `mt_fp4_float_qmm_mma`. ## Foundation + verification - **`quant::codec`** — pure-Rust element + scale codecs (E2M1/E4M3/E5M2, E8M0, FP32, FP16, symmetric intN), the single source of truth shared by packer, oracle, and kernels. - **`quant::format`** — the 30-format `QFormat` with `element()` / `scale_kind()` / `symmetric()` axis accessors; host pack/dequant with ragged-tail blocks + unified sub-byte bit-stream packing. - Every `(family × format × dtype)` is GPU-verified 1:1 against the oracle. Full workspace gate (`cargo test --workspace` + `clippy`) green; geometry audited byte-for-byte unchanged per family (no new freeze surface). The MPP/NAX/MoE-MPP cooperative kernels are runtime-gated to Apple10+ and auto-skip on the CI Paravirtual GPU.
0xClandestine
pushed a commit
that referenced
this pull request
Jun 3, 2026
…cross every weight-bearing kernel (#250) ## What this is PR 2 of the precision series: make fp4/fp8 **spec-conformant block-scaled** and bring the **full op × quant-format matrix to completion** — *every* weight-bearing kernel family supports **all 30 Track-1 quant formats** in fp16/bf16/fp32 activation, so the bench columns (GFLOP/s, %-peak, roofline) reveal the fastest precision per op. **No holes.** > Originally stacked on #247 (`ek/bench-metrics`); that has since merged into `dev`, so this branch is now precision-only on top of `dev`. ## Quantization is three orthogonal axes A quantized weight stores a small **code** per element + a **scale** that restores magnitude. Formats vary along three axes — **bit-width is not one of them**: 1. **Element** — a signed integer (`int2…int8`) *or* a micro-float codebook (E2M1 / E4M3 / E5M2). 2. **Scale** — how the per-block scale is stored: raw **FP32** per large group, a compact **E8M0** power-of-two per small block (OCP MX), an **E4M3** micro-scale × a tensor-wide FP32 global (NVFP4), or **FP16** (the memory-halving twin). 3. **Zero-point** — **symmetric** (no zero-point) or **asymmetric** (scale + bias). Track 1 (this PR) is the **symmetric** block-/group-scaled family; Track 2 is the pre-existing **asymmetric** affine integers (zero-point for MLX-checkpoint interop). The two are complementary, not a hierarchy. ## The 30 Track-1 formats | group | formats | element | block | scale | |---|---|---|---|---| | spec float | `nvfp4` | E2M1 | 16 | E4M3 + global FP32 | | | `mxfp4` / `mxfp8_e4m3` / `mxfp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | E8M0 pow-2 | | | `nvfp8` | E4M3 | 16 | FP32 | | legacy float | `fp4` / `fp8_e4m3` / `fp8_e5m2` | E2M1 / E4M3 / E5M2 | 32 | FP32 per-group | | **symmetric int** | `int2 / int3 / int4 / int5 / int6 / int8` | int N | 64 | FP32 per-group | | **MXINT** | `mxint2 / mxint3 / mxint4 / mxint5 / mxint6 / mxint8` | int N | 32 | E8M0 pow-2 | | **fp16-scale twins** | `*_f16` of every FP32-scaled format above | as twin | as twin | FP16 (2 B) | **30 formats** = 8 float + 6 symmetric int (FP32) + 6 MXINT (E8M0) + 10 fp16-scale twins. `mxint8` is OCP-ratified MXINT8; `mxint*` map to tensor-core block-scaling units for future NVIDIA / AMD targets. Sub-byte codes (2/3/4/5/6-bit) tight-bit-pack LSB-first into u32 words (a 4-bit stream is byte-identical to the classic nibble layout); 8-bit is one byte/code. All formats share the `quant::codec` bit primitives — one source of truth for the host packer, the CPU correctness oracle, and the in-kernel decode (the `e2m1_decode` / `e4m3_decode` / `e5m2_decode` / `int8_decode` intrinsics + the straddle-aware sub-byte bit-stream extraction). The integer + fp16 formats are emitted by a parameterized **(bit-width × scale-kind)** decode macro per family, so every variant reuses its family's proven dispatch geometry byte-for-byte (no new freeze surface). ## Coverage matrix — COMPLETE Every family below supports **all 30 formats**. Each `(family × format × dtype)` ships a `#[kernel]`, a 1:1 `#[test_kernel]` (GPU-verified vs `quant::format::dequant`), and a `#[bench]` with `.flops()`. | family | path | file(s) | |---|---|---| | dequant (standalone) | elementwise | `mlx/block_scaled_dequant.rs` | | qgemv (GEMV decode) | reduction | `mlx/block_scaled_matmul.rs` | | qmm (GEMM prefill) | reduction | `mlx/block_scaled_qmm.rs` | | qmm — simdgroup-MMA | simdgroup-matrix | `mlx/block_scaled_mma.rs` | | qmm — MPP | tensor engine | `mlx/block_scaled_qmm_mpp.rs` | | qmm — NAX | tensor engine | `mlx/block_scaled_qmm_nax.rs` | | MoE gather-qmm | reduction | `mlx/block_scaled_moe.rs` | | MoE gather — MPP (bm8/16/64) | tensor engine | `ffai/moe_mpp{,_bm8,_bm64}_block_scaled.rs` | | expert-indexed GEMV | reduction | `ffai/dequant_gemv_expert_indexed_block_scaled.rs` | | fused RMSNorm + GEMV | reduction | `ffai/rms_norm_block_scaled_qgemv.rs` | | fused gated-RMSNorm + GEMV | reduction | `ffai/gated_rms_norm_block_scaled_qgemv.rs` | | batched-Q/K/V qgemv + qmm | reduction | `ffai/batched_qkv_block_scaled_{qgemv,qmm}.rs` | | batched-4 qgemv + qmm | reduction | `ffai/batched_4_block_scaled_{qgemv,qmm}.rs` | | embedding gather | elementwise | `ffai/dequant_gather_block_scaled.rs` | | flash SDPA (block-scaled KV) | flash | `ffai/flash_block_scaled_sdpa.rs` — **d64/96/128/256/512** | | patch embed (linear projection) | reduction | `ffai/patch_embed_block_scaled.rs` | | patch embed (simdgroup-MMA) | simdgroup-matrix | `ffai/patch_embed_mma_block_scaled.rs` | | conv2d / conv3d (direct) | reduction | `ffai/{conv2d,conv3d}_block_scaled.rs` | | conv2d / conv3d (im2col simdgroup-MMA) | simdgroup-matrix | `ffai/{conv2d,conv3d}_mma_block_scaled.rs` | | depthwise conv2d | reduction | `ffai/depthwise_conv2d_block_scaled.rs` | | audio conv1d (STT) / fishspeech conv1d (TTS) | reduction | `ffai/{audio_conv1d,fishspeech_conv1d}_block_scaled.rs` | (`fp8_e4m3` / `fp8_e4m3_f16` reuse each family's `nvfp8` / `nvfp8_f16` kernel — identical 8-bit-E4M3 + scale shape.) **Flash KV covers every production head dim** (d64/96/128/256/512 × all 30 formats) — **no holes**. int8's group size (64) doesn't divide d96, so that case tiles with a **ragged trailing block** (a 64-block + a 32-block; `n_blocks = ceil(dim/block_size)`), the host packer and kernel rounding up identically so codes + scales stay self-consistent. ### The full integer matrix is a first-class citizen `int2/3/4/5/6/8` (FP32 group) + `mxint2/3/4/5/6/8` (E8M0 block) are present in **every** family above — including the fast tensor-engine paths (**simdgroup-MMA, MPP, NAX, MoE-MPP**), where integer throughput is highest on Apple GPUs / the ANE and the `mxint*` E8M0 layout maps to tensor-core block-scaling for future hardware. The core matmul / MoE / fused-norm / batched / KV-cache / attention families *also* carry the pre-existing **asymmetric affine** integers (scale + bias) for MLX-checkpoint interop. ### Symmetric block-scaled int4 now exists — affine int4 is still required, by design Block-scaled `int4` / `mxint4` are now present (symmetric). MLX 4-bit checkpoints are **asymmetric** affine (`w = scale·q + bias`), which a symmetric format cannot represent — so **affine int4 (Track 2, with its zero-point) remains required to load every MLX 4-bit model**. That's a deliberate design boundary (the zero-point), not a missing format. ("Legacy" applies only to the float-scale `fp4`/`fp8` within Track 1, superseded by spec mxfp4/nvfp4/mxfp8/nvfp8 but kept as labeled comparison variants.) ## fp16-scale axis Every FP32-scaled format has an **FP16-scale twin** (`*_f16`): same element + block, but the per-block scale is stored as a 2-byte IEEE half — the layout real checkpoints use, halving scale-buffer traffic. The host encoder does correct round-to-nearest-even **with full subnormal support**: wide-range elements (E5M2's `element_max` is 57344) push block scales into f16's subnormal range, and flushing them to zero would wipe whole blocks. The Apple-GPU `half` load matches the host decode exactly (subnormals included), so the CPU oracle holds. ## fp4 simdgroup-MMA f32 fix (from earlier in the series) Making the `#[test_kernel]` harness non-vacuous (it had been iterating an empty set — `metaltile::harness::registry::all_tests()` → `metaltile_std::all_tests()`) exposed two intertwined fp4-MMA defects, both fixed: 1. **Undefined shift.** `fp_quantized_mma::mt_fp4_qmm_mma` hand-rolled the E2M1 magnitude as `(mantissa + 2) << (exp - 1)`, undefined at `exp == 0` → miscompiled on f32 (output tile left unwritten). Replaced with the `e2m1_decode` intrinsic. 2. **Kernel name collision.** The block-scaled fp4 MMA kernel was *also* named `mt_fp4_qmm_mma`, colliding in the MSL/pipeline cache → order-dependent wrong pipelines (the real source of the ~0.46 f32 anomaly, and of an apparent `ffai_sdpa_multi_d256_causal` failure that was a *victim* of the contamination). Renamed to `mt_fp4_float_qmm_mma`. ## Foundation + verification - **`quant::codec`** — pure-Rust element + scale codecs (E2M1/E4M3/E5M2, E8M0, FP32, FP16, symmetric intN), the single source of truth shared by packer, oracle, and kernels. - **`quant::format`** — the 30-format `QFormat` with `element()` / `scale_kind()` / `symmetric()` axis accessors; host pack/dequant with ragged-tail blocks + unified sub-byte bit-stream packing. - Every `(family × format × dtype)` is GPU-verified 1:1 against the oracle. Full workspace gate (`cargo test --workspace` + `clippy`) green; geometry audited byte-for-byte unchanged per family (no new freeze surface). The MPP/NAX/MoE-MPP cooperative kernels are runtime-gated to Apple10+ and auto-skip on the CI Paravirtual GPU.
ekryski
added a commit
that referenced
this pull request
Jun 3, 2026
…s dead-strip fix (#252) ## Summary This branch adds the GPU kernels FFAI (an Apple-Silicon LLM/VLM/audio inference engine built on metaltile) needs to move its remaining vision, audio, video, and TTS paths off CPU, plus a test-infrastructure fix that closes a silent gap in the kernel-correctness gate. Every kernel ships a paired `#[test_kernel]` GPU-correctness test against a CPU oracle in the same commit; all pass via `tile test`. > **Stacked on #247** (`ek/bench-metrics`). The first commit (`feat(bench): perf metrics …`) belongs to #247 and only appears here because the `.flops()` annotations below depend on it. Once #247 merges I'll rebase to drop it; review the commits *above* it. ## What changed **New FFAI kernels** (each with a `#[test_kernel]` GPU-correctness test + `#[bench]`): | Kernel | Purpose | |---|---| | `transpose_th` | token↔head-major transpose for GPU vision attention | | `clamp_scalar` | element-wise scalar clamp (Gemma-4 clipped linears) | | `resize_normalize` (+ `_bicubic`) | fused resize + per-channel normalize + interleaved→NCHW (VL preprocess) | | `depthwise_conv2d_nhwc` | channel-last depthwise conv (FastVLM) | | `sdpa_bidirectional_windowed` (d80) | block-diagonal windowed bidirectional SDPA (Qwen2.5-VL) | | `sdpa_rel_pos_conformer` (d128) | content-dependent Transformer-XL relative-position SDPA (Conformer audio) | | `sdpa_decode_sink_buf` | per-head attention-sink decode (GPT-OSS) | | `avg_pool2d_nhwc`, `im2col_patch`, `pos_emb_2d_add`, `frame_diff_luma` | vision/video front-end ops | | `mel_spectrogram_magnitude` | magnitude log-Mel front-end | | `lstm` (`ffai_lstm` full-sequence + `lstm_cell`), `adain1d` | TTS acoustic stack (StyleTTS2/Kokoro) | | `conv1d_dilated`, `conv1d_transpose` | HiFi-GAN-style vocoder convs | **Cleanups** - De-model-named the stdlib: `kokoro.rs` → `adain1d.rs`/`lstm.rs`, `fishspeech_conv1d.rs` → `conv1d_dilated_transpose.rs` (kernels named by what they compute, not the model). - Migrated the new kernels to the bare `#[kernel]` + `#[bench]` form after the macro refactor. - `.flops()` annotations on the compute-bound kernels (conformer/windowed/sink SDPA, depthwise conv, LSTM, conv1d) so #247's GFLOP/s + roofline columns populate; memory-bound kernels left blank. **Test-infrastructure fix** (`test(harness): anchor kernel TUs …`) - `kernel_tests_harness` only called `all_tests()`; with nothing referencing kernel code, macOS `-dead_strip` removed every metaltile-std kernel translation unit — and its `#[test_kernel]` inventory submissions — so the harness iterated **zero** entries and passed vacuously (~1.4 s, exercising nothing). It now references `all_kernels()` to anchor the TUs (the same anchor `kernel_registry_consistency` relies on) and asserts both registries are non-empty. The harness now actually runs all GPU `#[test_kernel]` checks (~24 s) under `cargo test --workspace`. **Docs** - `docs/KERNEL_ORGANIZATION.md` — proposed organization of the kernel/ops tree by kernel family, reconciled with the v4 CLI subprocess rewrite. For discussion. ## Update — complete the StyleTTS2 / Kokoro GPU chain The original TTS batch above (`adain1d`, `lstm`, `conv1d_dilated`/`conv1d_transpose`) shipped the *building blocks*. These follow-on commits add the remaining kernels so the **entire** Kokoro / StyleTTS2 text→waveform chain (PLBERT text encoder → prosody predictor → AdaIN decoder → iSTFTNet generator → vocoder) runs GPU-resident in FFAI, verified end-to-end against the CPU / mlx-audio reference. **New kernels** (each with a paired `#[test_kernel]` + `#[bench]`): | Kernel | Purpose | |---|---| | `leaky_relu` | runtime-slope LeakyReLU — the generator / AdaIN-resblock activation (slope 0.2 / 0.1 / 0.01) | | `snake1d` | `x + (1/α)·sin²(αx)` with per-channel learnable α — the BigVGAN / iSTFTNet periodic activation | | `upsample_nearest1d` | nearest 1-D up-sample — the `AdainResBlk1d` shortcut + F0/N predictor up-sampler | | `conv1d_transpose_depthwise` | grouped (depthwise, `groups == channels`) conv-transpose — the StyleTTS2 decoder `pool` (the dense `conv1d_transpose` sums over all in-channels) | | `gelu_erf` (+ a new `erf` prelude intrinsic) | exact `0.5·x·(1+erf(x/√2))` GELU — PLBERT's `nn.GELU()` default; the tanh-approx `mt_gelu` drifts enough over 12 ALBERT layers to shift the predictor's rounding-sensitive durations | **Correctness fix to an existing kernel** - `adain1d`: **clamp variance ≥ 0 before `rsqrt`.** Variance was computed as `E[x²] − E[x]²`, which suffers catastrophic f32 cancellation when a channel is near-constant over a long time axis (the Kokoro generator runs adain over ~7,800 frames): the result can go slightly negative → `rsqrt(negative)` = **NaN**, poisoning the whole utterance. The true variance is ≥ 0, so clamp it; this matches the CPU two-pass reference. Latent bug affecting **any** `adain1d` consumer, not just Kokoro. **LSTM coverage** - Added a `ffai_lstm` GPU-correctness test at the real Kokoro `predictor.shared` shape (hidden 256 = 8 full simdgroups, seqLen 65, input 640) — the prior tests capped at hidden 40 / seqLen 8, never exercising the cross-simdgroup threadgroup recurrence over a long sequence. Passes. (FFAI side carries a note: the Kokoro front-end LSTMs run on CPU as a workaround for a *context-dependent* NaN that only appears after the full chain's many prior GPU dispatches and is not reproducible standalone — the kernel is bit-exact in isolation.) ## Verification - `tile test` — every new kernel passes against its CPU oracle (e.g. `sdpa_decode_sink_buf` 6/6, `sdpa_rel_pos_conformer` 6/6, `ffai_lstm` (incl. hidden-256/seqLen-65), `conv1d_*` + `conv1d_transpose_depthwise`, `snake1d`, `leaky_relu`, `upsample_nearest1d`, `gelu_erf`, `resize_normalize_bicubic`). - `make test` — full workspace gate green (0 failures); `make emit-all` — every kernel emits valid MSL under `xcrun metal`; the empty-body MSL detector is clean. - `kernel_registry_consistency` + the fixed `kernel_tests_harness` (all GPU checks, ~24 s) pass. - **Downstream end-to-end:** FFAI regenerates these kernels and runs the full Kokoro chain GPU-resident — PLBERT (gemm/sdpa/`gelu_erf`/layerNorm) → predictor F0/N AdaIN curves → AdaIN decoder (`conv1d_transpose_depthwise` pool) → iSTFTNet generator (`leaky_relu`/`snake1d`/`conv1d_*`) → vocoder iSTFT — producing a waveform identical to the verified CPU path (energy 75.2; per-stage parity ~1e-5). Also: GPT-OSS-20B greedy decode through `sdpa_decode_sink_buf` produces coherent output (156/200 unique tokens). Representative bench (`tile bench -f sdpa_rel_pos_conformer`, M1 Max), showing the GFLOP/s column from the `.flops()` annotation: ``` Shape │ MT(µs) │ MT(GB/s) │ GFLOP/s │ ok N=512K f32 │ 3536.5 │ 2.4 │ 455.4 │ ✓ N=512K f16 │ 1792.9 │ 2.3 │ 898.3 │ ✓ N=512K bf16 │ 2605.0 │ 1.6 │ 618.3 │ ✓ ``` ## AI assistance This PR used AI assistance (Claude) under human direction and review — spanning kernel design, DSL authoring, CPU-oracle tests, the dead-strip root-cause analysis, the adain1d NaN diagnosis, the bench annotations, and this description. Each kernel was validated against its CPU reference before inclusion.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements
docs/BENCH_METRICS_SPEC.md(Phases 1–4).Why
tile benchreported only GB/s (bytes_moved ÷ time). That can't answer the two questions the suite exists to answer:The per-kernel latency was already measured and then discarded (
let (gbps, _stats) = …). This surfaces it and adds the missing metrics — additive only: the existing GB/s / ref-vs-MT /okcolumns and theref/mtJSON keys thatbaselines/*.jsondiffing consumes are unchanged.What's new
MT(µs)default column (theminsample, matching the GB/s convention);Ref(µs)under-v. This is the metric that makes "fastest precision" directly readable..flops(...)(dense-equivalent useful FLOPs, so int4/int8/f16 are comparable); memory-bound elementwise/reduction kernels leave it blank.-v) —%BW(÷ peak DRAM bandwidth),%FLOP(÷ peak compute — the M5 Neural-Accelerator FP16 ceiling where applicable, the SIMD pipe otherwise; bf16 stays on the SIMD pipe),AI(arithmetic intensity, FLOPs/byte). Peaks come from a per-device table (crates/metaltile-std/src/device_specs.rs, M1–M5 Max). Unknown device → blank, never fails (CI's paravirtual GPU is safe).-v) —memory-bound/compute-bound/occupancy-limited/register-limited/latency-bound, combining the roofline position (AI vs the device ridge point) with the occupancy/register signals. The CLI's separate occupancy pass is folded into the per-row path (profile.rs), so it's computed once for every row instead of in a parallel map.Sample (Apple M1 Max)
Default table:
-vadds the roofline + verdict (gemv readsmemory-bound: AI ≈ 0.5–1.0, %BW saturating, tiny %FLOP):A tiled GEMM (
steel_gemm/fused_naxf16) reads 6568 GFLOP/s (~63% of the M1 Max FP16 ceiling) while its GB/s looks misleadingly tiny — exactly the case GFLOP/s exists to expose.Compatibility
JSON is additive:
op/shape/metric/ref/mtare byte-stable (baseline diffing readsref/mt); the newlatency_us,gflops,pct_peak_bw,pct_peak_flops,arith_intensity,bottleneckkeys are appended. No table snapshot tests existed, so no snapshot churn.Tests
to_gflops/to_gbpsusemin, guard the off-GPU stub), latency surfacing (mt_us/ref_us),DeviceSpecs::lookup(known + unknown + NA-vs-SIMD dtype pick), andclassify_bottleneck(memory/compute/occupancy/register/latency cases) are unit-tested.kernel_registry_consistency(every annotated setup still codegens) + the cli JSON-schema tests pass.cargo test --workspace --exclude metaltile-runtimegreen;make fmt-check/make typosclean.Notes
make clippy/make teston macOS hits a pre-existing, unrelated failure inmetaltile-runtime's test-onlymod perf(a duplicatemod perf+ a leftover.lock().unwrap()on aparking_lotmutex from perf(runtime): swap PSO + MSL cache locks to parking_lot::Mutex #234). It's identical todevand untouched here — fixed in a separate small PR.