From 1fe6c08c9e7d4ddafe60e9cc477fbadb12e5625d Mon Sep 17 00:00:00 2001 From: TheTom Date: Wed, 3 Jun 2026 13:03:19 -0500 Subject: [PATCH] =?UTF-8?q?feat(ffai,mlx):=20Step-3=20/=20DeepSeek-V3=20de?= =?UTF-8?q?coder=20kernels=20=E2=80=94=20clamped=20SwiGLU,=20head-wise=20a?= =?UTF-8?q?ttn=20gate,=20sigmoid+bias=20router?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three #[kernel]-DSL ports for the architectural primitives the Step-3.5-Flash / Step-3.7-Flash decoders need that aren't already in the metaltile-std surface. Each is a self-contained file: kernel body + #[test_kernel] (GPU-vs-CPU-oracle) + #[bench]. - mt_clamped_swiglu (mlx/clamped_swiglu.rs): out[i] = clip(silu(gate[i]), max=L) * clip(up[i], -L, L) Per-layer f32 `limit` constexpr; `limit <= 0` collapses to plain SwiGLU (shipped on an equivalence test). gpt-oss + Step-3 activation clip pattern. - ffai_attn_head_gate (ffai/attn_head_gate.rs): out[h, d] = attn[h, d] * sigmoid(gate[h]) Per-head sigmoid gate applied between SDPA and o_proj. 1D elementwise; owning head = idx / head_dim. Tested at full-attn (64h) + SWA (96h). - ffai_moe_router_sigmoid_bias (ffai/moe_router_sigmoid_bias.rs): scores[e] = sigmoid(logits[e]) + bias[e] DeepSeek-V3 routing score path; the sister of moe_router_sqrtsoftplus. Top-k selection stays in the shared pipeline. Tested at 288 + 256 experts. All three use the proven DSL idioms (free-function `exp`, f32-suffixed literals, `select` over ambiguous min/max, `#[constexpr]` scalars, linear `tid`) — verified on M5 via `tile test`. Transcendental tols match the sqrtsoftplus sibling (2e-4) for cross-GPU fast-math headroom. --- .../metaltile-std/src/ffai/attn_head_gate.rs | 119 ++++++++++++++++ crates/metaltile-std/src/ffai/mod.rs | 2 + .../src/ffai/moe_router_sigmoid_bias.rs | 112 +++++++++++++++ .../metaltile-std/src/mlx/clamped_swiglu.rs | 128 ++++++++++++++++++ crates/metaltile-std/src/mlx/mod.rs | 1 + 5 files changed, 362 insertions(+) create mode 100644 crates/metaltile-std/src/ffai/attn_head_gate.rs create mode 100644 crates/metaltile-std/src/ffai/moe_router_sigmoid_bias.rs create mode 100644 crates/metaltile-std/src/mlx/clamped_swiglu.rs diff --git a/crates/metaltile-std/src/ffai/attn_head_gate.rs b/crates/metaltile-std/src/ffai/attn_head_gate.rs new file mode 100644 index 00000000..a5d9afd1 --- /dev/null +++ b/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( + attn: Tensor, + gate: Tensor, + out: Tensor, + #[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::(); + let s = 1.0f32 / (1.0f32 + exp(0.0f32 - g_raw)); + let a = load(attn[idx]).cast::(); + store(out[idx], (a * s).cast::()); +} + +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 = (0..n).map(|i| (i % 23) as f32 * 0.1 - 1.0).collect(); + let gate: Vec = (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 = 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) + } +} diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index e0b3cb05..2da88533 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/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; diff --git a/crates/metaltile-std/src/ffai/moe_router_sigmoid_bias.rs b/crates/metaltile-std/src/ffai/moe_router_sigmoid_bias.rs new file mode 100644 index 00000000..c7674aa9 --- /dev/null +++ b/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` kernel; the legacy +// `bench(...)` registration expects a generic-T shape it can't bind to +// a no-`` signature. The new declarative `#[bench]` on +// `kernel_benches::bench_router` below handles registration directly. +#[kernel] +pub fn ffai_moe_router_sigmoid_bias( + logits: Tensor, + bias: Tensor, + mut scores: Tensor, +) { + 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 = (0..n_experts).map(|i| (i % 31) as f32 * 0.1 - 1.5).collect(); + let bias: Vec = (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 = + 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) + } +} diff --git a/crates/metaltile-std/src/mlx/clamped_swiglu.rs b/crates/metaltile-std/src/mlx/clamped_swiglu.rs new file mode 100644 index 00000000..51d697de --- /dev/null +++ b/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( + 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(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) + } +} diff --git a/crates/metaltile-std/src/mlx/mod.rs b/crates/metaltile-std/src/mlx/mod.rs index 018d660f..a77ee2d0 100644 --- a/crates/metaltile-std/src/mlx/mod.rs +++ b/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;