perf(core): fused residual-add RMSNorm and fused RoPE + KV-append kernels by inureyes · Pull Request #971 · lablup/mlxcel

inureyes · 2026-07-30T15:53:18Z

Summary

Adds the two fused decode kernels issue #905 asks for, wires both into the decode paths that exist on main today, and lands them behind kill switches so the measure-then-keep decision is a one-line flip rather than a revert. No benchmark numbers are in this PR: the measurement pass is deliberately serialized separately so GPU contention does not pollute it. Both fusions default ON so the end-to-end sweep measures them; MLXCEL_FUSED_ADD_RMSNORM=0 and MLXCEL_FUSED_ROPE_APPEND=0 restore the graph paths.

What changed

FusedAddRMSNorm (src/lib/mlx-cpp/turbo/fused_norm.cpp, task item 1). A two-output custom kernel computing new_residual = x + residual and normed = rms_norm(new_residual) * (weight_bias + weight) in one dispatch, replacing the elementwise Add plus fast::rms_norm pair that every pre-norm block pays at each residual join. MLX arrays are immutable from the graph's perspective, so the in-place residual update a mutable-tensor runtime would perform is expressed as the second output. One threadgroup per row, N_READS = 4 sweep matching MLX's own rms_single_row access pattern, two-stage (SIMD group then threadgroup) sum-of-squares reduction, precise::rsqrt. weight_bias absorbs the Gemma (1 + w) convention with no second weight tensor: the offset is folded in the weight's own dtype, which is exactly how GemmaRMSNorm::new builds adjusted_weight, so the two agree bit-for-bit rather than approximately.

Fused q/k RoPE + KV-append (src/lib/mlx-cpp/turbo/fused_rope_append.cpp, task item 3). A three-output kernel that reads the row-contiguous fused-QKV projection output whole and emits q, k and v already in their consumers' layouts, replacing three trailing-axis slices, three reshape/transpose pairs and two fast::rope calls with one dispatch. Taking the whole qkv rather than three slices is what keeps it a single dispatch: a last-axis slice is not row-contiguous and ensure_row_contiguous would insert exactly the materializing copies the kernel exists to remove. The rotation is transcribed from MLX's rope.metal (inv_freq = exp2(-d * log2(base)), theta = (scale * position) * inv_freq, fast::cos / fast::sin) so it agrees with fast::rope including the traditional convention and partial rope_dims.

Why that kernel emits the append payload instead of writing the slab. The obvious shape for "fused RoPE + KV-append" is a kernel that writes k and v straight into the destination cache. MLX cannot express that, and the header says so at length: outputs of fast::metal_kernel / cuda_kernel are always freshly allocated (CustomKernel::eval_gpu calls allocator::malloc per output and has no donation path), while the existing slice_update append does donate, because copy_gpu donates the input buffer when it is uniquely referenced and the cache's slab always is. Fusing the write would therefore trade an O(1) donated update for an O(capacity) slab copy: at 4 K context, 8 KV heads and head_dim 128 that is 8 MiB per tensor per layer per token. So the kernel produces the payload in the destination's own layout and leaves the O(1) donated slice_update to do the store. This is the one place the implementation deliberately departs from the issue's literal wording, and it is the reason.

Adoption (task item 2). src/models/llama3.rs TransformerBlock::forward and forward_batched take the fused residual join (these also serve Qwen2 / Qwen2.5 and every VLM whose text backbone is Llama3Model / Qwen2Model), and Attention::forward takes the fused RoPE path ahead of the graph fallback. Unlike forward_split_rope, the new path takes traditional as a real parameter, so traditional-RoPE checkpoints reach it rather than being bypassed. src/models/gemma.rs takes the fused residual join with weight_bias = 1.0 through a FusedAddRmsNormSpec impl on its own GemmaRMSNorm. Gemma 2 / 3 / 4 are deliberately not adopted: they use sandwich norm (norm(attn_out) then clip(x + normed)), which is not an add-then-norm pair, so there is nothing for this kernel to fuse there.

Rust surface (src/lib/mlxcel-core/src/layers.rs). fused_add_rms_norm and graph_add_rms_norm over a FusedAddRmsNormSpec trait exposing both the raw weight (what the kernel folds weight_bias into) and the effective weight (what the graph fallback hands fast_rms_norm), so neither path pays an extra op for the Gemma convention. FusedQKVLinear::forward_fused_rope_append returns None whenever the fused path is off, unavailable, or outside the kernel's geometry contract, so the launcher's throw is unreachable from the decode path (an exception crossing the cxx boundary is not recoverable). Backend availability is cached in a OnceLock rather than re-asked per join.

Microbench (task item 4). examples/fused_norm_rope_microbench.rs sweeps hidden sizes 2048 / 4096 / 8192 against batch 1 / 4 / 8 for both fusions, each against the exact graph it replaced, and prints a table plus a CSV block. Its doc comment states that the op-level number is a lower bound: both fusions also remove a full-width MLX intermediate per call, which shows up as allocator and dependency-tracking pressure rather than as kernel time.

Defaults, and how to flip them

Switch	Default	Where the default lives
`MLXCEL_FUSED_ADD_RMSNORM`	on	`FUSED_ADD_RMSNORM_DEFAULT` in `src/lib/mlxcel-core/src/layers.rs`
`MLXCEL_FUSED_ROPE_APPEND`	on	`FUSED_ROPE_APPEND_DEFAULT` in `src/lib/mlxcel-core/src/layers.rs`

Both are single const bool declarations sitting next to each other. Setting one to false makes off the default and turns the variable into an opt-in (=1), leaving the kernel and its tests in place; that is the intended outcome if the measurement does not justify a fusion. An unset variable takes the compiled-in default and an unrecognised value also takes it, so a typo in a deployment script cannot quietly change the decode graph in either direction. Both are read once and cached for the process lifetime.

CUDA is written but UNVALIDATED

This host is Apple Silicon (M1 Ultra), Metal only: no CUDA device, no nvcc. Both CUDA JIT strings are structural transliterations of their Metal counterparts, kept deliberately line-for-line parallel so a reviewer with a CUDA box can diff the two rather than re-derive the algorithm, and each carries an UNVALIDATED banner naming the parity test to run first. They have never been compiled or executed. Nothing in this PR claims a CUDA result. The kernel-selection gate is the same metal::is_available() pattern paged_attention.cpp already uses, so a CUDA build reaches the CUDA string and an Apple build never does.

Test plan

All commands prefixed with export DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer.

cargo test --release -p mlxcel-core --lib --features metal,accelerate fused_norm_parity_tests: 8 passed. Per-dtype parity (f32 / f16 / bf16) against the add + fast_rms_norm graph across decode and prefill shapes; the residual output landing on the graph add's value; weight_bias = 1.0 reproducing Gemma's precomputed (1 + w); the GemmaRMSNorm layer agreeing with the raw weight_bias = 1.0 call; LoRA-scaled and surgery-scaled weight cases; the kill switch; and greedy-argmax parity over a self-feeding 8-layer, 48-step stack.
cargo test --release -p mlxcel-core --lib --features metal,accelerate fused_rope_parity_tests: 9 passed. Parity against the exact Llama3 attention fallback graph over offsets to 131071 (the large offsets matter: RingSlidingKVCache positions are absolute and climb past the window capacity), multi-token windows, batch > 1, the interleaved traditional convention, partial rope_dims with the untouched tail, both destination layouts against each other, a bit-identity check that no rotation leaks onto V, and a scaled-projection case.
Both files re-run with MLXCEL_FUSED_ADD_RMSNORM=0 MLXCEL_FUSED_ROPE_APPEND=0: 8 and 9 passed. The kill-switch test asserts byte-identical output to graph_add_rms_norm in that mode, not merely close.
cargo test --release -p mlxcel-core --lib --features metal,accelerate layers::tests: 64 passed, including the new flag-parser and FusedAddRmsNormSpec convention pins.
cargo test --release --lib --features metal,accelerate models::llama3::tests: greedy token parity on a pinned prompt through the tiny synthetic Llama, in both fused and unfused modes.
cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings and the microbench example.

Gap, stated plainly: greedy token parity is on a synthetic checkpoint, not a real one. There are no model checkpoints on this host (models/ is empty), so the real-checkpoint greedy parity the issue asks for per adopted family has not been run and should be part of the measurement pass. The synthetic Llama exercises the same TransformerBlock::forward and Attention::forward the real family uses, so it covers the wiring, not the scale.

What is deliberately not here

Benchmark results. Measurement is serialized separately; docs/benchmark_results/fused-norm-rope-<hw>-<date>.md is for that pass to write.
Paged-pool wiring. The kernel supports the paged block-pool destination layout and it is covered by parity tests, but no caller uses it. The batched paged decode path it belongs to is issue Route production batched paged decode through the fused v2 kernel, retiring the gather-then-SDPA hot path #899's, and restructuring it here would collide with that work. The seam is a FusedRopeDestLayout::PagedPool variant with the reason in its doc comment.
Gemma 2 / 3 / 4 adoption, for the sandwich-norm reason above.
A .metal reference copy of either kernel. sparse_v_sdpa and paged_attention ship one; turbo4_delegated_sdpa does not, so this follows the latter.

Closes #905

First slice of the fused decode kernels for issue #905. Adds `src/lib/mlx-cpp/turbo/fused_norm.cpp` and its header: a two-output custom-kernel launcher that computes `new_residual = x + residual` and `normed = rms_norm(new_residual) * (weight_bias + weight)` in one dispatch, replacing the elementwise Add plus `fast::rms_norm` pair that every pre-norm block pays twice per token. MLX arrays are immutable from the graph's perspective, so the in-place residual update a mutable-tensor runtime would do is expressed as a two-output kernel `(normed, new_residual)`. `weight_bias` absorbs the Gemma `(1 + w)` convention without materializing a second weight tensor: `0.0` is the standard case, `1.0` is Gemma, and the bias is folded in the weight's own dtype so it reproduces `GemmaRMSNorm`'s precomputed `(1 + w)` bit-for-bit. Numerics are pinned to MLX's own `rms_single_row` Metal kernel so the fused and unfused paths stay interchangeable: fp32 accumulation over the dtype-rounded sum, `precise::rsqrt`, and the `x * inv_mean` product rounded to the activation dtype before the weight multiply. The residual sum is rounded to the activation dtype before it is squared, matching what the unfused RMSNorm reads back out of global memory. The CUDA source string is a line-for-line transliteration of the Metal body and is marked UNVALIDATED: this host is Apple Silicon with no CUDA device and no nvcc, so it has never been compiled or run. Also declares the bridge entry points for both fused kernels in `mlx_cxx_bridge.h`; the implementations and the Rust side follow in subsequent commits. Refs #905

Adds the second #905 kernel and wires both into the decode paths that exist on main today. `src/lib/mlx-cpp/turbo/fused_rope_append.cpp` is a three-output custom kernel that reads the row-contiguous fused-QKV projection output whole, applies rotary embedding to the q and k blocks, and emits q, k and v already in their consumers' layouts. Taking the whole `qkv` rather than three trailing-axis slices is what keeps it a single dispatch: a last-axis slice is not row-contiguous, and `ensure_row_contiguous` would insert exactly the materializing copies the kernel exists to remove. The rotation math is transcribed from MLX's own `rope.metal` (`exp2(-d * log2(base))`, `fast::cos` / `fast::sin`) so the fused and graph rotations agree. The kernel emits the append payload instead of writing the cache slab in place, and the header says why: outputs of `fast::metal_kernel` are always freshly allocated (`CustomKernel::eval_gpu` has no donation path), while the existing `slice_update` append donates the slab buffer when it is uniquely referenced. Fusing the write would trade an O(1) donated update for an O(capacity) slab copy, which at 4 K context and head_dim 128 is 8 MiB per tensor per layer per token. `dest_layout` selects the dense `KVCache` slab order or the paged pool row order; the paged variant is implemented and tested but deliberately left unwired, because the batched paged decode path it belongs to is issue #899's. Rust side: `layers::fused_add_rms_norm` and `layers::graph_add_rms_norm` with the `FusedAddRmsNormSpec` trait, which exposes both the raw weight (what the kernel folds `weight_bias` into) and the effective weight (what the graph fallback hands `fast_rms_norm`), so neither path pays an extra op for the Gemma `(1 + w)` convention. `FusedQKVLinear::forward_fused_rope_append` returns `None` whenever the fused path is off, unavailable, or outside the kernel's geometry contract, so the launcher's `throw` is unreachable from the decode path. Adoption: `llama3.rs` `TransformerBlock::forward` and `forward_batched` (which also serve the Qwen2 family) take the fused residual join, and `Attention::forward` takes the fused RoPE path ahead of the graph fallback. Unlike `forward_split_rope`, the new path takes `traditional` as a real parameter, so traditional-RoPE checkpoints reach it too. `gemma.rs` takes the fused residual join with `weight_bias = 1.0` via a `FusedAddRmsNormSpec` impl on its own `GemmaRMSNorm`. Both fusions default ON with runtime kill switches `MLXCEL_FUSED_ADD_RMSNORM=0` and `MLXCEL_FUSED_ROPE_APPEND=0`. The defaults are single named constants in `layers.rs` (`FUSED_ADD_RMSNORM_DEFAULT`, `FUSED_ROPE_APPEND_DEFAULT`) so the measure-then-keep decision in issue #905 is a one-line flip. Validated with `cargo check -p mlxcel-core --features metal,accelerate` and `cargo check --lib --features metal,accelerate`; parity tests follow in the next commit. Refs #905

Adds the regression guard issue #905 requires, plus the op-level microbench its measure-then-keep policy needs. `fused_norm_parity_tests.rs`: per-dtype numeric parity (f32 / f16 / bf16) against the `add` + `fast_rms_norm` graph across the shapes the decode and prefill paths produce; an explicit check that `weight_bias = 1.0` reproduces Gemma's precomputed `(1 + w)` tensor rather than approximating it; a residual-output check that the fused sum lands on the graph `add`'s value, because that value is what the next join accumulates onto; scaled-weight cases standing in for LoRA fusion and the surgery tooling, which rewrite norm weights in place; a kill-switch test that asserts byte-identical output to the graph path when `MLXCEL_FUSED_ADD_RMSNORM=0`; and a greedy-argmax parity run over a self-feeding 8-layer, 48-step synthetic stack, which is the end-to-end statement the per-element tolerances are a proxy for. `fused_rope_parity_tests.rs`: parity against the exact graph the Llama3 attention fallback builds, swept over offsets up to 131071. The large offsets are the point: `RingSlidingKVCache` positions are absolute and keep climbing past the window capacity, so a kernel that folded a window origin into the position would pass every offset-zero test and decode nonsense on rotated caches. Also covers multi-token windows (token `t` must rotate at `positions_base + t`), batch > 1, the interleaved `traditional` convention, partial `rope_dims` with the untouched tail, both destination layouts against each other, a bit-identity check that no rotation leaks onto V, and a scaled-projection case for merged adapters. Tolerances are stated per dtype as normalized RMS and normalized max deviation and are sized at a small multiple of one ulp, because the fused and graph paths differ only in where the residual sum is rounded, not in the arithmetic. Both files skip on a CPU-only build, matching `fused_moe_parity_tests.rs`. `layers.rs` unit tests pin the flag parser: the #905 switches are default-ON, the inverse of the opt-in `MLXCEL_FUSED_QK_NORM`, and an unrecognised value keeps the compiled-in default rather than silently changing the decode graph. A second test pins the `FusedAddRmsNormSpec` conventions, since a mis-implemented impl is a silent wrong-math bug. `examples/fused_norm_rope_microbench.rs` sweeps hidden sizes 2048/4096/8192 against batch 1/4/8 for both fusions, each against the exact graph it replaced, and prints a table plus a CSV block. Its doc comment states that the op-level number is a lower bound: both fusions also remove a full-width MLX intermediate per call, which shows up as allocator pressure rather than kernel time. Documents both kill switches in `docs/environment-variables.md` and the measure-then-keep gate in `docs/benchmarks.md`, including where the compiled-in defaults live so flipping one is a single-line change. No benchmark numbers are recorded here; the measurement pass is run separately so GPU contention does not pollute it. Refs #905

`fused_add_rms_norm` and `forward_fused_rope_append` asked the FFI whether the backend has a custom-kernel JIT on every call. That answer reaches `metal::is_available()` / `cu::is_available()` in C++ and cannot change mid-process, so re-asking at every residual join of every layer of every token was pure overhead on exactly the path the fusion exists to make cheaper. Both are now a `OnceLock`. The eligibility check also short-circuits on `array_ndim(weight) != 1` before allocating any shape vectors, so the disqualifying case costs one FFI call instead of three vector allocations. Also drops a pointless probe allocation from the microbench: it was constructing a one-element array just to read back a dtype it had itself chosen a line earlier. Refs #905

Adds a greedy decode parity test to the Llama3 suite. It runs the tiny synthetic Llama, the same `TransformerBlock::forward` and `Attention::forward` the real family uses, for 24 steps from a pinned prompt and compares the token sequence against a baseline captured with both #905 kill switches set to `0`. Why a pinned list rather than an in-process A/B: both gates are read once and cached for the process lifetime, because they sit on the per-token hot path and must not re-read the environment. A single test process therefore cannot exercise both sides, so running this file in both modes is what closes the loop. Verified both ways: 12 passed with the fusions on (the default) and 12 passed with `MLXCEL_FUSED_ADD_RMSNORM=0 MLXCEL_FUSED_ROPE_APPEND=0`. The synthetic model settles into a short cycle after a few steps, so the discriminating tokens are the leading ones; the tail is kept anyway because a fusion bug that only compounds across steps would surface there and nowhere else. This is checkpoint-level coverage of the wiring, not of scale: no model checkpoints exist on this host, so real-checkpoint greedy parity per adopted family still has to be run as part of the measurement pass. Refs #905

