refactor(kernels): migrate kv_cache + audio + vision families

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,13 +66,7 @@ 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 kv_cache;
-pub mod kv_cache_update_many;
 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;
@@ -108,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;
@@ -123,9 +112,5 @@ 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;

crates/metaltile-std/src/ffai/lstm.rs → 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<T>(
+pub fn mt_lstm<T>(
     x: Tensor<T>,
     w_ih: Tensor<T>,
     w_hh: Tensor<T>,
@@ -133,9 +133,9 @@ pub fn ffai_lstm<T>(
 /// 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<T>(
     x: Tensor<T>,
@@ -210,7 +210,7 @@ pub fn lstm_cell<T>(
 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))

crates/metaltile-std/src/ffai/mel_spectrogram.rs → crates/metaltile-std/src/kernels/audio/mel_spectrogram.rs

@@ -45,7 +45,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn mel_spectrogram<T>(
+pub fn mt_mel_spectrogram<T>(
     audio: Tensor<T>,
     window: Tensor<T>,
     mel_weight: Tensor<T>,
@@ -98,18 +98,18 @@ pub fn mel_spectrogram<T>(
 // ─────────────────────────────────────────────────────────────────────────
 // 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<T>(
 /// 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<T>(
+pub fn mt_mel_stft_window<T>(
     audio: Tensor<T>,
     window: Tensor<T>,
     mut out_re: Tensor<T>,
@@ -128,7 +128,7 @@ pub fn mel_stft_window<T>(
     #[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<T>(
 /// mel] = log(Σ_{k<n_freq} mel_weight[mel,k]·(re²+im²) + log_eps)`, where
 /// `re`/`im` are `fft_re`/`fft_im` from `mt_fft_n{n_fft}`. One thread per
 /// `(frame, mel)`; dispatch flat over `n_frames * n_mels`. Output is
-/// bit-identical in form to `mel_spectrogram` — only the spectrum source
+/// bit-identical in form to `mt_mel_spectrogram` — only the spectrum source
 /// (FFT vs in-thread DFT) differs.
 #[kernel]
-pub fn mel_filterbank<T>(
+pub fn mt_mel_filterbank<T>(
     fft_re: Tensor<T>,
     fft_im: Tensor<T>,
     mel_weight: Tensor<T>,
@@ -159,7 +159,7 @@ pub fn mel_filterbank<T>(
     #[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<T>(
 }
 
 /// **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<T>(
+pub fn mt_mel_spectrogram_magnitude<T>(
     audio: Tensor<T>,
     window: Tensor<T>,
     mel_weight: Tensor<T>,
@@ -198,7 +198,7 @@ pub fn mel_spectrogram_magnitude<T>(
     #[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<T>(
 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))

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;

crates/metaltile-std/src/ffai/snake1d.rs → crates/metaltile-std/src/kernels/audio/snake1d.rs

@@ -27,7 +27,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn ffai_snake1d<T>(
+pub fn mt_snake1d<T>(
     input: Tensor<T>,
     alpha: Tensor<T>,
     mut out: Tensor<T>,
@@ -43,12 +43,12 @@ pub fn ffai_snake1d<T>(
     store(out[i], y.cast::<T>());
 }
 
-/// 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))