Repository navigation
CUDA backend: keep wide attention heads that are not a power of 2 off the Triton SDPA kernel - #24427
shoumikhin wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24427
Note: Links to docs will display an error until the docs builds have been completed. ⏳ 4 Pending, 1 Unrelated FailureAs of commit 82b00ec with merge base ab68460 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 15m 52s —— View job The new head-size bound is correct and nothing blocks merging. The five issues are minor and mostly about tests: heads of 1024 are said to still use Triton, but nothing compiles the kernel at that size, and no test checks wide heads on ROCm. Detailed reviewPR Review: #24427SummaryThis PR makes Code Quality
Testing
API Design
Backward Compatibility
RecommendationApprove: none of the above blocks the fix. Before the cherry-pick, the most important steps are compiling a head of 1024 and checking
Not done: no GPU was available, so no CUDA tests were run, and |
…PA kernel The Triton SDPA kernel for head sizes that are not a power of 2 loads the whole head in one tile, and that tile is at most 256 wide. For a wider head, such as 640 in the mid-block attention of a video VAE encoder, it reads only the first 256 values of each query and key, and it never writes the rest of the output, which then holds whatever was in memory. The program runs and returns wrong numbers. The shape check that the kernel and the replacement pass share now rejects such heads, so the pass leaves the call to the regular CUDA lowering. Heads that are a power of 2 use the other kernel, which covers the whole head, so they still use Triton. test_sdpa_head_sizes lowers attention over head sizes on both sides of the limit and compares each result with eager. Two smaller tests check that the pass still gives heads of 512 and 1024 to Triton, and that a direct call with a wide head raises.
The check now compares the head size with the tile that _next_power_of_2 picks, so it stays right if that cap changes. The docstring of triton::sdpa states the limit for direct callers. A head of 640 is now exported through the regular lowering in the tests, which also checks the logged reason. The hand-built support test keeps only the heads that the kernel still takes, since some CI GPUs lack the shared memory to compile a head of 1024.
03c6ed4 to
47603ba
Compare
|
Thanks. The new commit fixes three of these: the check now follows the tile size from Two I left as they are:
|
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 15m 11s —— View job The new head-size check is correct and nothing blocks merging; the one minor issue is that the new Detailed reviewPR Review: #24427SummaryThis PR makes API Design
RecommendationApprove: the issue above is minor and predates this PR, and these heads fail with an error rather than giving wrong numbers.
Not done: this runner allows neither Python nor GPU commands, so no tests were run. |
…string The docstring said any power-of-2 head works, but the power-of-2 kernel has its own shared-memory limits for very wide heads. It now states only the limit for heads that are not a power of 2, and the support test is named for what it checks: which heads the pass gives to the kernel.
|
Thanks. The docstring now states only the limit this change adds: a head that is not a power of 2 must be at most 256. The support test is renamed to say what it checks, which is the heads the pass gives to the kernel. Rejecting power-of-2 heads above 1024 is a separate change. Those heads cannot give wrong numbers today: when they do not fit, compiling stops with an error. The right limit for heads of 1024 and below depends on how much shared memory each GPU has, so I am keeping it for a follow-up. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 13m 29s —— View job The head-size check is correct and nothing blocks merging. One minor note: the heads it rejects now use the MATH fallback, which gets slower with the square of the sequence length, and the timing in the description covers only 64 tokens. Detailed reviewPR Review: #24427SummaryThis PR makes Performance
RecommendationApprove: the issue above is minor, and the change replaces wrong numbers with correct ones.
Not done: no tests were run. |
|
Agreed. The timing is for the shape we measured only, which is 64 tokens. For long sequences the regular lowering builds the full score matrix, so it can need much more time and memory. The description now says so. Keeping these heads on Triton, for example by padding them with zeros to a power of 2, is a follow-up. |
What is wrong today
The CUDA backend replaces
scaled_dot_product_attentionwith its own Triton kernel when the kernel accepts the call. For head sizes that are not a power of 2, that kernel holds the whole head in one tile of at most 256 values. Wider heads of that kind, such as 288, 320, 384, 640 or 768, are still accepted. The kernel then uses only the first 256 values of each query and key, and it never writes the rest of the output. The program runs and returns wrong numbers, with no error. For example, the mid-block attention of the LingBot-VA video VAE encoder has heads of 640, and its lowered encoder drifts far from eager.What this change does
The shape check that the kernel and the replacement pass share now rejects head sizes that are not a power of 2 and are wider than that one tile. The pass already leaves calls the kernel rejects to the regular CUDA lowering, so such attention in an exported model now gives the right result. Code that calls the Triton kernel directly, such as the off-graph KV lowering, now gets an error instead of wrong numbers. Heads that are a power of 2, including 512 and 1024, use the other Triton kernel, which covers the whole head, so they still use Triton.
Because this fixes wrong results, it should also go to
release/1.6.What was tested
On an A100 with torch 2.14:
test_sdpa_head_sizeslowers attention with heads of 96, 192, 256, 512, 288, 320, 384, 640 and 768, runs it, and compares each result with eager. Without this change the last five fail.test_heads_the_kernel_takes_are_supportedchecks that the pass still gives heads of 192, 256, 512 and 1024 to the Triton kernel.test_wide_head_that_is_not_a_power_of_2_is_not_replacedchecks that a head of 640 exports through the regular lowering and logs why.test_wide_non_pow2_head_rejectedchecks that a direct call with a head of 640 raises.test_sdpa_splitk_replacement.pystill pass.