[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 training paths. FP8 paths are explicitly excluded via a @no_torch_dynamo()-wrapped eager island. Four softmax kernel families and the THD↔BSHD layout-conversion helpers are refactored from torch.autograd.Function subclasses into torch.library.custom_ops with fake implementations and register_autograd bindings. C++ backward kernels now allocate fresh output buffers to satisfy the no-aliasing contract, and the .item() call that produced an unbacked SymInt from cu_seqlens[-1] is replaced by a shape read captured before any layout transformation.

• softmax.py / softmax.cpp: Each of the four fused softmax kernels is wrapped as a pair of custom_ops (forward + backward), with autograd linked via register_autograd; C++ backward functions allocate fresh output tensors instead of writing in-place into the gradient input.
• utils.py / backends.py: ConvertTHDtoBSHD and ConvertBSHDtoTHD become thin .apply() wrappers around new te_attention::convert_* custom ops; total_tokens_q is hoisted from the cu_seqlens[-1].item() call to a shape read before any layout conversion, fixing a NameError on inference THD→BSHD paths and avoiding an unbacked-SymInt GPU sync.
• test_torch_compile.py: A new test_unfused_dpa_torch_compile parametrized test covers 5 QKV layouts with fullgraph=True + reduce-overhead, validating forward and backward through CUDA-graph replay.

Confidence Score: 5/5

Safe to merge; the non-FP8 training paths are well-covered by the new fullgraph compile test, and FP8 is correctly fenced behind a dynamo-disabled eager island.

The custom-op refactor is mechanically sound: fake impls correctly reflect output shapes, autograd hooks use register_autograd with setup_context (the compile-compatible pattern), and the C++ backward allocations remove input aliasing. The total_tokens_q hoisting is logically correct — it reads query_layer.shape[0] before any layout transformation, so the value is right for both the thd training path and the thd_2bshd inference path. No correctness or runtime-error issues were found in the new code.

No files require special attention; the only note is a stale docstring in forward_fused_softmax that still names the removed autograd.Function classes.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/softmax.py	Replaces four autograd.Function softmax classes with torch.library custom_op pairs (fwd + bwd), each with a fake impl and register_autograd binding; adds float() coercion on the scale parameter; docstring in forward_fused_softmax still names the removed classes.
transformer_engine/pytorch/csrc/extensions/softmax.cpp	Each of the four backward kernel wrappers now allocates a fresh output buffer (torch::empty) instead of writing in-place into output_grads, satisfying the custom_op no-aliasing contract.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	ConvertTHDtoBSHD and ConvertBSHDtoTHD are refactored from autograd.Function subclasses to thin .apply() wrappers around te_attention custom ops; the custom ops carry their own fake impls and register_autograd hooks.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Splits forward into a public dispatch shim and _forward; FP8 paths routed to @no_torch_dynamo _forward_eager; total_tokens_q is captured from query_layer.shape[0] before any layout transformation, avoiding the unbacked-SymInt .item() call and fixing NameError on thd_2bshd inference paths.
tests/pytorch/test_torch_compile.py	Adds test_unfused_dpa_torch_compile parametrized over 5 QKV layouts with fullgraph=True + reduce-overhead; tests fwd+bwd CUDA-graph capture and replay; dtype is hardcoded to bfloat16 (fp16 not exercised).

Sequence Diagram

sequenceDiagram
    participant Caller
    participant forward as UnfusedDPA.forward
    participant _fwd as UnfusedDPA._forward
    participant _fwd_eager as UnfusedDPA._forward_eager (@no_torch_dynamo)
    participant thd2bshd as te_attention::convert_thd_to_bshd (custom_op)
    participant softmax as te_softmax::scaled_*_softmax_fwd (custom_op)
    participant bshd2thd as te_attention::convert_bshd_to_thd (custom_op)

    Caller->>forward: call(fp8, fp8_output, args...)
    alt "fp8=True or fp8_output=True"
        forward->>_fwd_eager: eager island (graph break)
        _fwd_eager->>_fwd: delegate
    else non-FP8 compile-compatible path
        forward->>_fwd: direct call (fully traced)
    end
    alt "qkv_format == thd (training)"
        _fwd->>thd2bshd: ConvertTHDtoBSHD.apply(q/k/v, cu_seqlens, max_seqlen)
    end
    _fwd->>softmax: FusedScaleMaskSoftmax
    alt "q_format == thd"
        _fwd->>bshd2thd: ConvertBSHDtoTHD.apply(ctx, cu_seqlens_q, total_tokens_q)
        Note over _fwd,bshd2thd: total_tokens_q from query_layer.shape[0] before layout conversion
    end
    _fwd-->>Caller: context_layer

Loading

_{Reviews (3): Last reviewed commit: "[PyTorch] Fix UnboundLocalError on thd i..." | Re-trigger Greptile}

[cyanguwa]

LGTM but please take a look at Greptile's few comments. Thanks.

[pggPL]

/te-ci pytorch