Repository navigation
CUDA backend: lower fused attention variants - #24426
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24426
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 4a201ca 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. |
|
@claude /pr-review |
meta-cla
Bot
added
the
CLA Signed
|
Claude finished @shoumikhin's task in 27m 57s —— View job No blockers; the main gaps are that the tests neither check what the pass emits nor go through the CUDA decomposition this PR fixes, that the causal length guard keeps calls PR Review: #24426SummaryThis PR adds Code Quality
Testing
Performance
RecommendationApprove. Nothing blocks merging. The two test gaps are the ones worth fixing first: as written, no test would catch a rewrite with wrong arguments on ROCm, or a mismatch with the graph the CUDA decomposition really produces. |
c8fd243 to
be8d587
Compare
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 29m 51s —— View job No blockers. The one gap is that no test checks that the rewritten calls reach the Triton kernels, so a change that sends them to the slower math path would still pass every test. PR Review: #24426SummaryThis PR adds Testing
RecommendationApprove. Nothing blocks merging. Closing the test gap takes a few lines in |
Decomposing scaled_dot_product_attention on a CUDA device turns it into
one of its fused variants: _scaled_dot_product_flash_attention without a
mask, or _scaled_dot_product_efficient_attention with one. torch.export's
default decompositions do this, and so does Torch-TensorRT for an
attention it leaves to the CUDA backend, for example one whose keys are a
constant. The CUDA backend had no kernel for these ops:
RuntimeError: Method forward missing fallback kernels (1 total):
- at::_ops::_scaled_dot_product_flash_attention::call
A new pass, ReplaceFusedSdpaWithSdpaPass, turns the flash, efficient and
cuDNN variants back into scaled_dot_product_attention with the same
query, key, value, bias, dropout, causal flag and scale. It sets
enable_gqa unless the query and key head counts are known to be equal,
since the fused variants group fewer key heads the same way. It leaves
out a scale of None, because the Triton SDPA kernels take a float scale.
The Triton SDPA kernel or Inductor's math decomposition then lowers the
attention, as it does any other. The pass also looks inside control flow
branches.
The pass rewrites a node only when scaled_dot_product_attention computes
the same result. Only its attention output may be used, which is the case
for inference. A causal call must have no bias, since
scaled_dot_product_attention takes no mask together with is_causal. Its
query and key lengths must also be known to be equal: flash attention
aligns a causal mask to the bottom right, the other variants and
scaled_dot_product_attention to the top left, and the Triton decode
kernel for one query ignores it. Otherwise the node stays, the pass logs
why, and the error above still names it.
test_fused_sdpa decomposes attention on a CUDA device into a causal flash
attention with fewer key heads and a custom scale, an efficient attention
with a mask and a flash attention of one query over 256 keys. It lowers
the program, checks that the first and last calls get the Triton kernels,
runs it, and compares the outputs with eager. The tests in
passes/tests/test_replace_fused_sdpa.py check on the CPU which calls the
pass rewrites and with which keyword arguments (cuDNN, causal with equal
static or dynamic lengths, fewer key heads, inside a torch.cond branch),
which it keeps (causal with lengths that differ or may differ, causal with
a bias, a used log-sum-exp), that the program stays valid, and that the
pass runs before the Triton kernel replacement. They fail without this
change.
be8d587 to
4a201ca
Compare
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 25m 36s —— View job No blockers. The main gap is that the CPU test of the new pass only checks the keyword arguments of the call it emits. So on ROCm, and for cuDNN and PR Review: #24426SummaryThis PR adds
Nothing was run, because Python needs approval in this job. Statements about which variant PyTorch picks for a given attention come from PyTorch's SDPA dispatch, which I couldn't check here. Code Quality
Testing
RecommendationApprove. Nothing blocks merging. Fix the test gap first. |
|
Thanks for the review. Answers to the two items: Causal calls with different lengths in the efficient and cuDNN variants. I agree that these two variants align the causal mask to the top left, like Positional arguments in the CPU test. The CPU test checks what the pass decides: which calls it rewrites, the scale, and |
|
@pytorchbot cherry-pick --onto cherry-pick-24420-by-pytorch_bot_bot_ -c critical |
Cherry picking #24426The cherry pick PR is at #24455 and it is recommended to link a critical cherry pick PR with an issue. Details for Dev Infra teamRaised by workflow job |
What is wrong today
Decomposing
scaled_dot_product_attentionon a CUDA device turns it into one of its fused variants:_scaled_dot_product_flash_attentionwithout a mask, or_scaled_dot_product_efficient_attentionwith one.torch.export's default decompositions do this. So does Torch-TensorRT, for an attention it leaves to the CUDA backend, for example one whose keys are a constant. The CUDA backend has no kernel for these ops:What this change does
A new pass,
ReplaceFusedSdpaWithSdpaPass, turns the flash, efficient and cuDNN variants back intoscaled_dot_product_attention, with the same query, key, value, bias, dropout, causal flag and scale. The Triton SDPA kernel or Inductor's math decomposition then lowers it, as it lowers any other attention. The pass also looks inside control flow branches.enable_gqaunless the query and key head counts are known to be equal, because the fused variants group fewer key heads the same way.None, because the Triton SDPA kernels take a float scale.The pass rewrites a node only when
scaled_dot_product_attentioncomputes the same result:scaled_dot_product_attentiontakes no mask together withis_causal.scaled_dot_product_attentionto the top left, and the Triton decode kernel for one query ignores it.Otherwise the node stays, the pass logs why, and the error above still names it.
What was tested
With torch 2.14 on an A100:
test_fused_sdpadecomposes attention on a CUDA device into a causal flash attention with fewer key heads and a custom scale, an efficient attention with a mask, and a flash attention of one query over 256 keys. It then lowers the program, checks that the first and last calls get the Triton kernels, runs it, and compares each output with eager.test_replace_fused_sdpa.pychecks on the CPU which calls the pass rewrites and with which keyword arguments, which calls it keeps, that the program stays valid, and that the pass runs before the Triton kernel replacement.They fail without this change. The other CUDA tests in the same CI job pass.