[PyTorch] torch.compile support for FusedAttention

[pggPL]

Description

Follow-up to #3053 (custom-op framework, Linear) and #3286 (DotProductAttention under torch.compile for the FlashAttention / unfused backends). This PR makes the FusedAttention (cuDNN) backend trace under torch.compile(fullgraph=True) as well, so DotProductAttention no longer graph-breaks at FusedAttention.forward -- on Hopper/Blackwell that is the default backend, so until now the whole attention block was an eager island there.

The cuDNN fused_attn_fwd / fused_attn_bwd pair is registered as a torch custom op through dynamo/custom_op.py, exactly like Linear: one FusedAttnFwdArgs / FusedAttnBwdArgs dataclass each, a data-free fake forward/backward (TensorSpecs for the output, the softmax stats / rng_state aux tensors and the grads) and the real implementations calling the existing cpp_extensions.fused_attn bindings. The eager FusedAttnFunc is untouched.

Scope is F16/BF16 attention. Configurations the op does not cover keep working as before -- FusedAttention.forward now carries @no_torch_dynamo(when=...) and falls back to an eager island with a warning for: FP8 attention, context parallelism, score_mod, KV caching (inference_params), NVTE_FUSED_ATTN_USE_FAv2_BWD and CPU activation offloading.

Notes for reviewers:

• The cuDNN aux pack hands back the bias / softmax_offset inputs as aux tensors; a custom op may not return its inputs, so those are re-attached from the forward args in setup_context and only stats / max / rng_state cross the op boundary.
• For packed layouts (bs3hd, bsh3d, bs2hd, ...) the backward returns dq/dk/dv as views of one dQKV buffer, which torch.library rejects as aliasing outputs; they are made contiguous (a copy only in the packed case).
• The softmax-stats shape for thd depends on the cuDNN version and SM ([tq, h, 1] ragged vs [b, h, sq, 1]); the fake mirrors the C++ condition using the assume_constant_result version/capability helpers. I could only test on an RTX Ada (sm89), where cuDNN attention with thd is not available, so the thd branch of the fake is untested here -- the existing thd cases in test_dpa_torch_compile cover it on Hopper/Blackwell.

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

• New attention/dot_product_attention/fused_attn_op.py: FusedAttnFwdArgs / FusedAttnBwdArgs, real + fake forward/backward, register_custom_op(op_name="fused_attn").
• backends.py: FusedAttention.forward decorated with @no_torch_dynamo(when=_needs_eager_fused_attention); under torch.compile it dispatches to the custom op instead of FusedAttnFunc.apply.
• tests/pytorch/test_torch_compile.py: the fused backend now runs through test_dpa_torch_compile, test_dpa_torch_compile_cudagraphs and test_dpa_torch_compile_eager_fallback; test_dpa_torch_compile_around_fused keeps covering the eager-island path by disabling the op.

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