perf(KERNEL-QUANT-CIQ-GEMM-ROCM): cooperative-tile WMMA kernels — Shared rejected, BigTile accepted - #3036
Open
joral wants to merge 6 commits into
Open
Conversation
…asure, reject (mudler#3032) Issue mudler#3032's same-tool rocprofv3 trace found KQuantGemmKWmmaQ4K/Q6K 4.9x-10.6x slower per-kernel than llama.cpp's mul_mat_q on the identical tensor-core mechanism, with llama.cpp launching double the warps per block and 16-48x fewer, bigger blocks. Each warp here already owns a fully independent output tile with its own shared-memory slice, so widening the block from 4 to 8 warps was the obvious first lever to test: does packing more of that already-independent work into fewer, bigger launches close any of the gap. Templated both kernels on WarpsPerBlock and added an 8-warp instantiation behind VT_ROCM_QUANT_WMMA_WIDE=1 (default off, same same-binary-A/B posture VT_ROCM_QUANT_WMMA and VT_ROCM_Q6K_SMALL_PRIVATE already ship with). Hardware-verified correct on both configs: ctest -R rocm|cross_device, 46/46 cases, 84066/84066 assertions, zero regression. The answer is no. Op-level A/B (quant-gemm-bench, RX 9060 XT, best-of-4, idle host) across the six Q4_K/Q6_K prefill shapes: geomean -4.3%, a net regression. Each warp shares no loaded or dequantized data with any other warp in its block, so widening only spreads the same LDS-capped occupancy budget over fewer, bigger blocks rather than doing more work per warp -- the axis this measurement rules out, not the one that explains llama.cpp's actual advantage. The sharper, still-open hypothesis this points at -- cross-warp data reuse in llama.cpp's own kernel -- is recorded in the spec's Owed section, unread and unconfirmed, as the next traceable step. The toggle stays in the tree, default off, as a ready-made A/B for re-checking this specific axis on different hardware or a future toolchain revision, rather than making the next person re-derive it from scratch. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
…thesis by reading llama.cpp's source The wide-block experiment's spec entry left this as an unread inference. Reading ggml-cuda/mmq.cuh, mmq-config-rdna4.cuh, and mmq-load-tiles.cuh confirms it: llama.cpp's Q4_K block cooperatively loads and shares one 128x128 tile across all warps, while this kernel gives each warp its own independent 16x16 tile with no sharing at all, even between warps that share the same M-rows. That is the mechanism gap the rejected wider-block experiment could not reach. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
…e for the WMMA kernels (mudler#3034) The wider-block A/B (mudler#3032/mudler#3033) measured a net regression, and reading llama.cpp's mul_mat_q source found the actual mechanism it uses: one cooperatively-loaded shared tile per block, not just more warps. Tracing this kernel's own tile/warp index math (tile = blockIdx.x * WarpsPerBlock + threadIdx.y; it = tile / n_tiles; jt = tile % n_tiles) shows every warp in one block already shares the same activation rows (it), since n_tiles is always far larger than WarpsPerBlock. Each warp reads those identical rows independently via its own load_matrix_sync call, every superblock -- the precise, confirmed redundancy, and exactly why widening the block alone (mudler#3033) did not help: it made that same redundant pattern wider without adding any sharing to fix it. This spec scopes a smaller, targeted fix instead of matching llama.cpp's full 128x128 cooperative tile: reuse the already-landed 8-warp block (kQuantWmmaWideWarpsPerBlock, mudler#3033, currently unused on its own), and add cooperative activation staging so the shared superblock's BlockQ8_K bytes are loaded into shared memory once per block instead of once per warp. LDS budget checked: ~55 KiB total against the 64 KiB limit. Pull request shape: separate spec and implementation pull requests, same recorded preference this row already uses. This pull request lands the spec only. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
…e rejected, BigTile accepted (mudler#3034) Two experiments in one wave, because the second only makes sense in light of the first's result. Shared (KQuantGemmKWmmaQ6K/Q4KShared): stages the 8-warp block's shared activation rows once instead of once per warp, exactly the redundancy the spec traced through this kernel's own tile/warp index math. Three copy implementations tried (naive byte copy, per-row loop, flat word loop); the result barely moved between them, which is itself the finding -- the cost is the staging mechanism (one more shared-memory round trip and sync per superblock), not a coding inefficiency. Geomean -16% vs the shipping default, -11.5% vs the already-rejected wide-block-alone arm (mudler#3033). Rejected. Reading llama.cpp's actual shared-memory formula (ggml_cuda_mmq_get_nbytes_shared_x, mmq.cuh:415-419) explains why its 128x128 tile fits 64KiB LDS when a naive scale-up of Shared's own 8-row-reuse does not (86.5KiB): not a smaller encoding -- its block_q8_1_mmq (144B/128 elements) and our BlockQ8_K (292B/256 elements) have nearly identical density -- but a narrower K-chunk (32 and 128 elements at a time, against our 256-wide superblock), restaged up to 8x more often in exchange for a much smaller peak footprint. BigTile (KQuantGemmKWmmaQ6K/Q4KBigTile<OutT, WarpsPerBlock, ItGroup>): keeps this row's existing 256-wide staging granularity and only widens the reuse width, choosing ItGroup=3 (48 activation rows, ~24x reuse) as the largest that still fits alongside the existing ~50KiB per-warp weight-side footprint at 8 warps. 2D grid (blockIdx.x fixes each warp's jt as before; blockIdx.y gives the whole block a shared it_base spanning ItGroup, looped inside each warp), because a 2D grid needs no n_tiles%WarpsPerBlock==0 precondition on the it/M axis (kept on jt/N: a warp whose jt falls outside n_tiles must never exist in a launched block, since its early return would desync the block's later __syncthreads() calls from warps that keep running). Weight dequant and Q4_K's scale/min unpack now run once per superblock and serve all ItGroup iterations, reuse Shared did not have either. Geomean +16.8% (quant-gemm-bench, best-of-4), holding on the ragged non-16-aligned tail shape too (+19-20%). Real-model confirmation (Ornith-1.5-9B-Q4_K_M.gguf, isolated prefill, rocprofv3): total prefill kernel time -14.0%; the isolated quant-GEMM gap against llama.cpp's mul_mat_q narrows from 8.92x to 6.98x slower. Accepted, still default off pending review. Along the way: found that this row's existing WMMA correctness tests (N=48, n_tiles=3) never satisfy either arm's n_tiles%8==0 precondition, so every earlier green result for Shared was real but hollow -- those tests structurally never reached it, only its fallback. Added per-variant dispatch counters and two new tests at N=128/M=80 (m_tiles not a multiple of ItGroup, exercising the ragged last block) that prove the specific kernel launched, not infer it from shapes. Also read the compiled GCN assembly on both sides (ours via -save-temps; llama.cpp's own .so unbundled per-ggml_type from its concatenated clang-offload bundles) to rule out register allocation as an explanation: llama.cpp uses MORE registers than us in both formats (247 vs 192 for Q4_K, 202 vs 142 for Q6_K) with zero spilling. Our one real spill (Q4_K, 60B) traced to one-time setup/epilogue code outside the per-superblock loop -- cosmetic, not a meaningful cost. A handover document (.agents/specs/kernel-quant-ciq-gemm-rocm-handover.md) indexes this whole investigation -- three PRs, six measured hypotheses, exact reproduction recipes -- for continuing on a fresh context. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
…o valid rows (review finding, mudler#3036) Both KQuantGemmKWmmaQ6KBigTile and KQuantGemmKWmmaQ4KBigTile staged ItGroup*16 activation rows unconditionally, before the it_i bounds check that guards every consumer. qact (EnsureQuantScratch's scratch buffer) holds exactly m_tiles*16 rows with no padding, so the ragged last blockIdx.y block read past the buffer's logical end whenever m_tiles % ItGroup != 0. This never corrupted output, because the unread rows were also never consumed, but it was still an out-of-bounds device read. Clamp the staged row count in both kernel bodies to min(ItGroup*16, (m_tiles - it_base)*16), so the loop only copies rows that are actually valid for the current it_base. Also correct the dispatch-site comment for the n_tiles % wpb precondition: mutation testing shows it guards a zero-sized grid_x launch for narrow N, not a mid-block __syncthreads desync from an out-of-range jt. Note BigTile's precedence over SHARE_ACT when a caller sets both toggles at once. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
…view/fix cycle (mudler#3036) The handover document was written before the fresh reviewer found a real out-of-bounds device read in BigTile's activation-staging loop for the ragged M-tail, and before the fix landed. Records PR mudler#3036's current head, the finding, the fix commit, and the operator's independent gate rerun, plus a gotcha about green NMSE tests not proving the absence of an OOB read on hardware with no device-side ASan for RDNA/gfx12. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code]
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.
Issue #3034. Row:
KERNEL-QUANT-CIQ-GEMM-ROCM(alreadyACTIVE). Base:row/KERNEL-QUANT-CIQ-GEMM-ROCM-COOPTILE(#3035, spec-only, base-reachable,not yet merged to
main).Two experiments, because the second only makes sense in light of the
first's result.
Shared (
KQuantGemmKWmmaQ6K/Q4KShared): stages the 8-warp block'sshared activation rows once instead of once per warp — exactly the
redundancy the spec traced through this kernel's own tile/warp index math.
Three copy implementations tried; the result barely moved between them,
which is itself the finding — the cost is the staging mechanism (one more
shared-memory round trip and sync per superblock), not a coding
inefficiency. Geomean -16% vs the shipping default, -11.5% vs the
already-rejected wide-block-alone arm (#3033). Rejected.
Reading llama.cpp's actual shared-memory formula
(
ggml_cuda_mmq_get_nbytes_shared_x,mmq.cuh:415-419) explains why its128x128 tile fits 64 KiB LDS when a naive scale-up of Shared's own 8-row
reuse does not (86.5 KiB): not a smaller encoding — its
block_q8_1_mmq(144B/128 elements) and our
BlockQ8_K(292B/256 elements) have nearlyidentical density — but a narrower K-chunk (32 and 128 elements at a time,
against our 256-wide superblock), restaged up to 8x more often in exchange
for a much smaller peak footprint.
BigTile (
KQuantGemmKWmmaQ6K/Q4KBigTile<OutT, WarpsPerBlock, ItGroup>):keeps this row's existing 256-wide staging granularity and only widens the
reuse width, choosing
ItGroup=3(48 activation rows, ~24x reuse) as thelargest that still fits alongside the existing ~50 KiB per-warp weight-side
footprint at 8 warps. 2D grid (
blockIdx.xfixes each warp'sjtasbefore;
blockIdx.ygives the whole block a sharedit_basespanningItGroup, looped inside each warp) — a 2D grid needs non_tiles % WarpsPerBlock == 0precondition on theit/M axis (kept onjt/N: a warp whosejtfalls outsiden_tilesmust never exist in alaunched block, since its early
returnwould desync the block's later__syncthreads()calls from warps that keep running). Weight dequant andQ4_K's scale/min unpack now run once per superblock and serve all
ItGroupiterations — reuse Shared did not have either.Geomean +16.8% (
quant-gemm-bench, best-of-4), holding on the raggednon-16-aligned tail shape too (+19-20%). Real-model confirmation
(
Ornith-1.5-9B-Q4_K_M.gguf, isolated prefill,rocprofv3): total prefillkernel time -14.0%; the isolated quant-GEMM gap against llama.cpp's
mul_mat_qnarrows from 8.92x to 6.98x slower. Accepted, stilldefault off pending review.
Test-coverage gap found and fixed: this row's existing WMMA correctness
tests (N=48,
n_tiles=3) never satisfy either arm'sn_tiles % 8 == 0precondition, so every earlier green result for Sharedwas real but hollow — those tests structurally never reached it, only its
fallback. Added per-variant dispatch counters and two new tests at
N=128/M=80 (
m_tilesnot a multiple ofItGroup, exercising the raggedlast block) that prove the specific kernel launched, not infer it from
shapes.
ISA-level check: read the compiled GCN assembly on both sides (ours via
-save-temps; llama.cpp's own.sounbundled per-ggml_typefrom itsconcatenated clang-offload bundles) to rule out register allocation as an
explanation: llama.cpp uses MORE registers than us in both formats (247 vs
192 for Q4_K, 202 vs 142 for Q6_K) with zero spilling. Our one real spill
(Q4_K, 60B) traced to one-time setup/epilogue code outside the
per-superblock loop — cosmetic, not a meaningful cost.
Gated on
isravale(RX 9060 XT, gfx1200, ROCm 7.2.3):ctest -R rocm|cross_device, 48/48 cases, 84084/84084 assertions, zero regression(three remaining preflight failures —
test_tower_skip_rss_arm, thetoolssuite'stest_drop_file_cache,test_cpu_x86_llamacpp_floor— arepre-existing environment artifacts unrelated to this change).
A handover document (
.agents/specs/kernel-quant-ciq-gemm-rocm-handover.md)indexes this whole investigation — three PRs, six measured hypotheses,
exact reproduction recipes — for continuing this line of work on a fresh
context.
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:claude-sonnet-5 [Claude Code]