[2/2] cuda: TQ4_1S MMQ prefill on RDNA3 and RDNA2 - #337
Merged
TheTom merged 6 commits intoSep 3, 2026
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 MMQ entry points were gated on amd_mfma_available(), which is CDNA only, so RDNA3 fell back to dequantising to f16 and calling hipBLAS. Nothing in the kernel actually needs MFMA: the vec_dot is the generic q8_0_16 one already shared with Q3_K, and upstream MMQ dispatches WMMA on RDNA3 by itself. The only missing piece was the per-arch tile configuration, which existed solely in mmq-config-cdna.cuh. Mirror Q3_K's twelve RDNA3 entries for TQ4_1S, which is already the shared-memory layout it borrows, and let the gate accept WMMA. 7900XTX: prefill 759 -> 2533 t/s (3.34x), decode unchanged, 269 of 269 tq4_1s cases pass, perplexity 3.2595 with MMQ against 3.2698 without. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
ggml_cuda_mmq_load_tiles_tq4_1s wrote a single hardcoded tile layout and was commented "TQ4_1S MMQ is MFMA-only". Every other loader branches on AMD_MFMA_AVAILABLE / TURING_MMA_AVAILABLE / AMD_WMMA_AVAILABLE, because dp4a consumes a different shared-memory layout: rows are padded by one int, and the scales start after that padded quant region rather than at a fixed 2*MMQ_TILE_NE_K offset. mmq_get_dp4a_tile_x_sizes already mapped TQ4_1S to MMQ_DP4A_TXS_Q8_0_16, so the sizing was right all along; only the loader ignored it, writing quants and scales where the vec_dot never read them. TQ4_1S was also missing from the dp4a traits block, so it picked up the wrong vec_dot. With both fixed, TQ4_1S MMQ runs on hardware with no matrix cores at all. V620 (gfx1030): prefill 424 -> 1984 t/s (4.68x), decode unchanged, 269 of 269 tq4_1s cases pass, perplexity 3.2660 with MMQ against 3.2705 without. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
…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.
Owner
|
Reviewed. Your two commits are clean: the dp4a tile layout matches The only build failure is the Plan: land #336 when its CI is green, then push the guard commit here so this PR re-runs, then merge. One follow-up for either PR: |
Owner
getenv() ran on every mul_mat dispatch. It is short-circuited away on NVIDIA but is a libc call per node on the AMD path. Read it once through a static, like the other knobs. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NxP6x5bmUDYFvmouceN2mR
TheTom
merged commit Sep 3, 2026
ed0816b
into
TheTom:feature/turboquant-kv-cache
6 of 21 checks passed
Author
|
Heads-up: the "read GGML_TQ_MMQ once" change I pushed here (439fe67) initializes its static from a call to itself, so on AMD the first TQ mul_mat with n >= 8 deadlocks on the static's guard. Every prefill hangs with GGML_TQ_MMQ=1 on RDNA and CDNA; NVIDIA never calls the function. One-line fix with MI210 results in #348. Sorry for the trouble. |
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
Extends TQ4_1S MMQ prefill from CDNA to RDNA3 and RDNA2. Very little was actually missing — the vec_dot is the generic
q8_0_16one already shared with Q3_K, and MMQ already dispatches WMMA on RDNA3 by itself. Three things stood in the way.1. The gate. The entry points tested
amd_mfma_available(cc), which is CDNA-only. Widened to accept WMMA and RDNA2.2. Missing tile configs. TQ4_1S CASE entries existed only in
mmq-config-cdna.cuh, soggml_cuda_mmq_get_config_rdna3()returnedGGML_TYPE_COUNTfor every J and the J-selection loop aborted withJ_best=0. Added entries mirroring Q3_K, whose shared-memory layout TQ4_1S already borrows.3. A real bug in the tile loader, which is what RDNA2 needed.
ggml_cuda_mmq_load_tiles_tq4_1swrote a single hardcoded tile layout and carried the comment "TQ4_1S MMQ is MFMA-only". Every other loader branches onAMD_MFMA_AVAILABLE / TURING_MMA_AVAILABLE / AMD_WMMA_AVAILABLE, because dp4a consumes a different shared-memory layout: rows padded by one int, and the scales based after that padded quant region rather than at a fixed2*MMQ_TILE_NE_Koffset.mmq_get_dp4a_tile_x_sizesalready mapped TQ4_1S toMMQ_DP4A_TXS_Q8_0_16, so the sizing was right all along — only the loader ignored it, writing quants and scales where the vec_dot never read them. TQ4_1S was also missing from the dp4a traits block and so picked up the wrong vec_dot.Additional information
Qwen3.6-35B-A3B-TQ4max, ROCm 7.2.3, each card measured against itself with the path disabled:
Decode is unchanged on both. The V620 number is the interesting one: 4.68x prefill on hardware with no matrix cores at all, reaching about 86% of an MI210's prefill.
The two commits are split so each is independently bisectable — RDNA3 (configs + WMMA gate) does not depend on the loader work; RDNA2 (loader + traits + configs + gate) does.
GGML_TQ_LPRwas also swept 2→32 on both cards and is flat within ±0.2, so the wave64-tuned default of 16 is safe on wave32 and no per-arch tuning is needed.Requirements
J_best=0abort and the dp4a layout bug, wrote the fix and commit messages, and ran the correctness and perplexity verification on the 7900XTX and V620. I reviewed the change and am responsible for every line of it.