cuda: TurboQuant TQ4_1S decode optimisations (TQ-only half of #338) - #342
Merged
TheTom merged 7 commits intoSep 3, 2026
Merged
Conversation
The MoE decode dispatch gated the int8 dp4a path behind !GGML_CUDA_CC_IS_AMD, which excluded CDNA even though gfx90a has the dot4 instructions. Enable it there, keeping RDNA on the scalar path. Decode 94.3 -> 95.4 t/s on an MI210 with Qwen3.6-35B-A3B-TQ4max. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR (cherry picked from commit 5a662f4)
The MoE gate and up projections consume the same normed activation, so the forward-WHT + q8_1 rotation ran once per projection. Cache it per graph eval (same key as the mmvq shared-quantize cache: tensor identity, data pointer, size, epoch; main stream only) and return the cached buffer. Removes ~40 kernel dispatches per token on a 40-layer MoE. At the measured ~4.4us cost of a small dispatch on MI210 (1.3us hardware floor plus unhidable cold-miss latency for a kernel too small to fill the GPU) that is worth about 0.18ms/token. GGML_TQ_ROTCACHE=0 disables. (cherry picked from commit 7ef6f21)
…thmetic tq4_cents8_reg expanded 8 nibble indices to 8 int8 centroids with ~100 ALU ops of shifts and selects, about 400 per 32-weight block. That made the TQ decode kernels ALU-bound rather than bandwidth-bound: the MoE matvec was sustaining only ~24% of MI210 peak bandwidth while the dense Q8_0 matvecs on the same card reached ~72%. Use get_int_from_table_16 (__builtin_amdgcn_perm, the IQ4_NL/MXFP4 path) to do the lookup in hardware: about 14 ops instead of 100. The shift fallback existed because HIP __byte_perm() with a runtime selector lowers through a dynamically indexed byte union that PromoteAlloca turns into a 32KB LDS staging area. __builtin_amdgcn_perm has no such problem and maps 1:1 to the instruction - this is the same idiom already validated in the TQ4_1S MMQ tile loader. (cherry picked from commit 080626d)
… and MUSA tq4_cents8_reg() is on the dp4a decode path that NVIDIA also compiles, so the unguarded AMD builtin broke the CUDA and MUSA builds. Keep v_perm_b32 on HIP and use the equivalent __byte_perm selectors (0x5140 / 0x7362) elsewhere. (cherry picked from commit 949de16)
One output row was spread over all 32 lanes, so each lane covered only blocks_per_row/32 blocks and then paid a full 5-round warp reduction. The reduction (5 dependent ds_bpermute, ~200 cycles) cost more than the work (~150 cycles), which is why every attempt to add parallelism made things worse: split-K was monotonically bad (115 -> 92 t/s at KS=8), 64-lane rows lost 7%, and extra per-lane ILP was a wash. Give each row 16 lanes instead: 4 reduction rounds, twice the work per lane, and the warp still reads only two contiguous weight regions (going below 16 scatters reads across more rows and loses more to coalescing than it saves - 8 and 4 both measured worse). Never assign more lanes than there are blocks, so no lane idles on narrow projections; previously the down projection (16 blocks) left half the warp idle. Paired A/B, alternating, 3 pairs: +3.1%, +2.6%, +3.0%. Output is byte-identical to the old mapping under greedy decode. GGML_TQ_LPR=N overrides. (cherry picked from commit 625028a)
nvcc treats the unreferenced local as an error under -Werror all-warnings; the cached pre-rotation helper computes the block count itself now.
Owner
|
Pushed a one-liner: nvcc rejected the now-unused |
Owner
|
GB10 (sm_121, driver 580.173.02), head fa870a1, nvcc build with So the |
The small-batch TQ MUL_MAT_ID path is CUDA-graph capturable and a context can hold several captured graphs, so freeing the cache buffer when it is outgrown could leave a sibling graph replaying against pool memory that has been handed out again. Push outgrown buffers (including the cross-device case, which used to leak) onto a retire list drained at context teardown, the same shape as the q8 quantize cache. Also key the allocation off ctx.device throughout, clamp GGML_TQ_LPR to the instantiated lane counts, floor the computed default at 2, and update the AGENTS.md pitfalls for the byte-permute LUT and graph capture.
github-actions
Bot
added
the
documentation
Owner
|
Pushed 3df06b1 with the last review items so this can land without another round trip:
Re-running the GB10 sweeps on this head; merging when those and CI are green. |
TheTom
merged commit Sep 3, 2026
c26baf1
into
TheTom:feature/turboquant-kv-cache
8 of 23 checks passed
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
2 participants
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.
Overview
Four patches on the TurboQuant decode path, all TQ4_1S-specific and all exercised by
test-backend-ops:v_perm_b32instead of shift arithmeticPatch 2 of #338 (the deep TQ4_1S MoE down-proj + weighted-sum fuse) is dropped entirely: the generic reduce tail in the companion PR measured faster than it (100.0 vs 99.0 tg128), so there is no reason to carry the slower kernel behind a knob.
Changes relative to #338 from the review:
GGML_CUDA_CC_IS_AMD(cc) ? 16 : 32), since every measurement behind it is MI210; NVIDIA keeps 32.tq4_cents8_reg()(949de16) is carried over.Additional information
MI210 (gfx90a), ROCm 7.2.3, Qwen3.6-35B-A3B-TQ4max: numbers per patch are as measured in #338; the combined figures for this subset are in the first comment below once CI is green.
test-backend-ops -o MUL_MAT -p type_a=tq4_1sand-o MUL_MAT_ID: 149/149 and 38/38 on the MI210 (gfx90a, ROCm 7.2.3), withGGML_TQ_MMQ=1set so the native path is the one exercised. The full-o MUL_MATsuite is 1697/1697 on the final run of this branch; an earlier run showed theq5_crConvRot case flaking, as it does on the pristine base on random data.Requirements
🤖 Generated with Claude Code
https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR