Fused residual-add RMSNorm and fused RoPE + KV-append decode kernels

[inureyes]

Background

In the per-token decode loop, residual add, RMSNorm, RoPE application, and KV-cache append each run as separate MLX graph nodes, paying kernel-launch and intermediate-tensor costs every token. The repo has direct evidence this class of overhead matters: the MoE decode-gap investigation counted ~773 AsType ops per token on Metal before the CUDA dtype patch reduced that path to zero (docs/benchmark_results/moe-decode-gap-investigation.md, single-dtype-decode-astype-gb10-2026-07-10.md), and #824/#829 showed RMSNorm kernels are on the correctness- and performance-critical path.

The candidate fusions, standard in serving-oriented kernel libraries:

• FusedAddRMSNorm: compute residual' = input + residual and out = rmsnorm(residual') * (weight_bias + weight) in one pass. Stash the fp32 sum once, run a two-stage (warp then block) sum-of-squares reduction, apply rsqrt(mean + eps). A weight_bias scalar absorbs the Gemma (1 + w) weight convention (weight_bias = 1.0); weight_bias = 0.0 is the standard case.
• Fused RoPE + append: apply rotary embedding to q and k and write k/v into the cache in the same kernel (cos/sin from a precomputed cache or computed on the fly), removing separate rope ops and the append's read-modify-write.
• Gated activation: silu(x[:d]) * x[d:2d] style fused gate with 16-byte vectorized loads.

One MLX-specific constraint: MLX arrays are lazy and immutable from the graph's perspective, so an in-place residual update must be expressed as a two-output kernel (normed, new_residual) rather than true in-place mutation. mlx::core::fast::metal_kernel/cuda_kernel support multiple outputs.

Task

1. FusedAddRMSNorm kernel (Metal + CUDA JIT strings in a new src/lib/mlx-cpp/turbo/fused_norm.cpp, following the paged_attention.cpp dual-source pattern): inputs (x, residual, weight, eps, weight_bias), outputs (normed, new_residual). Expose through mlxcel-core as a layer helper with a graph-composed fallback.
2. Adopt in decode paths: wire into the Llama3-family shared decode layer first (src/models/llama3.rs, which also serves Qwen2-family), then Gemma families (using weight_bias = 1.0). Gate with env MLXCEL_FUSED_ADD_RMSNORM=0 kill switch. Prefill may adopt it too if parity holds (same kernel, larger row count).
3. Fused q/k RoPE + KV-append (CUDA first, Metal follow-up): one kernel applying rope to q and k and writing k/v into the destination cache slab (dense KVCache slab and paged pool block layouts both supported; start with whichever the batched decode path uses after Route production batched paged decode through the fused v2 kernel, retiring the gather-then-SDPA hot path #899 lands, coordinate to avoid conflicts). Fallback and kill switch MLXCEL_FUSED_ROPE_APPEND=0.
4. Measure-then-keep policy: each fusion lands only if the microbench shows a win on at least one backend and no loss on the other; otherwise the kernel stays available but unwired (record the numbers either way).

Performance validation (mandatory)

• Op-level microbench (add to benchmarks/ or an example binary): per-op latency for (add+rmsnorm) and (rope+append) fused vs unfused across hidden sizes {2048, 4096, 8192} and batch {1, 4, 8}, Metal and CUDA.
• End-to-end: decode tok/s on one dense model (Llama3-class 4bit) and one MoE model, batch 1 and 4, before/after, Metal and GB10 when available.
• Required outcomes: fused ops beat unfused at op level on the adopting backend; end-to-end decode shows no regression anywhere and the measured gain is reported honestly (small single-digit percent is the expectation; that is acceptable for landing).
• Record in docs/benchmark_results/fused-norm-rope-<hw>-<date>.md.

Regression guard (mandatory)

• Numerical parity: fused vs unfused outputs within documented per-dtype tolerance (fp16/bf16 activations, fp32 accumulation); Gemma (1+w) convention covered by an explicit test; rope parity across position offsets including rotated/ring caches (RingSlidingKVCache positions are absolute).
• Greedy token parity on pinned prompts for each adopted family vs unfused baseline.
• Kill switches restore unfused behavior; full cargo test -p mlxcel-core and affected model tests pass; families not yet adopted show zero change.
• LoRA and surgery paths (which rewrite weights) still see consistent results (one test each with a scaled weight).

References

• src/models/llama3.rs
• src/lib/mlxcel-core/src/layers.rs
• docs/benchmark_results/moe-decode-gap-investigation.md

Line numbers are indicative as of current main; search for the named symbols.

---

Part of epic #909.