diff --git a/crates/metaltile-std/src/ffai/batched_qkv_qgemv.rs b/crates/metaltile-std/src/ffai/batched_qkv_qgemv.rs index 2896a691..e527b23a 100644 --- a/crates/metaltile-std/src/ffai/batched_qkv_qgemv.rs +++ b/crates/metaltile-std/src/ffai/batched_qkv_qgemv.rs @@ -20,7 +20,7 @@ //! `[ceil(max(out_q,out_k,out_v)/8), 1, 3]`, TPG = 64 (2 simdgroups × //! 32 lanes). Uses `mt_qmv`'s mask-without-shift trick + algebraic-split //! accumulator (`s*q_dot + b*xs`) — identical inner loop to -//! `ffai_rms_norm_qgemv_fast` but without the RMSNorm phase. +//! `mt_rms_norm_qgemv_fast` but without the RMSNorm phase. //! out_q, out_k, out_v must each be multiples of 8; in_dim must be a //! multiple of 512; group_size must be 64. //! @@ -147,7 +147,7 @@ pub fn ffai_batched_qkv_qgemv( /// `simd_id` selects the simdgroup (0 or 1); each simdgroup independently /// computes 4 output rows (row0..row3). Uses `mt_qmv`'s mask-without-shift /// trick + algebraic-split accumulator — identical inner loop to -/// `ffai_rms_norm_qgemv_fast` but without the RMSNorm phase. +/// `mt_rms_norm_qgemv_fast` but without the RMSNorm phase. /// /// Grid: `[ceil(max(out_q,out_k,out_v)/8), 1, 3]`. /// out_q, out_k, out_v must be multiples of 8; in_dim must be a multiple diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index e687f6bf..983ebd9a 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -16,7 +16,6 @@ //! counterpart lands in mainline at a future pin, the file moves to //! `mlx/`. -pub mod adain1d; pub mod arg_reduce; pub mod attn_head_gate; pub mod aura_dequant_rotated; @@ -62,9 +61,6 @@ pub mod gated_delta_prep; pub mod gated_delta_prep_chunk; pub mod gated_delta_replay; pub mod gated_delta_wy; -pub mod gated_rms_norm_block_scaled_qgemv; -pub mod gated_rms_norm_qgemv; -pub mod gated_rmsnorm; pub mod gather; pub mod gelu_erf; pub mod gemm; @@ -128,10 +124,6 @@ pub mod patch_embed_mma_block_scaled; pub mod patch_unfold_qwen; pub mod pos_emb_2d_add; pub mod resize_normalize; -pub mod rms_norm_block_scaled_qgemv; -pub mod rms_norm_qgemv; -pub mod rms_norm_residual; -pub mod rms_norm_rope; pub mod sampling; pub mod sdpa_bidirectional; pub mod sdpa_bidirectional_d128_relpos; diff --git a/crates/metaltile-std/src/kernels/mod.rs b/crates/metaltile-std/src/kernels/mod.rs index daa7dc7e..bff78df2 100644 --- a/crates/metaltile-std/src/kernels/mod.rs +++ b/crates/metaltile-std/src/kernels/mod.rs @@ -8,4 +8,5 @@ //! consumer is regenerated from the new inventory after each family lands. pub mod convolution; +pub mod norm; pub mod rope; diff --git a/crates/metaltile-std/src/ffai/adain1d.rs b/crates/metaltile-std/src/kernels/norm/adain1d.rs similarity index 91% rename from crates/metaltile-std/src/ffai/adain1d.rs rename to crates/metaltile-std/src/kernels/norm/adain1d.rs index 7ad5b473..a5a49e13 100644 --- a/crates/metaltile-std/src/ffai/adain1d.rs +++ b/crates/metaltile-std/src/kernels/norm/adain1d.rs @@ -23,7 +23,7 @@ use metaltile::kernel; #[kernel] -pub fn adain1d( +pub fn mt_adain1d( x: Tensor, gamma: Tensor, beta: Tensor, @@ -70,7 +70,7 @@ pub fn adain1d( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::adain1d; + use super::mt_adain1d; use crate::utils::{pack_f32, unpack_f32}; fn ramp(n: usize, period: usize, amp: f32, start: f32) -> Vec { @@ -108,7 +108,7 @@ pub mod kernel_tests { let gamma = unpack_f32(&pack_f32(&gamma_f, dt), dt); let beta = unpack_f32(&pack_f32(&beta_f, dt), dt); let expected = naive_adain(&x, &gamma, &beta, rows, length, eps); - TestSetup::new(adain1d::kernel_ir_for(dt)) + TestSetup::new(mt_adain1d::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("x", pack_f32(&x_f, dt), dt)) .input(TestBuffer::from_vec("gamma", pack_f32(&gamma_f, dt), dt)) @@ -121,6 +121,11 @@ pub mod kernel_tests { } // (batch*channels) rows; non-128-aligned length to exercise the strided path. + // f32 tol is 1e-3 (looser than the 1e-4 house default) on purpose: variance + // uses the one-pass `E[x^2] - E[x]^2` form over a length-300 reduce_sum, which + // is cancellation-prone, and the GPU tree/simd reduction accumulates in a + // different order than the oracle's sequential sum, so the per-element diff + // exceeds 1e-4. Do not tighten without switching to a two-pass variance. #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 5e-2])] fn test_adain1d(dt: DType) -> TestSetup { adain_setup(8, 300, dt) } } @@ -129,13 +134,13 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::adain1d; + use super::mt_adain1d; #[bench(dtypes = [f32, f16, bf16])] fn bench_adain1d(dt: DType) -> BenchSetup { let (batch, channels, length) = (4usize, 512usize, 1024usize); let rows = batch * channels; - BenchSetup::new(adain1d::kernel_ir_for(dt)) + BenchSetup::new(mt_adain1d::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("x", rows * length, dt)) .buffer(BenchBuffer::random("gamma", rows, dt)) diff --git a/crates/metaltile-std/src/ffai/gated_rms_norm_block_scaled_qgemv.rs b/crates/metaltile-std/src/kernels/norm/gated_rms_norm_block_scaled_qgemv.rs similarity index 100% rename from crates/metaltile-std/src/ffai/gated_rms_norm_block_scaled_qgemv.rs rename to crates/metaltile-std/src/kernels/norm/gated_rms_norm_block_scaled_qgemv.rs diff --git a/crates/metaltile-std/src/ffai/gated_rms_norm_qgemv.rs b/crates/metaltile-std/src/kernels/norm/gated_rms_norm_qgemv.rs similarity index 97% rename from crates/metaltile-std/src/ffai/gated_rms_norm_qgemv.rs rename to crates/metaltile-std/src/kernels/norm/gated_rms_norm_qgemv.rs index 7b610ff0..fd419e89 100644 --- a/crates/metaltile-std/src/ffai/gated_rms_norm_qgemv.rs +++ b/crates/metaltile-std/src/kernels/norm/gated_rms_norm_qgemv.rs @@ -11,7 +11,7 @@ //! `y` is fp32 (GDN recurrence accumulates in fp32), `z` / `w` / //! `inner` are model dtype `T`. //! -//! 2. `ffai_dequant_gemv_int4` (the GDN out projection): +//! 2. `mt_dequant_gemv_int4` (the GDN out projection): //! `out[o] = sum_i (q[o, i] * scale + bias) * inner_flat[i]` //! where `inner_flat[r * Dv + d] = inner[r, d]` and `i in [0, Hv*Dv)`. //! @@ -73,7 +73,7 @@ //! //! ```text //! inner = mt_gated_rmsnorm(y, z, w, eps) // [Hv, Dv] -//! out = ffai_dequant_gemv_int4(inner, Wq, S, B) // [out_dim] +//! out = mt_dequant_gemv_int4(inner, Wq, S, B) // [out_dim] //! ``` //! //! Pinned by the in-source `#[test_kernel]`s. @@ -88,7 +88,7 @@ use metaltile::kernel; /// TPG = 64. See module doc for invariants. #[kernel] #[allow(clippy::too_many_arguments)] -pub fn ffai_gated_rms_norm_qgemv_int4_fast( +pub fn mt_gated_rms_norm_qgemv_int4_fast( y: Tensor, z: Tensor, norm_weight: Tensor, @@ -144,7 +144,7 @@ pub fn ffai_gated_rms_norm_qgemv_int4_fast( let zv = load(z[idx]).cast::(); let wv = load(norm_weight[d]).cast::(); // silu(z) = z / (1 + exp(-z)), inline fp32 - same form as - // `ffai_gated_rmsnorm` / `moe_down_swiglu_accum`. + // `mt_gated_rmsnorm` / `moe_down_swiglu_accum`. let gate = zv / (1.0f32 + exp(0.0f32 - zv)); let inner = yv * inv_rms * wv * gate; threadgroup_store("tg_inner", idx, inner); @@ -155,7 +155,7 @@ pub fn ffai_gated_rms_norm_qgemv_int4_fast( // ── Phase 2: 8-row int4 GEMV against `tg_inner` ──────────────────── // - // Mirrors `ffai_rms_norm_qgemv_fast` Phase 2 verbatim, except the + // Mirrors `mt_rms_norm_qgemv_fast` Phase 2 verbatim, except the // 16-element X stripe per lane is loaded from `tg_inner` (fp32) in // place of the on-the-fly `x[xi] * norm_weight[xi] * inv_rms` fuse. // @@ -381,7 +381,7 @@ pub mod kernel_tests { use metaltile::{test::*, test_kernel}; use super::{ - ffai_gated_rms_norm_qgemv_int4_fast, + mt_gated_rms_norm_qgemv_int4_fast, oracle::{naive, quantize_int4_row, round, source, u32_bytes}, }; use crate::utils::pack_f32; @@ -425,7 +425,7 @@ pub mod kernel_tests { eps, ); - TestSetup::new(ffai_gated_rms_norm_qgemv_int4_fast::kernel_ir_for(dt)) + TestSetup::new(mt_gated_rms_norm_qgemv_int4_fast::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("y", pack_f32(&y, DType::F32), DType::F32)) .input(TestBuffer::from_vec("z", pack_f32(&z, dt), dt)) @@ -449,13 +449,13 @@ pub mod kernel_tests { fn test_gated_rms_norm_qgemv_int4_fast(dt: DType) -> TestSetup { setup(4, 128, 512, 64, dt) } } -/// New-syntax benchmark for `ffai_gated_rms_norm_qgemv_int4_fast` at the +/// New-syntax benchmark for `mt_gated_rms_norm_qgemv_int4_fast` at the /// Qwen3.6-A3B production shape (hv=16, dv=128, in_dim=2048, out_dim=2048). /// MLX-less reduction kernel (`class=GenericEmpty`), so `Ref(GB/s)` blank. pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_gated_rms_norm_qgemv_int4_fast; + use super::mt_gated_rms_norm_qgemv_int4_fast; #[bench(dtypes = [f32, f16, bf16])] fn bench_gated_rms_norm_qgemv_int4_fast(dt: DType) -> BenchSetup { @@ -463,7 +463,7 @@ pub mod kernel_benches { let in_dim = hv * dv; let u32_per_row = in_dim / 8; let n_groups = in_dim / group_size; - BenchSetup::new(ffai_gated_rms_norm_qgemv_int4_fast::kernel_ir_for(dt)) + BenchSetup::new(mt_gated_rms_norm_qgemv_int4_fast::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("y", in_dim, DType::F32)) .buffer(BenchBuffer::random("z", in_dim, dt)) diff --git a/crates/metaltile-std/src/ffai/gated_rmsnorm.rs b/crates/metaltile-std/src/kernels/norm/gated_rmsnorm.rs similarity index 94% rename from crates/metaltile-std/src/ffai/gated_rmsnorm.rs rename to crates/metaltile-std/src/kernels/norm/gated_rmsnorm.rs index ae09e769..0db26c61 100644 --- a/crates/metaltile-std/src/ffai/gated_rmsnorm.rs +++ b/crates/metaltile-std/src/kernels/norm/gated_rmsnorm.rs @@ -51,7 +51,7 @@ use metaltile::kernel; /// /// `y` is fp32 (the GDN recurrence output); `z`, `w`, `out` are `T`. #[kernel] -pub fn ffai_gated_rmsnorm( +pub fn mt_gated_rmsnorm( y: Tensor, z: Tensor, w: Tensor, @@ -106,7 +106,7 @@ pub fn ffai_gated_rmsnorm( } } -/// New-syntax correctness for `ffai_gated_rmsnorm` (Reduction mode, one +/// New-syntax correctness for `mt_gated_rmsnorm` (Reduction mode, one /// threadgroup per row, `tpg = n/4` — `n` a multiple of 128, `n ≤ 4096`). /// fp32-in / `T`-out split: `y` is always packed f32; `z` / `w` / `out` use /// `dt`. Per-row oracle: `out_i = w_i · y_i · rsqrt(mean(y²)+eps) · silu(z_i)`, @@ -114,7 +114,7 @@ pub fn ffai_gated_rmsnorm( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_gated_rmsnorm; + use super::mt_gated_rmsnorm; use crate::utils::{pack_f32, unpack_f32}; fn setup(rows: usize, n: usize, dt: DType) -> TestSetup { @@ -140,7 +140,7 @@ pub mod kernel_tests { y.extend_from_slice(&yr); z.extend_from_slice(&zr_raw); } - TestSetup::new(ffai_gated_rmsnorm::kernel_ir_for(dt)) + TestSetup::new(mt_gated_rmsnorm::kernel_ir_for(dt)) .mode(KernelMode::Reduction) // `y` is fp32 regardless of T. .input(TestBuffer::from_vec("y", pack_f32(&y, DType::F32), DType::F32)) @@ -154,20 +154,20 @@ pub mod kernel_tests { } #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 2e-2, 8e-2])] - fn test_ffai_gated_rmsnorm(dt: DType) -> TestSetup { setup(4, 512, dt) } + fn test_mt_gated_rmsnorm(dt: DType) -> TestSetup { setup(4, 512, dt) } } -/// New-syntax benchmark for `ffai_gated_rmsnorm` (fused GDN post-step, fp32 `y`, +/// New-syntax benchmark for `mt_gated_rmsnorm` (fused GDN post-step, fp32 `y`, /// `T` gate / weight / output, n=4096, tpg=1024 — the Apple TPG cap). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_gated_rmsnorm; + use super::mt_gated_rmsnorm; #[bench(dtypes = [f32, f16, bf16])] fn bench_gated_rmsnorm(dt: DType) -> BenchSetup { let (rows, n) = (4096usize, 4096usize); - BenchSetup::new(ffai_gated_rmsnorm::kernel_ir_for(dt)) + BenchSetup::new(mt_gated_rmsnorm::kernel_ir_for(dt)) .mode(KernelMode::Reduction) // `y` is always fp32 (the GDN recurrence output). .buffer(BenchBuffer::random("y", rows * n, DType::F32)) diff --git a/crates/metaltile-std/src/mlx/layer_norm.rs b/crates/metaltile-std/src/kernels/norm/layer_norm.rs similarity index 100% rename from crates/metaltile-std/src/mlx/layer_norm.rs rename to crates/metaltile-std/src/kernels/norm/layer_norm.rs diff --git a/crates/metaltile-std/src/kernels/norm/mod.rs b/crates/metaltile-std/src/kernels/norm/mod.rs new file mode 100644 index 00000000..9b8984fb --- /dev/null +++ b/crates/metaltile-std/src/kernels/norm/mod.rs @@ -0,0 +1,23 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Normalization kernels — the norm family (see +//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`). RMSNorm and its fused forms +//! (residual add, RoPE, gated, and the quantized RMSNorm→GEMV fusions), +//! LayerNorm, and 1-D adaptive instance norm. Kernels are named for the +//! operation; the quantized fusions carry the weight format as a name axis +//! (`mt__rms_norm_qgemv`). +//! +//! Migrated from the legacy `mlx/` (`rms_norm`, `layer_norm`) + `ffai/` split. +//! The `*_block_scaled_qgemv` files hold the per-format matrix and will fold +//! into the base `*_qgemv` files as a format axis in a later pass (plan §7). + +pub mod adain1d; +pub mod gated_rms_norm_block_scaled_qgemv; +pub mod gated_rms_norm_qgemv; +pub mod gated_rmsnorm; +pub mod layer_norm; +pub mod rms_norm; +pub mod rms_norm_block_scaled_qgemv; +pub mod rms_norm_qgemv; +pub mod rms_norm_residual; +pub mod rms_norm_rope; diff --git a/crates/metaltile-std/src/mlx/rms_norm.rs b/crates/metaltile-std/src/kernels/norm/rms_norm.rs similarity index 100% rename from crates/metaltile-std/src/mlx/rms_norm.rs rename to crates/metaltile-std/src/kernels/norm/rms_norm.rs diff --git a/crates/metaltile-std/src/ffai/rms_norm_block_scaled_qgemv.rs b/crates/metaltile-std/src/kernels/norm/rms_norm_block_scaled_qgemv.rs similarity index 99% rename from crates/metaltile-std/src/ffai/rms_norm_block_scaled_qgemv.rs rename to crates/metaltile-std/src/kernels/norm/rms_norm_block_scaled_qgemv.rs index 7773ab89..c1d4240e 100644 --- a/crates/metaltile-std/src/ffai/rms_norm_block_scaled_qgemv.rs +++ b/crates/metaltile-std/src/kernels/norm/rms_norm_block_scaled_qgemv.rs @@ -6,7 +6,7 @@ //! integer family (int2/3/4/5/6 + MXINT2..6 + MXINT8). //! //! `y = qmatmul(rms_norm(x) · norm_weight, W_q)` in one dispatch — the int4 -//! fusion of `ffai/rms_norm_qgemv.rs` (`ffai_rms_norm_qgemv`, the simple +//! fusion of `kernels/norm/rms_norm_qgemv.rs` (`mt_rms_norm_qgemv`, the simple //! one-row-per-TG variant) with the block-scaled weight decode of //! `mlx/block_scaled_matmul.rs`. The normalized activation never leaves //! registers between the RMSNorm reduce and the matvec. diff --git a/crates/metaltile-std/src/ffai/rms_norm_qgemv.rs b/crates/metaltile-std/src/kernels/norm/rms_norm_qgemv.rs similarity index 96% rename from crates/metaltile-std/src/ffai/rms_norm_qgemv.rs rename to crates/metaltile-std/src/kernels/norm/rms_norm_qgemv.rs index 41c9f382..9e9a39ef 100644 --- a/crates/metaltile-std/src/ffai/rms_norm_qgemv.rs +++ b/crates/metaltile-std/src/kernels/norm/rms_norm_qgemv.rs @@ -8,7 +8,7 @@ //! //! Two variants: //! -//! **`ffai_rms_norm_qgemv`** — one output row per TG (the original port). +//! **`mt_rms_norm_qgemv`** — one output row per TG (the original port). //! Reduction-mode: one threadgroup per output row. Phase 1 reduces //! `sum(x²)` across the threadgroup → `inv_rms`; phase 2 is a //! pack-strided int4 GEMV that feeds on @@ -16,7 +16,7 @@ //! the normalized activation never leaves registers. Grid: `[out_dim, 1, 1]`, //! TPG ≥ 32. //! -//! **`ffai_rms_norm_qgemv_fast`** — 8 output rows per TG, mirroring +//! **`mt_rms_norm_qgemv_fast`** — 8 output rows per TG, mirroring //! `mt_qmv`'s geometry. Phase 1 (SSQ → `inv_rms`) is shared across all //! 8 rows — the TG-wide reduce amortizes the RMSNorm over 8 outputs. //! Phase 2 uses the `mt_qmv` mask-without-shift trick (X pre-scaled by @@ -49,7 +49,7 @@ use metaltile::kernel; /// with `inv_rms = rsqrt(mean(x²) + eps)`, weights int4-packed. /// One output row per threadgroup (original correctness-first variant). #[kernel] -pub fn ffai_rms_norm_qgemv( +pub fn mt_rms_norm_qgemv( x: Tensor, norm_weight: Tensor, weight: Tensor, @@ -124,7 +124,7 @@ pub fn ffai_rms_norm_qgemv( /// Grid: `[out_dim/8, 1, 1]`. out_dim must be a multiple of 8; /// in_dim must be a multiple of 512; group_size must be 64. #[kernel] -pub fn ffai_rms_norm_qgemv_fast( +pub fn mt_rms_norm_qgemv_fast( x: Tensor, norm_weight: Tensor, weight: Tensor, @@ -426,11 +426,11 @@ pub fn ffai_rms_norm_qgemv_fast( } } -// ─── ffai_rms_norm_qgemv_int8_fast ─────────────────────────────────────────── +// ─── mt_rms_norm_qgemv_int8_fast ─────────────────────────────────────────── // // Fused RMSNorm + int8-quantized GEMV — 8-row-per-TG perf variant. // -// Mirrors `ffai_rms_norm_qgemv_fast` (int4, 8-row-per-TG, 2 SG × 32 lanes) +// Mirrors `mt_rms_norm_qgemv_fast` (int4, 8-row-per-TG, 2 SG × 32 lanes) // but replaces the int4 nibble-unpack with int8 byte-extract: // - 4 bytes per u32 (vals_per_pack = 4 vs 8 for int4) // - mask = 0xFF, shifts = 0 / 8 / 16 / 24 @@ -452,10 +452,10 @@ pub fn ffai_rms_norm_qgemv_fast( /// Perf-tuned fused RMSNorm + int8 GEMV — 8 output rows per TG. /// -/// int8 variant of `ffai_rms_norm_qgemv_fast`. Byte-extract (4 vals/pack), +/// int8 variant of `mt_rms_norm_qgemv_fast`. Byte-extract (4 vals/pack), /// algebraic-split accumulator. Grid: `[out_dim/8, 1, 1]`. #[kernel] -pub fn ffai_rms_norm_qgemv_int8_fast( +pub fn mt_rms_norm_qgemv_int8_fast( x: Tensor, norm_weight: Tensor, weight: Tensor, @@ -663,7 +663,7 @@ pub fn ffai_rms_norm_qgemv_int8_fast( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::{ffai_rms_norm_qgemv, ffai_rms_norm_qgemv_fast, ffai_rms_norm_qgemv_int8_fast}; + use super::{mt_rms_norm_qgemv, mt_rms_norm_qgemv_fast, mt_rms_norm_qgemv_int8_fast}; use crate::utils::{pack_f32, unpack_f32}; const EPS: f32 = 1e-5; @@ -714,7 +714,7 @@ pub mod kernel_tests { } /// Affine per-group int8 quantize of one weight row — 4 bytes per u32, - /// the byte-strided layout `ffai_rms_norm_qgemv_int8_fast` decodes. + /// the byte-strided layout `mt_rms_norm_qgemv_int8_fast` decodes. fn quantize_int8_row(row: &[f32], group_size: usize) -> (Vec, Vec, Vec) { let in_dim = row.len(); let n_groups = in_dim / group_size; @@ -846,7 +846,7 @@ pub mod kernel_tests { #[test_kernel(dtypes = [f32, f16, bf16], tol = 2e-1)] fn test_rms_norm_qgemv(dt: DType) -> TestSetup { let (in_dim, gs, out_dim) = (256usize, 64usize, 8usize); - setup(ffai_rms_norm_qgemv::kernel_ir_for(dt), dt, in_dim, gs, out_dim, false).grid_3d( + setup(mt_rms_norm_qgemv::kernel_ir_for(dt), dt, in_dim, gs, out_dim, false).grid_3d( out_dim as u32, 1, 1, @@ -858,7 +858,7 @@ pub mod kernel_tests { #[test_kernel(dtypes = [f32, f16, bf16], tol = 2e-1)] fn test_rms_norm_qgemv_fast(dt: DType) -> TestSetup { let (in_dim, gs, out_dim) = (512usize, 64usize, 16usize); - setup(ffai_rms_norm_qgemv_fast::kernel_ir_for(dt), dt, in_dim, gs, out_dim, false).grid_3d( + setup(mt_rms_norm_qgemv_fast::kernel_ir_for(dt), dt, in_dim, gs, out_dim, false).grid_3d( (out_dim / 8) as u32, 1, 1, @@ -870,7 +870,7 @@ pub mod kernel_tests { #[test_kernel(dtypes = [f32, f16, bf16], tol = 2e-1)] fn test_rms_norm_qgemv_int8_fast(dt: DType) -> TestSetup { let (in_dim, gs, out_dim) = (512usize, 64usize, 16usize); - setup(ffai_rms_norm_qgemv_int8_fast::kernel_ir_for(dt), dt, in_dim, gs, out_dim, true) + setup(mt_rms_norm_qgemv_int8_fast::kernel_ir_for(dt), dt, in_dim, gs, out_dim, true) .grid_3d((out_dim / 8) as u32, 1, 1, [64, 1, 1]) } } @@ -881,7 +881,7 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::{ffai_rms_norm_qgemv, ffai_rms_norm_qgemv_fast, ffai_rms_norm_qgemv_int8_fast}; + use super::{mt_rms_norm_qgemv, mt_rms_norm_qgemv_fast, mt_rms_norm_qgemv_int8_fast}; const EPS: f32 = 1e-5; @@ -913,22 +913,22 @@ pub mod kernel_benches { #[bench(dtypes = [f32, f16, bf16])] fn bench_rms_norm_qgemv(dt: DType) -> BenchSetup { let (in_dim, gs, out_dim) = (4096usize, 64usize, 4096usize); - let s = BenchSetup::new(ffai_rms_norm_qgemv::kernel_ir_for(dt)).mode(KernelMode::Reduction); + let s = BenchSetup::new(mt_rms_norm_qgemv::kernel_ir_for(dt)).mode(KernelMode::Reduction); buffers(s, in_dim, gs, out_dim, dt, 4).grid_3d(out_dim as u32, 1, 1, [128, 1, 1]) } #[bench(dtypes = [f32, f16, bf16])] fn bench_rms_norm_qgemv_fast(dt: DType) -> BenchSetup { let (in_dim, gs, out_dim) = (4096usize, 64usize, 4096usize); - let s = BenchSetup::new(ffai_rms_norm_qgemv_fast::kernel_ir_for(dt)) - .mode(KernelMode::Reduction); + let s = + BenchSetup::new(mt_rms_norm_qgemv_fast::kernel_ir_for(dt)).mode(KernelMode::Reduction); buffers(s, in_dim, gs, out_dim, dt, 4).grid_3d((out_dim / 8) as u32, 1, 1, [64, 1, 1]) } #[bench(dtypes = [f32, f16, bf16])] fn bench_rms_norm_qgemv_int8_fast(dt: DType) -> BenchSetup { let (in_dim, gs, out_dim) = (4096usize, 64usize, 4096usize); - let s = BenchSetup::new(ffai_rms_norm_qgemv_int8_fast::kernel_ir_for(dt)) + let s = BenchSetup::new(mt_rms_norm_qgemv_int8_fast::kernel_ir_for(dt)) .mode(KernelMode::Reduction); buffers(s, in_dim, gs, out_dim, dt, 8).grid_3d((out_dim / 8) as u32, 1, 1, [64, 1, 1]) } diff --git a/crates/metaltile-std/src/ffai/rms_norm_residual.rs b/crates/metaltile-std/src/kernels/norm/rms_norm_residual.rs similarity index 92% rename from crates/metaltile-std/src/ffai/rms_norm_residual.rs rename to crates/metaltile-std/src/kernels/norm/rms_norm_residual.rs index 217dd62c..ae7f7300 100644 --- a/crates/metaltile-std/src/ffai/rms_norm_residual.rs +++ b/crates/metaltile-std/src/kernels/norm/rms_norm_residual.rs @@ -29,7 +29,7 @@ use metaltile::kernel; /// `out[r, i] = residual[r, i] + w[i] * x[r, i] * rsqrt(mean(x[r]²) + eps)`. #[kernel] -pub fn ffai_rms_norm_residual( +pub fn mt_rms_norm_residual( x: Tensor, residual: Tensor, w: Tensor, @@ -72,14 +72,14 @@ pub fn ffai_rms_norm_residual( } } -/// New-syntax correctness for `ffai_rms_norm_residual` (Reduction mode, one +/// New-syntax correctness for `mt_rms_norm_residual` (Reduction mode, one /// threadgroup per row, `tpg = n/4` — `n` a multiple of 128, `n ≤ 4096`). /// Per-row oracle on dtype-rounded inputs: normalize `x` then add `residual`: /// `out_i = residual_i + x_i / sqrt(mean(x²) + eps) * w_i`. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_rms_norm_residual; + use super::mt_rms_norm_residual; use crate::utils::{pack_f32, unpack_f32}; fn setup(rows: usize, n: usize, dt: DType) -> TestSetup { @@ -104,7 +104,7 @@ pub mod kernel_tests { x.extend_from_slice(&row); residual.extend_from_slice(&res); } - TestSetup::new(ffai_rms_norm_residual::kernel_ir_for(dt)) + TestSetup::new(mt_rms_norm_residual::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("x", pack_f32(&x, dt), dt)) .input(TestBuffer::from_vec("residual", pack_f32(&residual, dt), dt)) @@ -117,20 +117,20 @@ pub mod kernel_tests { } #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 2e-2, 1e-1])] - fn test_ffai_rms_norm_residual(dt: DType) -> TestSetup { setup(4, 512, dt) } + fn test_mt_rms_norm_residual(dt: DType) -> TestSetup { setup(4, 512, dt) } } -/// New-syntax benchmark for `ffai_rms_norm_residual` (fused RMSNorm + residual +/// New-syntax benchmark for `mt_rms_norm_residual` (fused RMSNorm + residual /// add, hidden axis n=4096, tpg=1024 — the Apple TPG cap). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_rms_norm_residual; + use super::mt_rms_norm_residual; #[bench(dtypes = [f32, f16, bf16])] fn bench_rms_norm_residual(dt: DType) -> BenchSetup { let (rows, n) = (4096usize, 4096usize); - BenchSetup::new(ffai_rms_norm_residual::kernel_ir_for(dt)) + BenchSetup::new(mt_rms_norm_residual::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("x", rows * n, dt)) .buffer(BenchBuffer::random("residual", rows * n, dt)) diff --git a/crates/metaltile-std/src/ffai/rms_norm_rope.rs b/crates/metaltile-std/src/kernels/norm/rms_norm_rope.rs similarity index 94% rename from crates/metaltile-std/src/ffai/rms_norm_rope.rs rename to crates/metaltile-std/src/kernels/norm/rms_norm_rope.rs index 283a2234..c42dd2b7 100644 --- a/crates/metaltile-std/src/ffai/rms_norm_rope.rs +++ b/crates/metaltile-std/src/kernels/norm/rms_norm_rope.rs @@ -35,7 +35,7 @@ use metaltile::kernel; /// Fused RMSNorm + paired-layout RoPE for one Q/K head per threadgroup. #[kernel] -pub fn ffai_rms_norm_rope( +pub fn mt_rms_norm_rope( x: Tensor, w: Tensor, inv_freqs: Tensor, @@ -69,7 +69,7 @@ pub fn ffai_rms_norm_rope( store(out[rs + lid + half], (normed_a * sin_t + normed_b * cos_t).cast::()); } -/// New-syntax correctness for `ffai_rms_norm_rope` (Reduction mode, one +/// New-syntax correctness for `mt_rms_norm_rope` (Reduction mode, one /// threadgroup per row, `tpg = axis_size/2` — `axis_size` a multiple of 64, /// `axis_size ≤ 2048`). Per-row oracle replays the whole-row RMSNorm scale, /// the per-row position `pos = offset + (row / n_heads) mod seq_len`, and the @@ -77,7 +77,7 @@ pub fn ffai_rms_norm_rope( pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_rms_norm_rope; + use super::mt_rms_norm_rope; use crate::utils::{pack_f32, unpack_f32}; /// Small, deterministic inverse-frequency table (length `half`). @@ -86,7 +86,7 @@ pub mod kernel_tests { } #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 1e-1])] - fn test_ffai_rms_norm_rope(dt: DType) -> TestSetup { + fn test_mt_rms_norm_rope(dt: DType) -> TestSetup { let (axis, n_heads, seq_len, offset, eps) = (128usize, 4usize, 8usize, 5usize, 1e-5f32); let rows = n_heads * seq_len; // one batch let half = axis / 2; @@ -111,7 +111,7 @@ pub mod kernel_tests { expected[base + lid + half] = na * s + nb * c; } } - TestSetup::new(ffai_rms_norm_rope::kernel_ir_for(dt)) + TestSetup::new(mt_rms_norm_rope::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("x", pack_f32(&x_raw, dt), dt)) .input(TestBuffer::from_vec("w", pack_f32(&w_raw, dt), dt)) @@ -127,12 +127,12 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `ffai_rms_norm_rope` (fused RMSNorm + RoPE, one +/// New-syntax benchmark for `mt_rms_norm_rope` (fused RMSNorm + RoPE, one /// `(batch, seq, head)` row per threadgroup, axis_size=128, tpg=64). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_rms_norm_rope; + use super::mt_rms_norm_rope; fn f32_bytes(v: &[f32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } @@ -143,7 +143,7 @@ pub mod kernel_benches { let half = axis / 2; let inv_freqs: Vec = (0..half).map(|i| 1.0 / 10000.0_f32.powf(i as f32 / half as f32)).collect(); - BenchSetup::new(ffai_rms_norm_rope::kernel_ir_for(dt)) + BenchSetup::new(mt_rms_norm_rope::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("x", rows * axis, dt)) .buffer(BenchBuffer::random("w", axis, dt)) diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs index 09ab6977..7f1c4b8e 100644 --- a/crates/metaltile-std/src/mlx/mod.rs +++ b/crates/metaltile-std/src/mlx/mod.rs @@ -38,7 +38,6 @@ pub mod gemv_masked; pub mod hadamard; pub mod hadamard_m; pub mod indexing; -pub mod layer_norm; pub mod logsumexp; pub mod quantized; pub mod quantized_mma_dynamic_m; @@ -48,7 +47,6 @@ pub mod quantized_nax; pub mod quantized_nax_int8; pub mod random; pub mod reduce; -pub mod rms_norm; pub mod scaled_dot_product_attention; pub mod scan; pub mod scatter_axis; diff --git a/crates/metaltile-std/tests/dispatch_pinning_guard.rs b/crates/metaltile-std/tests/dispatch_pinning_guard.rs index 2b39813f..3ede044e 100644 --- a/crates/metaltile-std/tests/dispatch_pinning_guard.rs +++ b/crates/metaltile-std/tests/dispatch_pinning_guard.rs @@ -24,7 +24,7 @@ use metaltile::{ Context, core::{dtype::DType, ir::KernelMode}, }; -use metaltile_std::mlx::rms_norm::mt_rms_norm; +use metaltile_std::kernels::norm::rms_norm::mt_rms_norm; /// Buffers for an `n`-wide single-row `mt_rms_norm` dispatch (constexpr `n` /// is passed as a 4-byte buffer, matching the harness's dispatch convention). diff --git a/crates/metaltile-std/tests/rms_norm_per_head_gpu.rs b/crates/metaltile-std/tests/rms_norm_per_head_gpu.rs index c45aca74..8452e815 100644 --- a/crates/metaltile-std/tests/rms_norm_per_head_gpu.rs +++ b/crates/metaltile-std/tests/rms_norm_per_head_gpu.rs @@ -46,7 +46,7 @@ mod common; use common::{Dt, gpu_lock, pack_bytes, unpack_bytes}; use metaltile::Context; -use metaltile_std::mlx::rms_norm::{mt_rms_norm, mt_rms_norm_small}; +use metaltile_std::kernels::norm::rms_norm::{mt_rms_norm, mt_rms_norm_small}; fn cpu_rms_norm_reference(x: &[f32], w: &[f32], rows: usize, n: usize, eps: f32) -> Vec { let mut out = vec![0.0f32; rows * n]; diff --git a/docs/specs/KERNEL_AUDIT.md b/docs/specs/KERNEL_AUDIT.md index b2492d98..a8b0633e 100644 --- a/docs/specs/KERNEL_AUDIT.md +++ b/docs/specs/KERNEL_AUDIT.md @@ -76,8 +76,8 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | scan (prefix sum/prod/max/min) | ✓ | ✓ | ✓ | `mlx/scan.rs` → `mt_scan` + `mt_scan_exclusive` (sum), `mt_scan_prod` / `mt_scan_max` / `mt_scan_min` + exclusive variants. Sum pair uses hardware `simd_scan_exclusive`; the prod/max/min pairs use a `tgs[lsize]` threadgroup buffer for sequential cross-thread prefix reads. | | softmax | ✓ | ✓ | ✓ | `mlx/softmax.rs` → `mt_softmax` (looped + single-row collapsed). | | logsumexp | ✓ | ✓ | ✓ | `mlx/logsumexp.rs` → `mt_logsumexp`. | -| layer_norm | ✓ | ✓ | ✓ | `mlx/layer_norm.rs` → `mt_layer_norm`. | -| rms_norm | ✓ | ✓ | ✓ | `mlx/rms_norm.rs` → `mt_rms_norm` + `mt_rms_norm_small` (2-elem/thread, small-head_dim per-head q_norm/k_norm) + `mt_rms_norm_wide` (strided wide-row variant for `head_dim > 4096`, e.g. Gemma 4 31B hidden=5376). | +| layer_norm | ✓ | ✓ | ✓ | `kernels/norm/layer_norm.rs` → `mt_layer_norm`. | +| rms_norm | ✓ | ✓ | ✓ | `kernels/norm/rms_norm.rs` → `mt_rms_norm` + `mt_rms_norm_small` (2-elem/thread, small-head_dim per-head q_norm/k_norm) + `mt_rms_norm_wide` (strided wide-row variant for `head_dim > 4096`, e.g. Gemma 4 31B hidden=5376). | | rope (standard) | ✓ | ✓ | ✓ | `kernels/rope/base.rs` → `mt_rope`. | | rope (frequency-band scaled) | ✗ | ✗ | ✓ | `kernels/rope/rope_banded.rs` → `mt_rope_banded`. Per-row `positions` tensor + row grid axis (decode = T=1, prefill = all tokens in one dispatch); generic dtype, optional frequency-band scaling (Llama-3 / Qwen). | | sdpa_vector (prefill / generic) | ✓ | ✓ | ✓ | `mlx/scaled_dot_product_attention.rs` → `mt_sdpa`. Scalar SDPA for short sequences. | @@ -132,16 +132,16 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | conv1d_causal_step (depthwise SSM conv stream) | ✗ | partial | ✓ | `ffai/ssm.rs` → `conv1d_causal_step`. fp32 state recurrence. | | ssm_replay (sequential tape capture + replay) | ✗ | ✓ | ✓ | `ffai/ssm_replay.rs` → `ssm_step_record` (SSD forward + dA/dBx tape) + `ssm_replay` (re-fold first k entries). | | fused_gate_activation (silu/gelu × up gate) | ✗ | ✓ | ✓ | `mlx/fused_gate_activation.rs` → `mt_fused_gate_gelu` (gelu-tanh) + `mt_fused_gate_clipped_swiglu` (GPT-OSS: `[-7,7]` clamp, `sigmoid(1.702·g)` gate, `+1` up bias). The `silu` variant ships separately as `mlx/swiglu.rs`. | -| rms_norm_residual (RMSNorm + residual add fused) | ✗ | ✓ | ✓ | `ffai/rms_norm_residual.rs` → `ffai_rms_norm_residual`. Reduction-mode, `N = TPG*4`. ~90 saved dispatches/token on Gemma4-30. | -| rms_norm_rope (RMSNorm + RoPE fused) | ✗ | ✓ | ✓ | `ffai/rms_norm_rope.rs` → `ffai_rms_norm_rope`. Paired-layout RoPE; Q/K post-projection norm+rope in one dispatch. | -| rms_norm_qgemv (RMSNorm + quantized GEMV fused) | ✗ | ✓ | ✓ | `ffai/rms_norm_qgemv.rs` → `ffai_rms_norm_qgemv` (int4, one-row-per-TG correctness shape) + `ffai_rms_norm_qgemv_fast` (int4, 8-row-per-TG perf path, PR #154) + `ffai_rms_norm_qgemv_int8_fast` (int8, 8-row-per-TG, PR #157). | +| rms_norm_residual (RMSNorm + residual add fused) | ✗ | ✓ | ✓ | `kernels/norm/rms_norm_residual.rs` → `mt_rms_norm_residual`. Reduction-mode, `N = TPG*4`. ~90 saved dispatches/token on Gemma4-30. | +| 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. | | sampling (softmax + categorical inverse-CDF) | ✗ | ✗ | ✓ | `ffai/sampling.rs` → `softmax_categorical_sample`. Companion to `ffai_argmax` for `T > 0` decode. | | logits processors (temperature, repetition penalty, top-k / top-p / min-p masks) | ✗ | ✗ | ✓ | `ffai/logits_{processors,topk,top_p,min_p}.rs` — in-place decode-form sampler stages composed before `softmax_categorical_sample`. FFAI-only. | | 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. | -| gated_rmsnorm (fp32-in gated RMSNorm → activation dtype) | ✗ | ✗ | ✓ | `ffai/gated_rmsnorm.rs` → `ffai_gated_rmsnorm`. Fused Qwen3.5 / 3.6 GDN post-step `out = w·rmsNorm(y)·silu(z)`; `y` arrives fp32 (the `gated_delta` recurrence output). Closes the per-GDN-layer host-side CPU sync (~75 % of Qwen3.5/3.6 layers). | +| gated_rmsnorm (fp32-in gated RMSNorm → activation dtype) | ✗ | ✗ | ✓ | `kernels/norm/gated_rmsnorm.rs` → `mt_gated_rmsnorm`. Fused Qwen3.5 / 3.6 GDN post-step `out = w·rmsNorm(y)·silu(z)`; `y` arrives fp32 (the `gated_delta` recurrence output). Closes the per-GDN-layer host-side CPU sync (~75 % of Qwen3.5/3.6 layers). | | conv2d (vision patch conv — im2col + tiled GEMM) | ✓ | ✓ | ✓ | `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) | ✗ | ✗ | ✓ | `ffai/patch_embed.rs` → `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): `ffai/patch_embed_mma.rs` → `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. | @@ -255,8 +255,8 @@ each family's proven dispatch geometry verbatim (no new freeze surface). | MoE gather-qmm | reduction | `mlx/block_scaled_moe.rs` | int3–8 | | MoE gather — MPP (bm8/16/64) | MPP | `ffai/moe_mpp{,_bm8,_bm64}_block_scaled.rs` | int4, int8 | | expert-indexed GEMV | reduction | `ffai/dequant_gemv_expert_indexed_block_scaled.rs` | int4 | -| fused RMSNorm + GEMV | reduction | `ffai/rms_norm_block_scaled_qgemv.rs` | int4, int8-fast | -| fused gated-RMSNorm + GEMV | reduction | `ffai/gated_rms_norm_block_scaled_qgemv.rs` | int4 | +| fused RMSNorm + GEMV | reduction | `kernels/norm/rms_norm_block_scaled_qgemv.rs` | int4, int8-fast | +| fused gated-RMSNorm + GEMV | reduction | `kernels/norm/gated_rms_norm_block_scaled_qgemv.rs` | int4 | | batched-Q/K/V qgemv + qmm | reduction | `ffai/batched_qkv_block_scaled_{qgemv,qmm}.rs` | int4, int8-fast | | batched-4 qgemv + qmm | reduction | `ffai/batched_4_block_scaled_{qgemv,qmm}.rs` | int4 | | embedding gather | elementwise | `ffai/dequant_gather_block_scaled.rs` | int3–8 | diff --git a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md index 5ab67ecc..6bdc6d0c 100644 --- a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md +++ b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md @@ -12,11 +12,13 @@ Where they overlap (naming, per-file shape), defer to the style guide. > **Status:** in progress. `convolution/` is the proven exemplar — its > *structural* consolidation has landed (move, dedup, per-format macro DRY-ing), -> and `rope/` has landed too. Two optimization phases remain open for every -> family: replace the interim `macro_rules!` with `variants(...)`, and merge the -> `*_block_scaled` / `*_mma` files into the dimensionality file (§4). "✅ done" -> in §6 means the structural move is in; the optimization phases are tracked -> separately. The remaining families migrate one-per-PR per §6. +> and `rope/` + `norm/` have landed too. Two optimization phases remain open for +> every family: replace the interim `macro_rules!` with `variants(...)`, and +> merge the `*_block_scaled` / `*_mma` files into the dimensionality file (§4). +> "✅ done" in §6 means the structural move is in; the optimization phases are +> tracked separately. The `*_block_scaled_qgemv` format matrices moved with +> `norm/` but still await the format-axis fold (§7). The remaining families +> migrate one-per-PR per §6. ## 1. Why @@ -49,7 +51,7 @@ crates/metaltile-std/src/kernels/ 2pass/batched/sink) · multi(+d256/tree-mask) · prefill_mma · flash_quantized · aura_flash · steel/attn moe/ moe orchestration · mpp(bm8/bm64 × int8) · bgemm/gemv(q2k/iq2xxs) · block_scaled_moe - norm/ rms_norm(+residual/rope/qgemv/gated) · layer_norm · adain1d + norm/ ✅ DONE — rms_norm(+residual/rope/qgemv/gated) · layer_norm · adain1d rope/ ✅ DONE — rope · rope_2d · rope_banded · rope_yarn · partial_rope convolution/ ✅ DONE — conv1d/2d/3d · depthwise · winograd · steel_conv (see §4) ssm/ ssm(_replay) · gated_delta(+wy/prep/chunk) · mamba pregate-rmsnorm @@ -154,8 +156,8 @@ payoff last: | Wave | Families | Rationale | Payoff | |---|---|---|---| -| ✅ done | `convolution/`, `rope/` | exemplar + first wave-1 family | 24k → ~1.6k | -| 1 | `norm/`, `sampling/`, `ops/` | self-contained, mostly elementwise / few formats | small, sets the pattern | +| ✅ done | `convolution/`, `rope/`, `norm/` | exemplar + wave-1 families | 24k → ~1.6k | +| 1 | `sampling/`, `ops/` | self-contained, mostly elementwise / few formats | small, sets the pattern | | 2 | `gemm/`, `ssm/`, `audio/`, `vision/`, `kv_cache/` | 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 |