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

Merged

pggPL merged 47 commits into

NVIDIA:mainfrom

pggPL:dpa_torch_compile

Aug 17, 2026

Commits on Jul 29, 2026

[PyTorch][torch.compile] Support for DotProductAttention
pggPL
committed
f6ed27c
Delegate the no-FP8 case of DPA's init_fp8_metadata to the base class
pggPL
committed
3053d76
Run DPA eagerly for undeclared packed q/k/v instead of guessing the layout
pggPL
committed
d6b285e
Log backend selection in eager only; shorten comments
pggPL
committed
1cd6704
Fix lint: unused import and keyword-arg-before-vararg
pggPL
committed
42a9e6c
Compare against eager under CUDA graphs, and tighten the tolerances
pggPL
committed
af11f8f
Support one cu_seqlens tensor for both q and kv under torch.compile
pggPL
committed
0b64c96
Apply the shared-cu_seqlens workaround to the FlashAttention v4 path too
pggPL
committed
7f480c6
Run FP8 attention eagerly under torch.compile
pggPL
committed
adc96c4
Drive the compile tests off ModelConfig instead of a hand-picked intersection
pggPL
committed
737904f
Scale the unfused comparison to the tensor, and take tolerances from utils
pggPL
committed
9026a41
Merge the per-backend compile tests, and give padding masks real padding
pggPL
committed
bd2c506
Run eagerly when max_seqlen has to be derived from cu_seqlens
pggPL
committed
725931e
Share the run-and-compare code between the compile tests
pggPL
committed
b2ab7eb
Use the shared run-and-compare code in the remaining compile tests
pggPL
committed
50d1984
Read the predicate's arguments by name, not by hardcoded position
pggPL
committed
f19c52d
Move the eager-fallback decorator to jit.py
pggPL
committed
fbc7146
Skip FusedAttention rather than special-casing it in the compile tests
pggPL
committed
6b24630

Commits on Jul 30, 2026

Cover the compiled path around FusedAttention
pggPL
committed
9a204b8
Log backend selection through the no-op logger, as get_attention_backend does
pggPL
committed
2188f18
Return the fused sub-backend as an int, in eager as well
pggPL
committed
10d141b
Rename eager_under_compile_if to fallback_to_eager_when
pggPL
committed
93a201d
Fold the eager fallback into no_torch_dynamo as a `when` predicate
pggPL
committed
0866319
Cover the declared packed layouts beyond bs3hd
pggPL
committed
a80a41a
Let ONNX export keep its own path in get_qkv_layout
pggPL
committed
9eb1449
Guard the assumption the argument binding rests on
pggPL
committed
5c5b451
Check that CUDA graphs were actually captured, and generalize a docstring
pggPL
committed
95030f9
Skip lazy compilation while already tracing
pggPL
committed
b0e56d4
Merge upstream/main into dpa_torch_compile
pggPL
committed
6f7c818
Run context parallel attention eagerly
pggPL
committed
2cf695b

Commits on Aug 4, 2026

Keep DotProductAttention compilable under a CUDA RNG states tracker

pggPL
and
claude
committed
b89446f
Run checkpointed attention eagerly

pggPL
and
claude
committed
fb50154
Support KV caching under torch.compile

pggPL
and
claude
committed
4b1ebb2
Compile the fused sbh3d QKV split

pggPL
and
claude
committed
eb5ad55
Run FlashAttention 4 eagerly

pggPL
and
claude
committed
7da1a66
Pin the sequence lengths backend selection bakes in

pggPL
and
claude
committed
d1beb9b
Capture CUDA graphs with a KV cache

pggPL
and
claude
committed
4af8c30
Test generation against a KV cache

pggPL
and
claude
committed
bd94fff
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed
688e10c
Drop the imports the custom ops made unused

pggPL
and
claude
committed
43eba11

Commits on Aug 7, 2026

Rename _maybe_unshare_cu_seqlens to _unalias_cu_seqlens
pggPL
committed
c4a38bc
Split the get_full_cu_seqlens cache bypass per reason
pggPL
committed
1b495ce
Merge remote-tracking branch 'upstream/main' into dpa_merge
pggPL
committed
753cfeb

Commits on Aug 13, 2026

Commits on Aug 17, 2026

Skip the max_seqlen bake-in under ONNX export
pggPL
committed
075bce8
Make undeclared non-contiguous q/k/v contiguous under compile
pggPL
committed
b3e9949