Repository navigation
[CUDA] GroupQueryAttention: opt-in tensor-core split-KV decode kernel - #33759
Yuto Ashida (y-chan) wants to merge 5 commits into
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Yuto Ashida (y-chan)
changed the title
|
At larger batch × context the opt-in kernel is currently slower than XQA on A100 (up to 1.21x) and RTX PRO 6000 (up to 1.10x), as shown in the batch table. In that regime it reads the KV cache at a lower rate than XQA (about 1.56 vs 1.75 TB/s on A100). I'm looking into a faster KV-streaming path and will update this PR with the results. |
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Opt-in (default off) fp16 split-KV decode kernel for GQA (group 8, head 128). The eligibility gate is tight (fp16, seq 1, shared BNSH cache, no softcap/sink/window/smooth-softmax, RoPE non-interleaved or packed), the kernel numerics are right in my A100 run (max error 6e-5 vs an fp64 reference, tails included), and the kernel source compiles for sm_75 through sm_120.
Blocking: the partials/counters are persistent members of the kernel object (group_query_attention.h L64-72) and the mutex only covers allocation, not use. Concurrent Session::Run calls on one session share the node instance but run on different streams, so two overlapping runs share one scratch area and one last-CTA counter array (inline comment). I reproduced wrong/NaN output and leftover non-zero counters at kernel level; per-run scratch gave 0 failures. Fix: per-run stream-tagged GetScratchBuffer and zero the counters on the run's stream (capturable).
Suggestion: TestGQASplitDecode cannot detect a missing dispatch, because the debug info prints FLASH_ATTENTION for the split path too (inline). Also uncovered: position_ids, exact 64/128 length boundaries / full cache, concurrent runs and graph replay. The slowdown vs XQA at large batch x context on A100 is acknowledged by the author (acceptable for an opt-in switch).
Validation: CI: only license/cla at the head (fork PR); no CUDA/Windows/plugin pipeline ran. Local: A100 (sm_80) standalone harness built from the PR's .cu (sequential correctness; overlap on shared vs private scratch) plus a compile-only check for sm_75/80/86/89/90/100/120 (CUDA 13.0, Linux). Not run: full ORT build, test_gqa.py, end-to-end concurrent Run, CUDA-graph replay, runtime on sm_89/90/100/120, Windows.
|
Updated the PR:
Same CUDA-graph benchmark as in the description (32 query / 4 KV heads, head size 128, QK-Norm, RoPE, fp16; us per A100 80GB SXM4:
L40:
RTX PRO 6000 Blackwell Max-Q:
The end-to-end numbers in the description were measured with the first version and have not been re-run. |
For one-token decode with 8 query heads per KV head and head size 128 (fp16, shared BNSH cache, no softcap / sink / local window, RoPE over the full head or packed QKV), ORT_GQA_SPLIT_DECODE=1 runs a dedicated kernel in place of XQA, cuDNN SDPA and the Flash kernels. Off by default. - Prep kernel: one warp per head applies QK-Norm (separate q/k epsilons), half-split RoPE and appends K/V to the cache, replacing UnpackRoPEAppend. - Attention kernel: one CTA per (batch, kv head, split). Each warp streams 8-key steps of K and V from global memory into tensor-core mma fragments with the 8 query heads as the rows of one tile, so K/V are read once per group. Warps are combined in shared memory and the last CTA of a (batch, kv head) merges the split partials, so there is no combine kernel and no memset. - The launched split count depends only on batch, KV heads, buffer length and SM count, so the grid is fixed across CUDA graph replays; the splits actually used follow the context length on the device. A100, 32/4 heads, QK-Norm, 48 layers in one CUDA graph, per layer: XQA 22.8 / 26.9 / 36.9 us vs 15.5 / 18.8 / 25.7 us at past 512 / 2048 / 8192. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add TestGQASplitDecode: one-token decode with past lengths 1 to 8191, a KV buffer 8000 rows longer than the context, batch 1-4, 8 query heads per KV head, head size 128, with and without QK-Norm, rotary and packed QKV. Each case runs with ORT_GQA_SPLIT_DECODE=1 (and checks that the Flash backend, which hosts split decode, was selected) and with the default backends, against the PyTorch reference. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…c ring At large batch x context the opt-in split decode kernel streamed KV at ~1.56 TB/s on A100 (XQA ~1.75 TB/s). Two causes: * Fragment loads straight from global memory touch only 64 bytes of each 128-byte line of K per request, which caps an SM at ~28 GB/s regardless of prefetch depth. K and V now go through a per-warp cp.async ring in shared memory (4 steps on sm_80/sm_90, 3 where the opt-in shared memory is 99 KB), copied as whole 256-byte rows and read back with an XOR swizzle. The ring is 128-byte aligned: at the default offset after the static shared variables every row straddles two lines, which cost ~13% of the streaming rate. * The split count rounded the CTA count up past the SM count (e.g. batch 8: 4 splits x 32 pairs = 128 CTAs on 108 SMs), adding a short second wave. It now fills one wave without exceeding it, and with more pairs than SMs uses one split unless the last wave would be short. A context that is not split is written straight to the output, without partials, counter or merge. A100, CUDA graph, 32 query / 4 KV heads, head 128, fp16, QK-Norm + RoPE, us per layer (new / XQA): batch 64 / past 8192: ~600 / ~613; batch 32 / past 8192: 311 / 345; batch 16 / past 8192: 165 / 174; batch 1 / past 8192: 24.6 / 37. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…eck in tests Allocate the split partials and counters per run on the run's compute stream instead of keeping them on the kernel object: concurrent Run() calls of one session share the kernel but run on their own streams, and the shared buffers let one run's last-CTA election and merge see another run's partials. The prep kernel, which runs first on the same stream, now zeroes the counters, so no memset is added and the kernel no longer has to leave them at zero. Report the path in the attention kernel debug info (SplitDecode=1, NumSplits = launched splits) and assert it in TestGQASplitDecode, which could not tell split decode from the regular Flash path. Add cases for totals on the 64/128-key boundaries, a full cache and explicit position_ids, a CUDA graph replay with context lengths that change on the device only, and concurrent runs on one session. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
34a310f to
eabb034
Compare
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
The original per-run ownership race and dispatch-observability concerns are addressed: scratch is now per invocation and stream-tagged, prep resets the counters, and the tests distinguish split decode with SplitDecode=1 while covering additional boundary/replay cases.
Two localized follow-ups remain below: avoid allocating regular-Flash workspace that the split-decode early return never uses, and make worker exceptions fail the concurrent-run regression.
Fresh validation: an AST-based probe using the actual worker function reproduced four raised worker runs while the test's final failures == [] condition still passed. The Python file parses and the diff check passes. The current head reports only CLA; no fresh native ORT/CUDA build, concurrent Session.Run execution, or CUDA Graph replay was performed. No approval submitted.
|
|
||
| auto cuda_stream = Stream(context); | ||
| if (softmax_lse_accum_bytes > 0) { | ||
| if (softmax_lse_accum_bytes > 0 && !data.use_split_decode) { |
There was a problem hiding this comment.
[P2] Skip the regular-Flash allocations, not just their memsets, on split decode
When data.use_split_decode is true, the three GetScratchBuffer calls immediately above still reserve the normal Flash LSE and accumulator buffers in addition to the new partials/counters. FlashAttention() then returns from LaunchGqaSplitDecode before any regular-Flash consumer can use these buffers. This leaves unnecessary allocator work and peak device memory in the new performance path (including potentially sizeable split accumulators). Please guard the regular-Flash workspace sizing/allocation block with !data.use_split_decode, retaining only values actually required by the split path.
| if not torch.equal(out, expected): | ||
| failures.append((past_t.tolist(), (out.float() - expected.float()).abs().max().item())) | ||
|
|
||
| pool = [threading.Thread(target=loop, args=w) for w in work] |
There was a problem hiding this comment.
[P2] Propagate exceptions from the concurrent workers into the test result
threading.Thread.join() does not re-raise worker exceptions. If runner.step() raises a CUDA/ORT error, the worker exits before appending anything to failures, so the final assertEqual(failures, []) can pass even when every concurrent run failed. This file also supports direct execution through unittest.main(), which does not turn those thread exceptions into test errors. I reproduced this with the exact worker function: 4/4 injected run failures left failures empty. Please use futures and call result() for each worker, or collect worker exceptions and fail on them after joining.
…il the concurrent test on worker errors Split decode returns before the regular Flash kernels run, so do not size or allocate the softmax LSE and split accumulator buffers for it. In the concurrent-run regression, run the workers through a ThreadPoolExecutor and collect each future's result, so an exception raised by Session.Run in a worker fails the test instead of leaving the failure list empty. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Thanks for the review. Both points are fixed in 1fe927b: the split decode path no longer sizes or allocates the regular Flash LSE/accumulator buffers, and the concurrent-run test now runs the workers through a |
Description
Adds an opt-in kernel for single-token decode in GroupQueryAttention, enabled with
ORT_GQA_SPLIT_DECODE=1(off by default, so the default backend selection is unchanged).It applies to fp16 decode with 8 query heads per KV head and head size 128, a shared BNSH KV cache,
no softcap / smooth softmax / head sink / local window, and either RoPE over the full head
(non-interleaved) or packed QKV. When enabled and Flash can run, these shapes take the new kernel in
place of XQA, cuDNN SDPA and the Flash fast-decode path. Two kernels in
group_query_attention_split_decode.{h,cu}:and appends K/V to the cache; this replaces
UnpackRoPEAppendfor these shapes.cp.async ring in shared memory (whole 256-byte rows, XOR-swizzled; 4 stages, or 3 where the opt-in
shared memory is 99 KB) into tensor-core mma fragments (QK m16n8k16, PV m16n8k8), with the 8 query
heads of the KV head as the rows of one tile, so K and V are read once per group. The warps of a CTA
are combined in shared memory and the last CTA of a (batch, kv head) merges the split partials, so
there is no combine kernel and no memset; with a single split the CTA writes the output directly.
CUDA graphs: the launched split count depends only on batch size, KV head count, KV buffer length
and SM count (as many splits as fill one wave of one CTA per SM, at most 64), so the grid is the same
in every replay. The number of splits actually used follows the context length on the device (128
keys per split at least); the remaining CTAs exit at once. The partials and counters are scratch
buffers of the run, allocated on its compute stream; the prep kernel zeroes the counters.
Motivation and Context
For one decode token the attention itself is short, and much of the time goes to fixed costs around
it. With QK-Norm (Qwen3-family models) the default XQA path runs
UnpackRoPEAppend, a memset and theXQA kernel; the Flash path runs
UnpackRoPEAppend, two memsets, the split-KV kernel and a combinekernel. On A100 at past 2048, per layer in a CUDA graph (nsys): XQA kernel 15.6 us +
UnpackRoPEAppend 5.1 + memset 2.1 + sequence lengths 1.9, vs split decode 14.7 + prep 2.7 + sequence
lengths 1.8. At longer contexts the split count scales with the SM count and the gap grows.
Results
Current numbers (batch 1 to 64 on A100, L40 and RTX PRO 6000, after the cp.async ring change) are
in this comment. The measurements below are from the first version of the kernel.
Earlier measurements (first version of the kernel)
A100 80GB SXM4, CUDA 12.9, fp16, batch 1, 32 query / 4 KV heads, head size 128, QK-Norm, RoPE.
48 GQA nodes with independent KV caches in one CUDA graph, GPU time per replay / 48 (us per layer):
ORT_GQA_SPLIT_DECODE=1ORT_ENABLE_XQA=0)Without a CUDA graph (same model): 26.3 / 28.6 / 39.0 (default) vs 20.3 / 21.3 / 30.5 us.
Single op with the existing
profile_gqa.py(kernel time per call, nsys; the KV cache stays in L2):default 21.5 / 23.8 / 33.2 us + 2.3 us memset vs 17.1 / 19.5 / 24.0 us and no memset.
Without QK-Norm the default path is XQA (and Flash fast decode with
ORT_ENABLE_XQA=0); in a CUDAgraph with packed QKV the opt-in kernel was 15.0 / 19.4 / 24.6 us vs XQA 23.7 / 31.0 / 36.0 us.
End to end, on a non-public decoder with the Qwen3-30B-A3B architecture (48 layers, 32/4 heads,
QK-Norm, CUDA graph, batch 1, 300 tokens): 5.84 -> 5.00 ms/token (171.3 -> 200.1 tokens/s) with
ORT_GQA_SPLIT_DECODE=1, generated text unchanged (only timestamps in the transcript moved).Other GPUs, same CUDA-graph benchmark (us per layer, past 512 / 2048 / 8192). On SM >= 9 the
cuDNN SDPA column is what
ORT_ENABLE_XQA=0selects:ORT_GQA_SPLIT_DECODE=1Without a CUDA graph: L40 14.7 / 19.3 / 40.6 (default) vs 12.0 / 16.9 / 36.2 us; RTX PRO 6000
22.9 / 23.7 / 28.0 (default) vs 17.2 / 20.0 / 26.7 us (cuDNN SDPA 40.0 / 31.3 / 32.6).
Batch size (CUDA graph, same shapes; us per layer at past 512 / 2048 / 8192; below 100% means the
opt-in kernel is faster):
A100:
ORT_GQA_SPLIT_DECODE=1L40:
ORT_GQA_SPLIT_DECODE=1RTX PRO 6000:
ORT_GQA_SPLIT_DECODE=1The gain is largest for batch 1, where fixed costs dominate. With more batch x context the work is
bound by reading the KV cache, and there XQA reads it faster on A100 (batch 64 / past 8192: about
1.75 TB/s vs 1.56 TB/s for the opt-in kernel), so the opt-in kernel is up to 1.21x slower on A100 and
up to 1.10x on RTX PRO 6000 (where cuDNN SDPA is the fastest or tied at large batch). On L40 both
saturate its lower bandwidth and the opt-in kernel stays slightly ahead.
End to end on the other GPUs (same decoder architecture, 300 tokens, CUDA graph):
ORT_GQA_SPLIT_DECODE=1(*) A different export of the same architecture (block-wise QMoE scales, 4096-token KV buffer). On L40
the generated text was unchanged; on RTX PRO 6000 the output showed small numerical differences
(a few timestamps and punctuation marks), as expected from a different accumulation order.
Tests
TestGQASplitDecode(test_gqa.py), 19 parity cases: past lengths 1 to 8191 including totalson the 64 / 128 / 256 boundaries, full caches, a KV buffer 8000 rows longer than the context, batch
1 to 40, 16/2, 32/4 and 64/8 heads, with and without QK-Norm, rotary, packed QKV and explicit
position_ids. Each case runs withORT_GQA_SPLIT_DECODE=1and with the default backends againstthe PyTorch reference, and checks the debug info (
SplitDecode=1only with the flag). The existingGQA tests use other group sizes.
test_gqa_split_decode_cuda_graph_replay: one captured graph replayed while the context lengthschange on the device only (across split counts and split edges), compared with the default
backends at every step.
test_gqa_split_decode_concurrent_runs: one session run from four threads, each output comparedwith that thread's serial result.
TestGQASplitDecode21 passed (also on RTX PRO 6000,sm120).
test_gqa.pyon L40: 975passed, 4 failed, 1 skipped; all four also fail without this change: the two A100 failures listed
below, and two INT8 KV-cache cases ("No available GroupQueryAttention kernel supports the quantized
KV cache").
onnxruntime_provider_test --gtest_filter=*GroupQueryAttention*113 passed,onnxruntime_test_all31 passed.test_gqa.py:751 passed, 194 failed, 3 skipped;
the failures are unrelated and also occur without this change: 192 FP8 KV-cache XQA cases expect
XQA, which is not available on SM 8.0;
TestFlashGQA::test_gqa_decode_bidirectionalis rejected bythe input check;
TestGQAAttentionBias::test_windowed_cache_attention_bias_is_rejectedhits theKV-capacity check before the attention-bias check, so the expected message differs.
How to reproduce
🤖 Generated with Claude Code