[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]

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the eligible follow-up review scope.

Summary

• Introduces slotted argument dataclasses for forward and backward state.
• Extracts forward, context setup, backward, and QKV-layout restoration helpers.
• Reduces FusedAttnFunc to a thin autograd wrapper and updates FusedAttention.forward to construct the new argument object.

Diagram

sequenceDiagram
    participant FA as FusedAttention
    participant AF as FusedAttnFunc
    participant FW as _fused_attn_forward_impl
    participant CTX as _fused_attn_setup_ctx
    participant BW as _fused_attn_backward_impl
    FA->>AF: apply(q, k, v, bias, offset, fwd_args)
    AF->>FW: forward(fwd_args)
    FW-->>AF: output, max_logit, saved tensors, attributes
    AF->>CTX: populate backward arguments
    CTX-->>AF: tensors to save
    AF-->>FA: output
    Note over AF: Autograd invokes backward
    AF->>BW: backward(bwd_args)
    BW-->>AF: dq, dk, dv, d_bias, d_offset

Loading

_{Reviews (3) · Last reviewed commit: "[PyTorch] Preserve separate attention pa..."}

[cyanguwa]

Should we add if ctx is not None: here like in Linear:

TransformerEngine/transformer_engine/pytorch/module/linear.py

Line 1837 in 1634a5a

if ctx is not None:

[pggPL]

FusedAttnFunc.apply is the only way this backend is invoked, so this if will always be true