refactor(gemm): migrate the dense matmul family into kernels/gemm/

crates/metaltile-codegen/src/spirv/mod.rs

@@ -2098,7 +2098,7 @@ pub fn safe_glsl_ident(name: &str) -> String {
         "image3D",
         // GLSL reserves these for future use, even though they aren't
         // currently used; shaderc rejects them. The corpus hit them on
-        // conv* / depthwise_conv2d / aura_value_int4 / ffai_gemm kernels.
+        // conv* / depthwise_conv2d / aura_value_int4 / mt_gemm kernels.
         "input",
         "output",
         "texture",

crates/metaltile-std/src/ffai/ffai_dequant_q4.rs

@@ -7,7 +7,7 @@
 //! the inline-dequant `gemv_q4`, which is bandwidth-bound — perfect at
 //! batch 1. Prefill is COMPUTE-bound: many tokens reuse each weight, so
 //! the right move is to dequant the weight ONCE to f16 and feed the
-//! tensor-core `ffai_gemm`. This kernel is that one-time expansion.
+//! tensor-core `mt_gemm`. This kernel is that one-time expansion.
 //!
 //! ## Q4 block layout (matches `ffai_ops::quantize_q4`)
 //!

crates/metaltile-std/src/ffai/gemm_q4_mpp.rs

@@ -5,7 +5,7 @@
 //! bench's Q4 weight layout (signed 4-bit, per-32-block scale `amax/7` stored
 //! f16) instead of Q8_0. Compute-bound prefill projections (q/k/v/o, mamba
 //! in/out_proj, shared experts, lm_head) run on this instead of the f32
-//! scalar `ffai_gemm` (which sat at ~0.1% of the tensor-core peak).
+//! scalar `mt_gemm` (which sat at ~0.1% of the tensor-core peak).
 //!
 //! Same 64×64×32 coop_tile geometry as `ffai_gemm_q8_mpp` (4 simdgroups,
 //! 2×2 warp grid, 128 threads/tg). Only the weight-dequant block differs.

crates/metaltile-std/src/ffai/gemm_q8.rs

@@ -10,10 +10,10 @@
 //!
 //! Weight is the resident Q8 split (`qs` int8 packed 4/u32 + per-32-block
 //! `d` scale), laid out as `[out_dim, in_dim]` (row-major over values).
-//! Mirrors `ffai_gemm`'s geometry exactly so the dispatch wrapper is
+//! Mirrors `mt_gemm`'s geometry exactly so the dispatch wrapper is
 //! identical apart from the two extra weight buffers.
 //!
-//! ## DISPATCH INVARIANTS (same as ffai_gemm)
+//! ## DISPATCH INVARIANTS (same as mt_gemm)
 //! - TPG = 1024 (32×32). Grid: (out_dim/32) × (n_rows/32) threadgroups.
 //! - `in_dim % 16 == 0` (K-tile) AND `in_dim % 32 == 0` (Q8 block).
 //! - Row/col edges handled in-kernel (clamp loads to 0, skip OOB stores).

crates/metaltile-std/src/ffai/mod.rs

@@ -59,11 +59,9 @@ pub mod gated_delta_prep_chunk;
 pub mod gated_delta_replay;
 pub mod gated_delta_wy;
 pub mod gelu_erf;
-pub mod gemm;
 pub mod gemm_q4_mpp;
 pub mod gemm_q8;
 pub mod gemm_q8_mpp;
-pub mod gemv_axpy_inplace;
 pub mod gemv_q8;
 pub mod gguf_dequant_iq2_xxs;
 pub mod gguf_dequant_iq2_xxs_raw;
@@ -108,9 +106,7 @@ pub mod moe_mpp_int8;
 pub mod moe_mpp_shared;
 pub mod moe_router_sigmoid_bias;
 pub mod moe_router_sqrtsoftplus;
-pub mod patch_embed;
 pub mod patch_embed_block_scaled;
-pub mod patch_embed_mma;
 pub mod patch_embed_mma_block_scaled;
 pub mod patch_unfold_qwen;
 pub mod pos_emb_2d_add;

crates/metaltile-std/src/ffai/ssm.rs

@@ -361,8 +361,8 @@ pub mod kernel_tests {
         };
 
         let kernels: Vec<(&str, metaltile::core::Kernel)> = vec![
-            ("ffai_gemm_batched", {
-                let mut k = super::super::gemm::ffai_gemm_batched::kernel_ir_for(DType::F32);
+            ("mt_gemm_batched", {
+                let mut k = crate::kernels::gemm::dense::mt_gemm_batched::kernel_ir_for(DType::F32);
                 k.mode = KernelMode::Reduction;
                 k
             }),

crates/metaltile-std/src/ffai/gemm.rs → crates/metaltile-std/src/kernels/gemm/dense.rs

@@ -34,7 +34,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn ffai_gemm<T>(
+pub fn mt_gemm<T>(
     weight: Tensor<T>,
     input: Tensor<T>,
     out: Tensor<T>,
@@ -91,7 +91,7 @@ pub fn ffai_gemm<T>(
 //
 // Computes, for each batch `z ∈ [0, batch)`:
 //   out[z][r, o] = Σ_k weight[z][o, k] · input[z][r, k]
-// i.e. the SAME `out = input · weightᵀ` contraction as `ffai_gemm`, replicated
+// i.e. the SAME `out = input · weightᵀ` contraction as `mt_gemm`, replicated
 // over a batch axis. Each operand is a contiguous stack of per-batch matrices
 // with element strides `w_stride` / `x_stride` / `o_stride` (counted in T
 // elements, NOT bytes) — the wrapper passes `out_dim*in_dim`, `n_rows*in_dim`,
@@ -100,15 +100,15 @@ pub fn ffai_gemm<T>(
 // This is what the Mamba2 SSD chunked-matmul prefill scan runs its 4 batched
 // GEMMs on (batch = n_chunks·n_heads) on Apple/HIP/Vulkan hardware MMA, where
 // cuBLAS strided-batched is unavailable. Same 32×32 / 16-K tiling as
-// `ffai_gemm`; the batch index rides the 3rd grid axis (`tgid_z`).
+// `mt_gemm`; the batch index rides the 3rd grid axis (`tgid_z`).
 //
-// ## DISPATCH INVARIANTS (identical to `ffai_gemm` plus the batch axis)
+// ## DISPATCH INVARIANTS (identical to `mt_gemm` plus the batch axis)
 // - TPG = 1024 (BM·BN). `in_dim % 16 == 0`.
 // - Grid: `((out_dim+31)/32, (n_rows+31)/32, batch)` Reduction-mode TGs.
 // - All operands row-major within each batch slice; out-of-range row/col edges
 //   handled in-kernel (clamp-load to 0, skip-store).
 #[kernel]
-pub fn ffai_gemm_batched<T>(
+pub fn mt_gemm_batched<T>(
     weight: Tensor<T>,
     input: Tensor<T>,
     out: Tensor<T>,
@@ -163,7 +163,7 @@ pub fn ffai_gemm_batched<T>(
     }
 }
 
-/// New-syntax correctness tests for `ffai_gemm` — the multi-row 32×32-tiled
+/// New-syntax correctness tests for `mt_gemm` — the multi-row 32×32-tiled
 /// GEMM `out[r, :] = weight · input[r, :]`. Reduction-mode (threadgroup-memory
 /// tiles + barriers).
 ///
@@ -179,7 +179,7 @@ pub fn ffai_gemm_batched<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::ffai_gemm;
+    use super::mt_gemm;
     use crate::utils::{pack_f32, unpack_f32};
 
     /// Triple-loop reference: out[r, o] = Σ_k weight[o, k] · input[r, k].
@@ -211,7 +211,7 @@ pub mod kernel_tests {
         let w = unpack_f32(&pack_f32(&weight_f, dt), dt);
         let x = unpack_f32(&pack_f32(&input_f, dt), dt);
         let expected = gemm_oracle(&w, &x, n_rows, in_dim, out_dim);
-        TestSetup::new(ffai_gemm::kernel_ir_for(dt))
+        TestSetup::new(mt_gemm::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("weight", pack_f32(&weight_f, dt), dt))
             .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt))
@@ -231,8 +231,8 @@ pub mod kernel_tests {
     #[test_kernel(dtypes = [f32, f16, bf16], tol = [5e-3, 5e-2, 2e-1])]
     fn test_gemm_edge(dt: DType) -> TestSetup { gemm_setup(20, 48, 100, dt) }
 
-    // ── ffai_gemm_batched ────────────────────────────────────────────────
-    use super::ffai_gemm_batched;
+    // ── mt_gemm_batched ────────────────────────────────────────────────
+    use super::mt_gemm_batched;
 
     fn gemm_batched_setup(
         batch: usize,
@@ -257,7 +257,7 @@ pub mod kernel_tests {
             let eb = gemm_oracle(wb, xb, n_rows, in_dim, out_dim);
             expected[b * n_rows * out_dim..(b + 1) * n_rows * out_dim].copy_from_slice(&eb);
         }
-        TestSetup::new(ffai_gemm_batched::kernel_ir_for(dt))
+        TestSetup::new(mt_gemm_batched::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("weight", pack_f32(&weight_f, dt), dt))
             .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt))
@@ -283,19 +283,19 @@ pub mod kernel_tests {
     fn test_gemm_batched_edge(dt: DType) -> TestSetup { gemm_batched_setup(2, 20, 48, 100, dt) }
 }
 
-/// New-syntax benchmark for `ffai_gemm`. Nemotron-class block-diffusion shape:
+/// New-syntax benchmark for `mt_gemm`. Nemotron-class block-diffusion shape:
 /// a 32-row block projected through a `[out_dim, in_dim]` weight (hidden 4096).
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::ffai_gemm;
+    use super::mt_gemm;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_gemm(dt: DType) -> BenchSetup {
         let (n_rows, in_dim, out_dim) = (32usize, 4096usize, 4096usize);
         let sz = dt.size_bytes();
         let bytes = out_dim * in_dim * sz + n_rows * in_dim * sz + n_rows * out_dim * sz;
-        BenchSetup::new(ffai_gemm::kernel_ir_for(dt))
+        BenchSetup::new(mt_gemm::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("weight", out_dim * in_dim, dt))
             .buffer(BenchBuffer::random("input", n_rows * in_dim, dt))

crates/metaltile-std/src/ffai/gemv_axpy_inplace.rs → crates/metaltile-std/src/kernels/gemm/gemv_axpy_inplace.rs

@@ -20,7 +20,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn ffai_gemv_axpy_inplace<T>(
+pub fn mt_gemv_axpy_inplace<T>(
     mat: Tensor<T>,
     vec: Tensor<T>,
     mut accum: Tensor<T>,
@@ -39,7 +39,7 @@ pub fn ffai_gemv_axpy_inplace<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::ffai_gemv_axpy_inplace;
+    use super::mt_gemv_axpy_inplace;
     use crate::utils::{pack_f32, unpack_f32};
 
     fn setup(m: usize, k: usize, weight: f32, dt: DType) -> TestSetup {
@@ -55,7 +55,7 @@ pub mod kernel_tests {
                 accum_dt[r] + weight * dot
             })
             .collect();
-        TestSetup::new(ffai_gemv_axpy_inplace::kernel_ir_for(dt))
+        TestSetup::new(mt_gemv_axpy_inplace::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("mat", pack_f32(&mat, dt), dt))
             .input(TestBuffer::from_vec("vec", pack_f32(&vec, dt), dt))
@@ -84,12 +84,12 @@ pub mod kernel_tests {
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::ffai_gemv_axpy_inplace;
+    use super::mt_gemv_axpy_inplace;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_gemv_axpy(dt: DType) -> BenchSetup {
         let (m, k) = (4096usize, 2048usize);
-        BenchSetup::new(ffai_gemv_axpy_inplace::kernel_ir_for(dt))
+        BenchSetup::new(mt_gemv_axpy_inplace::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("mat", m * k, dt))
             .buffer(BenchBuffer::random("vec", k, dt))

crates/metaltile-std/src/kernels/gemm/mod.rs

@@ -0,0 +1,20 @@
+//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric
+//! SPDX-License-Identifier: Apache-2.0
+//! Matrix-multiply kernels — the gemm family (see
+//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): dense GEMM / GEMV (+ masked,
+//! +axpy), the batched QKV / 4-way projection forms, patch-embed (im2col +
+//! matmul), and the MLX `steel/` tiled-GEMM templates.
+//!
+//! This is the **dense** half of the family, migrated from the legacy `mlx/` +
+//! `ffai/` split. The quantized matmuls (`*_q8`/`q4`/`block_scaled_qmm`/
+//! `quantized_*`/`fp_quantized_*` and the batched `*_qgemv`/`*_qmm` forms) are
+//! the quantized form of these ops and land here in a follow-up pass; the
+//! format-axis fold (plan §7) is deferred.
+
+pub mod dense;
+pub mod gemv;
+pub mod gemv_axpy_inplace;
+pub mod gemv_masked;
+pub mod patch_embed;
+pub mod patch_embed_mma;
+pub mod steel;

crates/metaltile-std/src/ffai/patch_embed.rs → crates/metaltile-std/src/kernels/gemm/patch_embed.rs

@@ -14,7 +14,7 @@
 //! the image and dots them with one weight row, no intermediate buffer.
 //!
 //! It differs from `conv2d` in layout, not arithmetic — `conv2d` keeps
-//! the NCHW image convention and writes NCHW output; `patch_embed` takes
+//! the NCHW image convention and writes NCHW output; `mt_patch_embed` takes
 //! the same NCHW image but treats the weight as a flat linear matrix
 //! `[hidden, patch_dim]` and writes transformer-token output
 //! `[num_patches, hidden]`, which is what a ViT block consumes directly.
@@ -46,7 +46,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn patch_embed<T>(
+pub fn mt_patch_embed<T>(
     image: Tensor<T>,
     weight: Tensor<T>,
     bias: Tensor<T>,
@@ -93,7 +93,7 @@ pub fn patch_embed<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::patch_embed;
+    use super::mt_patch_embed;
     use crate::utils::{pack_f32, unpack_f32};
 
     fn ramp(n: usize, period: usize, amp: f32) -> Vec<f32> {
@@ -164,7 +164,7 @@ pub mod kernel_tests {
         let bias = unpack_f32(&pack_f32(&bias_f, dt), dt);
         let expected =
             naive_patch_embed(&image, &weight, &bias, in_ch, in_h, in_w, patch_h, patch_w, hidden);
-        TestSetup::new(patch_embed::kernel_ir_for(dt))
+        TestSetup::new(mt_patch_embed::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .input(TestBuffer::from_vec("image", pack_f32(&image_f, dt), dt))
             .input(TestBuffer::from_vec("weight", pack_f32(&weight_f, dt), dt))
@@ -199,12 +199,12 @@ pub mod kernel_tests {
     }
 }
 
-/// New-syntax bench for `patch_embed` (ViT-L SigLIP stem shape).
+/// New-syntax bench for `mt_patch_embed` (ViT-L SigLIP stem shape).
 /// Grid3D, `grid_1d(num_patches * hidden, 256)`; bytes_moved = output stream.
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::patch_embed;
+    use super::mt_patch_embed;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_patch_embed(dt: DType) -> BenchSetup {
@@ -215,7 +215,7 @@ pub mod kernel_benches {
         let num_patches = (in_h / patch_h) * (in_w / patch_w);
         let patch_dim = in_ch * patch_h * patch_w;
         let n_out = num_patches * hidden;
-        BenchSetup::new(patch_embed::kernel_ir_for(dt))
+        BenchSetup::new(mt_patch_embed::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .buffer(BenchBuffer::random("image", in_ch * in_h * in_w, dt))
             .buffer(BenchBuffer::random("weight", hidden * patch_dim, dt))

crates/metaltile-std/src/ffai/patch_embed_mma.rs → crates/metaltile-std/src/kernels/gemm/patch_embed_mma.rs

@@ -72,7 +72,7 @@ use metaltile::kernel;
 /// Correctness pinned by the in-source `#[test_kernel]`s.
 #[kernel]
 #[allow(clippy::too_many_arguments)]
-pub fn patch_embed_mma<T>(
+pub fn mt_patch_embed_mma<T>(
     image: Tensor<T>,
     weight: Tensor<T>,
     bias: Tensor<T>,
@@ -292,7 +292,7 @@ pub fn patch_embed_mma<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::patch_embed_mma;
+    use super::mt_patch_embed_mma;
     use crate::utils::{pack_f32, unpack_f32};
 
     fn ramp(n: usize, period: usize, amp: f32) -> Vec<f32> {
@@ -367,7 +367,7 @@ pub mod kernel_tests {
         let expected = naive_patch_embed_mma(
             &image, &weight, &bias, in_ch, in_h, in_w, patch_h, patch_w, hidden,
         );
-        TestSetup::new(patch_embed_mma::kernel_ir_for(dt))
+        TestSetup::new(mt_patch_embed_mma::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("image", pack_f32(&image_f, dt), dt))
             .input(TestBuffer::from_vec("weight", pack_f32(&weight_f, dt), dt))
@@ -392,12 +392,12 @@ pub mod kernel_tests {
     fn test_patch_embed_mma_4x4(dt: DType) -> TestSetup { mma_setup(8, 32, 32, 4, 4, 32, dt) }
 }
 
-/// New-syntax bench for `patch_embed_mma` (ViT-L-ish 8×8 patch, hidden 1024).
+/// New-syntax bench for `mt_patch_embed_mma` (ViT-L-ish 8×8 patch, hidden 1024).
 /// Reduction mode, `grid_3d(hidden/32, num_patches/32, 1, [128,1,1])`.
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::patch_embed_mma;
+    use super::mt_patch_embed_mma;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_patch_embed_mma(dt: DType) -> BenchSetup {
@@ -408,7 +408,7 @@ pub mod kernel_benches {
         let num_patches = (in_h / patch_h) * (in_w / patch_w);
         let patch_dim = in_ch * patch_h * patch_w;
         let n_out = num_patches * hidden;
-        BenchSetup::new(patch_embed_mma::kernel_ir_for(dt))
+        BenchSetup::new(mt_patch_embed_mma::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("image", in_ch * in_h * in_w, dt))
             .buffer(BenchBuffer::random("weight", hidden * patch_dim, dt))

crates/metaltile-std/src/kernels/mod.rs

@@ -8,6 +8,7 @@
 //! consumer is regenerated from the new inventory after each family lands.
 
 pub mod convolution;
+pub mod gemm;
 pub mod norm;
 pub mod ops;
 pub mod rope;

crates/metaltile-std/src/mlx/mod.rs

@@ -25,8 +25,6 @@ pub mod fft;
 pub mod fp_quantized;
 pub mod fp_quantized_mma;
 pub mod fp_quantized_nax;
-pub mod gemv;
-pub mod gemv_masked;
 pub mod quantized;
 pub mod quantized_mma_dynamic_m;
 pub mod quantized_mpp;

crates/metaltile-std/src/mlx/steel/mod.rs

@@ -2,4 +2,3 @@
 //! SPDX-License-Identifier: Apache-2.0
 pub mod attn;
 pub use crate::kernels::convolution::steel_conv as conv;
-pub mod gemm;

crates/metaltile-std/tests/gemm_q8_correctness.rs

@@ -32,7 +32,7 @@ fn run_case(dt: Dt, in_dim: usize, out_dim: usize, n_rows: usize, tol: f32) {
         (0..n_rows * in_dim).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect();
     // Round inputs to the kernel's dtype so the f32 oracle sees the same
     // values the kernel loads (f16 input rounding, amplified by cancellation
-    // over in_dim terms, otherwise dominates the diff). Matches ffai_gemm's test.
+    // over in_dim terms, otherwise dominates the diff). Matches mt_gemm's test.
     if matches!(dt, Dt::F16) {
         input = unpack_bytes(&pack_bytes(&input, Dt::F16), Dt::F16);
     }