cuda: TurboQuant TQ4_1S decode optimisations - #338
Conversation
TQ4_1S weights previously had no MMQ path, so prefill dequantised them to f16 and went through hipBLAS. That threw away the format's whole point: the 5bpw footprint was paid for at load time and then not used at compute time. Add a TQ4_1S tile loader, MMQ type traits, and dispatch, and pre-rotate the activation so the block-local turbo WHT cancels against the weights. The weights stay in their rotated-domain int8 centroid form, so the stock MFMA-i8 MMQ kernel consumes them directly. The type reuses the Q3_K shared-memory layout, since its per-16 scales have the same shape, and so takes the same tile geometry Q3_K and IQ2_XS use on CDNA (occupancy 1, I=128); other geometries do not match what that layout expects and write out of bounds. Measured on an MI210 with Qwen3.6-35B-A3B-TQ4max, against this same branch with the path disabled: prefill 971 -> 2124 t/s (2.19x) decode 95 -> 93 t/s (-2.6%, this is a prefill-only path) ppl 3.2653 -> 3.2653 (identical) It also keeps roughly 13GB less VRAM than converting the experts to Q8_0. All 147 tq4_1s MUL_MAT cases pass. The path is env-gated behind GGML_TQ_MMQ so it can be A/B'd against the hipBLAS fallback. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
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
At decode the MoE tail ran the down-projection, then a separate MUL by the expert weights, then a reduction over the expert slots. Fold all three into the matvec so each expert's contribution is scaled and accumulated in registers. Two things were needed to make it fire in practice: the pattern match has to follow the graph shape the allocator actually produces, and the weights and ids have to be staged in pool scratch so the fusion applies on every layer instead of only some. About +1% decode on an MI210. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
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
The fused expert-reduce tail replaces MUL(expert weights) -> PERMUTE -> CONT -> SUM_ROWS after a MUL_MAT_ID. It reads the f32 MUL_MAT_ID output directly, so nothing in it is specific to TQ weights: make it weight-type agnostic and use it as the default tail everywhere. It accumulates into pool scratch and copies to the graph tensor last, which resolves the allocator aliasing hazards without having to veto the fusion. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
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.
…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.
Every transformer layer ends its attention and FFN blocks with MUL_MAT -> ADD(residual), which ran as two kernels. The mmvq fused path already indexes x_bias exactly like dst (per row and per column), so a full same-shape tensor works there, not just a broadcast bias - it was simply never wired up for a bare MUL_MAT -> ADD. About 80 dispatches per token on a 40-layer model. Decode only (should_fuse_mul_mat_vec_q requires ncols_dst == 1) and placed after the gate/up/GLU patterns so it cannot preempt them. GGML_CUDA_FUSE_RESIDUAL=0 disables.
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.
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.
…and MUSA tq4_1s_qs_to_int8x4() called the AMD-only byte-permute builtin unconditionally, so the CUDA and MUSA CI builds failed to compile mmq-load-tiles.cuh. Keep v_perm_b32 on HIP and use the equivalent __byte_perm selector elsewhere. The TQ4_1S MMQ path itself is still gated to AMD at dispatch time; this only makes the header compile on the other toolchains.
… 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.
|
Pushed two commits here: a merge of the #336 nvcc/MUSA guard, plus 949de16 which guards the two Review notes, in priority order. Patches 1, 5, 6, 8 (the TQ-specific ones) look sound and are the ones
Suggestion: split this into a TQ-only PR (patches 1, 5, 6, 8) and a generic-fusions PR (3, 4, 7, 9). I'll take the TQ half as soon as #336 lands and CI is green on it. The generic half needs the fixes above plus some kind of multi-op test before it goes default-on. Two nits: |
|
Thanks for the thorough pass, and for the two guard commits. Agreed on all six points and on the split. I'll open a TQ-only PR with patches 1, 5, 6 and 8 (with the lanes-per-row default gated to AMD, the duplicated comment removed and the On a personal note: I had surgery this week. It went well, but it will be a while before I can type properly, which is why this series landed in your queue all at once. I wanted what I had been working on in front of you in case things had not gone so well. Replies from me may be slow and terse for a bit. |
|
Split as discussed:
Closing this one in favour of the two above. |
cuda: TurboQuant TQ4_1S decode optimisations (TQ-only half of #338)
|
Just saw the note about the surgery. Glad it went well, and I hope the recovery is quick. No rush on any of these, health first. The series was in good shape when it landed, so nothing here needs you typing before you're ready. Slow and terse is completely fine. |
cuda: shared-quantize cache, residual and elementwise-chain fusions (generic half of #338)
2 participants
Overview
Nine patches on the TurboQuant decode path. They are independent of each other and each is separately measured, so they can be taken in whole or in part:
v_perm_b32instead of shift arithmeticPatch 1 is arguably a plain bug fix: the int8 dp4a decode path was gated behind
!GGML_CUDA_CC_IS_AMD, which excluded CDNA despite gfx90a having the dot4 instructions.Patch 6 is the significant one. TurboQuant decode was ALU-bound rather than bandwidth-bound — the MoE matvec sustained only ~24% of MI210 peak bandwidth while dense Q8_0 matvecs on the same card reached ~72%, because the nibble-to-centroid expansion burned roughly 100 ALU ops per call. Replacing it with a hardware byte-permute removes that. Note it must use
__builtin_amdgcn_perm; HIP's__byte_permwith a runtime selector lowers via PromoteAlloca into a 32KB LDS staging area and comes out about 3x slower.Patches 3, 4, 7 and 9 are not TurboQuant-specific and help any quantized type.
Additional information
MI210 (gfx90a), ROCm 7.2.3, Qwen3.6-35B-A3B-TQ4max, measured against this PR's own base (#336) with
GGML_TQ_MMQunset on both, so this is exactly the delta these nine patches add:+28% decode, and +21% on the non-MMQ prefill path as a side effect of the fusion work (with
GGML_TQ_MMQ=1prefill is unchanged at ~2120, since MMQ handles it).Patch 9 was also checked on the other two AMD architectures and ports cleanly: +1.8% on a 7900XTX (gfx1100) and +1.4% on a V620 (gfx1030). Patch 8's
GGML_TQ_LPRwas swept 2→32 on both and is flat within ±0.2, so its wave64-tuned default of 16 is safe on wave32.Verification
Patches 2, 4, 6 and 8 all live in the MoE decode matvec, which is why #334 and #335 exist: on AMD the
MUL_MAT_IDsuite aborted before reaching any TurboQuant case, so that path had no exercised coverage at all. With those two in place it is 979/979 including cases at real model shapes (256 experts, 8 used). I would rather land those first than have this reviewed on the strength of llama-bench numbers alone.Requirements