Skip to content

cuda: make the ADD+RMS_NORM+MUL fusion reachable - #209

Merged
bri-prism merged 2 commits into
PrismML-Eng:prismfrom
cklxx:cuda-add-rms-norm-fusion-all-nvidia
Oct 24, 2026
Merged

bri-prism merged 2 commits into
PrismML-Eng:prismfrom
cklxx:cuda-add-rms-norm-fusion-all-nvidia

Conversation

@cklxx

@cklxx cklxx commented Sep 19, 2026

Copy link
Copy Markdown

Summary

The fused add_rms_norm_f32 path added in d8d96cf (#135) never fires outside GB10 — and, I believe, doesn't fire on GB10 either for the standard residual pattern. Two independent gates block it:

  1. Arch gate. cc == GGML_CUDA_CC_DGX_SPARK excludes every other NVIDIA device, but add_rms_norm_f32 uses no architecture-specific instructions (plain loads, block_reduce, rsqrtf). Relaxed to GGML_CUDA_CC_IS_NVIDIA(cc).

  2. Memory-range check. With the arch gate relaxed, every candidate still failed ggml_cuda_check_fusion_memory_ranges. Instrumenting it shows why — for this graph ggml's allocator produces exact aliases:

    add=0x..400000  a=0x..40b480  b=0x..400000  rms=0x..405480  mul=0x..405480   # add is in-place over b
    add=0x..416480  a=0x..400000  b=0x..416480  rms=0x..400000  mul=0x..400000   # mul output reuses a
    

    Same base pointer, type, shape and contiguous layout, which the check treats like a partial overlap and rejects.

Since that aliasing comes from the graph allocator rather than the device, the same rejection should happen on GB10 — worth a quick check on your side (a temporary print in the range check is enough).

Change

  • ggml_cuda_check_fusion_memory_ranges gains an opt-in allow_exact_alias (default false, existing callers unchanged). An exact alias — identical data, type, shape, contiguous — is accepted; partial overlaps are still rejected for everyone.
  • The ADD+RMS_NORM+MUL fusion opts in. It is safe because the kernel is elementwise within a row: each thread reads a[col], b[col] before writing sum[col], and block_reduce syncs the block before dst[col] is written.
  • The norm weight is additionally required not to alias either output (it's broadcast over rows).

Verification (RTX 4070 Ti SUPER, sm_89, Ternary-Bonsai-2-27B PQ2_0)

  • test-backend-ops -o ADD_RMS_NORM: 25/25; RMS_NORM_MUL_ADD, ADD, RMS_NORM, MUL all pass.

  • nsys: fusion now fires — 96 add_rms_norm_f32 launches per 32 decode tokens, and rms_norm_f32<1024> launches drop by exactly 96 (258 → 162).

  • Greedy generation output coherent.

  • Throughput: no measurable change (interleaved base/patched, 3 rounds × 4 reps):

    tg128 (t/s) pp512 (t/s)
    base 64.29 / 63.80 / 63.56 1712.6 / 1707.2 / 1692.8
    patched 64.44 / 63.67 / 63.55 1713.5 / 1696.9 / 1692.9

    Expected: at batch 1 the fusion saves one n_embd-row read (~20 KB) per firing, 3 firings per token, against ~6.7 GB of weight traffic per token. The value here is that the kernel actually runs, not a speedup on this workload.

Not changed

While instrumenting the range check I also saw the upstream MUL_MAT, MUL_MAT, GLU fusion rejected 48/128 times on this model. That one is a real partial overlap (GLU output, 69632 B, starts at the same address as the 20480 B activation input) and unsafe to relax — the check is doing its job. Noting it in case it's useful; only fixable at the allocator level.

https://claude.ai/code/session_01UpAMXrmoyeC1gjM77dReJo

The fused add_rms_norm_f32 path (d8d96cf) was unreachable outside GB10 for
two independent reasons:

1. The gate `cc == GGML_CUDA_CC_DGX_SPARK` excluded every other NVIDIA
   device, although the kernel uses no architecture-specific instructions.

2. Even with that gate relaxed, ggml_cuda_check_fusion_memory_ranges
   rejects the fusion for the residual pattern ggml's allocator actually
   produces: the residual ADD is in-place (add->data == add->src[1]->data)
   and the norm output reuses the other input's buffer
   (mul->data == add->src[0]->data). Both are exact aliases - same base
   pointer, type, shape and contiguous layout - which the range check
   treats like a partial overlap.

The kernel is elementwise within a row: every thread reads a[col] and
b[col] before writing sum[col], and block_reduce synchronizes the block
before dst[col] is written. Exact aliases are therefore safe for it, so
the range check gains an opt-in `allow_exact_alias` that this fusion
passes. Partial overlaps are still rejected for all callers, and the
default is unchanged for existing callers.

The norm weight is additionally required not to alias either output, as
it is broadcast across rows.

Verified on RTX 4070 Ti SUPER (sm_89) with Ternary-Bonsai-2-27B PQ2_0:
- test-backend-ops -o ADD_RMS_NORM: 25/25
- fusion now fires (96 add_rms_norm_f32 launches per 32 decode tokens,
  rms_norm_f32<1024> launches drop by the same 96)
- throughput unchanged within noise (tg128 63.88 vs 63.89 t/s, pp512
  1704 vs 1701, interleaved 3x4 runs): at batch 1 the fusion saves one
  n_embd-row read per firing, which is negligible against weight traffic.

Claude-Session: https://claude.ai/code/session_01UpAMXrmoyeC1gjM77dReJo

@bri-prism bri-prism left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agent review: posted by the maintainer's coding agent at their request.

No findings in this source pass. I traced the exact-alias predicate through the fused kernel's row accesses and reduction, and checked the exclusion for the broadcast norm weight.

CUDA execution was not performed. The remaining validation is a fused-versus-unfused comparison for in-place residual input, norm output reusing an input buffer, and rejected partial overlap, including a non-GB10 NVIDIA device newly enabled by this patch.

Reviewed commit: b989116e1ceef747a659d146b1977709e97342fc.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

The RMS output can still be misused as an unwritten weight, and the fused alias path lacks regression coverage.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 High severity · 1 Medium severity

Open (2)

Comment thread ggml/src/ggml-cuda/ggml-cuda.cu Resolved
Comment thread ggml/src/ggml-cuda/ggml-cuda.cu Resolved
@bri-prism

Copy link
Copy Markdown
Collaborator

Agent benchmark follow-up, posted at the maintainer's request.

Pinned head b989116 against merge-base 9a9394a. Same public PTQ1 model on both arms, full GPU offload, flash attention, q4_0 K/V, batch/microbatch 512, 8 CPU threads. Three alternating baseline/candidate pairs, three repetitions per invocation, short/empty starting context.

GPU pp512 before → after, tok/s Paired change tg128 before → after, tok/s Paired change
RTX 3090 751.19 → 754.86 +0.5% 59.78 → 59.88 +0.2%
RTX 4090 1574.64 → 1573.04 -0.1% 88.53 → 89.27 +0.8%
H100 SXM 1229.30 → 1232.92 +0.3% 88.51 → 88.68 +0.2%
RTX 5090 1879.01 → 1883.52 +0.2% 119.59 → 119.75 +0.1%

Selected CPU-reference backend checks passed on the compared arms.

The selected ADD_RMS_NORM checks were included alongside the matrix-multiplication checks. These are short-context throughput measurements.

The percentages describe these paired runs; small changes should not be interpreted as established improvements. No long-context, multi-slot serving, or end-to-end logit-parity claim is made.

PQ2 control measurements (paired pp512 / tg128 changes):

  • RTX 3090: +0.5% / +0.5%.
  • RTX 4090: +0.3% / +1.7%.
  • H100 SXM: +0.6% / +0.3%.
  • RTX 5090: +0.8% / +0.2%.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
@bri-prism

Copy link
Copy Markdown
Collaborator

@cklxx @sb32445, #209 and #310 both change ggml_cuda_check_fusion_memory_ranges, so whichever lands second needs a small rebase. @cklxx, please also add a test-backend-ops case for the exact ADD -> RMS_NORM -> MUL pattern with both outputs live. @sb32445, #310 relaxes ggml_cuda_should_fuse_mul_mat_vec_q for every mmvq fusion at 2 to 4 columns, not only SwiGLU. Please require a contiguous bias when ncols > 1, and commit the 2 to 4 column fusion test cases.

professorpalmer pushed a commit to professorpalmer/llama.cpp-ada-mirai that referenced this pull request Oct 15, 2026
From PrismML-Eng/llama.cpp#209 by @cklxx (commits b989116e1 and 37fe65676, squashed).

The fused add_rms_norm_f32 path ran only on GB10 (cc == DGX_SPARK), and the fusion memory-range check rejected the
exact aliases that the graph allocator produces for the residual pattern. The kernel has no architecture-specific
code, so the gate is now any NVIDIA device. The range check gets an opt-in allow_exact_alias (same data, type, shape,
contiguous); partial overlaps are still rejected. The norm weight must not alias any output.

@bri-prism bri-prism left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agent review: posted by the maintainer's coding agent at their request.

CUDA validation on H200 (sm_90, CUDA 12.8, Release), this diff applied to prism 73bfe6c. This is the fused-versus-unfused check my earlier review asked for, on a non-GB10 NVIDIA device that the patch newly enables.

  • test-backend-ops: ADD_RMS_NORM 25/25, RMS_NORM_MUL_ADD 30/30, RMS_NORM 51/51.
  • Local cases (not part of the PR): I added 20 cases for ADD → RMS_NORM → MUL(weight) with the residual sum kept for a later consumer, at widths 1 to 16896, with 1 and 5 rows, both with a normal ADD and with ggml_add_inplace. All 20 pass, both fused and with GGML_CUDA_DISABLE_FUSION=1.
  • End to end: Ternary-Bonsai-2-9B PQ2_0-MTP on llama-server, 4 prompts × 128 greedy tokens, 2 rounds. Generated tokens are identical to stock prism without speculation and with MTP depth 3. Decode is +0.5% to +1.2% in both modes.

Reviewed commit: diff as of 2026-10-23.

@bri-prism
bri-prism merged commit 0493f95 into PrismML-Eng:prism Oct 24, 2026
1 check passed
bri-prism added a commit that referenced this pull request Oct 25, 2026
…rm at both norms of every layer, bit-identical, +4.8-5.4% decode on H200 (#391)

Fuse the residual ADD + RMS_NORM + MUL + signed FWHT into one launch at each FFN input, plus a layer-start look-ahead for the lone RMS_NORM before the Hadamard transform (skipping the GDN conv-state nodes in between). Bit-identical. When it does not fire, the ADD+RMS_NORM+MUL fusion from #209 handles those nodes as before.

On H200: 9B plain decode +4.8-5.4%, MTP3 +1.2-1.7%, DFlash7 +1.3-2.5%; 27B neutral. Validated together with #383 (GGML_CUDA_GDN_FUSION on and off): greedy identical, MUL_MAT 1666/1666, both fusions fire.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants