refactor(norm): migrate normalization family into kernels/norm/

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<T>(
 /// `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

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;

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;

crates/metaltile-std/src/ffai/adain1d.rs → crates/metaltile-std/src/kernels/norm/adain1d.rs

@@ -23,7 +23,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn adain1d<T>(
+pub fn mt_adain1d<T>(
     x: Tensor<T>,
     gamma: Tensor<T>,
     beta: Tensor<T>,
@@ -70,7 +70,7 @@ pub fn adain1d<T>(
 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<f32> {
@@ -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))

crates/metaltile-std/src/ffai/gated_rms_norm_qgemv.rs → 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<T>(
+pub fn mt_gated_rms_norm_qgemv_int4_fast<T>(
     y: Tensor<f32>,
     z: Tensor<T>,
     norm_weight: Tensor<T>,
@@ -144,7 +144,7 @@ pub fn ffai_gated_rms_norm_qgemv_int4_fast<T>(
             let zv = load(z[idx]).cast::<f32>();
             let wv = load(norm_weight[d]).cast::<f32>();
             // 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<T>(
 
     // ── 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,21 +449,21 @@ 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 {
         let (hv, dv, out_dim, group_size) = (16usize, 128usize, 2048usize, 64usize);
         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))

crates/metaltile-std/src/ffai/gated_rmsnorm.rs → 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<T>(
+pub fn mt_gated_rmsnorm<T>(
     y: Tensor<f32>,
     z: Tensor<T>,
     w: Tensor<T>,
@@ -106,15 +106,15 @@ pub fn ffai_gated_rmsnorm<T>(
     }
 }
 
-/// 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)`,
 /// with `silu(z) = z / (1 + exp(-z))`, on dtype-rounded `z` / `w`.
 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))

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_<format>_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;

crates/metaltile-std/src/ffai/rms_norm_block_scaled_qgemv.rs → 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.