Whitespace only, from `cargo fmt`. No behavior change. Refs #905

…r the microbench The op-level microbench crashed on its first configuration. Its consume-the-outputs step reshaped q to [B, n_heads, 1, D] and added it to k/v at [B, n_kv_heads, 1, D], which cannot broadcast once n_heads differs from n_kv_heads, and the harness always uses a GQA ratio of 4. Reshaping q into its KV grouping, [B, n_kv, ratio, D], broadcasts correctly and keeps the step identical in the fused and unfused paths. The crash was in the harness only; both parity suites and the model tests were unaffected. With the harness running, neither fusion demonstrated a win on Apple M1 Ultra, so both defaults move to off and the two env vars become opt-in. Measured over three repetitions at hidden {2048, 4096, 8192} by batch {1, 4, 8}: add-RMSNorm scattered 0.77x to 1.14x with no structure and its worst cell re-measured at 1.02x and 1.03x, while fused RoPE + append hovered around parity except at hidden 8192 batch 1, which came in at 0.94x, 0.90x and 0.89x, the only reproducible signal in the sweep and one that points the wrong way. Per-iteration time stays near 280-410us across a 4x change in hidden size, so the harness is dominated by fixed dispatch and synchronization cost rather than by the op, and it cannot resolve these kernels on this host. Issue #905's measure-then-keep policy lands a fusion unwired unless the microbench shows a win, so the kernels, their parity tests and the microbench all stay while only the wiring defaults change. This matches how MLXCEL_FUSED_QK_NORM (#326) shipped after the same outcome. Both constants sit adjacent in layers.rs and are one edit away from being enabled if a quiet host, an amortizing harness, or the CUDA backend shows a win. Both kill-switch tests derived their expectation as the negation of the disable flag, which silently hard-coded default-on and failed for the wrong reason once the constants flipped. They now assert the documented precedence instead: an explicit truthy or falsey value wins, and unset or unrecognised keeps the compiled-in default. That required widening the two constants to pub(crate). Both suites pass with the defaults off and with the switches explicitly on.

…d direction `fused_905_flags_default_on_and_respect_disable_values` asserted that an unset environment variable means enabled, which re-pinned default-on in a third place and failed once the measurement moved both constants to off. It now compares the unset case against `FUSED_ADD_RMSNORM_DEFAULT` and `FUSED_ROPE_APPEND_DEFAULT`, matching how the same test already handled unrecognised values, and the name no longer asserts a direction. The explicit-value coverage is unchanged: truthy and falsey spellings still win over the default in both directions.

inureyes added status:review Under review type:performance Performance improvements priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers status:done Completed and removed status:review Under review labels Jul 30, 2026

inureyes added 8 commits July 31, 2026 01:29

style: apply rustfmt to the #905 kernel wiring and tests

596963b

Whitespace only, from `cargo fmt`. No behavior change. Refs #905

inureyes force-pushed the feature/issue-905-fused-norm-rope-kernels branch from 08c24bb to b67cf28 Compare July 30, 2026 16:41

inureyes merged commit f0ee104 into main Jul 30, 2026
5 checks passed

inureyes mentioned this pull request Aug 3, 2026

Epic: Fused paged decode, sorting-free sampling, and serving performance techniques #909

Closed

inureyes deleted the feature/issue-905-fused-norm-rope-kernels branch August 4, 2026 12:08

inureyes mentioned this pull request Aug 5, 2026

fix(cuda): fused_rope_append diverges from the graph RoPE reference at offsets near 2^17 #1049

Closed

5 tasks done

inureyes self-assigned this Aug 31, 2026

Navigation Menu

Repository navigation

perf(core): fused residual-add RMSNorm and fused RoPE + KV-append kernels - #971

perf(core): fused residual-add RMSNorm and fused RoPE + KV-append kernels#971
inureyes merged 8 commits into
mainfrom
feature/issue-905-fused-norm-rope-kernels

inureyes commented Jul 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented Jul 30, 2026

Summary

What changed

Defaults, and how to flip them

CUDA is written but UNVALIDATED

Test plan

What is deliberately not here

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant