CUDA: gated_delta_net - share per-token k/q via shared memory for long prefills - #54
Draft
roberteg16 wants to merge 1 commit into
Draft
CUDA: gated_delta_net - share per-token k/q via shared memory for long prefills#54roberteg16 wants to merge 1 commit into
roberteg16 wants to merge 1 commit into
Conversation
roberteg16
force-pushed
the
rogarcia.gdn-shared-kq
branch
from
July 16, 2026 13:51
e332179 to
5bb1972
Compare
This comment was marked as outdated.
This comment was marked as outdated.
roberteg16
force-pushed
the
rogarcia.gdn-shared-kq
branch
from
July 28, 2026 17:29
5bb1972 to
2619b60
Compare
GATED_DELTA_NET is 18% of prefill GPU time on Qwen3.5-35B-A3B / gfx1151. The existing kernel is a token-by-token scan with one warp per state column, which leaves the machine latency bound: k/q are re-fetched by all 128 warps of a head on every token, a 256x read amplification that puts MemUnitBusy at 94.5% while moving only 17.9 GB/s. Add a chunked delta rule (WY representation / UT transform) as a single fused kernel, one block per (head, seq), state resident in registers. Keeping a whole chunk in 64 KB of LDS needs two things: - CS = 32 rather than 64. At CS = 64 the chunk's U + W + kq alone are 80 KB. - Never materialising U and W. Instead of v_new = U - W S with U = Tinv(V beta) and W = Tinv(K beta exp(g_cs)), use the identical v_new = Tinv (V beta - (K beta exp(g_cs)) S), so one [CS][S] buffer holds V*beta, becomes Y, then becomes v_new in place. The in-place triangular product needs no second buffer because rows CS/2..CS-1 are accumulated into registers before any row is overwritten. The two triangular products (Gram and kq) are packed across all 1024 threads with a triangle-to-rectangle fold: 496 + 496 + 32 is exactly the block size. The natural (tid/CS, tid%CS) mapping makes a wave issue a full 128-MAC dot with most of its lanes masked off. Other things that mattered: __launch_bounds__(1024, 16) caps VGPRs at 96, which is what lets two 32-wave workgroups share a WGP; folding the lane-uniform exp(g_cs[i]) into the value before the butterfly halves the number of values to reduce; two state columns per thread so each staged K/Q read feeds two FMAs. That occupancy target is only reachable on wave32, where 1024 threads are 32 waves and two workgroups per WGP give 16 waves per SIMD. On wave64 the same workgroup is 16 waves and tops out at 4, which -Wpass-failed turns into a build error, so the request is conditioned on the physical warp size. The dispatch is also gated on wave32: the lane mapping, the 16-lane butterflies and the LDS padding are all tuned for RDNA3.5. Verified to build clean with -Werror -Wpass-failed for gfx908/90a/942/950, gfx1030, gfx1100-1103, gfx1150-1153 and gfx1200/1201. GDN op at pp2048: 220.9 -> 116.7 ms. End to end against 0f0db62, interleaved: +5.7% at pp1024, +10.5% at pp2048, +20.0% at pp4096 with -ub 4096. Note the op only ever sees n_ubatch tokens, so at the default -ub 512 the gain is ~+2.5%. The larger figures need a larger ubatch - which is itself the point: with the sequential kernel -ub 4096 is a pessimisation at pp4096 (1520 vs 1656 t/s at -ub 512) because the scan's serial length grows with n_tokens, while with this kernel it becomes the best configuration (1825 t/s). Scalar gate and K == 1 only, S_v == 128, n_tokens >= 128; everything else falls back to the existing kernel, which is untouched. Off by default; GGML_CUDA_GDN_CHUNKED=1 enables it, =2 forces it for any multi-chunk shape so the tests can reach it. Also add five test-backend-ops cases: no existing GDN case combined head_size=128 with more than 64 tokens, and none covered GQA (v_repeat > 1), which this model actually uses (H_k=16, H_v=32). Assisted-by: Claude Opus 4.7
roberteg16
force-pushed
the
rogarcia.gdn-shared-kq
branch
from
July 28, 2026 18:07
2619b60 to
fa2898d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
GATED_DELTA_NETis 18.3% of prefill GPU time on Qwen3.5-35B-A3B / gfx1151. The existing kernel isa token-by-token scan with one warp per state column, which is latency bound:
k_t/q_tarere-fetched by all 128 warps of a head on every token (256x read amplification, MemUnitBusy 94.5%
while moving only 17.9 GB/s).
This replaces it, for the scalar-gate case, with a chunked delta rule (WY representation / UT
transform) as a single fused kernel: one block per (head, seq), the 128x128 state resident in
registers, no scratch buffer. The chunk's intermediates never leave LDS.
Keeping a whole chunk in 64 KB of LDS needs two things:
U + W + kqalone are 80 KB.v_new = U - W SwithU = Tinv(V beta)andW = Tinv(K beta exp(g_cs)), use the identicalv_new = Tinv (V beta - (K beta exp(g_cs)) S), so one[CS][S]buffer holdsV*beta, becomesY, then becomesv_newin place.Result: 64 KB of global traffic per (chunk, head) instead of 360 KB, and MemUnitBusy drops from
94.5% to 38.6%.
Single commit, 439 lines, 4 files, additions only. The existing kernel is untouched apart from a
26-line dispatch branch.
Results
GDN op at pp2048: 220.9 -> 116.7 ms (1.89x).
End to end, Qwen3.5-35B-A3B Q4_K_M on gfx1151, interleaved A/B against
0f0db6292:-ub 4096)-ub 4096)-ub 4096)-ub 512)-ub 512)Decode is structurally untouched: the op's
n_tokensis tokens per sequence, so it is 1 duringgeneration and the
>= 128gate can never be reached.test-backend-ops -o GATED_DELTA_NET: 42/42 in all modes. Five cases added, because no existing GDNcase combined
head_size=128with more than 64 tokens, and none covered GQA (v_repeat > 1) at alldespite the model running
H_k=16, H_v=32.Caveat: this is gated by
-ub, not by prompt lengthThe op only ever sees
n_ubatchtokens. A 4096-token prompt at the default ubatch is eight separate512-token GDN calls, so the gain is capped at the pp512 level no matter how long the prompt is.
That is why the two
-ubrows above differ so much. Read it as a column, at pp4096:With the sequential kernel, raising the ubatch is a pessimisation (1520 vs 1656) because the
scan's serial length grows with
n_tokens. With this kernel it becomes the best configuration.So best-config to best-config the gain is 1655.6 -> 1824.9, +10.2%, and it is only reachable if
the caller raises
-ub, which costs larger activation buffers. At the default ubatch the PR isworth ~+2.5-3%.
Scope
Taken only when
!kda && K == 1 && S_v == 128 && n_tokens >= 128. KDA,K > 1snapshots, other headsizes and single-chunk shapes fall back to the existing kernel.
Off by default.
GGML_CUDA_GDN_CHUNKED=1enables it from 128 tokens,=2forces it for anymulti-chunk shape so the tests can reach it.