[NVIDIA] perf: pipeline the extend-attention KV loads on Hopper - #36102
[NVIDIA] perf: pipeline the extend-attention KV loads on Hopper#36102Hert4 wants to merge 5 commits into
Conversation
num_stages for the Triton extend-attention kernel has been hardcoded to 1 since the kernel was written, with no comment and no prior discussion. The KV loop body is a straight-line global -> shared load of the K and V tiles followed by two dots, which is the shape Triton can software pipeline. Measured on H200 (sm_90), no sliding window, n=16384, bf16: head_dim q/kv ns=1 ns=2 ns=3 best 64 16/8 4.231 3.982 3.587 1.18x 64 32/4 7.960 7.499 6.771 1.18x 128 16/8 5.809 5.750 5.199 1.12x 128 32/4 10.895 10.686 10.045 1.08x 256 16/8 18.227 16.228 14.500 1.26x 256 32/4 35.077 31.708 28.354 1.24x 512 16/8 45.960 31.760 OOM 1.45x 512 32/4 91.029 61.047 OOM 1.49x At Lq=512 three stages need 299008 bytes of shared memory against a 232448 byte limit, hence the cap at 2 there. Output is bit-identical across num_stages 1/2/3 at all 16 configurations tested (exact comparison of output fingerprints, not a tolerance check), as expected for a scheduling hint. HIP and pre-Hopper CUDA keep num_stages=1: unmeasured, and the gfx950 specialization already fights register spills (sgl-project#34741). The unified kernel is deliberately left alone for the same reason. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6a1891a89a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not _is_cuda or CUDA_CAPABILITY[0] < 9: | ||
| return 1 | ||
| return 3 if Lq <= 256 else 2 |
There was a problem hiding this comment.
On sm120 (CUDA_CAPABILITY[0] == 12), this condition also returns 3 even though the block-size selector uses BLOCK_N=128 specifically because that workstation Blackwell GPU has only 100 KiB of shared memory. For normal bf16/fp16 attention with Lq=Lv=128, the three staged 128x128 K and V tiles alone require 192 KiB, so Triton cannot compile or launch extend prefill on RTX Pro 6000. Limit this override to the measured Hopper architecture or use a safe sm120-specific stage count.
Useful? React with 👍 / 👎.
|
@BBuf you tuned the Hopper block sizes for this exact kernel in #26588, so this is probably in your area: this PR is the same kind of arch-gated tuning, one line further down. It raises Measured 1.08x to 1.49x on H200 across I do not have the |
|
@alexnails thanks for picking this up. Could you add the Everything else is ready for review: the change is one helper plus one line at the call site, no signature or kernel changes, and output is bit-identical to |
|
Small update on the CI here, since none of the red checks on this PR come from the change itself. The root cause is #35961 renamed that helper to I opened #36175 with the four-line fix. Once that lands, a rebase here should give a clean run. This PR touches exactly one file, |
Codex review caught a real problem with the >= 9 gate: it also covers sm120 and sm100, neither of which was measured. sm120 picks BLOCK_N=128 for Lq<=128 precisely because it only has 100 KiB of shared memory. Three stages of 128x128 K and V tiles is 192 KiB, so extend prefill would fail to compile on RTX Pro 6000. sm100 uses its own block sizes to dodge PTX register exhaustion, and deeper pipelining has not been checked against that either. Narrowed the condition to == 9, which is exactly the architecture the numbers in this PR come from. Everything else keeps num_stages=1. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
bd524b0 to
1efcb1d
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1efcb1d73d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| """ | ||
| if not _is_cuda or CUDA_CAPABILITY[0] != 9: | ||
| return 1 | ||
| return 3 if Lq <= 256 else 2 |
There was a problem hiding this comment.
For Hopper's supported bf16 Lq=576, Lv=512 configuration (used by test_kimi_k3_absorbed_mla_shape), this selects two stages solely from Lq. With BLOCK_N=64, each stage must hold the 576-wide K tile (512 plus the 64 DPE columns) and the 512-wide V tile: 64 * (576 + 512) * 2 = 139,264 bytes; two stages require 278,528 bytes, exceeding the 232,448-byte Hopper limit documented above. Triton will therefore reject compilation/launch for extend prefill on this supported Kimi shape; include Lv/DPE in the limit or retain one stage for this case.
Useful? React with 👍 / 👎.
|
@alexnails when you have a spare moment, would you mind taking a look at this one? You wrote #34462 on this exact kernel, so you are probably the person best placed to judge whether the architecture gate here is drawn in the right place. The change itself is one line at the call site plus a helper: Two things have changed since the Codex review:
On the red checks: none of them are test failures. All 14 are No rush at all, and thanks for adding the label earlier. |
2 participants
Motivation
num_stagesfor the Triton extend-attention kernel has been hardcoded to1since the kernel was written, with no comment and no prior discussion that I could find. The KV loop body is a straight-line global-to-shared load of the K and V tiles followed by two dots, which is exactly the shape Triton can software pipeline. At one stage the global loads never overlap with compute, so the kernel waits on HBM.decode_attention.pyright next door already carries tuned buckets withnum_stages=2.There is precedent for architecture-gated tuning in this same file: #26588 (Hopper block sizes for
Lq=129..256), c1067f8 (gfx950 block sizes forhead_dim > 128), #34741 (gfx950 tile shape on Triton 3.7).Modifications
One new helper next to
_get_block_sizes_for_extend_attention, and one line changed at the call site inextend_attention_fwd. No signature changes, no kernel changes.The condition order matters:
CUDA_CAPABILITYis only defined inside theif _is_cuda:branch at the top of the module, so the short circuit is load-bearing.HIP and pre-Hopper CUDA are untouched and keep
num_stages=1. I have no ROCm hardware to measure on, and the gfx950 specialization is already fighting register spills per #34741, so raising the pipeline depth there is not something I want to guess at.extend_attention_fwd_unifiedalso has a hardcodednum_stages = 1. I left it alone deliberately, since I did not benchmark that path.Accuracy Tests
num_stagesis a scheduling hint and must not change results, so I checked with exact comparison rather than a tolerance. For every configuration I took an output fingerprint (sum, max absolute value, and the sum over a strided slice) atnum_stages1, 2, and 3 and compared them for equality.All 16 configurations produced identical fingerprints. Zero mismatches.
I was not able to run
test/registered/attention/locally: those tests pull a tokenizer from HuggingFace and the machine I benchmark on has no outbound internet access. CI should cover that once this is labeled.Speed Tests and Profiling
NVIDIA H200 (sm_90),
shared_memory_per_block_optin = 232448, bf16,n = 16384, fresh prompt with no prefix, 4 warmup and 12 timed iterations. Times in milliseconds.The OOM at
Lq=512, num_stages=3isRequired: 299008, Hardware limit: 232448. That is the reason the helper caps at 2 aboveLq=256rather than going to 3 everywhere.A note on the sliding-window path
I also measured all eight shapes above with
sliding_window_size=1024. There the gain is 1.00x to 1.03x, essentially nothing. That is expected rather than disappointing: on currentmainthe stage-2 loop still walks from block 0 and relies onSKIP_TILEto skip the body, so most iterations load nothing at all and there is nothing to pipeline.Once the loop is bounded by the window (#34462), that path becomes load-bound again and
num_stagesstarts paying there too. With that bound applied I measured a further 1.17x fromnum_stages=2athead_dim=256,window=1024,n=27000. The two changes compose; this one stands on its own for the non-windowed case.What this does not cover
Lq <= 256threshold is the boundary I measured, not something derived from a shared-memory formula. Other Hopper parts share the same limit so it should hold, but I only have one kind of card to say that with.CI States
Latest PR Test (Base): ❌ Run #32804741550
Latest PR Test (Extra): ❌ Run #32804741369
Latest PR Test (AMD ROCm 7.2): ❌ Run #32804741500