Skip to content

[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends #3286

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pggPL wants to merge 30 commits into
base: main
Choose a base branch
from
Open

[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends #3286

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
f6ed27c
[PyTorch][torch.compile] Support for DotProductAttention
pggPL Jul 27, 2026
3053d76
Delegate the no-FP8 case of DPA's init_fp8_metadata to the base class
pggPL Jul 28, 2026
d6b285e
Run DPA eagerly for undeclared packed q/k/v instead of guessing the l…
pggPL Jul 28, 2026
1cd6704
Log backend selection in eager only; shorten comments
pggPL Jul 28, 2026
42a9e6c
Fix lint: unused import and keyword-arg-before-vararg
pggPL Jul 29, 2026
af11f8f
Compare against eager under CUDA graphs, and tighten the tolerances
pggPL Jul 29, 2026
0b64c96
Support one cu_seqlens tensor for both q and kv under torch.compile
pggPL Jul 29, 2026
7f480c6
Apply the shared-cu_seqlens workaround to the FlashAttention v4 path too
pggPL Jul 29, 2026
adc96c4
Run FP8 attention eagerly under torch.compile
pggPL Jul 29, 2026
737904f
Drive the compile tests off ModelConfig instead of a hand-picked inte…
pggPL Jul 29, 2026
9026a41
Scale the unfused comparison to the tensor, and take tolerances from …
pggPL Jul 29, 2026
bd2c506
Merge the per-backend compile tests, and give padding masks real padding
pggPL Jul 29, 2026
725931e
Run eagerly when max_seqlen has to be derived from cu_seqlens
pggPL Jul 29, 2026
b2ab7eb
Share the run-and-compare code between the compile tests
pggPL Jul 29, 2026
50d1984
Use the shared run-and-compare code in the remaining compile tests
pggPL Jul 29, 2026
f19c52d
Read the predicate's arguments by name, not by hardcoded position
pggPL Jul 29, 2026
fbc7146
Move the eager-fallback decorator to jit.py
pggPL Jul 29, 2026
6b24630
Skip FusedAttention rather than special-casing it in the compile tests
pggPL Jul 29, 2026
9a204b8
Cover the compiled path around FusedAttention
pggPL Jul 30, 2026
2188f18
Log backend selection through the no-op logger, as get_attention_back…
pggPL Jul 30, 2026
10d141b
Return the fused sub-backend as an int, in eager as well
pggPL Jul 30, 2026
93a201d
Rename eager_under_compile_if to fallback_to_eager_when
pggPL Jul 30, 2026
0866319
Fold the eager fallback into no_torch_dynamo as a `when` predicate
pggPL Jul 30, 2026
a80a41a
Cover the declared packed layouts beyond bs3hd
pggPL Jul 30, 2026
9eb1449
Let ONNX export keep its own path in get_qkv_layout
pggPL Jul 30, 2026
5c5b451
Guard the assumption the argument binding rests on
pggPL Jul 30, 2026
95030f9
Check that CUDA graphs were actually captured, and generalize a docst…
pggPL Jul 30, 2026
b0e56d4
Skip lazy compilation while already tracing
pggPL Jul 30, 2026
6f7c818
Merge upstream/main into dpa_torch_compile
pggPL Jul 30, 2026
2cf695b
Run context parallel attention eagerly
pggPL Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading