[PyTorch] torch.compile support for FusedAttention

[pggPL]

Description

Depends on #3480 (the FusedAttnFunc restructuring); the first commit here is that PR. Review the second commit only.

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 attention forward/backward is registered as a torch custom op through dynamo/custom_op.py, exactly like Linear: the FusedAttnFwdArgs / FusedAttnBwdArgs dataclasses and the _fused_attn_forward_impl / _fused_attn_setup_ctx / _fused_attn_backward_impl helpers from #3480 are shared between the eager FusedAttnFunc and the op, so this PR only adds the data-free fake forward/backward (TensorSpecs for the output, the softmax stats / rng_state aux tensors and the grads), the registration, the eager-fallback gate and the dispatch.

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, NVTE_FUSED_ATTN_USE_FAv2_BWD and CPU activation offloading. KV caching (inference_params, paged and non-paged) goes through the op.

Notes for reviewers:

• The cuDNN aux pack hands back the bias / softmax_offset inputs as aux tensors, and the saved q/k/v/out are the inputs / output themselves; a custom op may not return its inputs, so the forward leaves those slots None and names their source in ctx_attrs["saved_from"] (from [PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers #3480), and only stats / 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 on the op path only (a copy only in the packed case; eager keeps returning views).
• 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 neither is cuDNN attention with a KV cache -- so the thd branch of the fake and the KV-cache cases (kv_cache_bshd-fused, test_dpa_torch_compile_kv_cache_decoding[*-fused]) are untested here and rely on the Hopper/Blackwell CI runners.

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

• backends.py: _fused_attn_forward_fake / _fused_attn_backward_fake, _fused_attn_backward_op_impl, register_custom_op(op_name="fused_attn"); FusedAttention.forward decorated with @no_torch_dynamo(when=_needs_eager_fused_attention) and, under torch.compile, dispatching 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, test_dpa_torch_compile_kv_cache_decoding 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