From e4657b38d320b88a3716d57abeaae9f848834ff4 Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Fri, 12 Jun 2026 15:56:53 -0600 Subject: [PATCH 1/3] refactor(ops): migrate core/elementwise primitives into kernels/ops/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes wave 1: moves the core/elementwise primitive kernels out of the legacy mlx/ + ffai/ split into kernels/ops/ per the consolidation plan. - mlx/{arange,arg_reduce,binary,binary_two,copy,gather_axis,hadamard, hadamard_m,indexing,logsumexp,random,reduce,scan,scatter_axis,strided, ternary,unary} -> kernels/ops/ (already mt_*; inventory-registered, so no import sites to update beyond one snapshot test) - ffai/{axpy_scalar_inplace,clamp_scalar,gather} + mt_vector_add.rs (renamed vector_add.rs) -> kernels/ops/ with ffai_ -> mt_ prefixes - delete ffai/arg_reduce.rs: ffai_argmax was a byte-for-byte duplicate of mt_argmax (verified) -> real dedup - fence.rs moves but stays intentionally undeclared (device/system fences have no DSL representation — the '1 out of scope' op) Updated the steel_msl_snapshots import, the op_group/name-filter test fixtures, the plan (ops done, wave 1 complete) and KERNEL_AUDIT rows. cargo build + workspace clippy clean; all ops kernel tests + steel_msl_snapshots pass. --- crates/metaltile-core/src/protocol.rs | 2 +- crates/metaltile-std/src/ffai/arg_reduce.rs | 174 ------------------ crates/metaltile-std/src/ffai/mod.rs | 5 - crates/metaltile-std/src/kernels/mod.rs | 1 + .../src/{mlx => kernels/ops}/arange.rs | 0 .../src/{mlx => kernels/ops}/arg_reduce.rs | 0 .../ops}/axpy_scalar_inplace.rs | 10 +- .../src/{mlx => kernels/ops}/binary.rs | 0 .../src/{mlx => kernels/ops}/binary_two.rs | 0 .../src/{ffai => kernels/ops}/clamp_scalar.rs | 14 +- .../src/{mlx => kernels/ops}/copy.rs | 0 .../src/{mlx => kernels/ops}/fence.rs | 0 .../src/{ffai => kernels/ops}/gather.rs | 21 +-- .../src/{mlx => kernels/ops}/gather_axis.rs | 0 .../src/{mlx => kernels/ops}/hadamard.rs | 0 .../src/{mlx => kernels/ops}/hadamard_m.rs | 0 .../src/{mlx => kernels/ops}/indexing.rs | 0 .../src/{mlx => kernels/ops}/logsumexp.rs | 0 crates/metaltile-std/src/kernels/ops/mod.rs | 35 ++++ .../src/{mlx => kernels/ops}/random.rs | 0 .../src/{mlx => kernels/ops}/reduce.rs | 0 .../src/{mlx => kernels/ops}/scan.rs | 0 .../src/{mlx => kernels/ops}/scatter_axis.rs | 0 .../src/{mlx => kernels/ops}/strided.rs | 0 .../src/{mlx => kernels/ops}/ternary.rs | 0 .../src/{mlx => kernels/ops}/unary.rs | 0 .../ops/vector_add.rs} | 10 +- crates/metaltile-std/src/mlx/mod.rs | 17 -- .../tests/steel_msl_snapshots.rs | 12 +- crates/metaltile/src/runner/harness.rs | 2 +- docs/specs/KERNEL_AUDIT.md | 35 ++-- docs/specs/KERNEL_CONSOLIDATION_PLAN.md | 19 +- 32 files changed, 96 insertions(+), 261 deletions(-) delete mode 100644 crates/metaltile-std/src/ffai/arg_reduce.rs rename crates/metaltile-std/src/{mlx => kernels/ops}/arange.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/arg_reduce.rs (100%) rename crates/metaltile-std/src/{ffai => kernels/ops}/axpy_scalar_inplace.rs (87%) rename crates/metaltile-std/src/{mlx => kernels/ops}/binary.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/binary_two.rs (100%) rename crates/metaltile-std/src/{ffai => kernels/ops}/clamp_scalar.rs (89%) rename crates/metaltile-std/src/{mlx => kernels/ops}/copy.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/fence.rs (100%) rename crates/metaltile-std/src/{ffai => kernels/ops}/gather.rs (86%) rename crates/metaltile-std/src/{mlx => kernels/ops}/gather_axis.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/hadamard.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/hadamard_m.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/indexing.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/logsumexp.rs (100%) create mode 100644 crates/metaltile-std/src/kernels/ops/mod.rs rename crates/metaltile-std/src/{mlx => kernels/ops}/random.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/reduce.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/scan.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/scatter_axis.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/strided.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/ternary.rs (100%) rename crates/metaltile-std/src/{mlx => kernels/ops}/unary.rs (100%) rename crates/metaltile-std/src/{ffai/mt_vector_add.rs => kernels/ops/vector_add.rs} (89%) diff --git a/crates/metaltile-core/src/protocol.rs b/crates/metaltile-core/src/protocol.rs index bc405b85..a0a5b5be 100644 --- a/crates/metaltile-core/src/protocol.rs +++ b/crates/metaltile-core/src/protocol.rs @@ -323,7 +323,7 @@ mod tests { assert_eq!(op_group("mt_softmax_f32"), "softmax"); assert_eq!(op_group("mt_softmax_bf16"), "softmax"); assert_eq!(op_group("softmax"), "softmax"); - assert_eq!(op_group("ffai_vector_add_f16"), "ffai_vector_add"); + assert_eq!(op_group("mt_vector_add_f16"), "vector_add"); } #[test] diff --git a/crates/metaltile-std/src/ffai/arg_reduce.rs b/crates/metaltile-std/src/ffai/arg_reduce.rs deleted file mode 100644 index 0f33623e..00000000 --- a/crates/metaltile-std/src/ffai/arg_reduce.rs +++ /dev/null @@ -1,174 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Generic `argmax` with u32 index output — FFAI's decode-form -//! greedy-sampler workhorse. -//! -//! Adapted from `mt_argmax_f32` (in `mlx/arg_reduce.rs`) but generic -//! over input dtype and emitting a `u32` index rather than a float-cast -//! version. Decode-form samplers (greedy token pick) need an integer -//! token id; the f32-output upstream variant doesn't fit that contract. -//! -//! Tie-breaking: strict `>` on values, smallest index on ties — matches -//! NumPy / PyTorch / MLX `argmax` semantics. -//! -//! Codegen-only — there's no MLX argmax template with the same -//! u32-output signature. Correctness validated in FFAI integration -//! tests against reference decoder output. - -use metaltile::kernel; - -// Tree-reduction strides: 128 → 64 → 32 → 16 → 8 → 4 → 2. -// Each iteration: threads with `lid < stride` merge the upper half into -// the lower half (take higher value; on ties take smaller index — NumPy -// argmax semantics). Final stride-1 merge writes the result directly -// to `out[0]` and is kept inline below. -// -// Originally hand-unrolled via a `macro_rules! argmax_step!` invoked -// 7×; the proc-macro does not expand inner declarative macros, so the -// expansion silently produced no IR. A DSL `for` loop over the seven -// stages yields identical MSL and survives the proc-macro intact. - -#[kernel] -pub fn ffai_argmax(inp: Tensor, out: Tensor, #[constexpr] n: u32) { - let lid = tid; - let mut best_val = neg_infinity(); - let mut best_idx = lid - lid; - threadgroup_alloc("tg_vals", 256); - threadgroup_alloc("tg_idxs", 256); - let n_iters = (n + lsize - 1u32) / lsize; - for _r in range(0u32, n_iters, 1u32) { - let pos = _r * lsize + lid; - if pos < n { - let v = load(inp[pos]).cast::(); - let better = v > best_val; - if better { - best_val = v; - best_idx = pos; - } - } - } - threadgroup_store("tg_vals", lid, best_val); - threadgroup_store("tg_idxs", lid, best_idx); - threadgroup_barrier(); - // 7-stage power-of-two halving reduction over the 256-thread group. - for _stage in range(0u32, 7u32, 1u32) { - let stride = 128u32 >> _stage; - if lid < stride { - let ov = threadgroup_load("tg_vals", lid + stride); - let oi = threadgroup_load("tg_idxs", lid + stride); - let tv = threadgroup_load("tg_vals", lid); - let ti = threadgroup_load("tg_idxs", lid); - let bet = (ov > tv) | ((ov == tv) & (oi < ti)); - threadgroup_store("tg_vals", lid, select(bet, ov, tv)); - threadgroup_store("tg_idxs", lid, select(bet, oi, ti)); - } - threadgroup_barrier(); - } - // Final stride-1 merge writes result directly to output. - if lid == 0u32 { - let ov = threadgroup_load("tg_vals", 1u32); - let oi = threadgroup_load("tg_idxs", 1u32); - let tv = threadgroup_load("tg_vals", 0u32); - let ti = threadgroup_load("tg_idxs", 0u32); - let bet = (ov > tv) | ((ov == tv) & (oi < ti)); - let final_idx = select(bet, oi, ti); - store(out[0], final_idx); - } -} - -/// New-syntax correctness for `ffai_argmax`. -/// -/// This is a **reduction-mode, MLX-less (ffai)** kernel — exactly the case the -/// legacy harness left unbenched via `class=GenericEmpty`. The new builders -/// express it with `.mode(KernelMode::Reduction)` (so codegen emits the -/// threadgroup reduction) and a one-threadgroup grid; the u32 index output is -/// compared exactly through the integer-aware `pack_f32`/`unpack_f32` path. -pub mod kernel_tests { - use metaltile::{test::*, test_kernel}; - - use super::ffai_argmax; - use crate::utils::pack_f32; - - /// One 256-lane threadgroup reduces all `n` elements (the kernel loops over - /// `n` in chunks of `lsize`). `max_idx` carries a lone `2.0` spike in a - /// field of `1.0` — an unambiguous argmax in every dtype (no rounding ties). - fn setup(n: usize, max_idx: usize, dt: DType) -> TestSetup { - let mut inp = vec![1.0f32; n]; - inp[max_idx] = 2.0; - TestSetup::new(ffai_argmax::kernel_ir_for(dt)) - .mode(KernelMode::Reduction) - .input(TestBuffer::from_vec("inp", pack_f32(&inp, dt), dt)) - .input(TestBuffer::zeros("out", 1, DType::U32)) - .constexpr("n", n as u32) - .expect(TestBuffer::from_vec( - "out", - pack_f32(&[max_idx as f32], DType::U32), - DType::U32, - )) - .grid_3d(1, 1, 1, [256, 1, 1]) - } - - // Integer index output — tol < 1 means the index must match exactly. - #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.5)] - fn test_ffai_argmax_interior(dt: DType) -> TestSetup { setup(256, 37, dt) } - - // n > lsize exercises the chunked outer loop (4 chunks of 256). - #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.5)] - fn test_ffai_argmax_multi_chunk(dt: DType) -> TestSetup { setup(1000, 813, dt) } - - /// Tie-break test: many positions share the maximum value; argmax must - /// return the SMALLEST such index (NumPy/PyTorch/MLX semantics). The - /// kernel encodes this in both the per-lane scan (strict `>`) and every - /// tree-merge step (`(ov > tv) | (ov == tv & oi < ti)`); a regression to - /// `>=` would return a larger index. `tie_lo`/`tie_hi` bracket a run of - /// equal maxima that straddles lane boundaries (positions chosen so the - /// run spans more than one 256-lane chunk + crosses tree-reduction - /// strides), so the expected answer is exactly `tie_lo`. - fn tie_setup(n: usize, tie_lo: usize, tie_hi: usize, dt: DType) -> TestSetup { - assert!(tie_lo < tie_hi && tie_hi < n); - let mut inp = vec![0.5f32; n]; - for v in inp.iter_mut().take(tie_hi + 1).skip(tie_lo) { - *v = 1.0; // a plateau of equal maxima - } - TestSetup::new(ffai_argmax::kernel_ir_for(dt)) - .mode(KernelMode::Reduction) - .input(TestBuffer::from_vec("inp", pack_f32(&inp, dt), dt)) - .input(TestBuffer::zeros("out", 1, DType::U32)) - .constexpr("n", n as u32) - .expect(TestBuffer::from_vec("out", pack_f32(&[tie_lo as f32], DType::U32), DType::U32)) - .grid_3d(1, 1, 1, [256, 1, 1]) - } - - // Plateau of equal maxima spanning lanes 100..400 (crosses the 256-lane - // chunk boundary + several tree strides) — must resolve to index 100. - #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.5)] - fn test_ffai_argmax_ties_take_smallest(dt: DType) -> TestSetup { tie_setup(512, 100, 400, dt) } - - // Production vocab size (Qwen ~152k): one threadgroup loops ~595 chunks. - // Spike near the end exercises the full chunked scan at realistic scale. - #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.5)] - fn test_ffai_argmax_vocab_152k(dt: DType) -> TestSetup { setup(152_064, 151_900, dt) } -} - -/// New-syntax benchmark for `ffai_argmax` — an MLX-less reduction kernel that -/// previously produced no bench rows (`class=GenericEmpty`). It now reports -/// real GB/s in `tile bench` with `Ref(GB/s)` blank (no MLX counterpart). -pub mod kernel_benches { - use metaltile::{bench, test::*}; - - use super::ffai_argmax; - - // Vocab-sized argmax (greedy decode): one threadgroup reduces 256K elements. - // Read-dominated, so bytes_moved counts the input. - #[bench(dtypes = [f32, f16, bf16])] - fn bench_argmax(dt: DType) -> BenchSetup { - let n = 256 * 1024usize; - BenchSetup::new(ffai_argmax::kernel_ir_for(dt)) - .mode(KernelMode::Reduction) - .buffer(BenchBuffer::random("inp", n, dt)) - .buffer(BenchBuffer::zeros("out", 1, DType::U32).output()) - .constexpr("n", n as u32) - .grid_3d(1, 1, 1, [256, 1, 1]) - .bytes_moved((n * dt.size_bytes()) as u64) - } -} diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index c54b3dbd..1b200326 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 arg_reduce; pub mod attn_head_gate; pub mod aura_dequant_rotated; pub mod aura_encode; @@ -26,7 +25,6 @@ pub mod aura_flash_sdpa; pub mod aura_score; pub mod aura_value; pub mod avg_pool2d_nhwc; -pub mod axpy_scalar_inplace; pub mod batched_4_block_scaled_qgemv; pub mod batched_4_block_scaled_qmm; pub mod batched_4_qgemv; @@ -35,7 +33,6 @@ pub mod batched_qkv_block_scaled_qgemv; pub mod batched_qkv_block_scaled_qmm; pub mod batched_qkv_qgemv; pub mod batched_qkv_qmm; -pub mod clamp_scalar; pub mod dequant_gather; pub mod dequant_gather_block_scaled; pub mod dequant_gemv; @@ -61,7 +58,6 @@ pub mod gated_delta_prep; pub mod gated_delta_prep_chunk; pub mod gated_delta_replay; pub mod gated_delta_wy; -pub mod gather; pub mod gelu_erf; pub mod gemm; pub mod gemm_q4_mpp; @@ -112,7 +108,6 @@ pub mod moe_mpp_int8; pub mod moe_mpp_shared; pub mod moe_router_sigmoid_bias; pub mod moe_router_sqrtsoftplus; -pub mod mt_vector_add; pub mod patch_embed; pub mod patch_embed_block_scaled; pub mod patch_embed_mma; diff --git a/crates/metaltile-std/src/kernels/mod.rs b/crates/metaltile-std/src/kernels/mod.rs index 28aba0ba..89cac5c7 100644 --- a/crates/metaltile-std/src/kernels/mod.rs +++ b/crates/metaltile-std/src/kernels/mod.rs @@ -9,5 +9,6 @@ pub mod convolution; pub mod norm; +pub mod ops; pub mod rope; pub mod sampling; diff --git a/crates/metaltile-std/src/mlx/arange.rs b/crates/metaltile-std/src/kernels/ops/arange.rs similarity index 100% rename from crates/metaltile-std/src/mlx/arange.rs rename to crates/metaltile-std/src/kernels/ops/arange.rs diff --git a/crates/metaltile-std/src/mlx/arg_reduce.rs b/crates/metaltile-std/src/kernels/ops/arg_reduce.rs similarity index 100% rename from crates/metaltile-std/src/mlx/arg_reduce.rs rename to crates/metaltile-std/src/kernels/ops/arg_reduce.rs diff --git a/crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs b/crates/metaltile-std/src/kernels/ops/axpy_scalar_inplace.rs similarity index 87% rename from crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs rename to crates/metaltile-std/src/kernels/ops/axpy_scalar_inplace.rs index 0969f03c..78c45d5d 100644 --- a/crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs +++ b/crates/metaltile-std/src/kernels/ops/axpy_scalar_inplace.rs @@ -11,7 +11,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_axpy_scalar_inplace(src: Tensor, mut accum: Tensor, #[constexpr] scalar: f32) { +pub fn mt_axpy_scalar_inplace(src: Tensor, mut accum: Tensor, #[constexpr] scalar: f32) { let i = tid; let s = load(src[i]).cast::(); let a = load(accum[i]).cast::(); @@ -21,7 +21,7 @@ pub fn ffai_axpy_scalar_inplace(src: Tensor, mut accum: Tensor, #[const pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_axpy_scalar_inplace; + use super::mt_axpy_scalar_inplace; use crate::utils::{pack_f32, unpack_f32}; fn setup(n: usize, dt: DType) -> TestSetup { @@ -32,7 +32,7 @@ pub mod kernel_tests { let accum_dt = unpack_f32(&pack_f32(&accum_in, dt), dt); let expected: Vec = accum_dt.iter().zip(&src_dt).map(|(a, s)| a + scalar * s).collect(); - TestSetup::new(ffai_axpy_scalar_inplace::kernel_ir_for(dt)) + TestSetup::new(mt_axpy_scalar_inplace::kernel_ir_for(dt)) .input(TestBuffer::from_vec("src", pack_f32(&src, dt), dt)) .input(TestBuffer::from_vec("accum", pack_f32(&accum_in, dt), dt)) .constexpr("scalar", scalar) @@ -47,12 +47,12 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_axpy_scalar_inplace; + use super::mt_axpy_scalar_inplace; #[bench(dtypes = [f32, f16, bf16])] fn bench_axpy_scalar(dt: DType) -> BenchSetup { let n = 4096usize; - BenchSetup::new(ffai_axpy_scalar_inplace::kernel_ir_for(dt)) + BenchSetup::new(mt_axpy_scalar_inplace::kernel_ir_for(dt)) .buffer(BenchBuffer::random("src", n, dt)) .buffer(BenchBuffer::random("accum", n, dt).output()) .constexpr("scalar", 0.5f32) diff --git a/crates/metaltile-std/src/mlx/binary.rs b/crates/metaltile-std/src/kernels/ops/binary.rs similarity index 100% rename from crates/metaltile-std/src/mlx/binary.rs rename to crates/metaltile-std/src/kernels/ops/binary.rs diff --git a/crates/metaltile-std/src/mlx/binary_two.rs b/crates/metaltile-std/src/kernels/ops/binary_two.rs similarity index 100% rename from crates/metaltile-std/src/mlx/binary_two.rs rename to crates/metaltile-std/src/kernels/ops/binary_two.rs diff --git a/crates/metaltile-std/src/ffai/clamp_scalar.rs b/crates/metaltile-std/src/kernels/ops/clamp_scalar.rs similarity index 89% rename from crates/metaltile-std/src/ffai/clamp_scalar.rs rename to crates/metaltile-std/src/kernels/ops/clamp_scalar.rs index de433eb5..cc03b06e 100644 --- a/crates/metaltile-std/src/ffai/clamp_scalar.rs +++ b/crates/metaltile-std/src/kernels/ops/clamp_scalar.rs @@ -27,7 +27,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_clamp_scalar(input: Tensor, out: Tensor, lo: Tensor, hi: Tensor) { +pub fn mt_clamp_scalar(input: Tensor, out: Tensor, lo: Tensor, hi: Tensor) { let i = program_id::<0>(); let lo_v = load(lo[0]); let hi_v = load(hi[0]); @@ -39,12 +39,12 @@ pub fn ffai_clamp_scalar(input: Tensor, out: Tensor, lo: Tensor, h store(out[i], clamped.cast::()); } -/// New-syntax correctness for `ffai_clamp_scalar`. Grid3D, grid `[n,1,1]`, +/// New-syntax correctness for `mt_clamp_scalar`. Grid3D, grid `[n,1,1]`, /// tpg `[1,1,1]`. Oracle clamps each element to `[lo, hi]`. pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_clamp_scalar; + use super::mt_clamp_scalar; use crate::utils::{pack_f32, unpack_f32}; fn f32_bytes(v: f32) -> Vec { v.to_le_bytes().to_vec() } @@ -58,7 +58,7 @@ pub mod kernel_tests { let input_f: Vec = (0..n).map(|i| (i as f32 - 128.0) * 0.1).collect(); let input = unpack_f32(&pack_f32(&input_f, dt), dt); let exp: Vec = input.iter().map(|&x| x.max(lo).min(hi)).collect(); - TestSetup::new(ffai_clamp_scalar::kernel_ir_for(dt)) + TestSetup::new(mt_clamp_scalar::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("input", pack_f32(&input_f, dt), dt)) .input(TestBuffer::from_vec("lo", f32_bytes(lo), DType::F32)) @@ -69,16 +69,16 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `ffai_clamp_scalar`. +/// New-syntax benchmark for `mt_clamp_scalar`. pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_clamp_scalar; + use super::mt_clamp_scalar; #[bench(dtypes = [f32, f16, bf16])] fn bench_clamp_scalar(dt: DType) -> BenchSetup { let n = 576 * 768usize; // SigLIP patch-grid activation size - BenchSetup::new(ffai_clamp_scalar::kernel_ir_for(dt)) + BenchSetup::new(mt_clamp_scalar::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("input", n, dt)) .buffer(BenchBuffer::random("lo", 1, DType::F32)) diff --git a/crates/metaltile-std/src/mlx/copy.rs b/crates/metaltile-std/src/kernels/ops/copy.rs similarity index 100% rename from crates/metaltile-std/src/mlx/copy.rs rename to crates/metaltile-std/src/kernels/ops/copy.rs diff --git a/crates/metaltile-std/src/mlx/fence.rs b/crates/metaltile-std/src/kernels/ops/fence.rs similarity index 100% rename from crates/metaltile-std/src/mlx/fence.rs rename to crates/metaltile-std/src/kernels/ops/fence.rs diff --git a/crates/metaltile-std/src/ffai/gather.rs b/crates/metaltile-std/src/kernels/ops/gather.rs similarity index 86% rename from crates/metaltile-std/src/ffai/gather.rs rename to crates/metaltile-std/src/kernels/ops/gather.rs index fbd076d5..0b4a54bd 100644 --- a/crates/metaltile-std/src/ffai/gather.rs +++ b/crates/metaltile-std/src/kernels/ops/gather.rs @@ -11,12 +11,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_gather( - table: Tensor, - indices: Tensor, - out: Tensor, - #[constexpr] dim: u32, -) { +pub fn mt_gather(table: Tensor, indices: Tensor, out: Tensor, #[constexpr] dim: u32) { let idx = program_id::<0>(); let token = idx / dim; let d = idx - token * dim; @@ -25,19 +20,19 @@ pub fn ffai_gather( store(out[idx], load(table[src])); } -/// New-syntax correctness for `ffai_gather` (Grid3D, one thread per output +/// New-syntax correctness for `mt_gather` (Grid3D, one thread per output /// element). Pure copy — expected output is the gathered rows /// `out[token, d] = table[indices[token], d]`, exact (tol 0). pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_gather; + use super::mt_gather; use crate::utils::{pack_f32, unpack_f32}; fn u32_bytes(v: &[u32]) -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() } #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.0)] - fn test_ffai_gather(dt: DType) -> TestSetup { + fn test_mt_gather(dt: DType) -> TestSetup { let (vocab, dim, n_tokens) = (17usize, 8usize, 6usize); // table[r, d] = r * 1000 + d — a wrong token/dim decomposition // cross-contaminates immediately. @@ -53,7 +48,7 @@ pub mod kernel_tests { table_dt[indices[token] as usize * dim + d] }) .collect(); - TestSetup::new(ffai_gather::kernel_ir_for(dt)) + TestSetup::new(mt_gather::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .input(TestBuffer::from_vec("table", pack_f32(&table, dt), dt)) .input(TestBuffer::from_vec("indices", u32_bytes(&indices), DType::U32)) @@ -64,12 +59,12 @@ pub mod kernel_tests { } } -/// New-syntax benchmark for `ffai_gather` (embedding-table lookup, Qwen-class +/// New-syntax benchmark for `mt_gather` (embedding-table lookup, Qwen-class /// dim, one thread per output element). pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_gather; + use super::mt_gather; fn u32_bytes(v: impl Iterator) -> Vec { v.flat_map(|x| x.to_le_bytes()).collect() @@ -79,7 +74,7 @@ pub mod kernel_benches { fn bench_gather(dt: DType) -> BenchSetup { let (vocab, dim, n_tokens) = (8192usize, 4096usize, 1024usize); let n_elems = n_tokens * dim; - BenchSetup::new(ffai_gather::kernel_ir_for(dt)) + BenchSetup::new(mt_gather::kernel_ir_for(dt)) .mode(KernelMode::Grid3D) .buffer(BenchBuffer::random("table", vocab * dim, dt)) .buffer(BenchBuffer::from_vec( diff --git a/crates/metaltile-std/src/mlx/gather_axis.rs b/crates/metaltile-std/src/kernels/ops/gather_axis.rs similarity index 100% rename from crates/metaltile-std/src/mlx/gather_axis.rs rename to crates/metaltile-std/src/kernels/ops/gather_axis.rs diff --git a/crates/metaltile-std/src/mlx/hadamard.rs b/crates/metaltile-std/src/kernels/ops/hadamard.rs similarity index 100% rename from crates/metaltile-std/src/mlx/hadamard.rs rename to crates/metaltile-std/src/kernels/ops/hadamard.rs diff --git a/crates/metaltile-std/src/mlx/hadamard_m.rs b/crates/metaltile-std/src/kernels/ops/hadamard_m.rs similarity index 100% rename from crates/metaltile-std/src/mlx/hadamard_m.rs rename to crates/metaltile-std/src/kernels/ops/hadamard_m.rs diff --git a/crates/metaltile-std/src/mlx/indexing.rs b/crates/metaltile-std/src/kernels/ops/indexing.rs similarity index 100% rename from crates/metaltile-std/src/mlx/indexing.rs rename to crates/metaltile-std/src/kernels/ops/indexing.rs diff --git a/crates/metaltile-std/src/mlx/logsumexp.rs b/crates/metaltile-std/src/kernels/ops/logsumexp.rs similarity index 100% rename from crates/metaltile-std/src/mlx/logsumexp.rs rename to crates/metaltile-std/src/kernels/ops/logsumexp.rs diff --git a/crates/metaltile-std/src/kernels/ops/mod.rs b/crates/metaltile-std/src/kernels/ops/mod.rs new file mode 100644 index 00000000..4231cbf7 --- /dev/null +++ b/crates/metaltile-std/src/kernels/ops/mod.rs @@ -0,0 +1,35 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Core / elementwise primitive kernels — the ops family (see +//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): binary / unary / ternary +//! elementwise, copy (+ strided), arange, random, reduce / arg_reduce, scan, +//! indexing, gather / scatter, hadamard, fence, logsumexp, clamp, axpy, and +//! vector add. Migrated from the legacy `mlx/` + `ffai/` split. +//! +//! `ffai/arg_reduce.rs` (`ffai_argmax`) was a byte-for-byte duplicate of +//! `mt_argmax` (in `arg_reduce.rs`) and was dropped, not moved. + +pub mod arange; +pub mod arg_reduce; +pub mod axpy_scalar_inplace; +pub mod binary; +pub mod binary_two; +pub mod clamp_scalar; +pub mod copy; +// `fence.rs` is intentionally NOT declared: device/system memory fences and +// atomics have no `#[kernel]` DSL representation (the "1 out of scope" op in the +// audit). The file is kept for the implementation notes in its header. +pub mod gather; +pub mod gather_axis; +pub mod hadamard; +pub mod hadamard_m; +pub mod indexing; +pub mod logsumexp; +pub mod random; +pub mod reduce; +pub mod scan; +pub mod scatter_axis; +pub mod strided; +pub mod ternary; +pub mod unary; +pub mod vector_add; diff --git a/crates/metaltile-std/src/mlx/random.rs b/crates/metaltile-std/src/kernels/ops/random.rs similarity index 100% rename from crates/metaltile-std/src/mlx/random.rs rename to crates/metaltile-std/src/kernels/ops/random.rs diff --git a/crates/metaltile-std/src/mlx/reduce.rs b/crates/metaltile-std/src/kernels/ops/reduce.rs similarity index 100% rename from crates/metaltile-std/src/mlx/reduce.rs rename to crates/metaltile-std/src/kernels/ops/reduce.rs diff --git a/crates/metaltile-std/src/mlx/scan.rs b/crates/metaltile-std/src/kernels/ops/scan.rs similarity index 100% rename from crates/metaltile-std/src/mlx/scan.rs rename to crates/metaltile-std/src/kernels/ops/scan.rs diff --git a/crates/metaltile-std/src/mlx/scatter_axis.rs b/crates/metaltile-std/src/kernels/ops/scatter_axis.rs similarity index 100% rename from crates/metaltile-std/src/mlx/scatter_axis.rs rename to crates/metaltile-std/src/kernels/ops/scatter_axis.rs diff --git a/crates/metaltile-std/src/mlx/strided.rs b/crates/metaltile-std/src/kernels/ops/strided.rs similarity index 100% rename from crates/metaltile-std/src/mlx/strided.rs rename to crates/metaltile-std/src/kernels/ops/strided.rs diff --git a/crates/metaltile-std/src/mlx/ternary.rs b/crates/metaltile-std/src/kernels/ops/ternary.rs similarity index 100% rename from crates/metaltile-std/src/mlx/ternary.rs rename to crates/metaltile-std/src/kernels/ops/ternary.rs diff --git a/crates/metaltile-std/src/mlx/unary.rs b/crates/metaltile-std/src/kernels/ops/unary.rs similarity index 100% rename from crates/metaltile-std/src/mlx/unary.rs rename to crates/metaltile-std/src/kernels/ops/unary.rs diff --git a/crates/metaltile-std/src/ffai/mt_vector_add.rs b/crates/metaltile-std/src/kernels/ops/vector_add.rs similarity index 89% rename from crates/metaltile-std/src/ffai/mt_vector_add.rs rename to crates/metaltile-std/src/kernels/ops/vector_add.rs index 4a6b0432..077a18c2 100644 --- a/crates/metaltile-std/src/ffai/mt_vector_add.rs +++ b/crates/metaltile-std/src/kernels/ops/vector_add.rs @@ -10,7 +10,7 @@ use metaltile::kernel; #[kernel] -pub fn ffai_vector_add(a: Tensor, b: Tensor, mut out: Tensor) { +pub fn mt_vector_add(a: Tensor, b: Tensor, mut out: Tensor) { let i = tid; let av = load(a[i]).cast::(); let bv = load(b[i]).cast::(); @@ -20,7 +20,7 @@ pub fn ffai_vector_add(a: Tensor, b: Tensor, mut out: Tensor) { pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::ffai_vector_add; + use super::mt_vector_add; use crate::utils::{pack_f32, unpack_f32}; fn setup(n: usize, dt: DType) -> TestSetup { @@ -29,7 +29,7 @@ pub mod kernel_tests { let a_dt = unpack_f32(&pack_f32(&a, dt), dt); let b_dt = unpack_f32(&pack_f32(&b, dt), dt); let expected: Vec = a_dt.iter().zip(&b_dt).map(|(x, y)| x + y).collect(); - TestSetup::new(ffai_vector_add::kernel_ir_for(dt)) + TestSetup::new(mt_vector_add::kernel_ir_for(dt)) .input(TestBuffer::from_vec("a", pack_f32(&a, dt), dt)) .input(TestBuffer::from_vec("b", pack_f32(&b, dt), dt)) .input(TestBuffer::zeros("out", n, dt)) @@ -47,12 +47,12 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::ffai_vector_add; + use super::mt_vector_add; #[bench(dtypes = [f32, f16, bf16])] fn bench_add(dt: DType) -> BenchSetup { let n = 4096usize; - BenchSetup::new(ffai_vector_add::kernel_ir_for(dt)) + BenchSetup::new(mt_vector_add::kernel_ir_for(dt)) .buffer(BenchBuffer::random("a", n, dt)) .buffer(BenchBuffer::random("b", n, dt)) .buffer(BenchBuffer::zeros("out", n, dt).output()) diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs index 2f898776..5da83380 100644 --- a/crates/metaltile-std/src/mlx/mod.rs +++ b/crates/metaltile-std/src/mlx/mod.rs @@ -14,10 +14,6 @@ //! expected to wire up eventually, it lives in `ffai/` until the //! comparison lands. -pub mod arange; -pub mod arg_reduce; -pub mod binary; -pub mod binary_two; pub mod block_scaled_dequant; pub mod block_scaled_matmul; pub mod block_scaled_mma; @@ -26,37 +22,24 @@ pub mod block_scaled_qmm; pub mod block_scaled_qmm_mpp; pub mod block_scaled_qmm_nax; pub mod clamped_swiglu; -pub mod copy; pub mod fft; pub mod fp_quantized; pub mod fp_quantized_mma; pub mod fp_quantized_nax; pub mod fused_gate_activation; -pub mod gather_axis; pub mod gemv; pub mod gemv_masked; -pub mod hadamard; -pub mod hadamard_m; -pub mod indexing; -pub mod logsumexp; pub mod quantized; pub mod quantized_mma_dynamic_m; pub mod quantized_mpp; pub mod quantized_mpp_int8; pub mod quantized_nax; pub mod quantized_nax_int8; -pub mod random; -pub mod reduce; pub mod scaled_dot_product_attention; -pub mod scan; -pub mod scatter_axis; pub mod sdpa_vector; pub mod sgload_smoke; pub mod steel; -pub mod strided; pub mod swiglu; -pub mod ternary; -pub mod unary; // `conv.rs` and `shared.rs` are placeholder/stale stubs left over from // the old `metaltile-bench` crate. They reference `crate::runner` which diff --git a/crates/metaltile-std/tests/steel_msl_snapshots.rs b/crates/metaltile-std/tests/steel_msl_snapshots.rs index 8f13771e..9493ba83 100644 --- a/crates/metaltile-std/tests/steel_msl_snapshots.rs +++ b/crates/metaltile-std/tests/steel_msl_snapshots.rs @@ -32,11 +32,13 @@ use metaltile::{ codegen::{MslGenerator, msl::MslConfig}, core::{dtype::DType, ir::KernelMode}, }; -use metaltile_std::mlx::{ - fp_quantized_nax, - hadamard_m, - quantized_nax, - steel::gemm::{steel_gemm_fused_nax, steel_gemm_gather_nax, steel_gemm_splitk_nax}, +use metaltile_std::{ + kernels::ops::hadamard_m, + mlx::{ + fp_quantized_nax, + quantized_nax, + steel::gemm::{steel_gemm_fused_nax, steel_gemm_gather_nax, steel_gemm_splitk_nax}, + }, }; /// Lower one of the NAX-family kernel IRs to MSL with its declared diff --git a/crates/metaltile/src/runner/harness.rs b/crates/metaltile/src/runner/harness.rs index 2f63802f..ebb815ac 100644 --- a/crates/metaltile/src/runner/harness.rs +++ b/crates/metaltile/src/runner/harness.rs @@ -1430,7 +1430,7 @@ mod name_filter_tests { })) .unwrap(); // kernel name hit - assert!(nf.matches("ffai_vector_add_f32", "bench_add", "")); + assert!(nf.matches("mt_vector_add_f32", "bench_add", "")); // legacy bench-fn-ident hit assert!(nf.matches("some_other_kernel_f32", "bench_vector_add", "")); assert!(!nf.matches("mt_exp_f32", "bench_exp", "")); diff --git a/docs/specs/KERNEL_AUDIT.md b/docs/specs/KERNEL_AUDIT.md index 20c24f60..ec599edf 100644 --- a/docs/specs/KERNEL_AUDIT.md +++ b/docs/specs/KERNEL_AUDIT.md @@ -60,22 +60,21 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | Op | MLX (upstream) | MLX (ekryski@alpha) | metaltile | Notes | |---|---|---|---|---| -| arange | ✓ | ✓ | ✓ | `mlx/arange.rs` → `mt_arange`. Generic `T`. | -| arg_reduce (argmax/argmin → float) | ✓ | ✓ | ✓ | `mlx/arg_reduce.rs` → `mt_argmax` + `mt_argmin`. Both generic over `T` (values widened to f32 for comparison); winning index emitted as `u32`; ties take the smallest index. | -| arg_reduce (argmax → u32 index) | ✗ | ✗ | ✓ | `ffai/arg_reduce.rs` → `ffai_argmax`. FFAI-only; integer-index sampler workhorse. | -| binary (elementwise add/sub/mul/div/min/max) | ✓ | ✓ | ✓ | `mlx/binary.rs` → 6 kernels. Generic `T`. | -| binary_two (fused two-output elementwise) | ✓ | ✓ | ✓ | `mlx/binary_two.rs` → `mt_binary_two`. | -| copy (contiguous) | ✓ | ✓ | ✓ | `mlx/copy.rs` → `mt_copy`. | -| copy (strided / general) | ✓ | ✓ | ✓ | `mlx/strided.rs` → `mt_strided_copy` (2-D padded) + `mt_strided_copy_nd` (arbitrary-rank). Each output element unravels its flat index against a runtime `shape` array and gathers `src[Σ coord_d · strides[d]]`. Covers padded copies, transposes, broadcasts (stride 0), and dilated slices in one kernel. | -| ternary (select) | ✓ | ✓ | ✓ | `mlx/ternary.rs` → `mt_select`. | -| unary (exp/log/sqrt/rsqrt/abs/silu/etc.) | ✓ | ✓ | ✓ | `mlx/unary.rs` → 7+ kernels including `mt_silu`. Plus `mt_scalar_fma_chain8` (fused 8-way scalar FMA) and `mt_add_rms_norm` (residual-add + RMSNorm fusion, Reduction mode). | +| arange | ✓ | ✓ | ✓ | `kernels/ops/arange.rs` → `mt_arange`. Generic `T`. | +| arg_reduce (argmax/argmin → u32 index) | ✓ | ✓ | ✓ | `kernels/ops/arg_reduce.rs` → `mt_argmax` + `mt_argmin`. Both generic over `T` (values widened to f32 for comparison); winning index emitted as `u32`; ties take the smallest index. (The former `ffai_argmax` was a byte-for-byte duplicate of `mt_argmax` and was dropped.) | +| binary (elementwise add/sub/mul/div/min/max) | ✓ | ✓ | ✓ | `kernels/ops/binary.rs` → 6 kernels. Generic `T`. | +| binary_two (fused two-output elementwise) | ✓ | ✓ | ✓ | `kernels/ops/binary_two.rs` → `mt_binary_two`. | +| copy (contiguous) | ✓ | ✓ | ✓ | `kernels/ops/copy.rs` → `mt_copy`. | +| copy (strided / general) | ✓ | ✓ | ✓ | `kernels/ops/strided.rs` → `mt_strided_copy` (2-D padded) + `mt_strided_copy_nd` (arbitrary-rank). Each output element unravels its flat index against a runtime `shape` array and gathers `src[Σ coord_d · strides[d]]`. Covers padded copies, transposes, broadcasts (stride 0), and dilated slices in one kernel. | +| ternary (select) | ✓ | ✓ | ✓ | `kernels/ops/ternary.rs` → `mt_select`. | +| unary (exp/log/sqrt/rsqrt/abs/silu/etc.) | ✓ | ✓ | ✓ | `kernels/ops/unary.rs` → 7+ kernels including `mt_silu`. Plus `mt_scalar_fma_chain8` (fused 8-way scalar FMA) and `mt_add_rms_norm` (residual-add + RMSNorm fusion, Reduction mode). | | swiglu (`silu(gate)·up` fused MLP activation) | ✗ | ✗ | ✓ | `mlx/swiglu.rs` → `mt_swiglu`. Standard modern-transformer MLP activation (Llama 4, Qwen3 dense + MoE, Gemma, Mistral). | -| random (key hash → u32) | ✓ | ✓ | ✓ | `mlx/random.rs` → `mt_random_hash`. | -| reduce (sum/prod/max/min — all + row + col) | ✓ | ✓ | ✓ | `mlx/reduce.rs` covers `all_reduce*`, `row_reduce*`, `col_reduce*`, and `seg_reduce*` (Grid3D one-thread-per-segment, contiguous fixed-length runs) — all four ops for each shape. | +| random (key hash → u32) | ✓ | ✓ | ✓ | `kernels/ops/random.rs` → `mt_random_hash`. | +| reduce (sum/prod/max/min — all + row + col) | ✓ | ✓ | ✓ | `kernels/ops/reduce.rs` covers `all_reduce*`, `row_reduce*`, `col_reduce*`, and `seg_reduce*` (Grid3D one-thread-per-segment, contiguous fixed-length runs) — all four ops for each shape. | | sort | ✓ | ✓ | ✓ | `kernels/sampling/sort.rs` → `mt_sort` (single-block bitonic) + `mt_merge` (multi-block merge) + `mt_sort_segmented` (per-row bitonic for `[batch, n]` matrices, `n ≤ 1024`, one TG per row). | -| 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. | +| scan (prefix sum/prod/max/min) | ✓ | ✓ | ✓ | `kernels/ops/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 | ✓ | ✓ | ✓ | `kernels/sampling/softmax.rs` → `mt_softmax` (looped + single-row collapsed). | -| logsumexp | ✓ | ✓ | ✓ | `mlx/logsumexp.rs` → `mt_logsumexp`. | +| logsumexp | ✓ | ✓ | ✓ | `kernels/ops/logsumexp.rs` → `mt_logsumexp`. | | 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`. | @@ -105,7 +104,7 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | gemv_masked | ✓ | ✓ | ✓ | `mlx/gemv_masked.rs` → `mt_gemv_masked`. | | quantized (affine_quantize / affine_dequantize) | ✓ | ✓ | ✓ | `mlx/quantized.rs` — quantize + dequantize for all widths: int2/int4/int8 (pack-aligned) + int3/int5/int6 (byte-stream). 12 kernels (`mt_affine_{quantize,dequantize}_int{2,3,4,5,6,8}`). int3/5/6 quantize uses bit-stream OR (lane 0 ORs codes into u32 words) to handle straddling — no atomics. | | quantized (affine_qmv / qvm / qmm — matvec / matmul) | ✓ | ✓ | ✓ | `mlx/quantized.rs` — **int4 perf**: `mt_qmv` (8-row-per-TG decode, mirrors MLX `qmv_fast`) + `mt_qmm` / `_bm2` / `_bm4` (M-batched prefill) + `mt_qmm_mma` / `_m16` (simdgroup-matrix MMA prefill) + `mt_qmm_mma_mpp` (MPP) + `mt_qmm_nax` (NAX). **int8 perf** (PR #154): `mt_qmv_int8_fast`, `mt_qmm_int8_fast` / `_bm2` / `_bm4`, `mt_qmm_mma_int8` / `_m16_int8`, `mt_qmm_mma_mpp_int8`, `mt_qmm_nax_int8` — pack-aligned (4 bytes/u32, byte-shift extract), closes the ~6–8× int8-vs-int4 perf gap. **Odd-bitwidth MMA** (PR #157): `mt_qmm_mma_b{3,5,6}` — straddle-aware two-word bit-stream dequant in the 4-SG MMA body. **All bit-widths × all dtypes**: `mt_{qmv,qvm,qmm}_b{3,4,5,6,8}` (correctness-first scalar family). **qvm perf**: `mt_qvm_int4_fast` (PR #154) — 8-col-per-TG, MLX `qvm_fast` shape. | -| quantized (gather_qmv / gather_qmm — gather variants) | ✓ | ✓ | ✓ | `ffai/moe.rs` → `mt_moe_gather_qmm_int4` (int4 affine grouped-gather) + `mt_moe_gather_qmm_b{3,5,6,8}` (all bit-widths, scalar). **int4 perf**: `mt_moe_gather_qmm_mma_int4{,_bm16}` + `_m8` (decode) + `_m{16,32}` (PR #157 short-prefill, hand-unrolled `acc0..accN` cells — the DSL doesn't lower runtime-indexed mutable arrays), MPP scale-ups `bm{8,16,64}_mpp` (`ffai/moe_mpp{,_bm8,_bm64}.rs`). **int8 perf** (PR #154): pack-aligned `mt_moe_gather_qmm_mma_int8` (1-SG MMA decode) + `_bm16_mpp` + `_bm8_mpp` (direct-input cooperative tensors, M=8 forbids coop-tensor) + `_bm64_mpp` (4-SG 2×2 long-context prefill). All MPP kernels stage bf16 through `half` cooperative tensors via `coop_stage(T)`. Bare-tensor `ffai/gather.rs` exists but is non-quantized. **Expert-indexed dequant GEMV** (PR #160): `dequant_gemv_int4_expert_indexed` — per-output-row expert selection for the gate/up FFN dispatch shape. | +| quantized (gather_qmv / gather_qmm — gather variants) | ✓ | ✓ | ✓ | `ffai/moe.rs` → `mt_moe_gather_qmm_int4` (int4 affine grouped-gather) + `mt_moe_gather_qmm_b{3,5,6,8}` (all bit-widths, scalar). **int4 perf**: `mt_moe_gather_qmm_mma_int4{,_bm16}` + `_m8` (decode) + `_m{16,32}` (PR #157 short-prefill, hand-unrolled `acc0..accN` cells — the DSL doesn't lower runtime-indexed mutable arrays), MPP scale-ups `bm{8,16,64}_mpp` (`ffai/moe_mpp{,_bm8,_bm64}.rs`). **int8 perf** (PR #154): pack-aligned `mt_moe_gather_qmm_mma_int8` (1-SG MMA decode) + `_bm16_mpp` + `_bm8_mpp` (direct-input cooperative tensors, M=8 forbids coop-tensor) + `_bm64_mpp` (4-SG 2×2 long-context prefill). All MPP kernels stage bf16 through `half` cooperative tensors via `coop_stage(T)`. Bare-tensor `kernels/ops/gather.rs` exists but is non-quantized. **Expert-indexed dequant GEMV** (PR #160): `dequant_gemv_int4_expert_indexed` — per-output-row expert selection for the gate/up FFN dispatch shape. | | moe (router top-k + permute + unpermute orchestration) | ✗ | ✓ | ✓ | `ffai/moe.rs` → `mt_moe_router_topk`, `mt_moe_permute`, `mt_moe_unpermute`. MoE expert-routing orchestration. The grouped quantized BGEMM that fuses per-expert FFN matmuls is counted under the `quantized (gather_*)` row. | | dequant_gather (quantized embedding-table gather) | ✗ | ✗ | ✓ | `ffai/dequant_gather.rs`. int{3,4,5,6,8} all bit-widths. FFAI-only. | | dequant_gemv (quantized GEMV, FFAI flavour) | ~ | ~ | ✓ | `ffai/dequant_gemv.rs` → `dequant_gemv_int{3,4,5,6,8}` (one-row-per-TG) + `dequant_gemv_int4_fast` (PR #154, 8-row-per-TG, mirrors MLX `qmv_fast`). The non-fast int4 kernel stays because FFAI's GPU-router opts into its indirect Swift wrapper. | @@ -114,10 +113,10 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | 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. | -| hadamard (hadamard_n + hadamard_m) | ✓ | ✓ | ✓ | `mlx/hadamard.rs` → `mt_hadamard_n{64,128,256,512,1024}` (FWHT, log2(N) butterfly passes). `mlx/hadamard_m.rs` → `mt_hadamard_m{12,20,28}` (non-pow2 M factor, Sloane-table bitmask accumulate). Generic over `T`. | +| 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) | ✓ | ✓ | ✓ | `ffai/gather.rs` → `ffai_gather`. FFAI's embedding-table gather. | -| indexing (scatter, scatter_axis, gather_axis, gather_front, masked_scatter) | ✓ | ✓ | ✓ | `mlx/gather_axis.rs` + `mlx/scatter_axis.rs` → `mt_gather_axis` / `mt_scatter_axis`; `mlx/indexing.rs` → `mt_gather_front`, `mt_scatter`, `mt_masked_scatter`. All one-thread-per-output Grid3D with bounds guards. | +| gather (bare-tensor embedding lookup) | ✓ | ✓ | ✓ | `kernels/ops/gather.rs` → `mt_gather`. FFAI's embedding-table gather. | +| indexing (scatter, scatter_axis, gather_axis, gather_front, masked_scatter) | ✓ | ✓ | ✓ | `kernels/ops/gather_axis.rs` + `kernels/ops/scatter_axis.rs` → `mt_gather_axis` / `mt_scatter_axis`; `kernels/ops/indexing.rs` → `mt_gather_front`, `mt_scatter`, `mt_masked_scatter`. All one-thread-per-output Grid3D with bounds guards. | | aura_encode (codebook quantize, fused) | ✗ | ✓ | ✓ | `ffai/aura_encode.rs`. Bit-widths 2/3/4/8. | | aura_dequant_rotated (bulk dequant to rotated codec space) | ✗ | ✓ | ✓ | `ffai/aura_dequant_rotated.rs`. bits ∈ {2,3,4,8}. | | aura_score (compressed-domain Q·K) | ✗ | ✓ | ✓ | `ffai/aura_score.rs`. bits ∈ {2,3,4,8}. Generic over `T`. | @@ -138,7 +137,7 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | 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) | ✗ | ✗ | ✓ | `kernels/sampling/categorical_sample.rs` → `mt_softmax_categorical_sample`. Companion to `ffai_argmax` for `T > 0` decode. | +| 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. | | 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). | diff --git a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md index d4fcd40d..27e635db 100644 --- a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md +++ b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md @@ -12,13 +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/` + `norm/` + `sampling/` 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. +> and all of wave 1 (`rope/`, `norm/`, `sampling/`, `ops/`) 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 +> `*_block_scaled_qgemv` format matrices moved with `norm/` but still await the +> format-axis fold (§7). Wave 2/3 families migrate one-per-PR per §6. ## 1. Why @@ -43,7 +43,7 @@ the FFAI emit path is unaffected). ``` crates/metaltile-std/src/kernels/ - ops/ elementwise/core primitives: binary · unary · ternary · copy · arange · + ops/ ✅ DONE — elementwise/core primitives: binary · unary · ternary · copy · arange · random · reduce · arg_reduce · scan · indexing · gather/scatter · hadamard · fence · clamp · logsumexp · vector_add · axpy · strided gemm/ gemm · gemv(_masked) · batched-projection (qkv / 4) · patch_embed · steel/gemm @@ -156,8 +156,7 @@ payoff last: | Wave | Families | Rationale | Payoff | |---|---|---|---| -| ✅ done | `convolution/`, `rope/`, `norm/`, `sampling/` | exemplar + wave-1 families | 24k → ~1.6k | -| 1 | `ops/` | self-contained, mostly elementwise / few formats | small, sets the pattern | +| ✅ done | `convolution/`, `rope/`, `norm/`, `sampling/`, `ops/` | exemplar + all of wave 1 | 24k → ~1.6k | | 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 | From 58ba46087b383ceec68d01876162ad11f9ca9889 Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Fri, 12 Jun 2026 18:26:19 -0600 Subject: [PATCH 2/3] refactor(ops): consolidate gated activations + rehome simdgroup-load probe Finishes rehoming the last mlx/ stragglers: - gated activations: fold mlx/{swiglu,clamped_swiglu,fused_gate_activation}.rs (4 kernels) into one kernels/ops/gated_activation.rs. Names unchanged (mt_swiglu, mt_clamped_swiglu, mt_fused_gate_gelu, mt_fused_gate_clipped_swiglu); ops/ is the activation home (mt_silu, which mt_swiglu cross-kernel-calls, lives in ops/unary.rs). 3 files -> 1. - mlx/sgload_smoke.rs -> probe/simdgroup_load_probe.rs; mt_sgload_smoke -> mt_simdgroup_load_probe. It's a HW-intrinsic validation kernel (sibling of probe/mma_layout_probe + mpp_matmul_smoke), so probe/ is its real home, not ops/. Renamed for what it does (validates the Op::SimdgroupLoad codegen path via a byte-exact round-trip). Test file renamed to match. Updated KERNEL_AUDIT (swiglu/fused rows) + the plan's ops listing. cargo build + clippy clean; gated-activation kernel tests + the probe integration test pass. --- .../src/kernels/ops/gated_activation.rs | 273 ++++++++++++++++++ crates/metaltile-std/src/kernels/ops/mod.rs | 6 +- .../metaltile-std/src/mlx/clamped_swiglu.rs | 128 -------- .../src/mlx/fused_gate_activation.rs | 176 ----------- crates/metaltile-std/src/mlx/mod.rs | 4 - crates/metaltile-std/src/mlx/swiglu.rs | 101 ------- crates/metaltile-std/src/probe/mod.rs | 1 + .../simdgroup_load_probe.rs} | 16 +- ...oke_gpu.rs => simdgroup_load_probe_gpu.rs} | 18 +- docs/specs/KERNEL_AUDIT.md | 4 +- docs/specs/KERNEL_CONSOLIDATION_PLAN.md | 2 +- 11 files changed, 298 insertions(+), 431 deletions(-) create mode 100644 crates/metaltile-std/src/kernels/ops/gated_activation.rs delete mode 100644 crates/metaltile-std/src/mlx/clamped_swiglu.rs delete mode 100644 crates/metaltile-std/src/mlx/fused_gate_activation.rs delete mode 100644 crates/metaltile-std/src/mlx/swiglu.rs rename crates/metaltile-std/src/{mlx/sgload_smoke.rs => probe/simdgroup_load_probe.rs} (92%) rename crates/metaltile-std/tests/{sgload_smoke_gpu.rs => simdgroup_load_probe_gpu.rs} (88%) diff --git a/crates/metaltile-std/src/kernels/ops/gated_activation.rs b/crates/metaltile-std/src/kernels/ops/gated_activation.rs new file mode 100644 index 00000000..9db1f7c7 --- /dev/null +++ b/crates/metaltile-std/src/kernels/ops/gated_activation.rs @@ -0,0 +1,273 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Gated MLP activations — `activation(gate) * up`. +//! +//! The SwiGLU-family fused activation used in modern transformer MLPs: given +//! the two halves of an MLP's gate/up projections (`w_gate·x`, `w_up·x`), apply +//! an activation to the gate and multiply by up in ONE elementwise pass — saves +//! a full-tensor read-modify-write versus separate `silu` + `mul` launches (the +//! intermediate stays in registers). One file covers the variants: +//! +//! - `mt_swiglu` — `silu(gate) · up` (Llama/Qwen/Gemma/Mistral) +//! - `mt_clamped_swiglu` — `clip(silu(gate), L) · clip(up, ±L)`; a per-layer +//! activation-clip limit (gpt-oss, StepFun Step-3). `limit <= 0` ⇒ plain SwiGLU +//! - `mt_fused_gate_gelu` — `gelu_approx(gate) · up` +//! - `mt_fused_gate_clipped_swiglu` — GPT-OSS clipped: clamp `±7`, +//! `sigmoid(1.702·g)` gate, `+1` up bias +//! +//! All are one-thread-per-output elementwise kernels computed in f32. The plain +//! ungated `silu` activation lives in `ops/unary.rs` (`mt_silu`); `mt_swiglu` +//! cross-kernel-calls it (inlined at codegen by `KernelInlinePass`). +//! +//! MLX reference: `mx.fast.swiglu` + `fused_gate_activation.metal` +//! (`apply_gate`; type 0 = silu = `mt_swiglu`). + +use metaltile::kernel; + +// ── SwiGLU: silu(gate) · up ─────────────────────────────────────────────────── + +#[kernel] +pub fn mt_swiglu(gate: Tensor, up: Tensor, out: Tensor) { + let idx = tid; + let g = load(gate[idx]).cast::(); + let u = load(up[idx]).cast::(); + // Cross-kernel call: KernelInlinePass splices mt_silu's scalar body here. + // mt_silu's input-param load is replaced by g (already f32), so silu runs in + // f32. Future fusion passes can identify the (silu, mul) → swiglu pattern. + let s = mt_silu(g); + store(out[idx], (s * u).cast::()); +} + +// ── Clamped SwiGLU: clip(silu(gate), L) · clip(up, ±L) ──────────────────────── + +// Bare `#[kernel]` with a runtime `limit: f32` constexpr. `limit <= 0` collapses +// to plain SwiGLU (the per-layer dispatch reaches for this only on marked layers). +#[kernel] +pub fn mt_clamped_swiglu( + gate: Tensor, + up: Tensor, + out: Tensor, + #[constexpr] limit: f32, +) { + let idx = tid; + let g = load(gate[idx]).cast::(); + let u = load(up[idx]).cast::(); + // silu(g) = g * sigmoid(g). Free-function `exp` + f32 literals so the binding + // isn't elided (the method form `(-g).exp()` nested in a larger expr drops + // out of codegen — see moe_router_sqrtsoftplus). + let sig = 1.0f32 / (1.0f32 + exp(0.0f32 - g)); + let s_raw = g * sig; + // Clip via `select`, not `min`/`max`: the DSL's min/max overloads are + // ambiguous on mixed int/float operands. silu's upper tail is clipped + // one-sided; `up` two-sided. + let active = limit > 0.0f32; + let neg = 0.0f32 - limit; + let s_clipped = select(active, select(s_raw > limit, limit, s_raw), s_raw); + let u_hi = select(u > limit, limit, u); + let u_lo = select(u_hi < neg, neg, u_hi); + let u_clipped = select(active, u_lo, u); + store(out[idx], (s_clipped * u_clipped).cast::()); +} + +// ── Fused gate-GELU: gelu_approx(gate) · up ─────────────────────────────────── + +/// GELU via the tanh approximation MLX uses — matches `gelu_approx_act` in +/// `fused_gate_activation.metal`. Computed in f32 so the cubic + tanh keep +/// precision regardless of `T`. +#[kernel] +pub fn mt_fused_gate_gelu(gate: Tensor, up: Tensor, out: Tensor) { + let idx = program_id::<0>(); + let g = load(gate[idx]).cast::(); + let u = load(up[idx]).cast::(); + // gelu_approx(x) = 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³))) + let x3 = g * g * g; + let inner = 0.7978845608f32 * (g + 0.044715f32 * x3); + let act = 0.5f32 * g * (1.0f32 + tanh(inner)); + store(out[idx], (act * u).cast::()); +} + +// ── Fused gate clipped-SwiGLU (GPT-OSS): clamp ±7, sigmoid(1.702·g), +1 up ───── + +/// The GPT-OSS variant: both halves clamped to `[-7, 7]`; gate uses +/// `sigmoid(1.702·g)`; up carries a `+1` bias: `g·sigmoid(1.702·g)·(u + 1)`. +/// Matches `clipped_swiglu` in `fused_gate_activation.metal`. Clamp is composed +/// from two `select`s (the DSL has no `clamp` builtin). +#[kernel] +pub fn mt_fused_gate_clipped_swiglu(gate: Tensor, up: Tensor, out: Tensor) { + let idx = program_id::<0>(); + let g_raw = load(gate[idx]).cast::(); + let u_raw = load(up[idx]).cast::(); + let g_hi = select(g_raw > 7.0f32, 7.0f32, g_raw); + let g = select(g_hi < (0.0f32 - 7.0f32), 0.0f32 - 7.0f32, g_hi); + let u_hi = select(u_raw > 7.0f32, 7.0f32, u_raw); + let u = select(u_hi < (0.0f32 - 7.0f32), 0.0f32 - 7.0f32, u_hi); + let sig = 1.0f32 / (1.0f32 + exp(0.0f32 - 1.702f32 * g)); + let act = g * sig * (u + 1.0f32); + store(out[idx], act.cast::()); +} + +/// Correctness for the gated activations. Oracles mirror each kernel exactly on +/// dtype-rounded inputs, computed in f32. +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::{mt_clamped_swiglu, mt_fused_gate_clipped_swiglu, mt_fused_gate_gelu, mt_swiglu}; + use crate::utils::{pack_f32, unpack_f32}; + + // ── mt_swiglu ───────────────────────────────────────────────────────────── + + fn swiglu_setup(n: usize, dt: DType) -> TestSetup { + let gate: Vec = (0..n).map(|i| (i % 17) as f32 * 0.35 - 3.0).collect(); + let up: Vec = (0..n).map(|i| (i % 13) as f32 * 0.2 - 1.0).collect(); + let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); + let u_dt = unpack_f32(&pack_f32(&up, dt), dt); + let expected: Vec = + g_dt.iter().zip(&u_dt).map(|(&g, &u)| (g / (1.0 + (-g).exp())) * u).collect(); + TestSetup::new(mt_swiglu::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt)) + .input(TestBuffer::from_vec("up", pack_f32(&up, dt), dt)) + .input(TestBuffer::zeros("out", n, dt)) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_mt_swiglu(dt: DType) -> TestSetup { swiglu_setup(1024, dt) } + + // ── mt_clamped_swiglu ───────────────────────────────────────────────────── + + fn clamped_setup(n: usize, limit: f32, dt: DType) -> TestSetup { + let gate: Vec = (0..n).map(|i| (i % 17) as f32 * 0.35 - 3.0).collect(); + let up: Vec = (0..n).map(|i| (i % 13) as f32 * 0.2 - 1.0).collect(); + let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); + let u_dt = unpack_f32(&pack_f32(&up, dt), dt); + let expected: Vec = g_dt + .iter() + .zip(&u_dt) + .map(|(&g, &u)| { + let s = g / (1.0 + (-g).exp()); // silu(g) = g * sigmoid(g) + let (s_c, u_c) = + if limit > 0.0 { (s.min(limit), u.max(-limit).min(limit)) } else { (s, u) }; + s_c * u_c + }) + .collect(); + TestSetup::new(mt_clamped_swiglu::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt)) + .input(TestBuffer::from_vec("up", pack_f32(&up, dt), dt)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("limit", limit) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_mt_clamped_swiglu_active(dt: DType) -> TestSetup { clamped_setup(1024, 7.0, dt) } + + /// `limit == 0` collapses to plain SwiGLU — equivalence with `mt_swiglu` is + /// the invariant the per-layer dispatch ships on. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_mt_clamped_swiglu_zero_limit_equals_plain(dt: DType) -> TestSetup { + clamped_setup(1024, 0.0, dt) + } + + // ── mt_fused_gate_{gelu,clipped_swiglu} ─────────────────────────────────── + + fn inputs(n: usize, dt: DType) -> (Vec, Vec, Vec, Vec) { + // Range spans beyond +/-7 so clipped_swiglu's clamp is exercised. + let gate: Vec = (0..n).map(|i| (i % 17) as f32 - 8.0).collect(); + let up: Vec = (0..n).map(|i| (i % 13) as f32 - 6.0).collect(); + let g = unpack_f32(&pack_f32(&gate, dt), dt); + let u = unpack_f32(&pack_f32(&up, dt), dt); + (gate, up, g, u) + } + + fn build( + kernel: metaltile::core::ir::Kernel, + gate: &[f32], + up: &[f32], + expected: &[f32], + dt: DType, + ) -> TestSetup { + TestSetup::new(kernel) + .input(TestBuffer::from_vec("gate", pack_f32(gate, dt), dt)) + .input(TestBuffer::from_vec("up", pack_f32(up, dt), dt)) + .input(TestBuffer::zeros("out", gate.len(), dt)) + .expect(TestBuffer::from_vec("out", pack_f32(expected, dt), dt)) + .grid_1d(gate.len(), 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 5e-2])] + fn test_fused_gate_gelu(dt: DType) -> TestSetup { + let (gate, up, g, u) = inputs(512, dt); + const C: f32 = 0.797_884_6; // sqrt(2/pi) + let expected: Vec = g + .iter() + .zip(&u) + .map(|(&g, &u)| 0.5 * g * (1.0 + (C * (g + 0.044715 * g * g * g)).tanh()) * u) + .collect(); + build(mt_fused_gate_gelu::kernel_ir_for(dt), &gate, &up, &expected, dt) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 5e-2])] + fn test_fused_gate_clipped_swiglu(dt: DType) -> TestSetup { + let (gate, up, g, u) = inputs(512, dt); + let expected: Vec = g + .iter() + .zip(&u) + .map(|(&g, &u)| { + let g = g.clamp(-7.0, 7.0); + let u = u.clamp(-7.0, 7.0); + g * (1.0 / (1.0 + (-1.702 * g).exp())) * (u + 1.0) + }) + .collect(); + build(mt_fused_gate_clipped_swiglu::kernel_ir_for(dt), &gate, &up, &expected, dt) + } +} + +/// Benchmarks for the gated activations (1D elementwise, 3 streams). +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{mt_clamped_swiglu, mt_fused_gate_clipped_swiglu, mt_fused_gate_gelu, mt_swiglu}; + + #[bench(dtypes = [f32, f16, bf16])] + fn bench_swiglu(dt: DType) -> BenchSetup { + let n = 1024 * 1024usize; + BenchSetup::new(mt_swiglu::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("gate", n, dt)) + .buffer(BenchBuffer::random("up", n, dt)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .grid_1d(n, 256) + .bytes_moved((3 * n * dt.size_bytes()) as u64) + } + + #[bench(dtypes = [f32, f16, bf16])] + fn bench_clamped_swiglu(dt: DType) -> BenchSetup { + let n = 1024 * 1024usize; + BenchSetup::new(mt_clamped_swiglu::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("gate", n, dt)) + .buffer(BenchBuffer::random("up", n, dt)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("limit", 7.0f32) + .grid_1d(n, 256) + .bytes_moved((3 * n * dt.size_bytes()) as u64) + } + + fn fb(kernel: metaltile::core::ir::Kernel, dt: DType) -> BenchSetup { + let n = 64 * 1024 * 1024usize; + BenchSetup::new(kernel) + .buffer(BenchBuffer::random("gate", n, dt)) + .buffer(BenchBuffer::random("up", n, dt)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .grid_1d(n, 256) + .bytes_moved((3 * n * dt.size_bytes()) as u64) + } + + #[bench(dtypes = [f32, f16, bf16])] + fn bench_gelu(dt: DType) -> BenchSetup { fb(mt_fused_gate_gelu::kernel_ir_for(dt), dt) } + + #[bench(dtypes = [f32, f16, bf16])] + fn bench_clipped(dt: DType) -> BenchSetup { + fb(mt_fused_gate_clipped_swiglu::kernel_ir_for(dt), dt) + } +} diff --git a/crates/metaltile-std/src/kernels/ops/mod.rs b/crates/metaltile-std/src/kernels/ops/mod.rs index 4231cbf7..b53b31d6 100644 --- a/crates/metaltile-std/src/kernels/ops/mod.rs +++ b/crates/metaltile-std/src/kernels/ops/mod.rs @@ -3,8 +3,9 @@ //! Core / elementwise primitive kernels — the ops family (see //! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): binary / unary / ternary //! elementwise, copy (+ strided), arange, random, reduce / arg_reduce, scan, -//! indexing, gather / scatter, hadamard, fence, logsumexp, clamp, axpy, and -//! vector add. Migrated from the legacy `mlx/` + `ffai/` split. +//! indexing, gather / scatter, hadamard, fence, logsumexp, clamp, axpy, +//! vector add, and the gated MLP activations (SwiGLU / GeGLU). Migrated from +//! the legacy `mlx/` + `ffai/` split. //! //! `ffai/arg_reduce.rs` (`ffai_argmax`) was a byte-for-byte duplicate of //! `mt_argmax` (in `arg_reduce.rs`) and was dropped, not moved. @@ -19,6 +20,7 @@ pub mod copy; // `fence.rs` is intentionally NOT declared: device/system memory fences and // atomics have no `#[kernel]` DSL representation (the "1 out of scope" op in the // audit). The file is kept for the implementation notes in its header. +pub mod gated_activation; pub mod gather; pub mod gather_axis; pub mod hadamard; diff --git a/crates/metaltile-std/src/mlx/clamped_swiglu.rs b/crates/metaltile-std/src/mlx/clamped_swiglu.rs deleted file mode 100644 index 55fbc9c5..00000000 --- a/crates/metaltile-std/src/mlx/clamped_swiglu.rs +++ /dev/null @@ -1,128 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Clamped SwiGLU activation — `clip(silu(gate), max=L) * clip(up, -L, L)`. -//! -//! Drop-in replacement for [`mt_swiglu`](super::swiglu::mt_swiglu) for -//! checkpoints that ship per-layer activation-clipping limits. Pattern: -//! -//! ```text -//! out[i] = clip(silu(gate[i]), max=L) * clip(up[i], -L, L) -//! ``` -//! -//! Two checkpoints in the wild use this shape — gpt-oss-120B and -//! StepFun's Step-3.7-Flash (per-layer `swiglu_limits` / `swiglu_limits_shared` -//! lists that are non-zero on a small subset of layers, zero elsewhere). -//! For layers whose limit is zero the caller should dispatch the plain -//! [`mt_swiglu`](super::swiglu::mt_swiglu); the clamped variant is the -//! one to reach for on the marked layers. -//! -//! Clipping happens on the f32 intermediates inside the kernel, before -//! the final cast back to `T`, so quant-stats fit the clipped range -//! regardless of activation dtype. -//! -//! ## ABI -//! -//! ```text -//! gate [N] T — w_gate · x -//! up [N] T — w_up · x -//! out [N] T — clipped SwiGLU output -//! limit f32 — constexpr; non-negative clip bound. `limit <= 0` -//! collapses to plain SwiGLU (no clip). -//! ``` -//! -//! Grid is 1D elementwise: one thread per output position. Caller -//! drives `grid_1d(n, 256)`. - -use metaltile::kernel; - -// Bare `#[kernel]` — the legacy `bench(...)` registration's `Binary` -// class can't represent the extra `limit: f32` runtime scalar; the new -// declarative `#[bench]` attribute on `kernel_benches::bench_clamped_swiglu` -// below registers the kernel for `tile bench` directly. -#[kernel] -pub fn mt_clamped_swiglu( - gate: Tensor, - up: Tensor, - out: Tensor, - #[constexpr] limit: f32, -) { - let idx = tid; - let g = load(gate[idx]).cast::(); - let u = load(up[idx]).cast::(); - // silu(g) = g * sigmoid(g). Free-function `exp` + f32 literals so the - // binding isn't elided (the method form `(-g).exp()` nested in a - // larger expr drops out of codegen — see moe_router_sqrtsoftplus). - let sig = 1.0f32 / (1.0f32 + exp(0.0f32 - g)); - let s_raw = g * sig; - // Clip via `select`, not `min`/`max`: the DSL's min/max overloads - // are ambiguous on mixed int/float operands (the dsv4_swiglu_limit - // sibling clamps the same way). `limit <= 0` collapses to plain - // SwiGLU. silu's upper tail is clipped one-sided; `up` two-sided. - let active = limit > 0.0f32; - let neg = 0.0f32 - limit; - let s_clipped = select(active, select(s_raw > limit, limit, s_raw), s_raw); - let u_hi = select(u > limit, limit, u); - let u_lo = select(u_hi < neg, neg, u_hi); - let u_clipped = select(active, u_lo, u); - store(out[idx], (s_clipped * u_clipped).cast::()); -} - -pub mod kernel_tests { - use metaltile::{test::*, test_kernel}; - - use super::mt_clamped_swiglu; - use crate::utils::{pack_f32, unpack_f32}; - - fn setup(n: usize, limit: f32, dt: DType) -> TestSetup { - let gate: Vec = (0..n).map(|i| (i % 17) as f32 * 0.35 - 3.0).collect(); - let up: Vec = (0..n).map(|i| (i % 13) as f32 * 0.2 - 1.0).collect(); - let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); - let u_dt = unpack_f32(&pack_f32(&up, dt), dt); - let expected: Vec = g_dt - .iter() - .zip(&u_dt) - .map(|(&g, &u)| { - let s = g / (1.0 + (-g).exp()); // silu(g) = g * sigmoid(g) - let (s_c, u_c) = - if limit > 0.0 { (s.min(limit), u.max(-limit).min(limit)) } else { (s, u) }; - s_c * u_c - }) - .collect(); - TestSetup::new(mt_clamped_swiglu::kernel_ir_for(dt)) - .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt)) - .input(TestBuffer::from_vec("up", pack_f32(&up, dt), dt)) - .input(TestBuffer::zeros("out", n, dt)) - .constexpr("limit", limit) - .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) - .grid_1d(n, 256) - } - - #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] - fn test_mt_clamped_swiglu_active(dt: DType) -> TestSetup { setup(1024, 7.0, dt) } - - /// `limit == 0` collapses to plain SwiGLU — equivalence with - /// [`mt_swiglu`](super::super::swiglu::mt_swiglu) is the - /// invariant we ship the per-layer dispatch on. - #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] - fn test_mt_clamped_swiglu_zero_limit_equals_plain(dt: DType) -> TestSetup { - setup(1024, 0.0, dt) - } -} - -pub mod kernel_benches { - use metaltile::{bench, test::*}; - - use super::mt_clamped_swiglu; - - #[bench(dtypes = [f32, f16, bf16])] - fn bench_clamped_swiglu(dt: DType) -> BenchSetup { - let n = 1024 * 1024usize; - BenchSetup::new(mt_clamped_swiglu::kernel_ir_for(dt)) - .buffer(BenchBuffer::random("gate", n, dt)) - .buffer(BenchBuffer::random("up", n, dt)) - .buffer(BenchBuffer::zeros("out", n, dt).output()) - .constexpr("limit", 7.0f32) - .grid_1d(n, 256) - .bytes_moved((3 * n * dt.size_bytes()) as u64) - } -} diff --git a/crates/metaltile-std/src/mlx/fused_gate_activation.rs b/crates/metaltile-std/src/mlx/fused_gate_activation.rs deleted file mode 100644 index c7faea5a..00000000 --- a/crates/metaltile-std/src/mlx/fused_gate_activation.rs +++ /dev/null @@ -1,176 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! Fused gate-activation — `activation(gate) * up`. -//! -//! The general form of the SwiGLU MLP activation. Given two -//! equally-sized inputs `gate` and `up` (the two halves of an MLP's -//! `w_gate · x` and `w_up · x` projections), produce -//! -//! ```text -//! out[i] = activation(gate[i]) * up[i] -//! ``` -//! -//! The `silu` activation is the standard modern-transformer case and -//! already ships as the dedicated [`mt_swiglu`](super::swiglu) kernel. -//! This module covers the other two activation variants MLX's -//! `fused_gate_activation.metal` carries: -//! -//! - **`gelu_approx`** — the tanh approximation of GELU -//! (`0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³)))`). Used by GELU-MLP -//! transformers that ship a fused gate/up projection. -//! - **`clipped_swiglu`** — the GPT-OSS clipped variant: both halves -//! are clamped to `[-7, 7]`, the gate side uses `sigmoid(1.702·g)`, -//! and the up side carries a `+1` bias before the multiply. -//! -//! All three are one-thread-per-output Grid3D kernels — no -//! cross-thread cooperation, so the reduction-mode dispatch hazards do -//! not apply. The `single_row` / `looped` dispatch split MLX uses is a -//! threadgroup-tiling perf detail; a flat element-parallel grid is the -//! tractable port and is what FFAI's elementwise activation path wants. -//! -//! MLX reference: `mlx/backend/metal/kernels/fused_gate_activation.metal` -//! — `apply_gate` with `activation_type ∈ {0,1,2}`. -//! `activation_type == 0` (silu) is `mt_swiglu`; this file ports `1` -//! and `2`. -//! -//! Codegen-only; correctness pinned by the in-source `#[test_kernel]`s. - -use metaltile::kernel; - -// Numeric constants are inlined as literals inside the kernel bodies -// below — the `#[kernel]` proc-macro parses the body as a token stream -// and does not substitute Rust `const` items. They are named here for -// documentation only: -// GELU_SQRT_2_OVER_PI = 0.7978845608 — √(2/π), the gelu-tanh -// inner scale. -// GELU_CUBIC_COEFF = 0.044715 — gelu-tanh cubic coeff. -// CLIPPED_SWIGLU_CLIP = 7.0 — GPT-OSS clamp bound. -// CLIPPED_SWIGLU_GATE_SCALE = 1.702 — GPT-OSS gate sigmoid scale. - -/// `out[i] = gelu_approx(gate[i]) * up[i]`. -/// -/// GELU via the tanh approximation MLX uses — matches -/// `gelu_approx_act` in `fused_gate_activation.metal`. Computed in f32 -/// regardless of `T` so the cubic + tanh keep their precision; the -/// result is cast back to `T` at the store. -#[kernel] -pub fn mt_fused_gate_gelu(gate: Tensor, up: Tensor, out: Tensor) { - let idx = program_id::<0>(); - let g = load(gate[idx]).cast::(); - let u = load(up[idx]).cast::(); - // gelu_approx(x) = 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³))) - let x3 = g * g * g; - let inner = 0.7978845608f32 * (g + 0.044715f32 * x3); - let act = 0.5f32 * g * (1.0f32 + tanh(inner)); - store(out[idx], (act * u).cast::()); -} - -/// `out[i] = clipped_swiglu(gate[i], up[i])` — the GPT-OSS variant. -/// -/// Both halves are clamped to `[-7, 7]`; the gate uses -/// `sigmoid(1.702·g)`; the up side has a `+1` bias before the -/// multiply: `g·sigmoid(1.702·g)·(u + 1)`. Matches `clipped_swiglu` -/// in `fused_gate_activation.metal`. The clamp is composed from two -/// `select`s (the DSL has no `clamp` builtin). -#[kernel] -pub fn mt_fused_gate_clipped_swiglu(gate: Tensor, up: Tensor, out: Tensor) { - let idx = program_id::<0>(); - let g_raw = load(gate[idx]).cast::(); - let u_raw = load(up[idx]).cast::(); - // clamp(x, -7, 7) = min(max(x, -7), 7), composed from select since - // the DSL exposes no `clamp` builtin (see ffai/rope_yarn.rs). - let g_hi = select(g_raw > 7.0f32, 7.0f32, g_raw); - let g = select(g_hi < (0.0f32 - 7.0f32), 0.0f32 - 7.0f32, g_hi); - let u_hi = select(u_raw > 7.0f32, 7.0f32, u_raw); - let u = select(u_hi < (0.0f32 - 7.0f32), 0.0f32 - 7.0f32, u_hi); - // gate side: g · sigmoid(1.702·g) - let sig = 1.0f32 / (1.0f32 + exp(0.0f32 - 1.702f32 * g)); - // up side carries a +1 bias before the multiply. - let act = g * sig * (u + 1.0f32); - store(out[idx], act.cast::()); -} - -/// New-syntax correctness for the fused gate-activation kernels (Grid3D, f32 -/// internal). Oracles mirror the kernels exactly on dtype-rounded inputs. -pub mod kernel_tests { - use metaltile::{test::*, test_kernel}; - - use super::{mt_fused_gate_clipped_swiglu, mt_fused_gate_gelu}; - use crate::utils::{pack_f32, unpack_f32}; - - fn inputs(n: usize, dt: DType) -> (Vec, Vec, Vec, Vec) { - // Range spans beyond +/-7 so clipped_swiglu's clamp is exercised. - let gate: Vec = (0..n).map(|i| (i % 17) as f32 - 8.0).collect(); - let up: Vec = (0..n).map(|i| (i % 13) as f32 - 6.0).collect(); - let g = unpack_f32(&pack_f32(&gate, dt), dt); - let u = unpack_f32(&pack_f32(&up, dt), dt); - (gate, up, g, u) - } - - fn build( - kernel: metaltile::core::ir::Kernel, - gate: &[f32], - up: &[f32], - expected: &[f32], - dt: DType, - ) -> TestSetup { - TestSetup::new(kernel) - .input(TestBuffer::from_vec("gate", pack_f32(gate, dt), dt)) - .input(TestBuffer::from_vec("up", pack_f32(up, dt), dt)) - .input(TestBuffer::zeros("out", gate.len(), dt)) - .expect(TestBuffer::from_vec("out", pack_f32(expected, dt), dt)) - .grid_1d(gate.len(), 256) - } - - #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 5e-2])] - fn test_fused_gate_gelu(dt: DType) -> TestSetup { - let (gate, up, g, u) = inputs(512, dt); - const C: f32 = 0.797_884_6; // sqrt(2/pi) - let expected: Vec = g - .iter() - .zip(&u) - .map(|(&g, &u)| 0.5 * g * (1.0 + (C * (g + 0.044715 * g * g * g)).tanh()) * u) - .collect(); - build(mt_fused_gate_gelu::kernel_ir_for(dt), &gate, &up, &expected, dt) - } - - #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-2, 5e-2])] - fn test_fused_gate_clipped_swiglu(dt: DType) -> TestSetup { - let (gate, up, g, u) = inputs(512, dt); - let expected: Vec = g - .iter() - .zip(&u) - .map(|(&g, &u)| { - let g = g.clamp(-7.0, 7.0); - let u = u.clamp(-7.0, 7.0); - g * (1.0 / (1.0 + (-1.702 * g).exp())) * (u + 1.0) - }) - .collect(); - build(mt_fused_gate_clipped_swiglu::kernel_ir_for(dt), &gate, &up, &expected, dt) - } -} - -/// New-syntax benchmarks for the fused gate-activation kernels. -pub mod kernel_benches { - use metaltile::{bench, test::*}; - - use super::{mt_fused_gate_clipped_swiglu, mt_fused_gate_gelu}; - - fn fb(kernel: metaltile::core::ir::Kernel, dt: DType) -> BenchSetup { - let n = 64 * 1024 * 1024usize; - BenchSetup::new(kernel) - .buffer(BenchBuffer::random("gate", n, dt)) - .buffer(BenchBuffer::random("up", n, dt)) - .buffer(BenchBuffer::zeros("out", n, dt).output()) - .grid_1d(n, 256) - .bytes_moved((3 * n * dt.size_bytes()) as u64) - } - - #[bench(dtypes = [f32, f16, bf16])] - fn bench_gelu(dt: DType) -> BenchSetup { fb(mt_fused_gate_gelu::kernel_ir_for(dt), dt) } - - #[bench(dtypes = [f32, f16, bf16])] - fn bench_clipped(dt: DType) -> BenchSetup { - fb(mt_fused_gate_clipped_swiglu::kernel_ir_for(dt), dt) - } -} diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs index 5da83380..eaffe58a 100644 --- a/crates/metaltile-std/src/mlx/mod.rs +++ b/crates/metaltile-std/src/mlx/mod.rs @@ -21,12 +21,10 @@ pub mod block_scaled_moe; pub mod block_scaled_qmm; pub mod block_scaled_qmm_mpp; pub mod block_scaled_qmm_nax; -pub mod clamped_swiglu; pub mod fft; pub mod fp_quantized; pub mod fp_quantized_mma; pub mod fp_quantized_nax; -pub mod fused_gate_activation; pub mod gemv; pub mod gemv_masked; pub mod quantized; @@ -37,9 +35,7 @@ pub mod quantized_nax; pub mod quantized_nax_int8; pub mod scaled_dot_product_attention; pub mod sdpa_vector; -pub mod sgload_smoke; pub mod steel; -pub mod swiglu; // `conv.rs` and `shared.rs` are placeholder/stale stubs left over from // the old `metaltile-bench` crate. They reference `crate::runner` which diff --git a/crates/metaltile-std/src/mlx/swiglu.rs b/crates/metaltile-std/src/mlx/swiglu.rs deleted file mode 100644 index b7fd23b4..00000000 --- a/crates/metaltile-std/src/mlx/swiglu.rs +++ /dev/null @@ -1,101 +0,0 @@ -//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric -//! SPDX-License-Identifier: Apache-2.0 -//! SwiGLU activation — `silu(gate) * up`. -//! -//! Fused element-wise activation used in every modern transformer MLP -//! (Llama 4, Qwen3 dense + MoE, Gemma, Mistral families): given two -//! equally-sized inputs `gate` and `up` (the two halves of an MLP's -//! `w_gate · x` and `w_up · x` outputs), produce -//! -//! ```text -//! out[i] = silu(gate[i]) * up[i] -//! = (gate[i] * sigmoid(gate[i])) * up[i] -//! ``` -//! -//! Existing baseline: two separate kernel launches — one applies -//! `silu(gate)` elementwise (`mt_silu` in `unary.rs`), the second -//! multiplies by `up` (`mt_binary` mul). Each load+store cycles the -//! intermediate `silu(gate)` value through device memory. -//! -//! Fusion saves one full-tensor RMW: the intermediate value stays in -//! registers, halving global memory traffic on the activation path. -//! At Qwen3-MoE expert intermediate=768 × prefill 512 tokens = -//! ~400KB per layer per expert; across 48 layers × 8 active experts -//! the saved bandwidth adds up. -//! -//! MLX reference: `mx.fast.swiglu` lives in -//! `mlx/mlx/backend/metal/kernels/fast.metal` as a single launch with -//! `silu(g) * u` in the body. We mirror that pattern. -//! -//! ## Cross-kernel calling -//! -//! `mt_swiglu` calls `mt_silu` via the DSL cross-kernel call syntax -//! (just the kernel name). `KernelInlinePass` splices the silu body -//! inline before MSL emission — no extra memory round-trip, same code -//! quality as a manual inline, with a clear compositional structure -//! that future fusion passes can reason about. -//! -//! Type-efficiency: `g` and `u` are loaded and cast to f32 before the -//! call. `KernelInlinePass` replaces `mt_silu`'s input-param load with -//! the actual f32 arg, so all arithmetic stays in f32 regardless of T. -//! No T→f32→T precision loss in the silu path. - -use metaltile::kernel; - -#[kernel] -pub fn mt_swiglu(gate: Tensor, up: Tensor, out: Tensor) { - let idx = tid; - let g = load(gate[idx]).cast::(); - let u = load(up[idx]).cast::(); - // Cross-kernel call: KernelInlinePass splices mt_silu's scalar body - // here. mt_silu's input-param load is replaced by g (already f32), - // so silu runs in f32. Future fusion passes can identify the - // (silu, mul) → swiglu composition pattern from this call site. - let s = mt_silu(g); - store(out[idx], (s * u).cast::()); -} - -/// New-syntax correctness for `mt_swiglu` (`silu(gate) * up`, computed in f32). -pub mod kernel_tests { - use metaltile::{test::*, test_kernel}; - - use super::mt_swiglu; - use crate::utils::{pack_f32, unpack_f32}; - - fn setup(n: usize, dt: DType) -> TestSetup { - let gate: Vec = (0..n).map(|i| (i % 17) as f32 * 0.35 - 3.0).collect(); - let up: Vec = (0..n).map(|i| (i % 13) as f32 * 0.2 - 1.0).collect(); - let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); - let u_dt = unpack_f32(&pack_f32(&up, dt), dt); - let expected: Vec = - g_dt.iter().zip(&u_dt).map(|(&g, &u)| (g / (1.0 + (-g).exp())) * u).collect(); - TestSetup::new(mt_swiglu::kernel_ir_for(dt)) - .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt)) - .input(TestBuffer::from_vec("up", pack_f32(&up, dt), dt)) - .input(TestBuffer::zeros("out", n, dt)) - .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) - .grid_1d(n, 256) - } - - #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] - fn test_mt_swiglu(dt: DType) -> TestSetup { setup(1024, dt) } -} - -/// New-syntax benchmark for `mt_swiglu` (vs MLX `mx.fast.swiglu`). -pub mod kernel_benches { - use metaltile::{bench, test::*}; - - use super::mt_swiglu; - - #[bench(dtypes = [f32, f16, bf16])] - fn bench_swiglu(dt: DType) -> BenchSetup { - // `idx = tid` (global) — keep within a comfortable single-grid size. - let n = 1024 * 1024usize; - BenchSetup::new(mt_swiglu::kernel_ir_for(dt)) - .buffer(BenchBuffer::random("gate", n, dt)) - .buffer(BenchBuffer::random("up", n, dt)) - .buffer(BenchBuffer::zeros("out", n, dt).output()) - .grid_1d(n, 256) - .bytes_moved((3 * n * dt.size_bytes()) as u64) - } -} diff --git a/crates/metaltile-std/src/probe/mod.rs b/crates/metaltile-std/src/probe/mod.rs index f5c25108..e779ec30 100644 --- a/crates/metaltile-std/src/probe/mod.rs +++ b/crates/metaltile-std/src/probe/mod.rs @@ -2,3 +2,4 @@ //! SPDX-License-Identifier: Apache-2.0 pub mod mma_layout_probe; pub mod mpp_matmul_smoke; +pub mod simdgroup_load_probe; diff --git a/crates/metaltile-std/src/mlx/sgload_smoke.rs b/crates/metaltile-std/src/probe/simdgroup_load_probe.rs similarity index 92% rename from crates/metaltile-std/src/mlx/sgload_smoke.rs rename to crates/metaltile-std/src/probe/simdgroup_load_probe.rs index 78dec11d..bd230fe3 100644 --- a/crates/metaltile-std/src/mlx/sgload_smoke.rs +++ b/crates/metaltile-std/src/probe/simdgroup_load_probe.rs @@ -36,7 +36,7 @@ //! Dispatch: grid `[1, 1, 1]`, tpg `[32, 1, 1]` (one simdgroup). //! //! Sample MSL the codegen produces (look for these in -//! `cargo run -p metaltile-cli -- inspect mt_sgload_smoke`): +//! `cargo run -p metaltile-cli -- inspect mt_simdgroup_load_probe`): //! //! ```text //! threadgroup T tg_tile[64]; @@ -58,7 +58,7 @@ use metaltile::kernel; /// - `dst`: `[64]` flat row-major 8×8 destination, written from /// the fragment in A/C lane convention #[kernel] -pub fn mt_sgload_smoke(src: Tensor, mut dst: Tensor) { +pub fn mt_simdgroup_load_probe(src: Tensor, mut dst: Tensor) { let lane = simd_lane; // A/C lane → frag-element mapping. Matches the probe kernel + // every MMA call site in `quantized.rs`. `fn0` / `fn1` are the @@ -108,14 +108,14 @@ pub fn mt_sgload_smoke(src: Tensor, mut dst: Tensor) { pub mod kernel_tests { use metaltile::{test::*, test_kernel}; - use super::mt_sgload_smoke; + use super::mt_simdgroup_load_probe; use crate::utils::{pack_f32, unpack_f32}; #[test_kernel(dtypes = [f32, f16], tol = 0.0)] - fn test_sgload_smoke(dt: DType) -> TestSetup { + fn test_simdgroup_load_probe(dt: DType) -> TestSetup { // Distinct, dtype-representable values so a dropped/garbled lane shows. let src_f: Vec = (0..64).map(|i| (i as f32 - 32.0) * 0.5).collect(); - TestSetup::new(mt_sgload_smoke::kernel_ir_for(dt)) + TestSetup::new(mt_simdgroup_load_probe::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .input(TestBuffer::from_vec("src", pack_f32(&src_f, dt), dt)) .input(TestBuffer::zeros("dst", 64, dt)) @@ -129,11 +129,11 @@ pub mod kernel_tests { pub mod kernel_benches { use metaltile::{bench, test::*}; - use super::mt_sgload_smoke; + use super::mt_simdgroup_load_probe; #[bench(dtypes = [f32, f16])] - fn bench_sgload_smoke(dt: DType) -> BenchSetup { - BenchSetup::new(mt_sgload_smoke::kernel_ir_for(dt)) + fn bench_simdgroup_load_probe(dt: DType) -> BenchSetup { + BenchSetup::new(mt_simdgroup_load_probe::kernel_ir_for(dt)) .mode(KernelMode::Reduction) .buffer(BenchBuffer::random("src", 64, dt)) .buffer(BenchBuffer::zeros("dst", 64, dt).output()) diff --git a/crates/metaltile-std/tests/sgload_smoke_gpu.rs b/crates/metaltile-std/tests/simdgroup_load_probe_gpu.rs similarity index 88% rename from crates/metaltile-std/tests/sgload_smoke_gpu.rs rename to crates/metaltile-std/tests/simdgroup_load_probe_gpu.rs index a2f81d59..f3968e86 100644 --- a/crates/metaltile-std/tests/sgload_smoke_gpu.rs +++ b/crates/metaltile-std/tests/simdgroup_load_probe_gpu.rs @@ -2,7 +2,7 @@ //! SPDX-License-Identifier: Apache-2.0 //! GPU smoke test for the `Op::SimdgroupLoad` HW intrinsic. //! -//! The kernel under test (`mt_sgload_smoke` in `mlx::sgload_smoke`) +//! The kernel under test (`mt_simdgroup_load_probe` in `probe::simdgroup_load_probe`) //! stages a flat-64 `src` into TG memory, runs **one** MSL //! `simdgroup_load(...)` instruction to land it into a //! `simdgroup_matrix` fragment, then scatters the fragment @@ -35,14 +35,14 @@ use metaltile::{ codegen::msl::{MslConfig, MslGenerator}, core::dtype::DType, }; -use metaltile_std::mlx::sgload_smoke::mt_sgload_smoke; +use metaltile_std::probe::simdgroup_load_probe::mt_simdgroup_load_probe; /// Sanity-check that the MSL the codegen produces actually contains /// a `simdgroup_load(...)` call — this is the whole point of having /// a smoke kernel for the primitive. #[test] -fn mt_sgload_smoke_emits_simdgroup_load_instruction_f32() { - let kernel = mt_sgload_smoke::kernel_ir_for(DType::F32); +fn mt_simdgroup_load_probe_emits_simdgroup_load_instruction_f32() { + let kernel = mt_simdgroup_load_probe::kernel_ir_for(DType::F32); let msl_gen = MslGenerator::new(MslConfig::default()); let msl = msl_gen.generate(&kernel).expect("MSL emit should succeed"); @@ -77,8 +77,8 @@ fn mt_sgload_smoke_emits_simdgroup_load_instruction_f32() { /// dtype-specialisation path of `kernel_ir_for` doesn't accidentally /// drop the primitive. #[test] -fn mt_sgload_smoke_emits_simdgroup_load_instruction_f16() { - let kernel = mt_sgload_smoke::kernel_ir_for(DType::F16); +fn mt_simdgroup_load_probe_emits_simdgroup_load_instruction_f16() { + let kernel = mt_simdgroup_load_probe::kernel_ir_for(DType::F16); let msl_gen = MslGenerator::new(MslConfig::default()); let msl = msl_gen.generate(&kernel).expect("MSL emit should succeed"); @@ -111,7 +111,7 @@ fn run_round_trip(dtype: Dt) { buffers.insert("dst".into(), vec![0u8; 64 * dtype.bytes()]); let ctx = Context::new().expect("Context::new on macOS"); - let kernel = mt_sgload_smoke::kernel_ir_for(dtype.to_dtype()); + let kernel = mt_simdgroup_load_probe::kernel_ir_for(dtype.to_dtype()); let result = ctx .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [1, 1, 1], [32, 1, 1]) @@ -133,7 +133,7 @@ fn run_round_trip(dtype: Dt) { } #[test] -fn mt_sgload_smoke_round_trips_8x8_tile_f32() { run_round_trip(Dt::F32); } +fn mt_simdgroup_load_probe_round_trips_8x8_tile_f32() { run_round_trip(Dt::F32); } #[test] -fn mt_sgload_smoke_round_trips_8x8_tile_f16() { run_round_trip(Dt::F16); } +fn mt_simdgroup_load_probe_round_trips_8x8_tile_f16() { run_round_trip(Dt::F16); } diff --git a/docs/specs/KERNEL_AUDIT.md b/docs/specs/KERNEL_AUDIT.md index ec599edf..def4be87 100644 --- a/docs/specs/KERNEL_AUDIT.md +++ b/docs/specs/KERNEL_AUDIT.md @@ -68,7 +68,7 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | copy (strided / general) | ✓ | ✓ | ✓ | `kernels/ops/strided.rs` → `mt_strided_copy` (2-D padded) + `mt_strided_copy_nd` (arbitrary-rank). Each output element unravels its flat index against a runtime `shape` array and gathers `src[Σ coord_d · strides[d]]`. Covers padded copies, transposes, broadcasts (stride 0), and dilated slices in one kernel. | | ternary (select) | ✓ | ✓ | ✓ | `kernels/ops/ternary.rs` → `mt_select`. | | unary (exp/log/sqrt/rsqrt/abs/silu/etc.) | ✓ | ✓ | ✓ | `kernels/ops/unary.rs` → 7+ kernels including `mt_silu`. Plus `mt_scalar_fma_chain8` (fused 8-way scalar FMA) and `mt_add_rms_norm` (residual-add + RMSNorm fusion, Reduction mode). | -| swiglu (`silu(gate)·up` fused MLP activation) | ✗ | ✗ | ✓ | `mlx/swiglu.rs` → `mt_swiglu`. Standard modern-transformer MLP activation (Llama 4, Qwen3 dense + MoE, Gemma, Mistral). | +| swiglu (`silu(gate)·up` fused MLP activation) | ✗ | ✗ | ✓ | `kernels/ops/gated_activation.rs` → `mt_swiglu` (+ `mt_clamped_swiglu`, per-layer activation-clip limit). Standard modern-transformer MLP activation (Llama 4, Qwen3 dense + MoE, Gemma, Mistral). | | random (key hash → u32) | ✓ | ✓ | ✓ | `kernels/ops/random.rs` → `mt_random_hash`. | | reduce (sum/prod/max/min — all + row + col) | ✓ | ✓ | ✓ | `kernels/ops/reduce.rs` covers `all_reduce*`, `row_reduce*`, `col_reduce*`, and `seg_reduce*` (Grid3D one-thread-per-segment, contiguous fixed-length runs) — all four ops for each shape. | | sort | ✓ | ✓ | ✓ | `kernels/sampling/sort.rs` → `mt_sort` (single-block bitonic) + `mt_merge` (multi-block merge) + `mt_sort_segmented` (per-row bitonic for `[batch, n]` matrices, `n ≤ 1024`, one TG per row). | @@ -130,7 +130,7 @@ Local verification of NAX kernels is the developer's responsibility on M4+ hardw | ssm_step (Mamba 2 SSD single-token decode) | ✗ | ✓ | ✓ | `ffai/ssm.rs` → `ssm_step`, `mt_ssm_step` (scalar `A`). 2D-`A_log` variant `ssm_step_a2d` (Jamba): per-(channel, state) `A_log`, moves Mamba 1 selective scan onto the GPU (previously host-side). | | 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`. | +| fused_gate_activation (silu/gelu × up gate) | ✗ | ✓ | ✓ | `kernels/ops/gated_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 in the same file as `mt_swiglu`. | | 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). | diff --git a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md index 27e635db..8f17ebe9 100644 --- a/docs/specs/KERNEL_CONSOLIDATION_PLAN.md +++ b/docs/specs/KERNEL_CONSOLIDATION_PLAN.md @@ -45,7 +45,7 @@ the FFAI emit path is unaffected). crates/metaltile-std/src/kernels/ ops/ ✅ DONE — elementwise/core primitives: binary · unary · ternary · copy · arange · random · reduce · arg_reduce · scan · indexing · gather/scatter · hadamard · - fence · clamp · logsumexp · vector_add · axpy · strided + fence · clamp · logsumexp · vector_add · axpy · strided · gated_activation gemm/ gemm · gemv(_masked) · batched-projection (qkv / 4) · patch_embed · steel/gemm sdpa/ ALL attention: bidirectional(+relpos/windowed/conformer) · decode(+d64..d512/ 2pass/batched/sink) · multi(+d256/tree-mask) · prefill_mma · flash_quantized · From 0c1a070660d88689d3e9a4daab1116583a78bbba Mon Sep 17 00:00:00 2001 From: Eric Kryski <599019+ekryski@users.noreply.github.com> Date: Fri, 12 Jun 2026 18:33:03 -0600 Subject: [PATCH 3/3] refactor(probe): rename mpp_matmul_smoke -> mpp_matmul_probe; document the probe convention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes the probe/ module consistent: all three validation kernels now use the _probe suffix (mma_layout_probe, simdgroup_load_probe, mpp_matmul_probe). - probe/mpp_matmul_smoke.rs -> mpp_matmul_probe.rs; mt_mpp_matmul_smoke -> mt_mpp_matmul_probe; integration test tests/mpp_matmul_smoke.rs renamed to match. - STYLE_GUIDE §1: document the convention — a kernel that validates a codegen path / HW intrinsic is a 'probe' (mt__probe, in probe/, outside kernels//). '_smoke' is reserved for *test* code (backend bring-up tests, codegen-sanity #[test]s), which stays test-side. - probe/mod.rs: add a module doc stating the same. No other kernels use the _smoke naming (the remaining *_smoke are legit tests: cuda/hip/vulkan backend bring-up + codegen-sanity test fns — left as-is). cargo build + clippy clean; probe integration + codegen unit tests pass. --- crates/metaltile-std/src/probe/mod.rs | 7 ++++++- .../{mpp_matmul_smoke.rs => mpp_matmul_probe.rs} | 8 ++++---- .../tests/moe_gather_qmm_mpp_bm64_correctness.rs | 2 +- .../{mpp_matmul_smoke.rs => mpp_matmul_probe.rs} | 14 +++++++------- docs/STYLE_GUIDE.md | 12 ++++++++++++ 5 files changed, 30 insertions(+), 13 deletions(-) rename crates/metaltile-std/src/probe/{mpp_matmul_smoke.rs => mpp_matmul_probe.rs} (96%) rename crates/metaltile-std/tests/{mpp_matmul_smoke.rs => mpp_matmul_probe.rs} (93%) diff --git a/crates/metaltile-std/src/probe/mod.rs b/crates/metaltile-std/src/probe/mod.rs index e779ec30..10e350dd 100644 --- a/crates/metaltile-std/src/probe/mod.rs +++ b/crates/metaltile-std/src/probe/mod.rs @@ -1,5 +1,10 @@ //! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric //! SPDX-License-Identifier: Apache-2.0 +//! Probe kernels — `#[kernel]`s whose purpose is to validate a codegen path or +//! HW intrinsic end-to-end (not production work). Named `mt__probe`; this +//! module sits at the crate root, outside the `kernels//` tree (see the +//! style guide §1). Distinct from `*_smoke` *tests*, which stay test-side. + pub mod mma_layout_probe; -pub mod mpp_matmul_smoke; +pub mod mpp_matmul_probe; pub mod simdgroup_load_probe; diff --git a/crates/metaltile-std/src/probe/mpp_matmul_smoke.rs b/crates/metaltile-std/src/probe/mpp_matmul_probe.rs similarity index 96% rename from crates/metaltile-std/src/probe/mpp_matmul_smoke.rs rename to crates/metaltile-std/src/probe/mpp_matmul_probe.rs index 3cd2e103..02b4684c 100644 --- a/crates/metaltile-std/src/probe/mpp_matmul_smoke.rs +++ b/crates/metaltile-std/src/probe/mpp_matmul_probe.rs @@ -30,14 +30,14 @@ use metaltile::core::{ }; use rustc_hash::FxHashMap; -/// Build the [`Kernel`] IR for `mt_mpp_matmul_smoke`. +/// Build the [`Kernel`] IR for `mt_mpp_matmul_probe`. /// /// Expresses the matmul via 6 primitive `CoopTile*` ops: /// `CoopTileSetup → CoopTileZero → CoopTileLoadA → CoopTileLoadB → CoopTileRun → CoopTileStoreC`. /// /// Dispatch geometry: 1 threadgroup × 32 threads = one simdgroup. pub fn kernel_ir() -> Kernel { - let mut k = Kernel::new("mt_mpp_matmul_smoke"); + let mut k = Kernel::new("mt_mpp_matmul_probe"); k.mode = KernelMode::Elementwise; // Params: A [M=16, K=32] fp16, B [K=32, N=16] fp16, C [M=16, N=16] fp32. @@ -149,7 +149,7 @@ mod tests { #[test] fn kernel_ir_constructs_and_has_three_params() { let k = kernel_ir(); - assert_eq!(k.name, "mt_mpp_matmul_smoke"); + assert_eq!(k.name, "mt_mpp_matmul_probe"); assert_eq!(k.params.len(), 3); assert_eq!(k.params[0].name, "A"); assert_eq!(k.params[1].name, "B"); @@ -171,7 +171,7 @@ mod tests { "MPP include missing from generated MSL:\n{msl}" ); assert!(msl.contains("mpp::tensor_ops::matmul2d_descriptor")); - assert!(msl.contains("kernel void mt_mpp_matmul_smoke")); + assert!(msl.contains("kernel void mt_mpp_matmul_probe")); } /// Developer aid — `cargo test -p metaltile-std --lib -- dump_generated_msl --nocapture` diff --git a/crates/metaltile-std/tests/moe_gather_qmm_mpp_bm64_correctness.rs b/crates/metaltile-std/tests/moe_gather_qmm_mpp_bm64_correctness.rs index e8d895d5..975bc205 100644 --- a/crates/metaltile-std/tests/moe_gather_qmm_mpp_bm64_correctness.rs +++ b/crates/metaltile-std/tests/moe_gather_qmm_mpp_bm64_correctness.rs @@ -54,7 +54,7 @@ fn pack_int4_row(weights: &[u32]) -> Vec { /// the kernel hits its pre-Metal-4 stub branch which writes zeros — the /// cosine assertion then fails with `cos = 0.0`. Skip rather than fail so /// CI's hosted Mac runner stays green; M5 Max + dev hardware still cover it. -/// Mirrors the gate added to `mpp_matmul_smoke.rs` + `qmm_mpp_correctness.rs`. +/// Mirrors the gate added to `mpp_matmul_probe.rs` + `qmm_mpp_correctness.rs`. fn skip_unless_apple10(test_name: &str) -> Option { let ctx = Context::new().expect("Context::new"); let family = ctx.chip_family(); diff --git a/crates/metaltile-std/tests/mpp_matmul_smoke.rs b/crates/metaltile-std/tests/mpp_matmul_probe.rs similarity index 93% rename from crates/metaltile-std/tests/mpp_matmul_smoke.rs rename to crates/metaltile-std/tests/mpp_matmul_probe.rs index 8fe32f49..17648da6 100644 --- a/crates/metaltile-std/tests/mpp_matmul_smoke.rs +++ b/crates/metaltile-std/tests/mpp_matmul_probe.rs @@ -2,7 +2,7 @@ //! SPDX-License-Identifier: Apache-2.0 //! GPU correctness oracle for the MPP `matmul2d` smoke kernel. //! -//! Dispatches `mt_mpp_matmul_smoke` (single threadgroup × single simdgroup, +//! Dispatches `mt_mpp_matmul_probe` (single threadgroup × single simdgroup, //! 16×32 fp16 @ 32×16 fp16 → 16×16 fp32) and validates against a naïve //! triple-loop CPU oracle. //! @@ -14,7 +14,7 @@ //! such toolchains this test will fail the correctness check, which is //! the intended signal. //! -//! Run: `cargo test --release -p metaltile-std --test mpp_matmul_smoke -- --nocapture` +//! Run: `cargo test --release -p metaltile-std --test mpp_matmul_probe -- --nocapture` #![cfg(target_os = "macos")] @@ -24,7 +24,7 @@ mod common; use common::gpu_lock; use metaltile::Context; -use metaltile_std::probe::mpp_matmul_smoke; +use metaltile_std::probe::mpp_matmul_probe; const M: usize = 16; const N: usize = 16; @@ -60,7 +60,7 @@ fn cpu_matmul_ref(a: &[f32], b: &[f32]) -> Vec { } #[test] -fn mpp_matmul_smoke_matches_cpu_reference() { +fn mpp_matmul_probe_matches_cpu_reference() { let _lock = gpu_lock(); let ctx = Context::new().expect("Context::new"); @@ -72,7 +72,7 @@ fn mpp_matmul_smoke_matches_cpu_reference() { // so coverage runners stay green; M5 Max + dev hardware still cover it. let family = ctx.chip_family(); if family.is_none_or(|lvl| lvl < 10) { - eprintln!("skip: mpp_matmul_smoke needs Apple10+ GPU (got chip_family={family:?})"); + eprintln!("skip: mpp_matmul_probe needs Apple10+ GPU (got chip_family={family:?})"); return; } @@ -99,12 +99,12 @@ fn mpp_matmul_smoke_matches_cpu_reference() { buffers.insert("B".into(), pack_f16_bytes(&b_q)); buffers.insert("C".into(), vec![0u8; M * N * 4]); - let kernel = mpp_matmul_smoke::kernel_ir(); + let kernel = mpp_matmul_probe::kernel_ir(); // 1 threadgroup × (32 threads = one simdgroup). let result = ctx .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [1, 1, 1], [32, 1, 1]) - .expect("dispatch mpp_matmul_smoke"); + .expect("dispatch mpp_matmul_probe"); let got = unpack_f32_bytes(result.outputs.get("C").expect("C buffer")); assert_eq!(got.len(), M * N); diff --git a/docs/STYLE_GUIDE.md b/docs/STYLE_GUIDE.md index 37105972..95702217 100644 --- a/docs/STYLE_GUIDE.md +++ b/docs/STYLE_GUIDE.md @@ -51,6 +51,18 @@ family — it folds into the op's file as a format axis (see §5). See family list and the migration state (the crate is mid-migration from the legacy `mlx/` + `ffai/` split). +### Probe kernels — `probe/`, not a `kernels/` family + +A kernel whose *purpose* is to validate a codegen path or HW intrinsic +end-to-end — not to do production work — is a **probe**. Probes live in +`crates/metaltile-std/src/probe/` (a crate-root module, like `utils`, **outside** +the `kernels//` tree) and are named **`mt__probe`**: +`mt_simdgroup_load_probe` (the `Op::SimdgroupLoad` round-trip), +`mt_mpp_matmul_probe`, `mt_mma_probe_*` (MMA layout). Use `_probe`, not `_smoke` +— "smoke" is reserved for *test* code (backend bring-up tests like +`tests/cuda_smoke.rs`, codegen-sanity `#[test] fn …_smoke()`), which stays +test-side and is not a kernel. + --- ## 2. File skeleton