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

[TheTom]

Summary

Three #[kernel]-DSL ports covering the architectural primitives the Step-3.5-Flash / Step-3.7-Flash decoder needs that aren't already in the metaltile-std surface. Each lands as a self-contained file: kernel body + #[test_kernel] (mod kernel_tests) + #[bench] (mod kernel_benches).

All three are validated on M5 via tile test — GPU output vs CPU oracle, across f32/f16/bf16.

Kernels

Kernel	File	Role
mt_clamped_swiglu	mlx/clamped_swiglu.rs	Drop-in for mt_swiglu with one-sided silu output clip + symmetric up clip — gpt-oss + Step-3 activation pattern
ffai_attn_head_gate	ffai/attn_head_gate.rs	out[h, d] = attn[h, d] * sigmoid(gate[h]) — head-wise sigmoid attention gate, applied between SDPA and o_proj
ffai_moe_router_sigmoid_bias	ffai/moe_router_sigmoid_bias.rs	DeepSeek-V3 routing score path: sigmoid(logits) + bias; top-k selection stays in the shared selection pipeline

1. mt_clamped_swiglu

    out[i] = clip(silu(gate[i]), max=L) * clip(up[i], -L, L)

Per-layer limit passed as an f32 constexpr. limit <= 0 collapses to plain SwiGLU — validated by test_mt_clamped_swiglu_zero_limit_equals_plain, the equivalence the per-layer dispatch ships on (most layers are limit=0; only a small subset carry non-zero limits). silu is g * sigmoid(g); clipping happens on the f32 intermediates before the cast back to T.

2. ffai_attn_head_gate

gate is [n_heads] — produced upstream by a Linear(hidden, n_heads). Sigmoid runs in f32 so the bf16 small-magnitude tail doesn't underflow. 1D elementwise over the flat [n_heads * head_dim] buffer; the owning head is idx / head_dim. Tests cover both full-attn (n_heads=64) and SWA (n_heads=96) shapes — the kernel doesn't bake the head count in.

3. ffai_moe_router_sigmoid_bias

Score-producer kernel only. The sigmoid+bias factorisation lets each expert score be evaluated independently (no softmax denominator coupling) — the sister of ffai_moe_router_sqrtsoftplus. Tested at both n_experts=288 and n_experts=256.

Why these three

The rest of the Step-3 decoder is a recombination of existing primitives — GQA SDPA, RoPE-llama-many with partial-rotary, per-head RMSNorm, MoE top-k routing, conv2d patch embedding, 2D RoPE, sdpa_bidirectional for the vision tower. The three above are the genuinely-new activation/gating/routing choices that don't fit any existing kernel.

Verification

• tile test (clamped_swiglu / attn_head_gate / router_sigmoid_bias) — all pass on M5, f32/f16/bf16
• full tile test — no regressions
• cargo clippy --all-targets --all-features -- -D warnings — clean
• cargo fmt --check — clean