feat(ffai,mlx): Step-3 / DeepSeek-V3 decoder kernels — clamped SwiGLU, head-wise attn gate, sigmoid+bias router

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

@@ -0,0 +1,119 @@
+//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric
+//! SPDX-License-Identifier: Apache-2.0
+//! Head-wise attention output gate — `attn_out[h, d] *= sigmoid(g[h])`.
+//!
+//! Applied between the SDPA output and `o_proj` for checkpoints that
+//! ship a per-head scalar gate (`use_head_wise_attn_gate` in upstream
+//! configs). StepFun's Step-3.7-Flash + Step-3.5-Flash are the current
+//! consumers; other gated-attention variants in the literature fit the
+//! same shape.
+//!
+//! The pattern in PyTorch reference impls is:
+//!
+//! ```python
+//!   g = self.g_proj(x)              # [batch, n_heads] — per-head scalar
+//!   attn = attn * sigmoid(g)[..., None]
+//! ```
+//!
+//! `g_proj` is a `Linear(hidden, n_heads)` projection of the layer's
+//! input. The fused kernel below does the `sigmoid(g) × attn` step in
+//! one pass; computing `g` itself stays as a separate gemv (it's just
+//! `hidden → n_heads` and reuses the existing dense matmul path).
+//!
+//! ## ABI
+//!
+//! ```text
+//!   attn    [n_heads, head_dim] T       — SDPA output, modified in place
+//!   gate    [n_heads]            T      — pre-sigmoid logits (g_proj(x))
+//!   out     [n_heads, head_dim] T       — gated SDPA output
+//!   head_dim u32                        — constexpr, distinguishes
+//!                                          per-head row size
+//! ```
+//!
+//! Grid is 1D elementwise over the flat `[n_heads * head_dim]` buffer;
+//! each thread derives its owning head as `idx / head_dim` and reuses
+//! that head's single sigmoid. Caller drives `grid_1d(n, 256)`.
+
+use metaltile::kernel;
+
+// Bare `#[kernel]` — the legacy `bench(...)` registration doesn't
+// fit this 3-input + 1-constexpr shape; the new declarative `#[bench]`
+// on `kernel_benches::bench_attn_head_gate` below handles registration.
+#[kernel]
+pub fn ffai_attn_head_gate<T>(
+    attn: Tensor<T>,
+    gate: Tensor<T>,
+    out: Tensor<T>,
+    #[constexpr] head_dim: u32,
+) {
+    // 1D elementwise over the flat [n_heads * head_dim] attn buffer; the
+    // owning head is `idx / head_dim` so the per-head gate is shared
+    // across its head_dim lanes.
+    let idx = tid;
+    let h = idx / head_dim;
+    // sigmoid(g[h]) is the per-head scalar. Computed in f32 so the
+    // small-magnitude tail (`g ≈ -10`) doesn't underflow on bf16.
+    // Free-function `exp` + f32 literals so the binding isn't elided.
+    let g_raw = load(gate[h]).cast::<f32>();
+    let s = 1.0f32 / (1.0f32 + exp(0.0f32 - g_raw));
+    let a = load(attn[idx]).cast::<f32>();
+    store(out[idx], (a * s).cast::<T>());
+}
+
+pub mod kernel_tests {
+    use metaltile::{test::*, test_kernel};
+
+    use super::ffai_attn_head_gate;
+    use crate::utils::{pack_f32, unpack_f32};
+
+    fn setup(n_heads: usize, head_dim: usize, dt: DType) -> TestSetup {
+        let n = n_heads * head_dim;
+        let attn: Vec<f32> = (0..n).map(|i| (i % 23) as f32 * 0.1 - 1.0).collect();
+        let gate: Vec<f32> = (0..n_heads).map(|h| (h % 7) as f32 * 0.5 - 1.5).collect();
+        let a_dt = unpack_f32(&pack_f32(&attn, dt), dt);
+        let g_dt = unpack_f32(&pack_f32(&gate, dt), dt);
+        let mut expected: Vec<f32> = Vec::with_capacity(n);
+        for h in 0..n_heads {
+            let s = 1.0_f32 / (1.0 + (-g_dt[h]).exp());
+            for d in 0..head_dim {
+                expected.push(a_dt[h * head_dim + d] * s);
+            }
+        }
+        TestSetup::new(ffai_attn_head_gate::kernel_ir_for(dt))
+            .input(TestBuffer::from_vec("attn", pack_f32(&attn, dt), dt))
+            .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt))
+            .input(TestBuffer::zeros("out", n, dt))
+            .constexpr("head_dim", head_dim as u32)
+            .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_ffai_attn_head_gate_step3_full(dt: DType) -> TestSetup { setup(64, 128, dt) }
+
+    /// SWA-layer shape variant: Step-3 uses 96 q-heads on
+    /// sliding-attention layers. Validates the kernel doesn't bake the
+    /// full-attn 64-head count in.
+    #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])]
+    fn test_ffai_attn_head_gate_step3_swa(dt: DType) -> TestSetup { setup(96, 128, dt) }
+}
+
+pub mod kernel_benches {
+    use metaltile::{bench, test::*};
+
+    use super::ffai_attn_head_gate;
+
+    #[bench(name = "ffai/attn_head_gate", dtypes = [f32, f16, bf16])]
+    fn bench_attn_head_gate(dt: DType) -> BenchSetup {
+        // Step-3 full-attn shape.
+        let (n_heads, head_dim) = (64usize, 128usize);
+        let n = n_heads * head_dim;
+        BenchSetup::new(ffai_attn_head_gate::kernel_ir_for(dt))
+            .buffer(BenchBuffer::random("attn", n, dt))
+            .buffer(BenchBuffer::random("gate", n_heads, dt))
+            .buffer(BenchBuffer::zeros("out", n, dt).output())
+            .constexpr("head_dim", head_dim as u32)
+            .grid_1d(n, 256)
+            .bytes_moved(((2 * n + n_heads) * dt.size_bytes()) as u64)
+    }
+}

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

