Skip to content

cuda: shared-quantize cache, residual and elementwise-chain fusions (generic half of #338) - #343

Merged
TheTom merged 5 commits into
TheTom:feature/turboquant-kv-cachefrom
jasstrong:pr/cuda-generic-fusions
Sep 3, 2026
Merged

cuda: shared-quantize cache, residual and elementwise-chain fusions (generic half of #338)#343
TheTom merged 5 commits into
TheTom:feature/turboquant-kv-cachefrom
jasstrong:pr/cuda-generic-fusions

Conversation

@jasstrong

@jasstrong jasstrong commented Sep 3, 2026

Copy link
Copy Markdown

The generic half of #338, split as requested in review, with the fixes from that review folded into the patches they belong to. Rebased onto the branch tip after #342 merged.

Overview

Two patches that are not TurboQuant-specific and help any quantized weight type:

# (in #338) patch effect
3 per-graph-eval shared-quantize cache for mmvq ~60% of quantize launches were duplicates
9 elementwise chain fusion (midi-kernel) +2.3%

plus one commit of test-backend-ops coverage. The MoE reduce-tail test is kept: it compares the tail against the CPU graph whether or not the backend fuses it.

Review fixes, by point

  1. Patch 7 is dropped. The base's bias fusion already takes a same-shape ADD after MUL_MAT (it checks ggml_are_same_shape on the ADD's inputs and goes through the n-uses check), so the residual case is covered without it; the earlier dispatch-count measurement predates the rebase that brought that fusion in. The residual test is kept and now exercises the base fusion, including ADD(mm, mm).
  2. Patch 9, overlap rule, numerics, and the elided-intermediate hole. An operand that is itself one of the nodes the chain elides is rejected (t0 = SILU(x); out = MUL(t0, t0) would otherwise read t0 unset; covered by the new self_mul test). The detector is its own static function with a fixed op array, reads op params with memcpy, and emits SQR. The chain vetoes any operand that overlaps the output unless it aliases it exactly (same pointer and size); ggml_cuda_check_fusion_memory_ranges is deliberately not used because it would veto that exact in-place case, which is the normal one. Exact in-place aliasing is what galloc produces for residual adds and is safe, because every thread reads index i before writing index i; a partial overlap at a nonzero offset is not, and neither is a broadcast operand inside the output. GELU calls ggml_cuda_op_gelu_single() and softplus uses the (x > 20) ? x : logf(1 + expf(x)) form, so fused and unfused paths agree bit for bit.
  3. Patch 4 is dropped. Upstream merged cuda: fuse MoE weighted expert reduction (cuda: fuse MoE weighted expert reduction ggml-org/llama.cpp#25952) on 2026-09-01: any weight type, k = 2..15, allocator-integrated, with its own test-backend-ops coverage. It only touches the reduction, so it is independent of which matmul produced the experts and would cover the TurboQuant MoE path too. One caveat for the rebase: cuda: fuse MoE weighted expert reduction ggml-org/llama.cpp#25952 matches the add-chain aggregation master builds (moe_out = ggml_add(moe_out, cur_experts[i])), while this branch's build_moe_ffn emits the PERMUTE -> CONT -> SUM_ROWS form (llama-graph.cpp:2287) that our tail matched. Whichever shape survives the rebase, carrying a competing tail here makes no sense; if the SUM_ROWS form stays, I can re-send the tail as a follow-up against it. For the record, the measurement Tom asked for: on Qwen3-30B-A3B Q4_K_M the fused tail gave identical perplexity (6.7883 ± 0.1208 either way) and, forced on, was 3% slower than the existing x_scale epilogue fusion for standard types (48.3 vs 49.9 tg128, measured while the card was shared; the ratio held on an idle card); on the TQ4max MoE it was worth +3.9% (121.1 ± 0.3 vs 116.6 ± 0.2 tg128, idle MI210). cuda: TurboQuant TQ4_1S decode optimisations #338's patch 2, the deep TQ4_1S fuse this tail had already beaten, is dropped for the same reason.
  4. Patch 3, cache lifetime. The cache is main-stream only, keyed like the rotation cache, and an outgrown buffer is retired rather than freed, then released at context teardown: a CUDA graph captured earlier can keep replaying against its buffer.
  5. Patch 8 landed in cuda: TurboQuant TQ4_1S decode optimisations (TQ-only half of #338) #342.
  6. Coverage. Whole-graph cases in test-backend-ops, compared against the CPU backend: MUL_MAT -> ADD(residual) including ADD(mm, mm), two MUL_MATs on one activation (the second quantize is a cache hit), SILU -> MUL -> ADD -> ADD(scalar) [-> GELU -> SOFTPLUS] -> SCALE -> CLAMP, and MUL_MAT_ID -> MUL -> PERMUTE -> CONT -> SUM_ROWS, for Q8_0, Q4_0 and TQ4_1S.

Nits: mmvq.cu includes <cstdlib>.

Additional information

MI210 (gfx90a), ROCm 7.2.3, GGML_TQ_MMQ=1 GGML_TQ_NATIVE=1:

suite result
-o MUL_MAT_RESIDUAL_FUSION 24/24
-o MUL_MAT_SHARED_SRC1 12/12
-o ELEM_CHAIN_FUSION 3/3 (incl. MUL(SILU(x), SILU(x)))
-o MOE_REDUCE_FUSION (unfused on this branch) 6/6
-o MUL_MAT -p type_a=tq4_1s 149/149
-o MUL_MAT_ID -p type_a=tq4_1s 38/38
-o MUL_MAT (all types) 1697/1697 on the final run; earlier runs showed 1692–1695 with the q5_cr / q8_cr ConvRot cases that also flake on the pristine base

The three quantized-graph tests use the same 5e-4 tolerance as test_mul_mat for quantized types; the f32 chain test keeps the default 1e-7, which is what proves the GELU and softplus forms now round like the unary kernels (they did not before the fix).

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: YES — Claude Code was used throughout this series: profiling and diagnosis, the kernel and fusion work, commit messages, the review fixes, the new tests and the verification runs.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR

Note on the base branch

While checking upstream for this split I noticed feature/turboquant-kv-cache is at a 2026-08-06 merge-base, 477 commits behind master. Besides ggml-org#25952, three landed changes bear directly on this decode work: CUDA: extend MOE fusion to specdec (41ef91f, the MoE GLU and top-k router fusions for more than one token, i.e. the speculative verify batches), CUDA: switch points per HW and quant type for the mvq->MMQ decode crossover (ggml-org#26079), and spec: fuse the DFlash encoder into the KV injection (ggml-org#27310). Your call when to rebase; happy to help carry the TQ-specific pieces across.

@TheTom

TheTom commented Sep 3, 2026

Copy link
Copy Markdown
Owner

Thanks for the split and for folding the fixes in; points 1 through 5 from the #338 review all check out against the code (the res != node guard, the exact-alias-only overlap rule, GELU via ggml_cuda_op_gelu_single() and the (x > 20) softplus form both bit-match unary.cu, the q8_cache retire list, and the single graph_epoch++).

CI: the cuda job fails on three things, all in ggml-cuda.cu:

  1. Lines 3659 and 3663: p0 = ((const float *) t->op_params)[0] in the chain_code lambda is a strict-aliasing error under gcc 13 with -Werror. Use memcpy(&p0, (const char *) t->op_params + 0*sizeof(float), sizeof(float)) (and the same for p1), which is what this file already does for max_bias further down.
  2. stl_algobase.h:437: __builtin_memcpy writing 20 bytes into a region of size 8 [-Werror=stringop-overflow=]. The inlining trace points at the pre-existing ops.insert(ops.end(), {...}) calls in the topk-moe block (3567, 3571, 3577), not your code; it is a gcc VRP false positive that the extra ~180 lines in ggml_cuda_try_fuse tipped over. Two things fix it together: ops.reserve(16) right after std::vector<ggml_op> ops;, and moving the chain detector into its own static function with ggml_op ops_ch[TQ_CHAIN_MAX_OPS] instead of the heap vector at 3736.
  3. windows (x64-openblas) also fails; I have not looked at that one.

Two blockers beyond the build:

Chain fusion can consume an elided intermediate. At 3718, if (a == prev) { other = b; } never checks b != prev. For t0 = SILU(x); out = MUL(t0, t0) the subgraph check passes (two in-subgraph uses count as fine, same reason the res != node guard was needed in patch 7), other becomes t0->data, and t0 is never written because the chain elides it. The alias veto does not save you: galloc usually gives t0 and out the same buffer, so exact() is true. Reject other if it is any node in cgraph->nodes[i .. j-1].

Patch 7 looks redundant. The base branch already has a generic MUL_MAT + ADD / ADD_ID bias fusion (the x_bias wiring around lines 4125 to 4150 on base, reaching both mul_mat_vec_f and mul_mat_vec_q), which is safe against x + x because it goes through ggml_can_fuse_ext with ggml_node_has_n_uses(..., 1). Your block at 3624 adds a contiguity check and a memory-range check but reaches fewer shapes. If there is a case the existing fusion misses that yours catches, please say which; otherwise drop patch 7 and let the existing one do the work.

Smaller:

  • The chain fusion skips ggml_cuda_check_fusion_memory_ranges, which every other fusion in this function calls.
  • The comment on the residual fusion says it is placed after the gate/up/GLU patterns; it is before them. No actual preemption today, but the comment is wrong.
  • alias_veto = overlaps(...) inside the loop should be |=; it is only correct because the loop condition short-circuits.
  • TQ_CHAIN_SQR is never emitted by chain_code.
  • The <cstdlib> include in mmvq.cu is harmless but common.cuh already provides it.

#342 is on its own track (one lifetime fix going in, see there). This one I'd hold until the chain operand check and the patch 7 question are settled.

jas and others added 3 commits September 3, 2026 08:00
Several matvecs in each decode layer consume the same normed activation,
and each was quantizing it to q8_1 independently: roughly 60% of the
quantize launches per token were duplicates. Cache the most recent
quantization, keyed on source pointer, size, type and graph epoch, and
reuse it within a graph evaluation.

The cache buffer is returned to the pool when the backend context is
destroyed, so it never outlives the pool it was taken from.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
(cherry picked from commit fd90c6b)
(cherry picked from commit 31002c7)
Collapse a run of consecutive same-shape elementwise ops into one kernel
that carries the value in a register across the whole run, instead of one
dispatch and one HBM round trip per op. Hybrid MoE graphs are full of these:
SIGMOID->MUL->ADD->ADD, ADD->SOFTPLUS->MUL, CLAMP->DIV, SILU->MUL, MUL->SCALE.

Supports ADD/MUL/DIV/SCALE/CLAMP and the common unary ops, chains up to 8
long, with same-shape or single-value-broadcast operands.

Aliasing: same-shape operands are read at index i and dst is written at
index i after those reads, so an in-place output (which the allocator
produces constantly for residual adds) is safe and must not be vetoed -
using the generic memory-range guard here rejected almost every candidate.
Only a broadcast operand, read at index 0 by every thread, needs the check.

1622 -> 1552 kernels/token on Qwen3.6-35B-A3B; +2.3% decode (paired A/B:
+3.5/+2.0/+1.6). Output byte-identical. GGML_CUDA_FUSE_CHAIN=0 disables.

(cherry picked from commit ec9bc5f)
(cherry picked from commit 68c0d7c)
test-backend-ops builds single-op graphs, so the residual-epilogue
fusion, the shared-quantize cache, the elementwise chain kernel and the
MoE expert-reduce tail never ran under CI. Add whole-graph cases for
each, compared against the CPU backend:

  MUL_MAT -> ADD(residual), including ADD(mm, mm), which must not fuse
  two MUL_MATs on one activation (second quantize is a cache hit)
  SILU -> MUL -> ADD -> ADD(scalar) [-> GELU -> SOFTPLUS] -> SCALE -> CLAMP
  MUL_MAT_ID -> MUL -> PERMUTE -> CONT -> SUM_ROWS

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
(cherry picked from commit 48b5e13)
@jasstrong
jasstrong force-pushed the pr/cuda-generic-fusions branch from 17e02d3 to 2fb1a70 Compare September 3, 2026 08:26
@jasstrong

Copy link
Copy Markdown
Author

Rebased onto the merged base and reworked per your list; new head pushed.

Patch 7 dropped. You're right: the base's bias fusion takes a same-shape ADD (it checks ggml_are_same_shape on the ADD's inputs and goes through the n-uses check), so it already covers the residual case. Our "~80 fewer dispatches" measurement was against the pre-rebase tree, which did not have that fusion yet. The residual test stays, since it now exercises the base fusion including the ADD(mm, mm) case.

Chain fusion, all in the reworked patch 9:

  • The elided-intermediate hole: an operand that is any node in nodes[i .. j-1] is rejected before it can be used. New test case MUL(SILU(x), SILU(x)) covers it and passes against the CPU graph.
  • Detector lifted into static int ggml_cuda_fuse_elem_chain(ctx, cgraph, i) with ggml_op ops_ch[TQ_CHAIN_MAX_OPS], and ops.reserve(16) after the vector in ggml_cuda_try_fuse, for the gcc 13 stringop-overflow report.
  • memcpy for the SCALE and CLAMP op params.
  • alias_veto |=, TQ_CHAIN_SQR now emitted for GGML_OP_SQR, <cstdlib> gone from mmvq.cu.
  • On ggml_cuda_check_fusion_memory_ranges: left out deliberately and now said so in a comment. It vetoes any overlap between the output and an input, and exact in-place aliasing is the normal case for this pattern (galloc hands residual adds their input's buffer); the chain's own rule allows only that exact case and vetoes every partial overlap and every broadcast operand inside the output.

The residual-fusion placement comment went with patch 7. Windows openblas was the fail-fast cancel, as you said.

Head 2fb1a70, MI210 (gfx90a), GGML_TQ_MMQ=1 GGML_TQ_NATIVE=1: MUL_MAT_RESIDUAL_FUSION 24/24, MUL_MAT_SHARED_SRC1 12/12, ELEM_CHAIN_FUSION 3/3 (the new MUL(SILU(x), SILU(x)) case included), MOE_REDUCE_FUSION 6/6. The tq4_1s MUL_MAT / MUL_MAT_ID suites passed 149/149 and 44/44 on the previous head of this branch; nothing in this round touches those kernels, and I'll re-run them on this head as soon as the card here is not shared with a live session.

The VMM pool asserts strict LIFO on free, and the q8 and TQ pre-rotation
caches hold long-lived pool buffers plus chronological retire lists. Freeing
each cache in turn, oldest retired entry first, trips that assert as soon as
two held buffers coexist (test-backend-ops -o MUL_MAT_SHARED_SRC1 passed 12/12
and then aborted in ~ggml_backend_cuda_context on a GB10). Gather every held
buffer from both caches and free them per device in descending address order,
which is reverse allocation order on the VMM pool and harmless on the legacy
pool.
@TheTom

TheTom commented Sep 3, 2026

Copy link
Copy Markdown
Owner

Second pass on 2fb1a70. The rework holds up: the elided-operand check covers every binary operand against all of nodes[i..j-1], the exact-alias rule uses ggml_nbytes on data pointers that already fold view offsets, every chain op is bit-identical to its unfused kernel, and the q8 cache key plus main-stream gate look right.

One real bug, found by running it on a GB10: test-backend-ops -o MUL_MAT_SHARED_SRC1 passes 12/12 and then aborts inside ~ggml_backend_cuda_context. The VMM pool (default on non-HIP CUDA) asserts strict LIFO on free, and the destructor freed q8_cache.ptr, then its retire list oldest-first, then the same for tq_rot_cache. With two or more held buffers that is the wrong order. The legacy pool does not care, which is why the MI210 runs were clean. Pushed f47832e: gather every held buffer from both caches and free them per device in descending address order, which is reverse allocation order on the VMM pool. This also covers the retire list that landed with #342, so it needs to go in with this PR or as its own follow-up if this one stalls.

Still checking one thing: the full test-backend-ops sweep on the GB10 segfaulted inside a ROLL case on this head. ROLL alone passes three times in a row, and this PR does not touch it, so I am running the same full sweep on the base branch to see whether it is a pre-existing GB10 issue before blaming the chain fusion. Will report either way.

Smaller items, none blocking:

  • The chain detector runs before the existing tuned fusions (multi-add, unary_mul, relu_sqr), so f32 contiguous same-shape cases now take the scalar chain kernel instead. Numerically identical, but the takeover is not mentioned anywhere; worth a comment and a quick check that ggml_cuda_op_fused_add is not faster for long ADD runs.
  • The four new tests compare numerics only. They would pass unchanged if the detector or the cache silently stopped firing. A fusion-hit counter or fusion_test_nodes() override would make them actually guard the feature.
  • SCALE with a nonzero bias is fused but never tested (ggml_scale(cur, 0.5f) only).
  • MOE_REDUCE_FUSION no longer tests a fusion in this PR now that patch 4 is gone; fine as MUL_MAT_ID coverage, the name is just misleading.
  • GGML_CUDA_FUSE_CHAIN and GGML_TQ_Q8CACHE use different prefixes and neither is documented.

@TheTom

TheTom commented Sep 3, 2026

Copy link
Copy Markdown
Owner

GB10 results on f47832e: the teardown abort is gone (MUL_MAT_SHARED_SRC1 12/12 clean exit, MUL_MAT 1697/1697, the four fusion suites all pass, and the full sweep matches base case for case plus 45 new cases). Two things still:

  1. test-backend-ops -o MUL_MAT_ID now aborts on the same VMM LIFO assert, but at runtime: backtrace is ggml_cuda_mul_mat_id -> ggml_cuda_pool_vmm::free. That function holds transient pool buffers (sorted src1/dst/ids) while the per-expert matmul allocates the persistent q8 cache on top of them, so the transients are no longer the top of the pool when they go out of scope. Sorting at teardown cannot fix that; a persistent buffer must not come from the LIFO pool at all while transients can be live. I'm moving both caches (q8 and the TQ pre-rotation one from cuda: TurboQuant TQ4_1S decode optimisations (TQ-only half of #338) #342) to a plain device allocation and keeping the retire-on-grow logic as is. Will push shortly and re-run the GB10 gates.

  2. The full sweep segfault at ROLL reproduces on the base branch and disappears with GGML_CUDA_DISABLE_GRAPHS=1 (18408/18408), so it is a pre-existing CUDA-graph bug on the GB10, not this PR. Filed separately.

…evice memory, not the pool

The VMM pool frees strict LIFO. Both caches allocated their long-lived buffer
from the pool while transient pool allocations were live (ggml_cuda_mul_mat_id
holds sorted src1/dst/ids buffers and then calls the per-expert matmul, which
allocates the q8 cache on top of them), so the transients were no longer the
top of the pool when they went out of scope and the free asserted. Sorting at
teardown cannot fix a runtime violation; a persistent buffer must not come
from the LIFO pool at all. Use ggml_cuda_device_malloc / cudaFree for both
caches and drop the teardown ordering logic. Retire-on-grow is unchanged.

CUDA graph capture on nvcc and on CDNA runs in relaxed mode, where an
allocation mid-capture is legal. The RDNA2 HIP path captures in ThreadLocal
mode; the pool's own growth had the same exposure there, and a grow forces a
re-capture anyway.
@TheTom

TheTom commented Sep 3, 2026

Copy link
Copy Markdown
Owner

Pushed ad5c1f1: both caches (q8 and the TQ pre-rotation one) now take their buffer from ggml_cuda_device_malloc and free with cudaFree, so the pool's LIFO discipline never sees them. Retire-on-grow is unchanged; the teardown sort from f47832e is gone since order no longer matters.

One thing for your AMD side when you're back at a keyboard: the RDNA2 HIP path captures graphs in ThreadLocal mode (nvcc and CDNA use Relaxed), where an allocation inside a capture errors. The pool's own growth had the same exposure there and a grow forces a re-capture anyway, so I don't expect a change in practice, but MI210 is CDNA and cannot tell us. No action needed now.

GB10 gates are running on this head; will merge when they and CI are green.

@TheTom

TheTom commented Sep 3, 2026

Copy link
Copy Markdown
Owner

GB10 (sm_121), head ad5c1f1, nvcc build:

-o MUL_MAT_ID                 979/979    (was aborting on the pool assert)
-o MUL_MAT_SHARED_SRC1        12/12
-o MUL_MAT_RESIDUAL_FUSION    24/24
-o ELEM_CHAIN_FUSION           3/3
-o MOE_REDUCE_FUSION           6/6
-o MUL_MAT                  1697/1697
full sweep, GGML_CUDA_DISABLE_GRAPHS=1   18453/18453   (base: 18408/18408, the delta is this PR's new cases)
full sweep, graphs on                     segfault at ROLL, same point as base (#347)

That is every CUDA gate I have. Merging once musa and the Windows re-run finish.

@TheTom
TheTom merged commit 1208c59 into TheTom:feature/turboquant-kv-cache Sep 3, 2026
17 of 35 checks passed
@jasstrong

Copy link
Copy Markdown
Author

Thanks for catching the pool-order bug and for running the GB10 gates; the device-allocation fix is the right shape and I'll carry ad5c1f1 over to our fleet branch, where the same caches (plus an fp16 variant for the MFMA path) still come from the pool. On the RDNA2 point: our V620 runs with HIP graphs on, so I'll check whether a cache grow inside a capture changes anything there and report back. The smaller items from your first pass (precedence comment and a timing check against ggml_cuda_op_fused_add, hit counters for the four tests, a SCALE bias case, the MOE_REDUCE test name, the env prefixes) I'll send as a follow-up PR rather than leave on the merged one.

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.

2 participants