[All] Refactor fused attention APIs with cuDNN-frontend support checks and opaque config/params handles

[cyanguwa]

Description

TE currently hand-maintains the fused-attention backend-selection logic in nvte_get_fused_attn_backend, duplicating cuDNN's support rules. This list drifts out of sync as cuDNN evolves, and the support check can disagree with what actually runs.

This PR replaces that logic with cuDNN-frontend's production-grade support checks. The new nvte_get_fused_attn_backend_v2 builds the same graph cuDNN executes at runtime, so the probe and execution can no longer diverge. It caches the graph on success and returns a diagnostic message on failure, giving users actionable guidance (e.g. adjust the config, GPU architecture, or cuDNN version).

This PR also reworks nvte_fused_attn_fwd / nvte_fused_attn_bwd into nvte_fused_attn_fwd_v2 / nvte_fused_attn_bwd_v2, which take opaque, attribute-based config/params handles instead of long flat argument lists — improving TE's API and ABI stability.

Legacy APIs are retained as deprecated shims that route through the v2 APIs, so existing callers keep working.

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

API rework (opaque config/params + v2 entry points)

• Opaque config/params handles (common/fused_attn/config_and_params.{h,cpp}, common/include/transformer_engine/fused_attn.h): new NVTEFusedAttnConfig / NVTEFusedAttnFwdParams / NVTEFusedAttnBwdParams with create/destroy/get/set attribute accessors, for better API/ABI stability. The cache key, probe, and execution now all originate from one place via make_config / derive / make_cache_key.
• v2 APIs (common/fused_attn/fused_attn*.{cpp,cu}): nvte_get_fused_attn_backend_v2, nvte_fused_attn_fwd_v2, and nvte_fused_attn_bwd_v2. The F16 and FP8 is_supported_* probes copy the config, set direction, derive(), and attempt a null-pointer graph build via check_support — i.e. the same graph cuDNN builds at runtime, so probe and execution can't diverge.
• Deprecated shims: legacy nvte_get_fused_attn_backend / nvte_fused_attn_fwd / nvte_fused_attn_bwd are retained, routed through the v2 APIs.
• Bindings updated to v2: PyTorch (csrc/extensions/attention.cpp) and JAX (jax/csrc/extensions/attention.cpp).

Correctness & backend selection

• Process-wide graph cache: cache is now process-wide (was thread-local) and guarded by a mutex, so a compiled graph is reused across threads instead of rebuilt per thread (still thread-safe).
• Bias-shape handling fix: applied consistently across common, PyTorch, and JAX.
• Per-step CP config checks: cp_per_step_configs probes each context-parallel step instead of only the global, non-CP config.
• log2(0) guard: avoids UB when casting -inf to size_t in get_max_batch_size / get_max_tokens.

Diagnostics

• NVTE_DEBUG / NVTE_DEBUG_LEVEL for JAX (parity with PyTorch): level 1 reports the selected backend; level 2 adds a diagnostic message explaining why fused attention was rejected.

Cleanup / removals

• Removed NVTE_FUSED_ATTN_BACKEND — the two remaining backends (F16, FP8) are mutually exclusive now that max512 is gone.
• Removed dead Q_ID/.../MASK_VAL_ID macros (used only by the max512 backend).
• Removed dead cudnn_frontend::xxx utility functions (used only by fp8_impl_v0 and max512).
• Unified include-guard names across fused_attn/ headers.

Tests

• Enabled previously skipped tests: 29 in L0 PyTorch (27 padding + post_scale_bias, 2 D256 bprop) and 528 in L0 JAX (420 SWA + dropout/post_scale_bias, 108 padding + post_scale_bias).
• No significant CI time increase (~5 min on GB200 for tests/pytorch/test_attention.py).

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

[greptile-apps]

Greptile Summary

The PR replaces hand-maintained fused-attention selection with cuDNN-frontend graph support checks and introduces opaque v2 configuration and parameter handles.

• Routes PyTorch and JAX bindings through the v2 fused-attention APIs.
• Adds a process-wide, device-keyed graph cache and cache diagnostics.
• Updates bias-shape, context-parallel probing, diagnostics, documentation, and tests.

Confidence Score: 3/5

The PR does not yet appear safe to merge because several previously reported compatibility and probe-versus-execution mismatches remain unverified at the current head.

The device-aware cache fix and null-pointer initialization are established, but the available evidence does not verify that the deprecated query supplies batch, training, output-format, and gradient-layout fields consistently or that JAX FP8 probing reflects the active recipe; those unresolved paths can still reject supported configurations or accept configurations that fail during execution.

Files Needing Attention: transformer_engine/common/fused_attn/fused_attn.cpp, transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu, transformer_engine/common/include/transformer_engine/fused_attn.h, transformer_engine/jax/cpp_extensions/attention.py

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/config_and_params.cpp	Implements opaque fused-attention objects, attribute accessors, configuration derivation, and device-aware cache-key normalization.
transformer_engine/common/fused_attn/fused_attn.cpp	Replaces backend selection and execution routing with v2 configuration-driven support checks while retaining deprecated shims.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Uses the shared process-wide graph cache for F16/BF16 support checks and execution.
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Adopts configuration-derived FP8 graph support checks and the shared device-aware cache.
transformer_engine/common/include/transformer_engine/fused_attn.h	Adds the public opaque-handle API, v2 entry points, attribute enums, and compatibility declarations.
transformer_engine/jax/cpp_extensions/attention.py	Extends JAX fused-attention backend queries with richer shape, bias, mask, and diagnostic configuration.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Updates PyTorch backend-selection configuration to match the v2 fused-attention support probe.

Sequence Diagram

sequenceDiagram
  participant Frontend as PyTorch/JAX
  participant Config as Opaque config/params
  participant Probe as Backend support probe
  participant Cache as Process-wide graph cache
  participant cuDNN as cuDNN frontend
  Frontend->>Config: Populate attention attributes
  Config->>Probe: Derive normalized configuration
  Probe->>Cache: Lookup device-keyed graph
  alt cache miss
    Cache->>cuDNN: Build and validate graph
    cuDNN-->>Cache: Cache supported graph
  end
  Probe-->>Frontend: Backend and diagnostic
  Frontend->>Cache: Execute through v2 API

Loading

_{Reviews (30): Last reviewed commit: "fix merge with torch.compile PRs" | Re-trigger Greptile}

[cyanguwa]

/te-ci L1

[cyanguwa]

/te-ci L1

[cyanguwa]

@greptileai

[cyanguwa]

Note: The explicit is_fused_attn_kernel_available() pre-check is no longer needed here:
FusedAttnRunner._check_configs() now runs the cuDNN-frontend support probe during
setup and skips with a diagnostic message when no fused-attn backend is available.

[cyanguwa]

/te-ci L1