@@ -17,6 +17,7 @@
 //! `mlx/`.
 
 pub mod arg_reduce;
+pub mod attn_head_gate;
 pub mod audio_conv1d;
 pub mod audio_conv1d_block_scaled;
 pub mod aura_dequant_rotated;
@@ -123,6 +124,7 @@ pub mod moe_mpp_bm8_block_scaled;
 pub mod moe_mpp_bm8_int8;
 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;

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

@@ -0,0 +1,112 @@
+//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric
+//! SPDX-License-Identifier: Apache-2.0
+//! MoE router — sigmoid + bias + (optional) top-k normalize + scale.
+//!
+//! The DeepSeek-V3 routing pattern, adopted by StepFun's Step-3 family:
+//!
+//! ```text
+//!   gates  = sigmoid(W_router · x) + router_bias        [n_experts]
+//!   topk   = argpartition(-gates, k)[:, :k]             [k]   (indices)
+//!   weights = gates.gather(topk)                        [k]
+//!   if norm_topk_prob:  weights = weights / weights.sum()
+//!   weights = weights * scaling_factor
+//! ```
+//!
+//! Distinct from the more common softmax-routing pattern Qwen3 / Llama /
+//! GPT-OSS use: the sigmoid + bias factorisation lets each expert score
+//! be evaluated independently (no softmax denominator coupling), and
+//! the per-expert bias is what the upstream "noisy gating with auxiliary
+//! load-balancing loss" line learns.
+//!
+//! This file ships the **score path only** — `sigmoid(logits) + bias`
+//! into a `[n_experts]` scores tensor. Top-k selection + normalize +
+//! scale are kept as a separate kernel because the same selection
+//! pipeline is shared with the (existing) softmax-routing path; only
+//! the score producer changes per family.
+//!
+//! ## ABI
+//!
+//! ```text
+//!   logits     [n_experts] f32      — pre-sigmoid router output
+//!                                       (`W_router · x`)
+//!   bias       [n_experts] f32      — per-expert routing bias
+//!   scores     [n_experts] f32      — out: `sigmoid(logits) + bias`
+//! ```
+//!
+//! Grid is 1D elementwise: one thread per expert. Caller drives
+//! `grid_1d(n_experts, 64)` — n_experts is typically 64-288 across
+//! the production checkpoints (DeepSeek-V3 = 256, Step-3 = 288).
+
+use metaltile::kernel;
+
+// Bare `#[kernel]` — non-generic, all-`Tensor<f32>` kernel; the legacy
+// `bench(...)` registration expects a generic-T shape it can't bind to
+// a no-`<T>` signature. The new declarative `#[bench]` on
+// `kernel_benches::bench_router` below handles registration directly.
+#[kernel]
+pub fn ffai_moe_router_sigmoid_bias(
+    logits: Tensor<f32>,
+    bias: Tensor<f32>,
+    mut scores: Tensor<f32>,
+) {
+    let idx = tid;
+    let l = load(logits[idx]);
+    let b = load(bias[idx]);
+    // Free-function `exp` + f32-suffixed literals: the method form
+    // `(-l).exp()` nested in a larger expression elides its binding in
+    // the DSL codegen mangler (same fix as moe_router_sqrtsoftplus).
+    let s = 1.0f32 / (1.0f32 + exp(0.0f32 - l));
+    store(scores[idx], s + b);
+}
+
+pub mod kernel_tests {
+    use metaltile::{test::*, test_kernel};
+
+    use super::ffai_moe_router_sigmoid_bias;
+    use crate::utils::{pack_f32, unpack_f32};
+
+    fn setup(n_experts: usize) -> TestSetup {
+        let dt = DType::F32;
+        let logits: Vec<f32> = (0..n_experts).map(|i| (i % 31) as f32 * 0.1 - 1.5).collect();
+        let bias: Vec<f32> = (0..n_experts).map(|i| (i % 7) as f32 * 0.05 - 0.15).collect();
+        let l_dt = unpack_f32(&pack_f32(&logits, dt), dt);
+        let b_dt = unpack_f32(&pack_f32(&bias, dt), dt);
+        let expected: Vec<f32> =
+            l_dt.iter().zip(&b_dt).map(|(&l, &b)| 1.0_f32 / (1.0 + (-l).exp()) + b).collect();
+        TestSetup::new(ffai_moe_router_sigmoid_bias::kernel_ir())
+            .input(TestBuffer::from_vec("logits", pack_f32(&logits, dt), dt))
+            .input(TestBuffer::from_vec("bias", pack_f32(&bias, dt), dt))
+            .input(TestBuffer::zeros("scores", n_experts, dt))
+            .expect(TestBuffer::from_vec("scores", pack_f32(&expected, dt), dt))
+            .grid_1d(n_experts, 64)
+    }
+
+    // tol 2e-4: sigmoid routes through `exp`, so the GPU↔CPU-oracle gap
+    // is dominated by transcendental fast-math and varies by GPU family
+    // — same rationale as the sister moe_router_sqrtsoftplus. Still far
+    // tighter than a logic error would ever land.
+    #[test_kernel(dtypes = [f32], tol = [2e-4])]
+    fn test_router_sigmoid_bias_step3(_dt: DType) -> TestSetup { setup(288) }
+
+    /// DeepSeek-V3 shape — same router pattern at a different expert count.
+    #[test_kernel(dtypes = [f32], tol = [2e-4])]
+    fn test_router_sigmoid_bias_dsv3(_dt: DType) -> TestSetup { setup(256) }
+}
+
+pub mod kernel_benches {
+    use metaltile::{bench, test::*};
+
+    use super::ffai_moe_router_sigmoid_bias;
+
+    #[bench(name = "ffai/moe_router_sigmoid_bias", dtypes = [f32])]
+    fn bench_router(_dt: DType) -> BenchSetup {
+        let dt = DType::F32;
+        let n_experts = 288usize;
+        BenchSetup::new(ffai_moe_router_sigmoid_bias::kernel_ir())
+            .buffer(BenchBuffer::random("logits", n_experts, dt))
+            .buffer(BenchBuffer::random("bias", n_experts, dt))
+            .buffer(BenchBuffer::zeros("scores", n_experts, dt).output())
+            .grid_1d(n_experts, 64)
+            .bytes_moved((3 * n_experts * dt.size_bytes()) as u64)
+    }
+}

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

@@ -0,0 +1,128 @@
+//! 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<T>(
+    gate: Tensor<T>,
+    up: Tensor<T>,
+    out: Tensor<T>,
+    #[constexpr] limit: f32,
+) {
+    let idx = tid;
+    let g = load(gate[idx]).cast::<f32>();
+    let u = load(up[idx]).cast::<f32>();
+    // 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::<T>());
+}
+
+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<f32> = (0..n).map(|i| (i % 17) as f32 * 0.35 - 3.0).collect();
+        let up: Vec<f32> = (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<f32> = 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(name = "mlx/clamped_swiglu", 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)
+    }
+}

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

@@ -25,6 +25,7 @@ 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 copy;
 pub mod fft;
 pub mod fp_quantized;