Repository navigation
Conversation
With too few heads for a wmma tile, a block scores 64 keys against 8 tokens: the keys are staged once in half precision, the queries one head at a time, and each thread owns one key for two tokens, so no dot product needs a cross thread reduction. Batches smaller than a token tile keep the vector kernel. test-backend-ops measures 4 heads.
github-actions
Bot
added
testing
|
On my Frankensetup (a 3060 connected via only PCIe3x4, and 128GB of system RAM at 2400MT/s): This branch: Master (1537a0a): So this branch is plausibly a bit faster but well within the margin of error. |
|
On the other hand the microbenchmark shows a huge improvement. From
to
raw benchmark outputvs |
|
Thanks for testing! With the experts on the CPU the run is bound by system RAM and PCIe, so an e2e number can't show a kernel that only speeds up the GPU side, while your microbench shows exactly that: 3.3x on the 3060. |
|
Also I didn't give At depth 65536 (otherwise all the same config) this branch speeds up pp2048 from:
to
Which is much more significant. |
|
Quite interesting results on 2x RTX PRO 6000 Blackwell Max-Q:
And with |
| const half2 k_val = k_shared[kl][c]; | ||
| #pragma unroll | ||
| for (int j = 0; j < TOKENS_PER_THREAD; ++j) { | ||
| ggml_cuda_mad(qk[j], k_val, q_shared[tl + j*TOKEN_GROUPS][c]); |
There was a problem hiding this comment.
from what I understand this will do f16 product, so is |q * k| > f16_max it will overflow
There was a problem hiding this comment.
Right, fixed: the queries now stay in float and every product is computed in float from the widened keys, which also made the kernel 10% faster on the RTX PRO 6000!
Address review from am17an: the half2 products overflow once a single q * k exceeds the f16 range. The queries stay in float in shared memory and each half2 of keys is widened once for both tokens, so every product and sum is computed in float.
|
@ServeurpersoCom I did a bot review and here are the findings. Please ignore the ones that don't make sense Possible speedups: |
|
Thanks! I'll go through the bot findings one by one and benchmark the kernel for each of them to see which ones actually pay off, every pass on this kernel speeds up Qwen3.8 Flash Next prefill a bit more. The first finding led to a good one: staging every head at once so each key is widened only once gives +17-20% on the kernel and +1.8% PP @ 128k e2e. The rest didn't pay off: the f32 key tile is 5-9% slower because the occupancy drops, and the direct F16 copy gains nothing measurable, I only kept it since it's free. |
The tile kernel stages the queries and weights of every head at once, so each key element is widened from half once and feeds all heads, with a single barrier. F16 keys are copied into the tile without a float round trip. Keeping the keys in float in shared memory measures slower, the occupancy drops.
|
Microbenchmark on the 3060 now hits |
|
I think the kernel is in a good spot now: 4.76 ms vs 16.4 ms for the vector kernel on the RTX PRO 6000 (3.4x), and the indexer went from 9.6% to around 3% of the GPU time at 128k. The last real lever would be tensor cores, running the token x head rows against the key tile through mma like the wmma kernel does for 32/64 heads, but with the indexer that small in the e2e it isn't worth the extra complexity here. |
|
can you ask check PPL master vs this branch before merging? Probably something at slightly larger context. |
|
Sure, running it now on wikitext-2 with my prod config at 32k and 64k context, master vs this branch: llama-perplexity -m Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf -f wiki.test.raw -c 32768 -b 4096 -ub 2048 -fa on -ctk q8_0 -ctv q8_0 -ngl 999 -lm dio -lzm off -ot per_layer_token_embd=CPU (+ same with -c 65536) PPL is unchanged, wikitext-2 with my prod config:
|
|
I'm testing HIP on a ROCm machine because the CI is struggling. |
Each thread of the tile kernel now scores two keys for a single token, so a warp shares its token and the query reads are broadcasts: six shared reads per element pair instead of nine for the same products. The inner loop is unrolled by 8, which keeps gfx908 at 63 VGPRs with no spill where the fully unrolled loop needed over a thousand, and makes the kernel 36x faster on an R9700 and slightly faster on CUDA.
|
The hip-quality-check failure was real: fully unrolled, the tile kernel needed 1110 VGPRs on gfx908 and spilled the same way on RDNA4, 610 ms vs 17 ms on an R9700 at kv 65536, nb 2048. Each thread now owns two keys for one token so the query reads are broadcasts, with the inner loop unrolled by 8: 63 VGPRs on gfx908 with no spill, 36x faster on the R9700, and slightly faster on CUDA too (+2.5-5% on the op). |
am17an
added
the
merge ready
|
I'm not sure why the normal MUSA build succeeds, but this broke MUSA Docker build: Edit: Ah, normal build only does arch cc/ @yeahdongcn |
This should fix it: the tile kernel now stages the queries in head passes, two heads per pass on MUSA (25 KB of shared memory under the 28 KB cap of archs 21/22) and a single pass elsewhere that folds to the merged kernel. Validated on an R9700 with the MUSA path forced on (192/192, CUDA and ROCm unchanged): 5068f0b @yeahdongcn a run on MUSA hardware would confirm it. I can follow-up |
Sure. I can test it on an MTT GPU when I'm back at work on October 8. |
…ml-org#29901) * cuda: tile the lightning indexer over keys and tokens for 4 heads With too few heads for a wmma tile, a block scores 64 keys against 8 tokens: the keys are staged once in half precision, the queries one head at a time, and each thread owns one key for two tokens, so no dot product needs a cross thread reduction. Batches smaller than a token tile keep the vector kernel. test-backend-ops measures 4 heads. * cuda: multiply the lightning indexer tile in float Address review from am17an: the half2 products overflow once a single q * k exceeds the f16 range. The queries stay in float in shared memory and each half2 of keys is widened once for both tokens, so every product and sum is computed in float. * cuda: widen each lightning indexer key once for all heads The tile kernel stages the queries and weights of every head at once, so each key element is widened from half once and feeds all heads, with a single barrier. F16 keys are copied into the tile without a float round trip. Keeping the keys in float in shared memory measures slower, the occupancy drops. * cuda: stop the lightning indexer tile from spilling registers on ROCm Each thread of the tile kernel now scores two keys for a single token, so a warp shares its token and the query reads are broadcasts: six shared reads per element pair instead of nine for the same products. The inner loop is unrolled by 8, which keeps gfx908 at 63 VGPRs with no spill where the fully unrolled loop needed over a thousand, and makes the kernel 36x faster on an R9700 and slightly faster on CUDA.
7 participants
Overview
Follow up to #29825, as suggested by @am17an: this ports the structure of the Vulkan shader by @0cc4m to CUDA for the 4 head case (qwen4exp), where the vector kernel reads every key again for each token.
A block scores 64 keys against 8 tokens with the keys staged once in half precision and half precision products accumulated in float, and batches smaller than a token tile keep the vector kernel so decoding is unchanged.
On an RTX PRO 6000 the kernel is 2.5x faster in prefill (6.4 vs 16.4 ms at kv 65536, nb 2048), and qwen4exp with my prod config goes from PP 3867 -> 4038 t/s @ 128k with TG unchanged, the indexer dropping from 9.6% to 4.0% of the GPU time in nsys.
Additional information
test-backend-ops passes and test-llama-archs passes on CUDA for every lightning indexer model (deepseek32, deepseek4, glm-dsa, dots3note, hy_v4, glm5-next, qwen4exp).
The dispatch follows the existing style of the file, one case per K type in each branch. Listing the types once in a macro could be an optional cleanup in a separate PR, since several dispatches across the codebase could be tidied the same way, out of scope here.
Older architectures will gain even more performance
The gap was found in the benchmark @0cc4m posted on #29825: at 4 heads the Vulkan shader runs at 9.32 TFLOPs on a 3090 against 2.34 for the CUDA vector kernel (4x), and 6.12 against 2.28 on a GB10. On my RTX PRO 6000 the vector kernel already reaches 8.5 TFLOPs, so the 2.5x measured here is the low end, and I expect a larger gain on Ampere and other older architectures. Numbers from someone with a 3090 would be welcome.
Requirements