From f87b9e6a10fcdc6052f1eb4ff331e884c30a1a45 Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Sat, 13 Jun 2026 10:17:47 -0600 Subject: [PATCH 1/5] refactor(kv_cache): migrate kv_cache family into kernels/kv_cache/ - mlx/fft.rs + ffai/{kv_cache,kv_cache_update_many} -> kernels/kv_cache/ (kv_cache.rs -> cache.rs to avoid module_inception). - mt_ prefixes on the bare-named kernels: mt_kv_cache_update(_many), mt_quantize_kv(_int4/_int8/_fp8_e4m3/_fp8_e5m2), mt_bulk_dequant_kv(...). fft already mt_*. - update the kv_cache_quant_roundtrip integration test imports + the conv1d doc-comment cross-ref. cargo build clean; kv_cache + fft + quantize_kv tests pass, integration test passes. --- crates/metaltile-std/src/ffai/mod.rs | 2 - .../conv1d_causal_step_silu_cast_many.rs | 2 +- .../kv_cache.rs => kernels/kv_cache/cache.rs} | 96 +++++++++---------- .../src/{mlx => kernels/kv_cache}/fft.rs | 0 .../metaltile-std/src/kernels/kv_cache/mod.rs | 11 +++ .../kv_cache/update_many.rs} | 12 +-- crates/metaltile-std/src/kernels/mod.rs | 1 + crates/metaltile-std/src/mlx/mod.rs | 1 - .../tests/kv_cache_quant_roundtrip_gpu.rs | 40 ++++---- 9 files changed, 87 insertions(+), 78 deletions(-) rename crates/metaltile-std/src/{ffai/kv_cache.rs => kernels/kv_cache/cache.rs} (93%) rename crates/metaltile-std/src/{mlx => kernels/kv_cache}/fft.rs (100%) create mode 100644 crates/metaltile-std/src/kernels/kv_cache/mod.rs rename crates/metaltile-std/src/{ffai/kv_cache_update_many.rs => kernels/kv_cache/update_many.rs} (95%) diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index 7ea71518..7b2df82e 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -70,8 +70,6 @@ pub mod gguf_dequant_q8_0; pub mod gguf_iq2_xxs_extract_qs; pub mod im2col_patch; pub mod im2col_patch_interleaved; -pub mod kv_cache; -pub mod kv_cache_update_many; pub mod leaky_relu; pub mod lstm; pub mod mel_spectrogram; diff --git a/crates/metaltile-std/src/kernels/convolution/conv1d_causal_step_silu_cast_many.rs b/crates/metaltile-std/src/kernels/convolution/conv1d_causal_step_silu_cast_many.rs index 5ac09c1e..87cb3ca0 100644 --- a/crates/metaltile-std/src/kernels/convolution/conv1d_causal_step_silu_cast_many.rs +++ b/crates/metaltile-std/src/kernels/convolution/conv1d_causal_step_silu_cast_many.rs @@ -22,7 +22,7 @@ //! On top of the bandwidth save the per-channel kernel collapses //! `T * (conv_step + silu_cast) = 2T` dispatches per layer into one — //! same dispatch-saving pattern as `mt_rope_banded` and -//! `kv_cache_update_many`. +//! `mt_kv_cache_update_many`. //! //! Layout: //! diff --git a/crates/metaltile-std/src/ffai/kv_cache.rs b/crates/metaltile-std/src/kernels/kv_cache/cache.rs similarity index 93% rename from crates/metaltile-std/src/ffai/kv_cache.rs rename to crates/metaltile-std/src/kernels/kv_cache/cache.rs index bb25dfa3..8c818688 100644 --- a/crates/metaltile-std/src/ffai/kv_cache.rs +++ b/crates/metaltile-std/src/kernels/kv_cache/cache.rs @@ -12,7 +12,7 @@ //! biases [n_kv_heads, max_seq, head_dim / group_size] T //! //! Affine-quant kernels use `#[kernel(variants(BITS = [4, 8], suffix = -//! "int{BITS}")]` so `quantize_kv_int4/8` and `bulk_dequant_kv_int4/8` +//! "int{BITS}")]` so `mt_quantize_kv_int4/8` and `mt_bulk_dequant_kv_int4/8` //! are generated from a single parameterised body. The fp8 kernels //! (e4m3 / e5m2) have float-literal format constants that cannot be //! expressed as integer variant parameters, so they remain explicit. @@ -29,7 +29,7 @@ use metaltile::kernel; // Dest layout: [n_kv_heads, max_seq, head_dim]. One thread per output // element (n_kv_heads * head_dim total threads). #[kernel] -pub fn kv_cache_update( +pub fn mt_kv_cache_update( src: Tensor, out: Tensor, #[constexpr] head_dim: u32, @@ -52,9 +52,9 @@ pub fn kv_cache_update( /// Affine KV-cache quantize — int4 and int8 variants. /// -/// Produces `quantize_kv_int4` and `quantize_kv_int8`. One thread per group. +/// Produces `mt_quantize_kv_int4` and `mt_quantize_kv_int8`. One thread per group. #[kernel(variants(BITS = [4, 8], suffix = "int{BITS}"))] -pub fn quantize_kv( +pub fn mt_quantize_kv( src: Tensor, mut out_w: Tensor, mut out_s: Tensor, @@ -114,10 +114,10 @@ pub fn quantize_kv( /// Affine KV-cache bulk dequant — int4 and int8 variants. /// -/// Produces `bulk_dequant_kv_int4` and `bulk_dequant_kv_int8`. +/// Produces `mt_bulk_dequant_kv_int4` and `mt_bulk_dequant_kv_int8`. /// One thread per output element. #[kernel(variants(BITS = [4, 8], suffix = "int{BITS}"))] -pub fn bulk_dequant_kv( +pub fn mt_bulk_dequant_kv( in_w: Tensor, in_s: Tensor, in_b: Tensor, @@ -183,7 +183,7 @@ pub fn bulk_dequant_kv( /// fp8 E4M3 KV-cache quantize — one thread per group. Stores the group amax /// as scale and packs fp8-quantized codes (4 per u32, 8 bits each). #[kernel] -pub fn quantize_kv_fp8_e4m3( +pub fn mt_quantize_kv_fp8_e4m3( src: Tensor, mut out_w: Tensor, mut out_s: Tensor, @@ -244,7 +244,7 @@ pub fn quantize_kv_fp8_e4m3( /// fp8 E5M2 KV-cache quantize — one thread per group. #[kernel] -pub fn quantize_kv_fp8_e5m2( +pub fn mt_quantize_kv_fp8_e5m2( src: Tensor, mut out_w: Tensor, mut out_s: Tensor, @@ -303,7 +303,7 @@ pub fn quantize_kv_fp8_e5m2( /// fp8 E4M3 KV-cache bulk dequant — one thread per output element. #[kernel] -pub fn bulk_dequant_kv_fp8_e4m3( +pub fn mt_bulk_dequant_kv_fp8_e4m3( in_w: Tensor, in_s: Tensor, mut out: Tensor, @@ -351,7 +351,7 @@ pub fn bulk_dequant_kv_fp8_e4m3( /// fp8 E5M2 KV-cache bulk dequant — one thread per output element. #[kernel] -pub fn bulk_dequant_kv_fp8_e5m2( +pub fn mt_bulk_dequant_kv_fp8_e5m2( in_w: Tensor, in_s: Tensor, mut out: Tensor, @@ -399,21 +399,21 @@ pub mod kernel_tests { use metaltile::{test::*, test_kernel}; use super::{ - bulk_dequant_kv_fp8_e4m3, - bulk_dequant_kv_fp8_e5m2, - bulk_dequant_kv_int4, - bulk_dequant_kv_int8, - kv_cache_update, - quantize_kv_fp8_e4m3, - quantize_kv_fp8_e5m2, - quantize_kv_int4, - quantize_kv_int8, + mt_bulk_dequant_kv_fp8_e4m3, + mt_bulk_dequant_kv_fp8_e5m2, + mt_bulk_dequant_kv_int4, + mt_bulk_dequant_kv_int8, + mt_kv_cache_update, + mt_quantize_kv_fp8_e4m3, + mt_quantize_kv_fp8_e5m2, + mt_quantize_kv_int4, + mt_quantize_kv_int8, }; use crate::utils::{pack_f32, unpack_f32}; fn u32_bytes(v: &[u32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } - // ── kv_cache_update ────────────────────────────────────────────── + // ── mt_kv_cache_update ────────────────────────────────────────────── #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.0)] fn test_kv_cache_update(dt: DType) -> TestSetup { let (n_kv_heads, head_dim, max_seq, position) = (4usize, 16usize, 8usize, 3usize); @@ -429,7 +429,7 @@ pub mod kernel_tests { expected[dst] = src_dt[h * head_dim + d]; } } - TestSetup::new(kv_cache_update::kernel_ir_for(dt)) + TestSetup::new(mt_kv_cache_update::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("src", pack_f32(&src, dt), dt)) .input(TestBuffer::from_vec("out", pack_f32(&cache, dt), dt)) @@ -440,7 +440,7 @@ pub mod kernel_tests { .grid_1d(n_kv_heads * head_dim, 256) } - // ── quantize_kv_int4 / int8 — scale + bias check ───────────────── + // ── mt_quantize_kv_int4 / int8 — scale + bias check ───────────────── // // group_size=8 is a multiple of both vals_per_pack (8 for int4, 4 for // int8) so the pack loop is well-formed for both bit-widths. @@ -511,15 +511,15 @@ pub mod kernel_tests { // shared setup helper retains a use site. #[allow(dead_code)] fn test_quantize_kv_int4(dt: DType) -> TestSetup { - quant_scale_bias_setup(quantize_kv_int4::kernel_ir_for(dt), 4, dt) + quant_scale_bias_setup(mt_quantize_kv_int4::kernel_ir_for(dt), 4, dt) } #[allow(dead_code)] // bench-only (see test_quantize_kv_int4 note) fn test_quantize_kv_int8(dt: DType) -> TestSetup { - quant_scale_bias_setup(quantize_kv_int8::kernel_ir_for(dt), 8, dt) + quant_scale_bias_setup(mt_quantize_kv_int8::kernel_ir_for(dt), 8, dt) } - // ── bulk_dequant_kv_int4 / int8 — exact dequant oracle ─────────── + // ── mt_bulk_dequant_kv_int4 / int8 — exact dequant oracle ─────────── // // scale=1, bias=0 → each output equals the unpacked quantized integer, // so the dequant is exact regardless of dtype rounding (small ints). @@ -578,10 +578,10 @@ pub mod kernel_tests { #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.0, variants(BITS = [4, 8], suffix = "int{BITS}"))] fn test_bulk_dequant_kv(dt: DType) -> TestSetup { - dequant_setup(bulk_dequant_kv_intBITS::kernel_ir_for(dt), BITS, dt) + dequant_setup(mt_bulk_dequant_kv_intBITS::kernel_ir_for(dt), BITS, dt) } - // ── quantize_kv_fp8_e4m3 / e5m2 — scale check ──────────────────── + // ── mt_quantize_kv_fp8_e4m3 / e5m2 — scale check ──────────────────── // // fp8 quant is scale-only: scale = group amax / fp8_max. Exact. fn quant_fp8_scale_setup( @@ -638,15 +638,15 @@ pub mod kernel_tests { // exact decode), so only the scale is pinned here. #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 1e-3, 1e-2])] fn test_quantize_kv_fp8_e4m3(dt: DType) -> TestSetup { - quant_fp8_scale_setup(quantize_kv_fp8_e4m3::kernel_ir_for(dt), 240.0, dt) + quant_fp8_scale_setup(mt_quantize_kv_fp8_e4m3::kernel_ir_for(dt), 240.0, dt) } #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 1e-3, 1e-2])] fn test_quantize_kv_fp8_e5m2(dt: DType) -> TestSetup { - quant_fp8_scale_setup(quantize_kv_fp8_e5m2::kernel_ir_for(dt), 57344.0, dt) + quant_fp8_scale_setup(mt_quantize_kv_fp8_e5m2::kernel_ir_for(dt), 57344.0, dt) } - // ── bulk_dequant_kv_fp8_e4m3 / e5m2 — exact fp8 decode oracle ───────── + // ── mt_bulk_dequant_kv_fp8_e4m3 / e5m2 — exact fp8 decode oracle ───────── // // Pack KNOWN fp8 bytes whose decoded magnitudes are exact in every dtype // (the fp8 grid is a subset of f16/bf16), with scale=1, so the dequant @@ -720,7 +720,7 @@ pub mod kernel_tests { (0xB8, -1.0), (0xC0, -2.0), ]; - dequant_fp8_setup(bulk_dequant_kv_fp8_e4m3::kernel_ir_for(dt), &palette, dt) + dequant_fp8_setup(mt_bulk_dequant_kv_fp8_e4m3::kernel_ir_for(dt), &palette, dt) } // e5m2 byte → exact value: 0x00→0, 0x34→0.25, 0x38→0.5, 0x3C→1.0, @@ -738,7 +738,7 @@ pub mod kernel_tests { (0xBC, -1.0), (0xB8, -0.5), ]; - dequant_fp8_setup(bulk_dequant_kv_fp8_e5m2::kernel_ir_for(dt), &palette, dt) + dequant_fp8_setup(mt_bulk_dequant_kv_fp8_e5m2::kernel_ir_for(dt), &palette, dt) } } @@ -748,15 +748,15 @@ pub mod kernel_benches { use metaltile::{bench, test::*}; use super::{ - bulk_dequant_kv_fp8_e4m3, - bulk_dequant_kv_fp8_e5m2, - bulk_dequant_kv_int4, - bulk_dequant_kv_int8, - kv_cache_update, - quantize_kv_fp8_e4m3, - quantize_kv_fp8_e5m2, - quantize_kv_int4, - quantize_kv_int8, + mt_bulk_dequant_kv_fp8_e4m3, + mt_bulk_dequant_kv_fp8_e5m2, + mt_bulk_dequant_kv_int4, + mt_bulk_dequant_kv_int8, + mt_kv_cache_update, + mt_quantize_kv_fp8_e4m3, + mt_quantize_kv_fp8_e5m2, + mt_quantize_kv_int4, + mt_quantize_kv_int8, }; fn u32_bytes(n: usize) -> Vec { vec![0u8; n * 4] } @@ -771,7 +771,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_kv_cache_update(dt: DType) -> BenchSetup { let elems = N_KV_HEADS * HEAD_DIM; - BenchSetup::new(kv_cache_update::kernel_ir_for(dt)) + BenchSetup::new(mt_kv_cache_update::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("src", elems, dt)) .buffer(BenchBuffer::zeros("out", N_KV_HEADS * MAX_SEQ * HEAD_DIM, dt).output()) @@ -805,7 +805,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16], variants(BITS = [4, 8], suffix = "int{BITS}"))] fn bench_quantize_kv(dt: DType) -> BenchSetup { - quant_bench(quantize_kv_intBITS::kernel_ir_for(dt), BITS, dt) + quant_bench(mt_quantize_kv_intBITS::kernel_ir_for(dt), BITS, dt) } fn dequant_bench(kernel: metaltile::core::ir::Kernel, bits: u32, dt: DType) -> BenchSetup { @@ -831,7 +831,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16], variants(BITS = [4, 8], suffix = "int{BITS}"))] fn bench_bulk_dequant_kv(dt: DType) -> BenchSetup { - dequant_bench(bulk_dequant_kv_intBITS::kernel_ir_for(dt), BITS, dt) + dequant_bench(mt_bulk_dequant_kv_intBITS::kernel_ir_for(dt), BITS, dt) } // fp8 quantize is scale-only (no out_b buffer). @@ -855,11 +855,11 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_quantize_kv_fp8_e4m3(dt: DType) -> BenchSetup { - quant_fp8_bench(quantize_kv_fp8_e4m3::kernel_ir_for(dt), dt) + quant_fp8_bench(mt_quantize_kv_fp8_e4m3::kernel_ir_for(dt), dt) } #[bench(dtypes = [f32, f16, bf16])] fn bench_quantize_kv_fp8_e5m2(dt: DType) -> BenchSetup { - quant_fp8_bench(quantize_kv_fp8_e5m2::kernel_ir_for(dt), dt) + quant_fp8_bench(mt_quantize_kv_fp8_e5m2::kernel_ir_for(dt), dt) } // fp8 dequant is scale-only (no in_b buffer). @@ -883,10 +883,10 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_bulk_dequant_kv_fp8_e4m3(dt: DType) -> BenchSetup { - dequant_fp8_bench(bulk_dequant_kv_fp8_e4m3::kernel_ir_for(dt), dt) + dequant_fp8_bench(mt_bulk_dequant_kv_fp8_e4m3::kernel_ir_for(dt), dt) } #[bench(dtypes = [f32, f16, bf16])] fn bench_bulk_dequant_kv_fp8_e5m2(dt: DType) -> BenchSetup { - dequant_fp8_bench(bulk_dequant_kv_fp8_e5m2::kernel_ir_for(dt), dt) + dequant_fp8_bench(mt_bulk_dequant_kv_fp8_e5m2::kernel_ir_for(dt), dt) } } diff --git a/crates/metaltile-std/src/mlx/fft.rs b/crates/metaltile-std/src/kernels/kv_cache/fft.rs similarity index 100% rename from crates/metaltile-std/src/mlx/fft.rs rename to crates/metaltile-std/src/kernels/kv_cache/fft.rs diff --git a/crates/metaltile-std/src/kernels/kv_cache/mod.rs b/crates/metaltile-std/src/kernels/kv_cache/mod.rs new file mode 100644 index 00000000..b20d50f3 --- /dev/null +++ b/crates/metaltile-std/src/kernels/kv_cache/mod.rs @@ -0,0 +1,11 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! KV-cache kernels — the kv_cache family (see +//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): the cache-update path +//! (single + batched), KV quantization / dequantization (incl. fp8), and the +//! FFT used by the STT front-end (`mt_fft` + Bluestein non-pow2 stages). +//! Migrated from the legacy `mlx/` (`fft`) + `ffai/` split. + +pub mod cache; +pub mod fft; +pub mod update_many; diff --git a/crates/metaltile-std/src/ffai/kv_cache_update_many.rs b/crates/metaltile-std/src/kernels/kv_cache/update_many.rs similarity index 95% rename from crates/metaltile-std/src/ffai/kv_cache_update_many.rs rename to crates/metaltile-std/src/kernels/kv_cache/update_many.rs index 50c17815..c9e5c6a3 100644 --- a/crates/metaltile-std/src/ffai/kv_cache_update_many.rs +++ b/crates/metaltile-std/src/kernels/kv_cache/update_many.rs @@ -45,7 +45,7 @@ use metaltile::kernel; #[kernel] -pub fn kv_cache_update_many( +pub fn mt_kv_cache_update_many( src: Tensor, positions: Tensor, out: Tensor, @@ -73,7 +73,7 @@ pub fn kv_cache_update_many( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::kv_cache_update_many; + use super::mt_kv_cache_update_many; use crate::utils::{pack_f32, unpack_f32}; fn u32_bytes(v: &[u32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } @@ -101,7 +101,7 @@ pub mod kernel_tests { } } let total = n_tokens * n_kv_heads * head_dim; - TestSetup::new(kv_cache_update_many::kernel_ir_for(dt)) + TestSetup::new(mt_kv_cache_update_many::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("src", pack_f32(&src, dt), dt)) .input(TestBuffer::from_vec("positions", u32_bytes(&positions), DType::U32)) @@ -114,12 +114,12 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `kv_cache_update_many` — a Qwen-class prefill +/// New-syntax benchmark for `mt_kv_cache_update_many` — a Qwen-class prefill /// batch appended in one dispatch (Grid3D, one thread per source element). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::kv_cache_update_many; + use super::mt_kv_cache_update_many; fn u32_bytes(v: impl Iterator) -> Vec { v.flat_map(|x| x.to_le_bytes()).collect() @@ -129,7 +129,7 @@ pub mod kernel_benches { fn bench_kv_cache_update_many(dt: DType) -> BenchSetup { let (n_tokens, n_kv_heads, head_dim, max_seq) = (512usize, 8usize, 128usize, 4096usize); let total = n_tokens * n_kv_heads * head_dim; - BenchSetup::new(kv_cache_update_many::kernel_ir_for(dt)) + BenchSetup::new(mt_kv_cache_update_many::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("src", total, dt)) .buffer(BenchBuffer::from_vec( diff --git a/crates/metaltile-std/src/kernels/mod.rs b/crates/metaltile-std/src/kernels/mod.rs index 18443bc0..c7eb5668 100644 --- a/crates/metaltile-std/src/kernels/mod.rs +++ b/crates/metaltile-std/src/kernels/mod.rs @@ -9,6 +9,7 @@ pub mod convolution; pub mod gemm; +pub mod kv_cache; pub mod norm; pub mod ops; pub mod rope; diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs index 181fdda3..ac6f6553 100644 --- a/crates/metaltile-std/src/mlx/mod.rs +++ b/crates/metaltile-std/src/mlx/mod.rs @@ -21,7 +21,6 @@ pub mod block_scaled_moe; pub mod block_scaled_qmm; pub mod block_scaled_qmm_mpp; pub mod block_scaled_qmm_nax; -pub mod fft; pub mod fp_quantized; pub mod fp_quantized_mma; pub mod fp_quantized_nax; diff --git a/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs b/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs index 643a9ac4..bfaf1f07 100644 --- a/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs +++ b/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs @@ -1,6 +1,6 @@ //! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric //! SPDX-License-Identifier: Apache-2.0 -//! GPU correctness for `quantize_kv_int4/int8` + `bulk_dequant_kv_int4/int8` +//! GPU correctness for `mt_quantize_kv_int4/int8` + `mt_bulk_dequant_kv_int4/int8` //! via raw → quantize → dequant round-trip. //! //! These four kernels ship in `ffai::kv_cache` but had no end-to-end @@ -11,7 +11,7 @@ //! //! Coverage rationale (mirrors the legacy //! `tests/kv_cache_update_gpu_correctness.rs`, removed in #240): -//! `quantize_kv_*` and `bulk_dequant_kv_*` are emitted from `macro_rules!` +//! `mt_quantize_kv_*` and `mt_bulk_dequant_kv_*` are emitted from `macro_rules!` //! shells (the proc-macro doesn't expand the inner declarative macro, //! so embedding kernel bodies in nested macros would silently produce //! empty kernels). The round-trip pins both the quantize geometry @@ -26,8 +26,8 @@ //! Each test: //! 1. Build random [n_kv_heads, head_dim] source slot (centred so //! group ranges are well-distributed) -//! 2. Dispatch `quantize_kv_*` → cache buffers at `position` -//! 3. Dispatch `bulk_dequant_kv_*` reading the whole [0..n_positions) +//! 2. Dispatch `mt_quantize_kv_*` → cache buffers at `position` +//! 3. Dispatch `mt_bulk_dequant_kv_*` reading the whole [0..n_positions) //! slice → reconstructed values //! 4. For the slot we just wrote, compare reconstructed vs source //! with a per-bits relative tolerance (int4: ±range/15 = one @@ -43,11 +43,11 @@ use std::collections::BTreeMap; use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes, unpack_u32_bytes}; use metaltile::{Context, core::ir::KernelMode}; -use metaltile_std::ffai::kv_cache::{ - bulk_dequant_kv_int4, - bulk_dequant_kv_int8, - quantize_kv_int4, - quantize_kv_int8, +use metaltile_std::kernels::kv_cache::cache::{ + mt_bulk_dequant_kv_int4, + mt_bulk_dequant_kv_int8, + mt_quantize_kv_int4, + mt_quantize_kv_int8, }; /// Shape parameters covering Qwen3-class K/V slots. @@ -140,12 +140,12 @@ fn roundtrip_int4(shape: &Shape, dt: Dt, source: &[f32]) -> Vec { buffers.insert("group_size".into(), (shape.group_size as u32).to_le_bytes().to_vec()); buffers.insert("position".into(), (shape.position as u32).to_le_bytes().to_vec()); - let mut qkernel = quantize_kv_int4::kernel_ir_for(dtype); + let mut qkernel = mt_quantize_kv_int4::kernel_ir_for(dtype); qkernel.mode = KernelMode::Grid3D; let (grid, tpg) = quantize_dispatch_grid(shape, bits); let q_out = ctx .dispatch_with_grid(&qkernel, &buffers, &BTreeMap::new(), grid, tpg) - .expect("quantize_kv_int4 dispatch"); + .expect("mt_quantize_kv_int4 dispatch"); let w_bytes = q_out.outputs.get("out_w").expect("out_w buffer").clone(); let s_bytes = q_out.outputs.get("out_s").expect("out_s buffer").clone(); @@ -163,12 +163,12 @@ fn roundtrip_int4(shape: &Shape, dt: Dt, source: &[f32]) -> Vec { dbuf.insert("group_size".into(), (shape.group_size as u32).to_le_bytes().to_vec()); dbuf.insert("n_positions".into(), (shape.n_positions as u32).to_le_bytes().to_vec()); - let mut dkernel = bulk_dequant_kv_int4::kernel_ir_for(dtype); + let mut dkernel = mt_bulk_dequant_kv_int4::kernel_ir_for(dtype); dkernel.mode = KernelMode::Grid3D; let (dgrid, dtpg) = dequant_dispatch_grid(shape); let d_out = ctx .dispatch_with_grid(&dkernel, &dbuf, &BTreeMap::new(), dgrid, dtpg) - .expect("bulk_dequant_kv_int4 dispatch"); + .expect("mt_bulk_dequant_kv_int4 dispatch"); let out_bytes = d_out.outputs.get("out").expect("out buffer"); unpack_bytes(out_bytes, dt) @@ -199,12 +199,12 @@ fn roundtrip_int8(shape: &Shape, dt: Dt, source: &[f32]) -> Vec { buffers.insert("group_size".into(), (shape.group_size as u32).to_le_bytes().to_vec()); buffers.insert("position".into(), (shape.position as u32).to_le_bytes().to_vec()); - let mut qkernel = quantize_kv_int8::kernel_ir_for(dtype); + let mut qkernel = mt_quantize_kv_int8::kernel_ir_for(dtype); qkernel.mode = KernelMode::Grid3D; let (grid, tpg) = quantize_dispatch_grid(shape, bits); let q_out = ctx .dispatch_with_grid(&qkernel, &buffers, &BTreeMap::new(), grid, tpg) - .expect("quantize_kv_int8 dispatch"); + .expect("mt_quantize_kv_int8 dispatch"); let w_bytes = q_out.outputs.get("out_w").expect("out_w buffer").clone(); let s_bytes = q_out.outputs.get("out_s").expect("out_s buffer").clone(); @@ -221,12 +221,12 @@ fn roundtrip_int8(shape: &Shape, dt: Dt, source: &[f32]) -> Vec { dbuf.insert("group_size".into(), (shape.group_size as u32).to_le_bytes().to_vec()); dbuf.insert("n_positions".into(), (shape.n_positions as u32).to_le_bytes().to_vec()); - let mut dkernel = bulk_dequant_kv_int8::kernel_ir_for(dtype); + let mut dkernel = mt_bulk_dequant_kv_int8::kernel_ir_for(dtype); dkernel.mode = KernelMode::Grid3D; let (dgrid, dtpg) = dequant_dispatch_grid(shape); let d_out = ctx .dispatch_with_grid(&dkernel, &dbuf, &BTreeMap::new(), dgrid, dtpg) - .expect("bulk_dequant_kv_int8 dispatch"); + .expect("mt_bulk_dequant_kv_int8 dispatch"); let out_bytes = d_out.outputs.get("out").expect("out buffer"); unpack_bytes(out_bytes, dt) @@ -340,7 +340,7 @@ fn kv_cache_int8_roundtrip_bf16() { // ── Cross-slot isolation ───────────────────────────────────────────── // -// `quantize_kv_*` writes only to its `position` slot — verify by +// `mt_quantize_kv_*` writes only to its `position` slot — verify by // pre-filling neighboring slots with a sentinel and checking they // survive a quantize+dequant cycle. Catches index formula regressions // (e.g. accidentally striding by head_dim instead of max_seq). @@ -378,12 +378,12 @@ fn kv_cache_int8_does_not_touch_other_slots_f32() { buffers.insert("group_size".into(), (shape.group_size as u32).to_le_bytes().to_vec()); buffers.insert("position".into(), (shape.position as u32).to_le_bytes().to_vec()); - let mut qkernel = quantize_kv_int8::kernel_ir_for(dtype); + let mut qkernel = mt_quantize_kv_int8::kernel_ir_for(dtype); qkernel.mode = KernelMode::Grid3D; let (grid, tpg) = quantize_dispatch_grid(&shape, bits); let q_out = ctx .dispatch_with_grid(&qkernel, &buffers, &BTreeMap::new(), grid, tpg) - .expect("quantize_kv_int8 dispatch"); + .expect("mt_quantize_kv_int8 dispatch"); let w_after = unpack_u32_bytes(q_out.outputs.get("out_w").expect("out_w")); let s_after = unpack_bytes(q_out.outputs.get("out_s").expect("out_s"), dt); From bc3c50b649012bac6f2d409e851f43d55997e668 Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Sat, 13 Jun 2026 10:22:32 -0600 Subject: [PATCH 2/5] refactor(audio): migrate audio/speech family into kernels/audio/ ffai/{lstm,mel_spectrogram,snake1d,upsample_nearest1d,vocoder} -> kernels/audio/ with mt_ prefixes: mt_lstm, mt_mel_spectrogram(+_magnitude/_stft_window/_filterbank), mt_snake1d, mt_upsample_nearest1d, mt_vocoder_istft. cargo build + clippy clean; mel/lstm/snake/vocoder/upsample tests pass. --- crates/metaltile-std/src/ffai/mod.rs | 5 -- .../src/{ffai => kernels/audio}/lstm.rs | 18 +++--- .../audio}/mel_spectrogram.rs | 60 +++++++++++-------- crates/metaltile-std/src/kernels/audio/mod.rs | 12 ++++ .../src/{ffai => kernels/audio}/snake1d.rs | 12 ++-- .../audio}/upsample_nearest1d.rs | 12 ++-- .../src/{ffai => kernels/audio}/vocoder.rs | 12 ++-- crates/metaltile-std/src/kernels/mod.rs | 1 + 8 files changed, 75 insertions(+), 57 deletions(-) rename crates/metaltile-std/src/{ffai => kernels/audio}/lstm.rs (97%) rename crates/metaltile-std/src/{ffai => kernels/audio}/mel_spectrogram.rs (93%) create mode 100644 crates/metaltile-std/src/kernels/audio/mod.rs rename crates/metaltile-std/src/{ffai => kernels/audio}/snake1d.rs (93%) rename crates/metaltile-std/src/{ffai => kernels/audio}/upsample_nearest1d.rs (92%) rename crates/metaltile-std/src/{ffai => kernels/audio}/vocoder.rs (97%) diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index 7b2df82e..b72caf17 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -71,8 +71,6 @@ pub mod gguf_iq2_xxs_extract_qs; pub mod im2col_patch; pub mod im2col_patch_interleaved; pub mod leaky_relu; -pub mod lstm; -pub mod mel_spectrogram; pub mod moe; pub mod moe_bgemm_iq2xxs_bm64; pub mod moe_bgemm_iq2xxs_mpp; @@ -121,9 +119,6 @@ pub mod sdpa_multi; pub mod sdpa_multi_d256; pub mod sdpa_prefill_d512_sink; pub mod sdpa_rel_pos_conformer; -pub mod snake1d; pub mod ssm; pub mod ssm_replay; pub mod transpose_th; -pub mod upsample_nearest1d; -pub mod vocoder; diff --git a/crates/metaltile-std/src/ffai/lstm.rs b/crates/metaltile-std/src/kernels/audio/lstm.rs similarity index 97% rename from crates/metaltile-std/src/ffai/lstm.rs rename to crates/metaltile-std/src/kernels/audio/lstm.rs index 10a8e40f..202e0991 100644 --- a/crates/metaltile-std/src/ffai/lstm.rs +++ b/crates/metaltile-std/src/kernels/audio/lstm.rs @@ -3,12 +3,12 @@ //! LSTM kernels — the recurrent building block style-vector TTS encoders / //! prosody / duration predictors need (no FFAI model used an LSTM before). //! This file holds two forms: -//! * [`ffai_lstm`] — runs the **whole sequence** recurrence on the GPU in +//! * [`mt_lstm`] — runs the **whole sequence** recurrence on the GPU in //! one dispatch (the "GPU from the start" form). //! * [`lstm_cell`] — ONE timestep, leaving the recurrence on the host (a //! per-step CPU↔GPU sync); use when per-step host control is wanted. //! -//! ## `ffai_lstm` +//! ## `mt_lstm` //! //! Runs the full sequence recurrence on the GPU in **one threadgroup**: //! thread `j` owns hidden unit `j`, with the hidden/cell state `h` / `c` @@ -49,7 +49,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_lstm( +pub fn mt_lstm( x: Tensor, w_ih: Tensor, w_hh: Tensor, @@ -133,9 +133,9 @@ pub fn ffai_lstm( /// one thread per `(b, j)` unit) — the per-step form, leaving the recurrence /// on the host (the host loops `t`, runs this forward `0..L` + backward /// `L-1..0` with separate weights for a bidirectional layer). Use this when -/// per-step host control is wanted; use [`ffai_lstm`] above to run the whole +/// per-step host control is wanted; use [`mt_lstm`] above to run the whole /// sequence on the GPU in one dispatch. Takes the `bias_ih` / `bias_hh` -/// split separately (vs. `ffai_lstm`'s precombined `bias`). +/// split separately (vs. `mt_lstm`'s precombined `bias`). #[kernel] pub fn lstm_cell( x: Tensor, @@ -210,7 +210,7 @@ pub fn lstm_cell( pub mod kernel_tests { use metaltile::{core::ir::Kernel, test::*, test_kernel}; - use super::{ffai_lstm, lstm_cell}; + use super::{lstm_cell, mt_lstm}; use crate::utils::{pack_f32, unpack_f32}; fn sigmoid(x: f32) -> f32 { 1.0 / (1.0 + (-x).exp()) } @@ -296,7 +296,7 @@ pub mod kernel_tests { 0, &mut expected, ); - TestSetup::new(ffai_lstm::kernel_ir_for(dt)) + TestSetup::new(mt_lstm::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("x", pack_f32(&x_f, dt), dt)) .input(TestBuffer::from_vec("w_ih", pack_f32(&w_ih_f, dt), dt)) @@ -443,13 +443,13 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::{ffai_lstm, lstm_cell}; + use super::{lstm_cell, mt_lstm}; #[bench(dtypes = [f32, f16, bf16])] fn bench_lstm(dt: DType) -> BenchSetup { let (seq_len, input_dim, hidden) = (200usize, 256usize, 256usize); let out_stride = hidden; - BenchSetup::new(ffai_lstm::kernel_ir_for(dt)) + BenchSetup::new(mt_lstm::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("x", seq_len * input_dim, dt)) .buffer(BenchBuffer::random("w_ih", 4 * hidden * input_dim, dt)) diff --git a/crates/metaltile-std/src/ffai/mel_spectrogram.rs b/crates/metaltile-std/src/kernels/audio/mel_spectrogram.rs similarity index 93% rename from crates/metaltile-std/src/ffai/mel_spectrogram.rs rename to crates/metaltile-std/src/kernels/audio/mel_spectrogram.rs index f4a647b7..a199357e 100644 --- a/crates/metaltile-std/src/ffai/mel_spectrogram.rs +++ b/crates/metaltile-std/src/kernels/audio/mel_spectrogram.rs @@ -45,7 +45,7 @@ use metaltile::kernel; #[kernel] -pub fn mel_spectrogram( +pub fn mt_mel_spectrogram( audio: Tensor, window: Tensor, mel_weight: Tensor, @@ -98,18 +98,18 @@ pub fn mel_spectrogram( // ───────────────────────────────────────────────────────────────────────── // FFT-routed STFT path. // -// `mel_spectrogram` does a direct DFT *inside every (frame, mel_bin) +// `mt_mel_spectrogram` does a direct DFT *inside every (frame, mel_bin) // thread* — so the full O(n_freq·n_fft) power spectrum is recomputed // `n_mels` times per frame. The FFT route splits it into three stages: // -// 1. `mel_stft_window` — extract + window each frame into FFT input +// 1. `mt_mel_stft_window` — extract + window each frame into FFT input // planes (real = windowed sample, imag = 0). // 2. `mt_fft_n{n_fft}` — one radix-2 FFT per frame (O(n_fft·log n_fft)). -// 3. `mel_filterbank` — power = re²+im², Mel-weight, log. +// 3. `mt_mel_filterbank` — power = re²+im², Mel-weight, log. // // The spectrum is now computed once per (frame, k) and the transform is // O(N log N) instead of O(N²). `n_fft` must be a power of two (the -// `mt_fft_n*` set). The single-kernel `mel_spectrogram` is kept for +// `mt_fft_n*` set). The single-kernel `mt_mel_spectrogram` is kept for // non-pow2 `n_fft` and single-dispatch callers. // ───────────────────────────────────────────────────────────────────────── @@ -118,7 +118,7 @@ pub fn mel_spectrogram( /// audio[frame*hop + t] · window[t]`, `out_im` zeroed. One thread per /// `(frame, t)`; dispatch flat over `n_frames * n_fft`. #[kernel] -pub fn mel_stft_window( +pub fn mt_mel_stft_window( audio: Tensor, window: Tensor, mut out_re: Tensor, @@ -128,7 +128,7 @@ pub fn mel_stft_window( #[constexpr] n_out: u32, ) { // `n_out = n_frames * n_fft`. Guard the threadgroup-rounded dispatch tail - // (see `mel_spectrogram`) from OOB `audio` reads / `out_re`/`out_im` writes. + // (see `mt_mel_spectrogram`) from OOB `audio` reads / `out_re`/`out_im` writes. let idx = program_id::<0>(); if idx < n_out { let t = idx % n_fft; @@ -144,10 +144,10 @@ pub fn mel_stft_window( /// mel] = log(Σ_{k( +pub fn mt_mel_filterbank( fft_re: Tensor, fft_im: Tensor, mel_weight: Tensor, @@ -159,7 +159,7 @@ pub fn mel_filterbank( #[constexpr] n_out: u32, ) { // `n_out = n_frames * n_mels`. Guard the threadgroup-rounded dispatch tail - // (see `mel_spectrogram`) from OOB `fft_re`/`fft_im` reads / `out` writes. + // (see `mt_mel_spectrogram`) from OOB `fft_re`/`fft_im` reads / `out` writes. let idx = program_id::<0>(); if idx < n_out { let mel_bin = idx % n_mels; @@ -180,12 +180,12 @@ pub fn mel_filterbank( } /// **Magnitude** log-Mel front-end — `|STFT| = sqrt(re²+im²)` through the -/// filterbank, vs the power (`re²+im²`) front-end of `mel_spectrogram` above. +/// filterbank, vs the power (`re²+im²`) front-end of `mt_mel_spectrogram` above. /// The amplitude-correct front-end the Gemma 4 audio encoder + several /// streaming-ASR models are trained on (feeding power degrades them). /// Direct-DFT, one thread per `(frame, mel_bin)`. #[kernel] -pub fn mel_spectrogram_magnitude( +pub fn mt_mel_spectrogram_magnitude( audio: Tensor, window: Tensor, mel_weight: Tensor, @@ -198,7 +198,7 @@ pub fn mel_spectrogram_magnitude( #[constexpr] n_out: u32, ) { // `n_out = n_frames * n_mels`. Guard the threadgroup-rounded dispatch tail - // (see `mel_spectrogram`) from OOB `audio`/`mel_weight` reads / `out` writes. + // (see `mt_mel_spectrogram`) from OOB `audio`/`mel_weight` reads / `out` writes. let idx = program_id::<0>(); if idx < n_out { let mel_bin = idx % n_mels; @@ -234,7 +234,12 @@ pub fn mel_spectrogram_magnitude( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::{mel_filterbank, mel_spectrogram, mel_spectrogram_magnitude, mel_stft_window}; + use super::{ + mt_mel_filterbank, + mt_mel_spectrogram, + mt_mel_spectrogram_magnitude, + mt_mel_stft_window, + }; use crate::utils::{pack_f32, unpack_f32}; const PI: f32 = std::f32::consts::PI; @@ -309,7 +314,7 @@ pub mod kernel_tests { // error — but only under low-precision *input* rounding (f16/bf16 quantize // the audio enough to sit a bin on the null; f32's finer grid does not). // The kernel is generic and correct — the math is identical across dtypes - // — so correctness is gated at f32; the post-FFT `mel_filterbank` test below + // — so correctness is gated at f32; the post-FFT `mt_mel_filterbank` test below // keeps f16/bf16 coverage on the path with no in-thread cancellation. See // the mel row in `specs/KERNEL_AUDIT.md`. #[test_kernel(dtypes = [f32], tol = [3e-3])] @@ -330,7 +335,7 @@ pub mod kernel_tests { &audio_dt, &window_dt, &mw_dt, n_fft, n_freq, n_mels, hop_length, n_frames, log_eps, ); let n_out = n_frames * n_mels; - TestSetup::new(mel_spectrogram::kernel_ir_for(dt)) + TestSetup::new(mt_mel_spectrogram::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("audio", pack_f32(&audio, dt), dt)) .input(TestBuffer::from_vec("window", pack_f32(&window, dt), dt)) @@ -364,7 +369,7 @@ pub mod kernel_tests { } } let exp_im = vec![0.0f32; n]; - TestSetup::new(mel_stft_window::kernel_ir_for(dt)) + TestSetup::new(mt_mel_stft_window::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("audio", pack_f32(&audio, dt), dt)) .input(TestBuffer::from_vec("window", pack_f32(&window, dt), dt)) @@ -405,7 +410,7 @@ pub mod kernel_tests { expected[frame * n_mels + mel_bin] = (acc + log_eps).ln(); } } - TestSetup::new(mel_filterbank::kernel_ir_for(dt)) + TestSetup::new(mt_mel_filterbank::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("fft_re", pack_f32(&fft_re, dt), dt)) .input(TestBuffer::from_vec("fft_im", pack_f32(&fft_im, dt), dt)) @@ -463,7 +468,7 @@ pub mod kernel_tests { // f32-only correctness gate, same direct-DFT cancellation-null reason as // `test_mel_spectrogram` (magnitude folds an extra `sqrt`, if anything - // sharpening the null sensitivity); the post-FFT `mel_filterbank` test keeps + // sharpening the null sensitivity); the post-FFT `mt_mel_filterbank` test keeps // f16/bf16 coverage. Looser f32 tol than the power sibling: the `sqrt` // amplifies the benign GPU↔CPU DFT accumulation-order difference. #[test_kernel(dtypes = [f32], tol = [1.5e-2])] @@ -484,7 +489,7 @@ pub mod kernel_tests { &audio_dt, &window_dt, &mw_dt, n_fft, n_freq, n_mels, hop_length, n_frames, log_eps, ); let n_out = n_frames * n_mels; - TestSetup::new(mel_spectrogram_magnitude::kernel_ir_for(dt)) + TestSetup::new(mt_mel_spectrogram_magnitude::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("audio", pack_f32(&audio, dt), dt)) .input(TestBuffer::from_vec("window", pack_f32(&window, dt), dt)) @@ -506,7 +511,12 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::{mel_filterbank, mel_spectrogram, mel_spectrogram_magnitude, mel_stft_window}; + use super::{ + mt_mel_filterbank, + mt_mel_spectrogram, + mt_mel_spectrogram_magnitude, + mt_mel_stft_window, + }; const N_FFT: usize = 400; const N_MELS: usize = 80; @@ -519,7 +529,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_mel_spectrogram(dt: DType) -> BenchSetup { let n_out = N_FRAMES * N_MELS; - BenchSetup::new(mel_spectrogram::kernel_ir_for(dt)) + BenchSetup::new(mt_mel_spectrogram::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("audio", n_samples(), dt)) .buffer(BenchBuffer::random("window", N_FFT, dt)) @@ -543,7 +553,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_mel_stft_window(dt: DType) -> BenchSetup { let n = N_FRAMES * N_FFT; - BenchSetup::new(mel_stft_window::kernel_ir_for(dt)) + BenchSetup::new(mt_mel_stft_window::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("audio", n_samples(), dt)) .buffer(BenchBuffer::random("window", N_FFT, dt)) @@ -559,7 +569,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_mel_filterbank(dt: DType) -> BenchSetup { let n_out = N_FRAMES * N_MELS; - BenchSetup::new(mel_filterbank::kernel_ir_for(dt)) + BenchSetup::new(mt_mel_filterbank::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("fft_re", N_FRAMES * N_FFT, dt)) .buffer(BenchBuffer::random("fft_im", N_FRAMES * N_FFT, dt)) @@ -580,7 +590,7 @@ pub mod kernel_benches { let n_freq = n_fft / 2 + 1; let n_samples = (n_frames - 1) * hop_length + n_fft; let n_out = n_frames * n_mels; - BenchSetup::new(mel_spectrogram_magnitude::kernel_ir_for(dt)) + BenchSetup::new(mt_mel_spectrogram_magnitude::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("audio", n_samples, dt)) .buffer(BenchBuffer::random("window", n_fft, dt)) diff --git a/crates/metaltile-std/src/kernels/audio/mod.rs b/crates/metaltile-std/src/kernels/audio/mod.rs new file mode 100644 index 00000000..0e5df71f --- /dev/null +++ b/crates/metaltile-std/src/kernels/audio/mod.rs @@ -0,0 +1,12 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Audio / speech front-end kernels — the audio family (see +//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): mel spectrogram (+ STFT window, +//! filterbank, magnitude), LSTM, the vocoder iSTFT, Snake1d activation, and +//! 1-D nearest upsampling. Migrated from the legacy `ffai/`. + +pub mod lstm; +pub mod mel_spectrogram; +pub mod snake1d; +pub mod upsample_nearest1d; +pub mod vocoder; diff --git a/crates/metaltile-std/src/ffai/snake1d.rs b/crates/metaltile-std/src/kernels/audio/snake1d.rs similarity index 93% rename from crates/metaltile-std/src/ffai/snake1d.rs rename to crates/metaltile-std/src/kernels/audio/snake1d.rs index 2213a5ca..1e986afe 100644 --- a/crates/metaltile-std/src/ffai/snake1d.rs +++ b/crates/metaltile-std/src/kernels/audio/snake1d.rs @@ -27,7 +27,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_snake1d( +pub fn mt_snake1d( input: Tensor, alpha: Tensor, mut out: Tensor, @@ -43,12 +43,12 @@ pub fn ffai_snake1d( store(out[i], y.cast::()); } -/// New-syntax correctness for `ffai_snake1d`. Grid3D, grid `[C·length,1,1]`, +/// New-syntax correctness for `mt_snake1d`. Grid3D, grid `[C·length,1,1]`, /// tpg `[1,1,1]`. Oracle applies snake per element with the channel's `α`. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_snake1d; + use super::mt_snake1d; use crate::utils::{pack_f32, unpack_f32}; #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 1e-2, 5e-2])] @@ -68,7 +68,7 @@ pub mod kernel_tests { x + (1.0 / (a + 1e-9)) * s * s }) .collect(); - TestSetup::new(ffai_snake1d::kernel_ir_for(dt)) + TestSetup::new(mt_snake1d::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::from_vec("alpha", pack_f32(&alpha_f, dt), dt)) @@ -83,13 +83,13 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_snake1d; + use super::mt_snake1d; #[bench(dtypes = [f32, f16, bf16])] fn bench_snake1d(dt: DType) -> BenchSetup { let (c, length) = (128usize, 7801usize); let n = c * length; - BenchSetup::new(ffai_snake1d::kernel_ir_for(dt)) + BenchSetup::new(mt_snake1d::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", n, dt)) .buffer(BenchBuffer::random("alpha", c, dt)) diff --git a/crates/metaltile-std/src/ffai/upsample_nearest1d.rs b/crates/metaltile-std/src/kernels/audio/upsample_nearest1d.rs similarity index 92% rename from crates/metaltile-std/src/ffai/upsample_nearest1d.rs rename to crates/metaltile-std/src/kernels/audio/upsample_nearest1d.rs index 3ae41b67..72062a92 100644 --- a/crates/metaltile-std/src/ffai/upsample_nearest1d.rs +++ b/crates/metaltile-std/src/kernels/audio/upsample_nearest1d.rs @@ -24,7 +24,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_upsample_nearest1d( +pub fn mt_upsample_nearest1d( input: Tensor, mut out: Tensor, #[constexpr] in_len: u32, @@ -37,13 +37,13 @@ pub fn ffai_upsample_nearest1d( store(out[i], load(input[c * in_len + t_in])); } -/// New-syntax correctness for `ffai_upsample_nearest1d`. Grid3D, grid +/// New-syntax correctness for `mt_upsample_nearest1d`. Grid3D, grid /// `[C·factor·in_len,1,1]`, tpg `[1,1,1]`. Oracle repeats each sample /// `factor` times. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_upsample_nearest1d; + use super::mt_upsample_nearest1d; use crate::utils::{pack_f32, unpack_f32}; fn setup(c: usize, in_len: usize, factor: usize, dt: DType) -> TestSetup { @@ -59,7 +59,7 @@ pub mod kernel_tests { input[ch * in_len + t_in] }) .collect(); - TestSetup::new(ffai_upsample_nearest1d::kernel_ir_for(dt)) + TestSetup::new(mt_upsample_nearest1d::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::zeros("out", n_out, dt)) @@ -82,13 +82,13 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_upsample_nearest1d; + use super::mt_upsample_nearest1d; #[bench(dtypes = [f32, f16, bf16])] fn bench_upsample_nearest1d(dt: DType) -> BenchSetup { let (c, in_len, factor) = (512usize, 130usize, 2usize); let n_out = c * in_len * factor; - BenchSetup::new(ffai_upsample_nearest1d::kernel_ir_for(dt)) + BenchSetup::new(mt_upsample_nearest1d::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", c * in_len, dt)) .buffer(BenchBuffer::zeros("out", n_out, dt).output()) diff --git a/crates/metaltile-std/src/ffai/vocoder.rs b/crates/metaltile-std/src/kernels/audio/vocoder.rs similarity index 97% rename from crates/metaltile-std/src/ffai/vocoder.rs rename to crates/metaltile-std/src/kernels/audio/vocoder.rs index 7bfe81bd..0bb8c70f 100644 --- a/crates/metaltile-std/src/ffai/vocoder.rs +++ b/crates/metaltile-std/src/kernels/audio/vocoder.rs @@ -56,7 +56,7 @@ use metaltile::kernel; #[kernel] -pub fn vocoder_istft( +pub fn mt_vocoder_istft( spec_re: Tensor, spec_im: Tensor, window: Tensor, @@ -126,7 +126,7 @@ pub fn vocoder_istft( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::vocoder_istft; + use super::mt_vocoder_istft; use crate::utils::{pack_f32, unpack_f32}; const PI: f32 = std::f32::consts::PI; @@ -224,7 +224,7 @@ pub mod kernel_tests { let im_dt = unpack_f32(&pack_f32(&im, dt), dt); let win_dt = unpack_f32(&pack_f32(&window, dt), dt); let expected = naive_istft(&re_dt, &im_dt, &win_dt, n_frames, n_fft, n_freq, hop); - TestSetup::new(vocoder_istft::kernel_ir_for(dt)) + TestSetup::new(mt_vocoder_istft::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("spec_re", pack_f32(&re, dt), dt)) .input(TestBuffer::from_vec("spec_im", pack_f32(&im, dt), dt)) @@ -239,19 +239,19 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `vocoder_istft` — a Kokoro-class iSTFTNet tail +/// New-syntax benchmark for `mt_vocoder_istft` — a Kokoro-class iSTFTNet tail /// over many frames (Grid3D, one thread per output sample). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::vocoder_istft; + use super::mt_vocoder_istft; #[bench(dtypes = [f32, f16, bf16])] fn bench_vocoder_istft(dt: DType) -> BenchSetup { let (n_frames, n_fft, hop) = (2048usize, 20usize, 5usize); let n_freq = n_fft / 2 + 1; let out_len = (n_frames - 1) * hop + n_fft; - BenchSetup::new(vocoder_istft::kernel_ir_for(dt)) + BenchSetup::new(mt_vocoder_istft::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("spec_re", n_frames * n_freq, dt)) .buffer(BenchBuffer::random("spec_im", n_frames * n_freq, dt)) diff --git a/crates/metaltile-std/src/kernels/mod.rs b/crates/metaltile-std/src/kernels/mod.rs index c7eb5668..8e9f5680 100644 --- a/crates/metaltile-std/src/kernels/mod.rs +++ b/crates/metaltile-std/src/kernels/mod.rs @@ -7,6 +7,7 @@ //! model-specific usage notes live in comments above the kernel. The FFAI emit //! consumer is regenerated from the new inventory after each family lands. +pub mod audio; pub mod convolution; pub mod gemm; pub mod kv_cache; From 0a7260441b763d33b08c2fc3be5e56711372e06a Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Sat, 13 Jun 2026 10:26:10 -0600 Subject: [PATCH 3/5] refactor(vision): migrate vision front-end family into kernels/vision/ ffai/{avg_pool2d_nhwc,frame_diff_luma,im2col_patch,im2col_patch_interleaved, patch_unfold_qwen,pos_emb_2d_add,resize_normalize,transpose_th} -> kernels/vision/ with mt_ prefixes. patch_unfold_qwen.rs -> patch_unfold.rs and ffai_patch_unfold_qwen -> mt_patch_unfold (Qwen stays as a usage note in the comment, not the name; test/bench fns de-modelled too). cargo build + clippy clean; resize/im2col/patch_unfold/transpose/avg_pool/ frame_diff/pos_emb tests pass. --- crates/metaltile-std/src/ffai/mod.rs | 8 ------- crates/metaltile-std/src/kernels/mod.rs | 1 + .../vision}/avg_pool2d_nhwc.rs | 12 +++++----- .../vision}/frame_diff_luma.rs | 10 ++++----- .../{ffai => kernels/vision}/im2col_patch.rs | 10 ++++----- .../vision}/im2col_patch_interleaved.rs | 10 ++++----- .../metaltile-std/src/kernels/vision/mod.rs | 16 ++++++++++++++ .../vision/patch_unfold.rs} | 16 +++++++------- .../vision}/pos_emb_2d_add.rs | 10 ++++----- .../vision}/resize_normalize.rs | 22 +++++++++---------- .../{ffai => kernels/vision}/transpose_th.rs | 16 +++++++------- 11 files changed, 70 insertions(+), 61 deletions(-) rename crates/metaltile-std/src/{ffai => kernels/vision}/avg_pool2d_nhwc.rs (96%) rename crates/metaltile-std/src/{ffai => kernels/vision}/frame_diff_luma.rs (96%) rename crates/metaltile-std/src/{ffai => kernels/vision}/im2col_patch.rs (96%) rename crates/metaltile-std/src/{ffai => kernels/vision}/im2col_patch_interleaved.rs (96%) create mode 100644 crates/metaltile-std/src/kernels/vision/mod.rs rename crates/metaltile-std/src/{ffai/patch_unfold_qwen.rs => kernels/vision/patch_unfold.rs} (95%) rename crates/metaltile-std/src/{ffai => kernels/vision}/pos_emb_2d_add.rs (95%) rename crates/metaltile-std/src/{ffai => kernels/vision}/resize_normalize.rs (96%) rename crates/metaltile-std/src/{ffai => kernels/vision}/transpose_th.rs (92%) diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index b72caf17..3ab9449d 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -24,7 +24,6 @@ pub mod aura_flash_pass2; pub mod aura_flash_sdpa; pub mod aura_score; pub mod aura_value; -pub mod avg_pool2d_nhwc; pub mod batched_4_block_scaled_qgemv; pub mod batched_4_block_scaled_qmm; pub mod batched_4_qgemv; @@ -51,7 +50,6 @@ pub mod dsv4_swiglu_limit; pub mod ffai_dequant_q4; pub mod flash_block_scaled_sdpa; pub mod flash_quantized_sdpa; -pub mod frame_diff_luma; pub mod gate_up_swiglu_fused; pub mod gated_delta; pub mod gated_delta_prep; @@ -68,8 +66,6 @@ pub mod gguf_dequant_iq2_xxs_raw; pub mod gguf_dequant_q2_k; pub mod gguf_dequant_q8_0; pub mod gguf_iq2_xxs_extract_qs; -pub mod im2col_patch; -pub mod im2col_patch_interleaved; pub mod leaky_relu; pub mod moe; pub mod moe_bgemm_iq2xxs_bm64; @@ -104,9 +100,6 @@ pub mod moe_router_sigmoid_bias; pub mod moe_router_sqrtsoftplus; pub mod patch_embed_block_scaled; pub mod patch_embed_mma_block_scaled; -pub mod patch_unfold_qwen; -pub mod pos_emb_2d_add; -pub mod resize_normalize; pub mod sdpa_bidirectional; pub mod sdpa_bidirectional_d128_relpos; pub mod sdpa_bidirectional_windowed; @@ -121,4 +114,3 @@ pub mod sdpa_prefill_d512_sink; pub mod sdpa_rel_pos_conformer; pub mod ssm; pub mod ssm_replay; -pub mod transpose_th; diff --git a/crates/metaltile-std/src/kernels/mod.rs b/crates/metaltile-std/src/kernels/mod.rs index 8e9f5680..0375c411 100644 --- a/crates/metaltile-std/src/kernels/mod.rs +++ b/crates/metaltile-std/src/kernels/mod.rs @@ -15,3 +15,4 @@ pub mod norm; pub mod ops; pub mod rope; pub mod sampling; +pub mod vision; diff --git a/crates/metaltile-std/src/ffai/avg_pool2d_nhwc.rs b/crates/metaltile-std/src/kernels/vision/avg_pool2d_nhwc.rs similarity index 96% rename from crates/metaltile-std/src/ffai/avg_pool2d_nhwc.rs rename to crates/metaltile-std/src/kernels/vision/avg_pool2d_nhwc.rs index 79f43716..e7e0d002 100644 --- a/crates/metaltile-std/src/ffai/avg_pool2d_nhwc.rs +++ b/crates/metaltile-std/src/kernels/vision/avg_pool2d_nhwc.rs @@ -27,7 +27,7 @@ use metaltile::kernel; #[kernel] -pub fn avg_pool2d_nhwc( +pub fn mt_avg_pool2d_nhwc( input: Tensor, out: Tensor, #[constexpr] batch: u32, @@ -75,7 +75,7 @@ pub fn avg_pool2d_nhwc( pub mod kernel_tests { use metaltile::{core::ir::Kernel, test::*, test_kernel}; - use super::avg_pool2d_nhwc; + use super::mt_avg_pool2d_nhwc; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, period: usize, amp: f32) -> Vec { @@ -173,13 +173,13 @@ pub mod kernel_tests { // Gemma 4-style 3×3 stride-3 non-overlapping pool (the pooling kernel). #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 5e-3, 3e-2])] fn test_avg_pool2d_nhwc_3x3_s3(dt: DType) -> TestSetup { - setup(avg_pool2d_nhwc::kernel_ir_for(dt), 1, 16, 24, 24, 3, 3, 3, 0, dt) + setup(mt_avg_pool2d_nhwc::kernel_ir_for(dt), 1, 16, 24, 24, 3, 3, 3, 0, dt) } // 2×2 stride-2 with padding-1 — exercises the partial-window divisor. #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 5e-3, 3e-2])] fn test_avg_pool2d_nhwc_2x2_pad(dt: DType) -> TestSetup { - setup(avg_pool2d_nhwc::kernel_ir_for(dt), 2, 8, 15, 15, 2, 2, 2, 1, dt) + setup(mt_avg_pool2d_nhwc::kernel_ir_for(dt), 2, 8, 15, 15, 2, 2, 2, 1, dt) } } @@ -187,7 +187,7 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::avg_pool2d_nhwc; + use super::mt_avg_pool2d_nhwc; #[bench(dtypes = [f32, f16, bf16])] fn bench_avg_pool2d_nhwc(dt: DType) -> BenchSetup { @@ -196,7 +196,7 @@ pub mod kernel_benches { let out_h = (in_h + 2 * pad - k) / stride + 1; let out_w = (in_w + 2 * pad - k) / stride + 1; let n_out = batch * out_h * out_w * ch; - BenchSetup::new(avg_pool2d_nhwc::kernel_ir_for(dt)) + BenchSetup::new(mt_avg_pool2d_nhwc::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", batch * in_h * in_w * ch, dt)) .buffer(BenchBuffer::zeros("out", n_out, dt).output()) diff --git a/crates/metaltile-std/src/ffai/frame_diff_luma.rs b/crates/metaltile-std/src/kernels/vision/frame_diff_luma.rs similarity index 96% rename from crates/metaltile-std/src/ffai/frame_diff_luma.rs rename to crates/metaltile-std/src/kernels/vision/frame_diff_luma.rs index 61c03565..b164b0c4 100644 --- a/crates/metaltile-std/src/ffai/frame_diff_luma.rs +++ b/crates/metaltile-std/src/kernels/vision/frame_diff_luma.rs @@ -28,7 +28,7 @@ use metaltile::kernel; #[kernel] -pub fn frame_diff_luma( +pub fn mt_frame_diff_luma( frame0: Tensor, frame1: Tensor, out: Tensor, @@ -73,7 +73,7 @@ pub fn frame_diff_luma( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::frame_diff_luma; + use super::mt_frame_diff_luma; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, step: f32, start: f32) -> Vec { @@ -116,7 +116,7 @@ pub mod kernel_tests { let f0 = unpack_f32(&pack_f32(&f0_f, dt), dt); let f1 = unpack_f32(&pack_f32(&f1_f, dt), dt); let expected = naive(&f0, &f1, in_w, out_h, out_w, ds); - TestSetup::new(frame_diff_luma::kernel_ir_for(dt)) + TestSetup::new(mt_frame_diff_luma::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("frame0", pack_f32(&f0_f, dt), dt)) .input(TestBuffer::from_vec("frame1", pack_f32(&f1_f, dt), dt)) @@ -142,14 +142,14 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::frame_diff_luma; + use super::mt_frame_diff_luma; #[bench(dtypes = [f32, f16, bf16])] fn bench_frame_diff_luma(dt: DType) -> BenchSetup { let (in_h, in_w, ds) = (720usize, 1280usize, 16usize); let out_h = in_h / ds; let out_w = in_w / ds; - BenchSetup::new(frame_diff_luma::kernel_ir_for(dt)) + BenchSetup::new(mt_frame_diff_luma::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("frame0", in_h * in_w * 3, dt)) .buffer(BenchBuffer::random("frame1", in_h * in_w * 3, dt)) diff --git a/crates/metaltile-std/src/ffai/im2col_patch.rs b/crates/metaltile-std/src/kernels/vision/im2col_patch.rs similarity index 96% rename from crates/metaltile-std/src/ffai/im2col_patch.rs rename to crates/metaltile-std/src/kernels/vision/im2col_patch.rs index 64c535fb..a48d9422 100644 --- a/crates/metaltile-std/src/ffai/im2col_patch.rs +++ b/crates/metaltile-std/src/kernels/vision/im2col_patch.rs @@ -28,7 +28,7 @@ use metaltile::kernel; #[kernel] -pub fn im2col_patch( +pub fn mt_im2col_patch( input: Tensor, out: Tensor, #[constexpr] channels: u32, @@ -58,7 +58,7 @@ pub fn im2col_patch( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::im2col_patch; + use super::mt_im2col_patch; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, period: usize, amp: f32) -> Vec { @@ -102,7 +102,7 @@ pub mod kernel_tests { let input_f = ramp(channels * in_h * in_w, 17, 5.0); let input = unpack_f32(&pack_f32(&input_f, dt), dt); let expected = naive(&input, channels, in_h, in_w, patch, grid_h, grid_w); - TestSetup::new(im2col_patch::kernel_ir_for(dt)) + TestSetup::new(mt_im2col_patch::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::zeros("out", n_out, dt)) @@ -128,7 +128,7 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::im2col_patch; + use super::mt_im2col_patch; #[bench(dtypes = [f32, f16, bf16])] fn bench_im2col_patch(dt: DType) -> BenchSetup { @@ -136,7 +136,7 @@ pub mod kernel_benches { let in_h = grid_h * patch; let in_w = grid_w * patch; let n_out = grid_h * grid_w * channels * patch * patch; - BenchSetup::new(im2col_patch::kernel_ir_for(dt)) + BenchSetup::new(mt_im2col_patch::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", channels * in_h * in_w, dt)) .buffer(BenchBuffer::zeros("out", n_out, dt).output()) diff --git a/crates/metaltile-std/src/ffai/im2col_patch_interleaved.rs b/crates/metaltile-std/src/kernels/vision/im2col_patch_interleaved.rs similarity index 96% rename from crates/metaltile-std/src/ffai/im2col_patch_interleaved.rs rename to crates/metaltile-std/src/kernels/vision/im2col_patch_interleaved.rs index 328ae056..d3250008 100644 --- a/crates/metaltile-std/src/ffai/im2col_patch_interleaved.rs +++ b/crates/metaltile-std/src/kernels/vision/im2col_patch_interleaved.rs @@ -32,7 +32,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_im2col_patch_interleaved( +pub fn mt_im2col_patch_interleaved( input: Tensor, out: Tensor, #[constexpr] channels: u32, @@ -71,7 +71,7 @@ pub fn ffai_im2col_patch_interleaved( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_im2col_patch_interleaved; + use super::mt_im2col_patch_interleaved; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, period: usize, amp: f32) -> Vec { @@ -141,7 +141,7 @@ pub mod kernel_tests { scale, bias, ); - TestSetup::new(ffai_im2col_patch_interleaved::kernel_ir_for(dt)) + TestSetup::new(mt_im2col_patch_interleaved::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::zeros("out", n_out, dt)) @@ -171,7 +171,7 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_im2col_patch_interleaved; + use super::mt_im2col_patch_interleaved; #[bench(dtypes = [f32, f16, bf16])] fn bench_interleaved(dt: DType) -> BenchSetup { @@ -180,7 +180,7 @@ pub mod kernel_benches { let in_h = grid_h * patch; let in_w = grid_w * patch; let n_out = grid_h * grid_w * patch_dim_padded; - BenchSetup::new(ffai_im2col_patch_interleaved::kernel_ir_for(dt)) + BenchSetup::new(mt_im2col_patch_interleaved::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", channels * in_h * in_w, dt)) .buffer(BenchBuffer::zeros("out", n_out, dt).output()) diff --git a/crates/metaltile-std/src/kernels/vision/mod.rs b/crates/metaltile-std/src/kernels/vision/mod.rs new file mode 100644 index 00000000..bcc674ed --- /dev/null +++ b/crates/metaltile-std/src/kernels/vision/mod.rs @@ -0,0 +1,16 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Vision front-end kernels — the vision family (see +//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): image resize+normalize (+bicubic), +//! im2col patch extraction (+ interleaved), patch unfold (non-square grids), +//! 2-D positional-embedding add, NHWC avg-pool, token-major transpose, and luma +//! frame differencing. Migrated from the legacy `ffai/`. + +pub mod avg_pool2d_nhwc; +pub mod frame_diff_luma; +pub mod im2col_patch; +pub mod im2col_patch_interleaved; +pub mod patch_unfold; +pub mod pos_emb_2d_add; +pub mod resize_normalize; +pub mod transpose_th; diff --git a/crates/metaltile-std/src/ffai/patch_unfold_qwen.rs b/crates/metaltile-std/src/kernels/vision/patch_unfold.rs similarity index 95% rename from crates/metaltile-std/src/ffai/patch_unfold_qwen.rs rename to crates/metaltile-std/src/kernels/vision/patch_unfold.rs index 9090ff62..bdc110e5 100644 --- a/crates/metaltile-std/src/ffai/patch_unfold_qwen.rs +++ b/crates/metaltile-std/src/kernels/vision/patch_unfold.rs @@ -38,7 +38,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_patch_unfold_qwen( +pub fn mt_patch_unfold( frames: Tensor, out: Tensor, #[constexpr] channels: u32, @@ -88,7 +88,7 @@ pub fn ffai_patch_unfold_qwen( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_patch_unfold_qwen; + use super::mt_patch_unfold; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, period: usize, amp: f32) -> Vec { @@ -179,7 +179,7 @@ pub mod kernel_tests { grid_t, is_image, ); - TestSetup::new(ffai_patch_unfold_qwen::kernel_ir_for(dt)) + TestSetup::new(mt_patch_unfold::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("frames", pack_f32(&frames_f, dt), dt)) .input(TestBuffer::zeros("out", n_out, dt)) @@ -198,11 +198,11 @@ pub mod kernel_tests { // Still image: 3 ch, patch 2, temporal 2 (frame reused), merge 2, 2 blocks // (side 4, img 8×8), grid_t 1. real_dim = 2·3·4 = 24, padded to 32. #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 1e-3, 1e-2])] - fn test_qwen_image(dt: DType) -> TestSetup { setup(dt, 3, 2, 2, 2, 2, 32, 1, true) } + fn test_patch_unfold_image(dt: DType) -> TestSetup { setup(dt, 3, 2, 2, 2, 2, 32, 1, true) } // Video: distinct frames, grid_t 2 (4 frames), same geometry. #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 1e-3, 1e-2])] - fn test_qwen_video(dt: DType) -> TestSetup { setup(dt, 3, 2, 2, 2, 2, 32, 2, false) } + fn test_patch_unfold_video(dt: DType) -> TestSetup { setup(dt, 3, 2, 2, 2, 2, 32, 2, false) } } /// New-syntax bench: Qwen2.5-VL image unfold (3 ch, patch 14, temporal 2, @@ -210,10 +210,10 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_patch_unfold_qwen; + use super::mt_patch_unfold; #[bench(dtypes = [f32, f16, bf16])] - fn bench_patch_unfold_qwen(dt: DType) -> BenchSetup { + fn bench_patch_unfold(dt: DType) -> BenchSetup { let (channels, patch, temporal_patch, merge, merge_blocks) = (3usize, 14usize, 2usize, 2usize, 8usize); let side = merge_blocks * merge; @@ -222,7 +222,7 @@ pub mod kernel_benches { let patch_dim_padded = patch_dim.div_ceil(16) * 16; let n_patches = side * side; let n_out = n_patches * patch_dim_padded; - BenchSetup::new(ffai_patch_unfold_qwen::kernel_ir_for(dt)) + BenchSetup::new(mt_patch_unfold::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("frames", channels * img_side * img_side, dt)) .buffer(BenchBuffer::zeros("out", n_out, dt).output()) diff --git a/crates/metaltile-std/src/ffai/pos_emb_2d_add.rs b/crates/metaltile-std/src/kernels/vision/pos_emb_2d_add.rs similarity index 95% rename from crates/metaltile-std/src/ffai/pos_emb_2d_add.rs rename to crates/metaltile-std/src/kernels/vision/pos_emb_2d_add.rs index 3efac4d6..6b6f3007 100644 --- a/crates/metaltile-std/src/ffai/pos_emb_2d_add.rs +++ b/crates/metaltile-std/src/kernels/vision/pos_emb_2d_add.rs @@ -31,7 +31,7 @@ use metaltile::kernel; #[kernel] -pub fn pos_emb_2d_add( +pub fn mt_pos_emb_2d_add( tokens: Tensor, pos_x: Tensor, pos_y: Tensor, @@ -53,7 +53,7 @@ pub fn pos_emb_2d_add( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::pos_emb_2d_add; + use super::mt_pos_emb_2d_add; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, step: f32, start: f32) -> Vec { @@ -87,7 +87,7 @@ pub mod kernel_tests { let pos_y = ramp(grid_h * hidden, 0.009, 0.1); let tokens = unpack_f32(&pack_f32(&tokens_f, dt), dt); let expected = naive(&tokens, &pos_x, &pos_y, hidden, grid_h, grid_w); - TestSetup::new(pos_emb_2d_add::kernel_ir_for(dt)) + TestSetup::new(mt_pos_emb_2d_add::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("tokens", pack_f32(&tokens_f, dt), dt)) .input(TestBuffer::from_vec("pos_x", pack_f32(&pos_x, DType::F32), DType::F32)) @@ -112,13 +112,13 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::pos_emb_2d_add; + use super::mt_pos_emb_2d_add; #[bench(dtypes = [f32, f16, bf16])] fn bench_pos_emb_2d_add(dt: DType) -> BenchSetup { let (hidden, grid_h, grid_w) = (1152usize, 28usize, 28usize); let n_patches = grid_h * grid_w; - BenchSetup::new(pos_emb_2d_add::kernel_ir_for(dt)) + BenchSetup::new(mt_pos_emb_2d_add::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("tokens", n_patches * hidden, dt)) .buffer(BenchBuffer::random("pos_x", grid_w * hidden, DType::F32)) diff --git a/crates/metaltile-std/src/ffai/resize_normalize.rs b/crates/metaltile-std/src/kernels/vision/resize_normalize.rs similarity index 96% rename from crates/metaltile-std/src/ffai/resize_normalize.rs rename to crates/metaltile-std/src/kernels/vision/resize_normalize.rs index d5cb0118..1c0750ed 100644 --- a/crates/metaltile-std/src/ffai/resize_normalize.rs +++ b/crates/metaltile-std/src/kernels/vision/resize_normalize.rs @@ -27,7 +27,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_resize_normalize( +pub fn mt_resize_normalize( input: Tensor, mean: Tensor, std: Tensor, @@ -100,7 +100,7 @@ pub fn ffai_resize_normalize( /// weights (no renormalize). Per-tap branch specialization is exact at /// boundaries (cubic1(1)=cubic2(1)=0). #[kernel] -pub fn ffai_resize_normalize_bicubic( +pub fn mt_resize_normalize_bicubic( input: Tensor, mean: Tensor, std: Tensor, @@ -202,12 +202,12 @@ pub fn ffai_resize_normalize_bicubic( store(out[(c * th + oy) * tw + ox], normed.cast::()); } -/// New-syntax correctness for `ffai_resize_normalize` vs a CPU bilinear +/// New-syntax correctness for `mt_resize_normalize` vs a CPU bilinear /// reference. Grid3D, grid `[target_w, target_h, 3]`, tpg `[1,1,1]`. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::{ffai_resize_normalize, ffai_resize_normalize_bicubic}; + use super::{mt_resize_normalize, mt_resize_normalize_bicubic}; use crate::utils::{pack_f32, unpack_f32}; fn u32_bytes(v: u32) -> Vec { v.to_le_bytes().to_vec() } @@ -261,7 +261,7 @@ pub mod kernel_tests { let src_f: Vec = (0..sh * sw * 3).map(|i| ((i % 17) as f32) / 17.0).collect(); let src = unpack_f32(&pack_f32(&src_f, dt), dt); let exp = cpu_ref(&src, sw, sh, tw, th, &mean, &std); - TestSetup::new(ffai_resize_normalize::kernel_ir_for(dt)) + TestSetup::new(mt_resize_normalize::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&src_f, dt), dt)) .input(TestBuffer::from_vec( @@ -346,7 +346,7 @@ pub mod kernel_tests { let src_f: Vec = (0..sh * sw * 3).map(|i| ((i % 23) as f32) / 23.0).collect(); let src = unpack_f32(&pack_f32(&src_f, dt), dt); let exp = cpu_ref_bicubic(&src, sw, sh, tw, th, &mean, &std); - TestSetup::new(ffai_resize_normalize_bicubic::kernel_ir_for(dt)) + TestSetup::new(mt_resize_normalize_bicubic::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&src_f, dt), dt)) .input(TestBuffer::from_vec( @@ -376,7 +376,7 @@ pub mod kernel_tests { let src_f: Vec = (0..sh * sw * 3).map(|i| ((i % 17) as f32) / 17.0).collect(); let src = unpack_f32(&pack_f32(&src_f, dt), dt); let exp = cpu_ref_bicubic(&src, sw, sh, tw, th, &mean, &std); - TestSetup::new(ffai_resize_normalize_bicubic::kernel_ir_for(dt)) + TestSetup::new(mt_resize_normalize_bicubic::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&src_f, dt), dt)) .input(TestBuffer::from_vec( @@ -399,17 +399,17 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `ffai_resize_normalize` at a representative VL +/// New-syntax benchmark for `mt_resize_normalize` at a representative VL /// preprocess shape (≈640×480 source → 448×448). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::{ffai_resize_normalize, ffai_resize_normalize_bicubic}; + use super::{mt_resize_normalize, mt_resize_normalize_bicubic}; #[bench(dtypes = [f32, f16, bf16])] fn bench_resize_normalize(dt: DType) -> BenchSetup { let (sw, sh, tw, th) = (640usize, 480usize, 448usize, 448usize); - BenchSetup::new(ffai_resize_normalize::kernel_ir_for(dt)) + BenchSetup::new(mt_resize_normalize::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", sh * sw * 3, dt)) .buffer(BenchBuffer::random("mean", 3, DType::F32)) @@ -440,7 +440,7 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_resize_normalize_bicubic(dt: DType) -> BenchSetup { let (sw, sh, tw, th) = (640usize, 480usize, 448usize, 448usize); - BenchSetup::new(ffai_resize_normalize_bicubic::kernel_ir_for(dt)) + BenchSetup::new(mt_resize_normalize_bicubic::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", sh * sw * 3, dt)) .buffer(BenchBuffer::random("mean", 3, DType::F32)) diff --git a/crates/metaltile-std/src/ffai/transpose_th.rs b/crates/metaltile-std/src/kernels/vision/transpose_th.rs similarity index 92% rename from crates/metaltile-std/src/ffai/transpose_th.rs rename to crates/metaltile-std/src/kernels/vision/transpose_th.rs index 5559d880..9b9db1ba 100644 --- a/crates/metaltile-std/src/ffai/transpose_th.rs +++ b/crates/metaltile-std/src/kernels/vision/transpose_th.rs @@ -3,7 +3,7 @@ //! A vision (or audio) tower's attention stage-1 — per-head RMSNorm plus //! `mt_rope_2d` — emits Q/K/V in **token-major** layout //! `[n_tokens, n_heads, head_dim]` (one contiguous head block per token). -//! But `ffai_sdpa_bidirectional` reads K/V **head-major** +//! But `mt_sdpa_bidirectional` reads K/V **head-major** //! `[n_heads, n_tokens, head_dim]` — its `kv_slab = kvh * kv_stride * //! head_dim` indexing walks one head's full token run contiguously. This //! kernel performs that physical reshape on the GPU so the whole attention @@ -37,7 +37,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_transpose_th( +pub fn mt_transpose_th( input: Tensor, out: Tensor, #[constexpr] n_tokens: u32, @@ -52,14 +52,14 @@ pub fn ffai_transpose_th( store(out[out_idx], load(input[in_idx])); } -/// New-syntax correctness for `ffai_transpose_th`. Grid3D, grid +/// New-syntax correctness for `mt_transpose_th`. Grid3D, grid /// `[n_tokens, n_heads, head_dim]`, tpg `[1,1,1]`. Oracle moves element /// `(token, head, d)` of the token-major input to `(head, token, d)` of /// the head-major output. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_transpose_th; + use super::mt_transpose_th; use crate::utils::{pack_f32, unpack_f32}; fn oracle(input: &[f32], n_tokens: usize, n_heads: usize, head_dim: usize) -> Vec { @@ -87,7 +87,7 @@ pub mod kernel_tests { // tol = 0. let input = unpack_f32(&pack_f32(&input_f, dt), dt); let exp = oracle(&input, n_tokens, n_heads, head_dim); - TestSetup::new(ffai_transpose_th::kernel_ir_for(dt)) + TestSetup::new(mt_transpose_th::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::zeros("out", n_tokens * n_heads * head_dim, dt)) @@ -99,17 +99,17 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `ffai_transpose_th` at the SigLIP production +/// New-syntax benchmark for `mt_transpose_th` at the SigLIP production /// shape (576 patches, 16 heads, head_dim 64). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_transpose_th; + use super::mt_transpose_th; #[bench(dtypes = [f32, f16, bf16])] fn bench_transpose_th(dt: DType) -> BenchSetup { let (n_tokens, n_heads, head_dim) = (576usize, 16usize, 64usize); - BenchSetup::new(ffai_transpose_th::kernel_ir_for(dt)) + BenchSetup::new(mt_transpose_th::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", n_tokens * n_heads * head_dim, dt)) .buffer(BenchBuffer::zeros("out", n_tokens * n_heads * head_dim, dt).output()) From 9d31a7f63a4e87e5631c042e075303ffb9c19f46 Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Sat, 13 Jun 2026 10:27:16 -0600 Subject: [PATCH 4/5] docs: mark kv_cache/audio/vision families done (plan + audit paths) --- docs/specs/KERNEL_AUDIT.md | 14 +++++++------- docs/specs/KERNEL_CONSOLIDATION_PLAN.md | 8 ++++---- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docs/specs/KERNEL_AUDIT.md b/docs/specs/KERNEL_AUDIT.md index e30c8145..34f8be5f 100644 --- a/docs/specs/KERNEL_AUDIT.md +++ b/docs/specs/KERNEL_AUDIT.md @@ -112,7 +112,7 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | fp_quantized_mma | ✗ | ✗ | ✓ | `mlx/fp_quantized_mma.rs` (PR #157) → `mt_fp4_qmm_mma` + `mt_fp8_e4m3_qmm_mma`. Simdgroup-matrix BM=BN=BK=32 MMA — same 4-SG 2×2 scaffold as `mt_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 `mlx/block_scaled_mma.rs` (its float-scale fp4 kernel is `mt_fp4_float_qmm_mma`). | | fp_quantized_nax | ✓ | ✓ | ✓ | `mlx/fp_quantized_nax.rs` → `mt_fp_qmm_nax`. 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 | ✓ | ✓ | ✓ | `mlx/quantized_nax.rs` → `mt_qmm_nax` (int4) + `mt_qmm_nax_int8` (int8, PR #154 in `mlx/quantized_nax_int8.rs`). MPP counterpart of `mt_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) | ✓ | ✓ | ✓ | `mlx/fft.rs` → `mt_fft_n{32,64,128,256,512,1024}` (iterative radix-2 Cooley–Tukey, forward + inverse via `inv` constexpr; complex via parallel real/imag planes). **Non-pow2 Bluestein** (PR #157): `mt_fft_bluestein_preprocess` + `mt_fft_bluestein_chirp_filter` + `mt_fft_bluestein_cmul` + `mt_fft_bluestein_postprocess` — 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. | +| fft (radix + readwrite + non-pow2) | ✓ | ✓ | ✓ | `kernels/kv_cache/fft.rs` → `mt_fft_n{32,64,128,256,512,1024}` (iterative radix-2 Cooley–Tukey, forward + inverse via `inv` constexpr; complex via parallel real/imag planes). **Non-pow2 Bluestein** (PR #157): `mt_fft_bluestein_preprocess` + `mt_fft_bluestein_chirp_filter` + `mt_fft_bluestein_cmul` + `mt_fft_bluestein_postprocess` — 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` → `mt_hadamard_n{64,128,256,512,1024}` (FWHT, log2(N) butterfly passes). `kernels/ops/hadamard_m.rs` → `mt_hadamard_m{12,20,28}` (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](#fence-ops--intentionally-out-of-scope). | | gather (bare-tensor embedding lookup) | ✓ | ✓ | ✓ | `kernels/ops/gather.rs` → `mt_gather`. FFAI's embedding-table gather. | @@ -135,8 +135,8 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | rms_norm_rope (RMSNorm + RoPE fused) | ✗ | ✓ | ✓ | `kernels/norm/rms_norm_rope.rs` → `mt_rms_norm_rope`. 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` → `mt_rms_norm_qgemv` (int4, one-row-per-TG correctness shape) + `mt_rms_norm_qgemv_fast` (int4, 8-row-per-TG perf path, PR #154) + `mt_rms_norm_qgemv_int8_fast` (int8, 8-row-per-TG, PR #157). | | batched_qkv_qgemv (Q/K/V 4-bit qGEMV → 1 dispatch) | ✗ | ✓ | ✓ | `ffai/batched_qkv_qgemv.rs` → `ffai_batched_qkv_qgemv` (one-row-per-TG) + `ffai_batched_qkv_qgemv_fast` (8-row-per-TG, GQA-guarded, PR #154). `program_id::<2>()` selects Q/K/V, output concatenated `[Q\|K\|V]`. | -| kv_cache_update (raw bf16/fp16 single-token append) | ✗ | ✗ | ✓ | `ffai/kv_cache.rs` → `kv_cache_update`. FFAI-only; raw cache append. | -| kv_cache (affine-quant int4/int8/fp8 quantize + bulk dequant) | ~ | ~ | ✓ | `ffai/kv_cache.rs` — `quantize_kv` + `bulk_dequant_kv` for int4/int8. **fp8** (PR #157): `quantize_kv_fp8_{e4m3,e5m2}` + `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. | +| kv_cache_update (raw bf16/fp16 single-token append) | ✗ | ✗ | ✓ | `kernels/kv_cache/cache.rs` → `mt_kv_cache_update`. FFAI-only; raw cache append. | +| kv_cache (affine-quant int4/int8/fp8 quantize + bulk dequant) | ~ | ~ | ✓ | `kernels/kv_cache/cache.rs` — `mt_quantize_kv` + `mt_bulk_dequant_kv` for int4/int8. **fp8** (PR #157): `mt_quantize_kv_fp8_{e4m3,e5m2}` + `mt_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` → `mt_softmax_categorical_sample`. Companion to `mt_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 `mt_softmax_categorical_sample`. | | sdpa_decode + learned attention sink (GPT-OSS-20B) | ✗ | ~ | ✓ | `ffai/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. | @@ -144,9 +144,9 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | conv2d (vision patch conv — im2col + tiled GEMM) | ✓ | ✓ | ✓ | `ffai/conv2d.rs` → `conv2d_patch14` / `conv2d_patch16` + `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` → `mt_patch_embed`. 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** (PR #157): `kernels/gemm/patch_embed_mma.rs` → `mt_patch_embed_mma` — 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` → `mt_rope_2d`. 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) | ✓ | ✓ | ✓ | `ffai/mel_spectrogram.rs` → `mel_spectrogram` (single-dispatch direct-DFT) + radix-FFT path `mel_stft_window` → `mt_fft_n{n_fft}` → `mel_filterbank` (three kernels, O(N log N)). Generic over `T` per PR #152. STT front-end. All four kernels are bounds-guarded (`idx < n_out`) for threadgroup-rounded dispatch. **Correctness of the two direct-DFT kernels (`mel_spectrogram`, `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 `mel_filterbank` (post-FFT, no in-thread cancellation). | +| mel_spectrogram (STFT + log-Mel filterbank) | ✓ | ✓ | ✓ | `kernels/audio/mel_spectrogram.rs` → `mt_mel_spectrogram` (single-dispatch direct-DFT) + radix-FFT path `mt_mel_stft_window` → `mt_fft_n{n_fft}` → `mt_mel_filterbank` (three kernels, O(N log N)). Generic over `T` per PR #152. STT front-end. All four kernels are bounds-guarded (`idx < n_out`) for threadgroup-rounded dispatch. **Correctness of the two direct-DFT kernels (`mt_mel_spectrogram`, `mt_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 `mt_mel_filterbank` (post-FFT, no in-thread cancellation). | | audio_conv1d (wide-stride 1D conv — STT patch embed) | ✓ | ✓ | ✓ | `ffai/audio_conv1d.rs` → `audio_conv1d`. Dense wide-stride multi-channel 1D conv (NCL); distinct from depthwise `conv1d_causal_step`. STT front-end. | -| vocoder / iSTFT (TTS waveform synthesis) | ✓ | ✓ | ✓ | `ffai/vocoder.rs` → `vocoder_istft`. Inverse-STFT overlap-add — one thread per output sample gathers every covering frame, inverse-DFTs with Hermitian symmetry, COLA-normalises. TTS waveform synthesis. | +| vocoder / iSTFT (TTS waveform synthesis) | ✓ | ✓ | ✓ | `kernels/audio/vocoder.rs` → `mt_vocoder_istft`. 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 @@ -210,7 +210,7 @@ subnormal range). The **asymmetric** integer track: **int2 / int3 / int4 / int5 / int6 / int8**, per-group (64) scale **+ bias** (zero-point), in `mlx/quantized.rs`, -`ffai/dequant_gemv.rs`, `ffai/dequant_gather.rs`, `ffai/kv_cache.rs`, and the +`ffai/dequant_gemv.rs`, `ffai/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. @@ -344,7 +344,7 @@ A few rows mix multiple `.metal` files into one op or split one file into multip 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`](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 intrinsics** — `mel_spectrogram`, `mt_softmax`, `mt_logsumexp`, `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. +- **`fast::` math intrinsics** — `mt_mel_spectrogram`, `mt_softmax`, `mt_logsumexp`, `mt_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/metaltile-codegen/src/passes/vectorize.rs`). **fp32 accumulators** are correctness-required across all production shapes; the f16/bf16-accumulator proposal was rejected. diff --git a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md index 873d3b91..8d7ed0b1 100644 --- a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md +++ b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md @@ -59,11 +59,11 @@ crates/metaltile-std/src/kernels/ ssm/ ssm(_replay) · gated_delta(+wy/prep/chunk) · mamba pregate-rmsnorm quant/ INFRA + the op×format matrix (§7): codec · format · gguf · block_scaled_* · quantized_* · fp_quantized_* · affine · aura codec stack · dequant_* - audio/ mel_spectrogram(+magnitude/stft/filterbank) · lstm · vocoder · snake1d · upsample - vision/ resize_normalize(+bicubic) · im2col · patch_unfold · pos_emb_2d · avg_pool2d · + audio/ ✅ DONE — mel_spectrogram(+magnitude/stft/filterbank) · lstm · vocoder · snake1d · upsample + vision/ ✅ DONE — resize_normalize(+bicubic) · im2col · patch_unfold · pos_emb_2d · avg_pool2d · transpose_th · frame_diff · broadcast_affine sampling/ ✅ DONE — logits_topk/top_p/min_p/processors · categorical_sample · softmax · sort - kv_cache/ kv_cache(_update_many) · fft + kv_cache/ ✅ DONE — kv_cache(_update_many) · fft primitives.rs cross-family decode/reduce ops (mt_decode_e2m1/e4m3/e5m2/e8m0, mt_unpack_nbit, …) mod.rs pub mod ops; pub mod gemm; pub mod sdpa; … ``` @@ -159,7 +159,7 @@ payoff last: | Wave | Families | Rationale | Payoff | |---|---|---|---| | ✅ done | `convolution/`, `rope/`, `norm/`, `sampling/`, `ops/` | exemplar + all of wave 1 | 24k → ~1.6k | -| 2 | `gemm/` (🔨 dense in; quantized matmuls next), `ssm/`, `audio/`, `vision/`, `kv_cache/` | moderate size, few cross-deps | medium | +| 2 | ✅ `audio/` `vision/` `kv_cache/` done; `gemm/` (🔨 dense in, quantized next); `ssm/` next | moderate size, few cross-deps | medium | | 3 | `sdpa/`, `moe/`, **`quant/`** | hardest axes (head-dim d64..d512; bm8/bm64×int8; the 30-format matrix) — most of the ~150k LOC | the bulk | ## 7. The `quant/` umbrella — collapsing the op × format matrix From 0e99ee392536a4898b31096e6a880b58bde17cd7 Mon Sep 17 00:00:00 2001 From: TheTom Date: Mon, 22 Jun 2026 16:25:34 -0500 Subject: [PATCH 5/5] test(kv_cache): fix stale ffai::kv_cache doc ref -> kernels::kv_cache Doc-comment only; the kernels moved to kernels/kv_cache/ in this PR. --- crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs b/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs index bfaf1f07..65f1f5f8 100644 --- a/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs +++ b/crates/metaltile-std/tests/kv_cache_quant_roundtrip_gpu.rs @@ -3,7 +3,7 @@ //! GPU correctness for `mt_quantize_kv_int4/int8` + `mt_bulk_dequant_kv_int4/int8` //! via raw → quantize → dequant round-trip. //! -//! These four kernels ship in `ffai::kv_cache` but had no end-to-end +//! These four kernels ship in `kernels::kv_cache` but had no end-to-end //! coverage before this file. They're how `AffineQuantizedKVCache` //! shrinks per-token K/V slots 4× (int4) or 2× (int8) at decode time — //! a wrong index formula in either direction would silently corrupt