cuda: make the ADD+RMS_NORM+MUL fusion reachable

[cklxx]

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

[bri-prism]

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]

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

[Copilot]

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 · 1

Open (2)

• Require distinct storage for MUL weight · New
• Add regression test for ADD-RMS_NORM-MUL aliasing · New

[bri-prism]

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%.

[bri-prism]

@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.