[PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers

[pggPL]

Description

Behavior-neutral restructuring of FusedAttnFunc, the cuDNN fused attention autograd.Function, mirroring what #2967 did for Linear ahead of #3053. Groundwork for #3472 (torch.compile support for the fused attention backend), which registers these helpers as a torch custom op through dynamo/custom_op.py.

• FusedAttnFwdArgs / FusedAttnBwdArgs dataclasses replace the 37 positional arguments of FusedAttnFunc.apply and the loose ctx.* attributes.
• Module-level _fused_attn_forward_impl, _fused_attn_setup_ctx and _fused_attn_backward_impl hold the actual logic; FusedAttnFunc is a thin wrapper taking the differentiable tensors (q, k, v, attn_bias, softmax_offset) plus one args object, so its backward returns 6 grads instead of 36 Nones.
• The cuDNN aux pack is stored in fixed slots (softmax_stats, rng_state, aux_bias, aux_softmax_offset) instead of a variable-length list.
• A saved tensor identical to a forward input or the output is not saved twice: the forward hands back None in that slot and names its source in ctx_attrs["saved_from"]; _fused_attn_setup_ctx re-attaches it. This is what a custom op needs (it may not return its own inputs) and keeps the FP8 paths -- where a None f16 slot means "not needed" -- exact.
• Dropped ctx attributes the backward never read (fp8_recipe, fp8_meta, is_output_fp8).

No functional change intended. Tested on an RTX Ada with the fused backend forced: tests/pytorch/attention/test_attention.py (test_dpa*, test_dot_product_attention, test_transformer_layer), test_kv_cache.py, test_cpu_offloading_v1.py and test_torch_compile.py all pass. FP8 attention and context parallelism could not be exercised on that GPU and rely on CI.

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

• attention/dot_product_attention/backends.py: FusedAttnFwdArgs, FusedAttnBwdArgs, _fused_attn_forward_impl, _fused_attn_setup_ctx, _fused_attn_backward_impl, _reload_qkv_layout; FusedAttnFunc reduced to a wrapper; FusedAttention.forward builds the args object.

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

[pggPL]

/te-ci pytorch L1

[greptile-apps]

Greptile Summary

The PR restructures the cuDNN fused-attention autograd implementation into typed forward/backward argument containers and module-level helpers without intending to change behavior.

• Replaces the large positional autograd interface with FusedAttnFwdArgs and FusedAttnBwdArgs.
• Moves forward, context setup, backward, and CPU-offload layout handling into focused helpers.
• Normalizes saved auxiliary tensors into fixed slots while preserving FP8, cache, bias, and softmax state.
• Updates FusedAttention.forward to construct the new argument container and invoke the thinner autograd wrapper.

Confidence Score: 5/5

The PR appears safe to merge with no changed-code defects identified.

The refactored paths preserve saved-tensor handling, native auxiliary-pack ordering, the supported shared page-table model, backward gradient arity, and module importability.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Refactors fused-attention forward and backward state into dataclasses and helper functions while preserving the examined eager, FP8, auxiliary-state, paging, and autograd contracts.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[FusedAttention.forward] --> B[Build FusedAttnFwdArgs]
    B --> C[FusedAttnFunc.apply]
    C --> D[_fused_attn_forward_impl]
    D --> E[_fused_attn_setup_ctx]
    E --> F[Save backward tensors and metadata]
    F --> G[FusedAttnFunc.backward]
    G --> H[Restore FusedAttnBwdArgs]
    H --> I[_fused_attn_backward_impl]
    I --> J[dQ, dK, dV, bias and offset gradients]

Loading

_{Reviews (1): Last reviewed commit: "[PyTorch] Match the fused attention aux ..." | Re-trigger Greptile}