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

[pggPL]

Description

DotProductAttention.forward no longer carries @no_torch_dynamo, so torch.compile captures the whole module -- input unpacking, qkv layout, backend selection and the backend itself -- with fullgraph=True on the FlashAttention and UnfusedDotProductAttention backends.

FusedAttention stays an eager island (@no_torch_dynamo on its own forward): the graph break moves from the entry of DotProductAttention down to the backend call, so everything around it is now compiled.

Follow-up to #3189 (traceable get_attention_backend), #3200 (declarative packed QKV), #3201 (unfused DPA) and #3268 (device-side padding mask).

Type of change

• New feature (non-breaking change which adds functionality)
• Bug fix (non-breaking change which fixes an issue)

Changes

Removing the decorator exposed several things that could not be traced. Each is fixed at the source rather than worked around:

• DotProductAttention.init_fp8_metadata called get_fp8_recipe() unconditionally, which builds a default Recipe when no autocast is active -- and constructing one is not traceable. With FP8 off it now returns through the base class, which is where the rest of that function ended up anyway.

• get_qkv_layout recognizes packed q/k/v from UntypedStorage.data_ptr and Tensor.storage_offset, neither of which dynamo can trace, and neither of which means anything on the fake tensors used while tracing. While tracing, the layout comes from the format instead; packed inputs are declared through qkv_layer/kv_layer ([PyTorch][torch.compile] DotProductAttention: declarative packed QKV/KV inputs #3200). ONNX export keeps its own path -- the exporter runs through dynamo, so it sets is_compiling() as well.

• get_indices, which packs bshd/sbhd inputs for FlashAttention's varlen entry point, built the index list with a Python comprehension over sequence lengths held in a GPU tensor: a device synchronization per sequence in eager, and data-dependent while tracing. It is now built on the device, verified exhaustively against the previous implementation (10304 cases, 0 mismatches).

• get_full_cu_seqlens keyed its cache on torch.is_inference_mode_enabled(), a fundamental graph break. The cache only saves an allocation, so while tracing the tensor is built directly.

• The fused sub-backend enum does not survive a graph break. _get_fused_attn_backend is decorated with @torch.compiler.assume_constant_result, and dynamo reconstructs such a value by re-emitting the call that produced it -- which is only valid inside the frame that made it. Crossing the break into FusedAttention, the resumed frame received the producing function instead of the enum member, and it reached fused_attn_fwd. get_attention_backend now returns the sub-backend as a plain int, which is baked into the graph as a literal; FusedAttnBackend.cast turns it back into a member, as fused_attn_fwd/bwd already did.

• Backend selection is logged through the no-op logger while tracing, as get_attention_backend does: logging.Logger methods graph-break.

• lazy_compile skips its nested torch.compile while tracing. Filling its closure cell is a side effect dynamo rejects inside a higher-order op, so compiling a module whose autograd.Function backward makes the first call to a jit_fuser'd function failed with Unsupported: HOP: Unsafe side effect. Here that is PackTensors.backward on the FlashAttention varlen path; an earlier eager run hides it by filling the cell beforehand. The nested torch.compile is inlined into the outer graph either way, so calling the function directly while tracing is equivalent -- and in eager nothing changes.

• no_torch_dynamo takes a when predicate, so a method can be an eager island for some calls only. The predicate receives the call's arguments keyed by parameter name and returns the reason a particular call cannot be traced.

Falling back to eager

Four kinds of call cannot be compiled and run as an eager island instead, each with a warning naming what to do differently:

	why	how to stay compiled
FP8 attention (fp8_dpa/fp8_mha)	quantizers, fp8_meta mutation and Float8Tensor crossing the graph boundary, none of it covered by tests	--
context parallelism (cp_group set)	the implementation is built around P2P communication and its own CUDA stream, neither of which dynamo traces	--
packed q/k/v passed as plain views	recognizing them takes the data pointers dynamo cannot read	declare the packing via qkv_layer/kv_layer
thd without max_seqlen_q/max_seqlen_kv	deriving it reads the sequence lengths off cu_seqlens: a device synchronization, and a data-dependent value while tracing	pass max_seqlen_q/max_seqlen_kv

The predicate is deliberately narrow for FP8: it asks the recipe for fp8_dpa/fp8_mha rather than whether an autocast is active, so FP8 GEMMs with attention in high precision -- the common training setup -- stay on the compiled path.

Without fullgraph, these calls work and warn. With fullgraph=True they raise, which is what asking for a full graph means.

Testing

tests/pytorch/test_torch_compile.py, 56 passed / 23 skipped on RTX Ada (sm89, flash-attn 2.7.4):

• the configuration matrix -- 14 configurations described with the same ModelConfig the eager attention tests use, against each backend that supports them, which get_available_attention_backends() decides. It covers bshd/sbhd/thd, causal / no_mask / padding / sliding window / arbitrary masks, GQA, cross attention, alibi, post-scale bias, MLA head dims, off-by-one softmax, one shared cu_seqlens tensor for q and kv, and all four declared packed layouts. Compiled with fullgraph=True, compared against eager in forward and every input gradient;
• CUDA graphs (mode="reduce-overhead") -- fresh inputs each iteration and compared against eager, since a replay reading stale buffers produces finite values and non-None gradients too; inductor's skip counter is asserted to be zero, so a declined capture cannot pass silently;
• the compiled path around FusedAttention -- it is not compiled itself, but everything before and after the break is, and what crosses that break has to survive it;
• the eager fallbacks -- that each fires exactly when it should, matches eager, and raises under fullgraph=True.

Backend selection is invalidated right before every compiled call, so it is traced rather than served from the cache the eager call populated.

Comparison is exact for backends whose kernel is the same either way. The unfused backend is built from PyTorch ops that inductor fuses and reassociates, so its absolute tolerance is taken from the tensor's own scale -- an off-by-one softmax differs by two bfloat16 roundings on individual elements, which no elementwise relative tolerance covers near zero.

attention/test_attention.py, attention/test_kv_cache.py, test_gqa.py and test_onnx_export.py show no new failures relative to main.

Not covered by this hardware and wanted from CI: FP8 and MXFP8 attention, the fused sub-backends, FlashAttention 3 and 4.

Known limitations

• FusedAttention is not compiled; it is skipped from the matrix, and removing that skip is all that is needed once its forward traces. Registering tex.fused_attn_fwd/bwd as custom ops -- which is what lets flash-attn's kernels sit inside a graph -- would be the way there.
• The unfused backend is compared to eager with a tolerance rather than exactly, for the reason above.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

The PR enables broader torch.compile coverage for DotProductAttention.

• Conditionally preserves eager execution for unsupported FP8, context-parallel, undeclared packed-view, and incomplete THD calls.
• Makes QKV layout handling, sequence-index generation, backend selection, fused-backend values, and lazy JIT helpers trace-compatible.
• Adds forward, backward, CUDA-graph, backend, and eager-fallback coverage across FlashAttention and unfused configurations.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Conditionally compiles DPA forward while retaining eager islands for unsupported configurations and making backend logging trace-safe.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Makes backend values, QKV layout handling, sequence indices, and cumulative-length allocation compatible with Dynamo tracing.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Avoids passing one cumulative-sequence tensor through two FlashAttention autograd inputs during compilation.
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Documents and normalizes fused-attention backend enum values across graph breaks.
transformer_engine/pytorch/jit.py	Adds conditional Dynamo disabling and avoids nested lazy compilation while an outer graph is being captured.
tests/pytorch/test_torch_compile.py	Adds broad compiled DPA correctness, CUDA-graph, fused-island, and fallback-path coverage.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[DotProductAttention call] --> B{Requires eager fallback?}
    B -- Yes --> C[Eager island with warning]
    B -- No --> D[Compile input handling and backend selection]
    D --> E{Selected backend}
    E -- FlashAttention --> F[Compiled FlashAttention path]
    E -- Unfused --> G[Compiled PyTorch attention path]
    E -- FusedAttention --> H[Eager FusedAttention island]
    F --> I[Compiled output handling]
    G --> I
    H --> I

Loading

_{Reviews (2): Last reviewed commit: "Run context parallel attention eagerly" | Re-trigger Greptile}