[PyTorch][torch.compile] Support for UnfusedDotProductAttention

[pggPL]

Description

Make the UnfusedDotProductAttention backend traceable by torch.compile(fullgraph=True, mode="reduce-overhead"), so the forward and backward can be captured into CUDA graphs without graph breaks.

Scope:

1. bf16/fp16 path (compile-supported): register the TE softmax kernels and THD<->BSHD conversion helpers as torch.library.custom_ops with fake impls and autograd bindings; remove an unbacked-SymInt .item() from the hot path of ConvertBSHDtoTHD.
2. FP8 is explicitly NOT supported under torch.compile: with fp8=True (emulation) and/or fp8_output=True (Float8Tensor output, a tensor subclass that cannot cross a graph boundary) the backend runs as an eager island — the forward dispatches to a torch._dynamo.disable'd wrapper, the same mechanism DotProductAttention and FusedAttention use module-wide. FP8 attention always involves delayed scaling regardless of the recipe: S and dP are produced inside the kernel, so their amax cannot be known before quantization and they use delayed-scaling quantizers even under Float8CurrentScaling (see DPA.init_fp8_metadata) — and delayed scaling (Float8Quantizer, tensor scale/amax state) is not supported under torch.compile.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• softmax.py / softmax.cpp: scaled_*_softmax_{forward,backward} as custom ops; C++ backward kernels allocate a fresh output buffer instead of writing in-place into output_grad (custom ops and cudagraph trees forbid input aliasing).
• utils.py: ConvertTHDtoBSHD / ConvertBSHDtoTHD as custom ops; num_tokens passed by the caller instead of cu_seqlens[-1].item().
• backends.py: UnfusedDotProductAttention.forward dispatches FP8 calls to an eager (dynamo-disabled) wrapper; the non-FP8 path is traced with no graph breaks.
• tests/pytorch/test_torch_compile.py: test_unfused_dpa_torch_compile (5 qkv layouts, fullgraph + reduce-overhead, fwd+bwd captured into CUDA graphs and replayed).

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

This PR makes UnfusedDotProductAttention traceable by torch.compile(fullgraph=True, mode="reduce-overhead") for bf16/fp16 paths by wrapping the TE softmax kernels and THD\u2194BSHD layout converters as torch.library.custom_ops with fake impls and autograd bindings, and by routing FP8 calls through a @no_torch_dynamo-decorated eager wrapper. The C++ softmax backward functions are also fixed to write into a fresh buffer instead of aliasing their input, which is required by the mutates_args=() contract.

• softmax.py: Four torch.autograd.Function classes replaced by custom_op + register_autograd pairs; FusedScaleMaskSoftmax updated to call the new ops via torch.ops.te_softmax.*.
• utils.py: ConvertTHDtoBSHD / ConvertBSHDtoTHD refactored to thin wrappers around new te_attention::convert_* custom ops; ConvertBSHDtoTHD.apply now requires a num_tokens arg (previously computed internally via cu_seqlens[-1].item()).
• backends.py: forward dispatches FP8 to a dynamo-disabled eager island; total_tokens_q is captured from query_layer.shape[0] before the THD\u2192BSHD conversion to avoid a GPU-syncing .item() call.

Confidence Score: 3/5

The training path through the new custom ops is well-structured, but the refactoring introduced a latent regression that crashes inference calls with THD query and non-THD KV tensors.

total_tokens_q is only set in the qkv_format == thd branch but consumed by q_format == thd, a condition also true for the thd_2bshd inference layout where total_tokens_q is never assigned, reverting a path that worked before this PR.

backends.py — specifically the if q_format == thd output block and its interaction with the thd_2bshd inference format.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Adds forward/eager-dispatch split for FP8 vs non-FP8 paths; introduces total_tokens_q to avoid .item() GPU sync — but leaves the thd_2bshd inference path with an unset total_tokens_q that will raise NameError.
transformer_engine/pytorch/attention/dot_product_attention/softmax.py	Replaces four torch.autograd.Function softmax classes with torch.library.custom_op + register_autograd pairs; FusedScaleMaskSoftmax updated to call the new ops; scale is now explicitly cast to float.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	THD↔BSHD conversions refactored from torch.autograd.Function to torch.library.custom_op with full fake/autograd registration; wrapper classes kept but ConvertBSHDtoTHD.apply arity changed from 2 to 3 args with misleading docstring.
transformer_engine/pytorch/csrc/extensions/softmax.cpp	All four softmax backward functions now allocate a fresh output buffer instead of writing in-place into output_grad, satisfying the mutates_args=() contract of torch.library.custom_op.
tests/pytorch/test_torch_compile.py	Adds test_unfused_dpa_torch_compile covering 5 qkv layouts with fullgraph=True + reduce-overhead; only tests bfloat16 despite the PR claiming fp16 support.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["UnfusedDotProductAttention.forward(fp8, fp8_output, ...)"] --> B{fp8 or fp8_output?}
    B -- Yes --> C["_forward_eager()\n@no_torch_dynamo — eager island"]
    B -- No --> D["_forward() — fully compiled"]
    C --> E["_forward() inside eager context"]
    D --> F{qkv_format?}
    F -- thd --> G["ConvertTHDtoBSHD.apply()\ncustom_op with fake + autograd"]
    F -- bshd/sbhd --> H["transpose / no-op"]
    G --> I["BSHD attention core"]
    H --> I
    I --> J["FusedScaleMaskSoftmax\ncustom_op with fake + autograd"]
    J --> K["bmm + context layer"]
    K --> L{q_format?}
    L -- thd --> M["ConvertBSHDtoTHD.apply(total_tokens_q)"]
    L -- sbhd/bshd --> N["reshape / transpose"]
    M --> O["Output tensor"]
    N --> O

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["UnfusedDotProductAttention.forward(fp8, fp8_output, ...)"] --> B{fp8 or fp8_output?}
    B -- Yes --> C["_forward_eager()\n@no_torch_dynamo — eager island"]
    B -- No --> D["_forward() — fully compiled"]
    C --> E["_forward() inside eager context"]
    D --> F{qkv_format?}
    F -- thd --> G["ConvertTHDtoBSHD.apply()\ncustom_op with fake + autograd"]
    F -- bshd/sbhd --> H["transpose / no-op"]
    G --> I["BSHD attention core"]
    H --> I
    I --> J["FusedScaleMaskSoftmax\ncustom_op with fake + autograd"]
    J --> K["bmm + context layer"]
    K --> L{q_format?}
    L -- thd --> M["ConvertBSHDtoTHD.apply(total_tokens_q)"]
    L -- sbhd/bshd --> N["reshape / transpose"]
    M --> O["Output tensor"]
    N --> O

Loading

Comments Outside Diff (2)

1. transformer_engine/pytorch/attention/dot_product_attention/backends.py, line 723-732 (link)

   NameError in THD→BSHD inference path

   total_tokens_q is only assigned inside the if qkv_format == "thd": block, but if q_format == "thd": (which consumes it) is entered whenever q_format is "thd" — including the "thd_2bshd" inference case where qkv_format != "thd" and the assignment block is skipped. Any inference call with a qkv_layout whose query portion is "thd" but whose KV portion differs (e.g. "thd_bshd_bshd") will raise NameError: name 'total_tokens_q' is not defined. The old ConvertBSHDtoTHD computed num_tokens = cu_seqlens[-1] internally, so this path worked before this PR. The fix is to capture the token count before the layout-specific transformation block, e.g. set total_tokens_q = query_layer.shape[0] at the top of the if qkv_format == "thd_2bshd": block (when query_layer still holds the original THD tensor), paralleling the existing assignment in the thd branch.

2. tests/pytorch/test_torch_compile.py, line 156 (link)

   Float16 path not exercised

   The PR description explicitly claims that both bf16 and fp16 paths are compile-compatible, but dtype is hardcoded to torch.bfloat16 and torch.float16 is never tested. Different precision paths can exercise different code branches in FusedScaleMaskSoftmax.is_kernel_available (specifically the self.input_in_fp16 guard) and different downcasting logic in forward_torch_softmax. A parametrize over [torch.bfloat16, torch.float16] would provide the coverage needed to verify the fp16 claim.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

API arity change not reflected in the class docstring

The docstring says "kept so that callsites can continue to use the .apply(...) syntax", but the signature changed from the old apply(bshd_tensor, cu_seqlens) (2 args) to the new apply(bshd_tensor, cu_seqlens, num_tokens) (3 args). Any caller not yet updated to pass num_tokens will silently get a TypeError at runtime. Since this class is not exported publicly this is contained, but the misleading comment increases maintenance risk.