[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.
• Fused attention graph cache debug NVTE_FUSED_ATTN_CACHE_DEBUG: opt-in instrumentation that reports cuDNN graph build-vs-execution counts and per-stage cudnn-frontend build timings, so cache hit/miss/build/exec behaviors and graph build time can be inspected. Off by default; available for both PyTorch and Jax.

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: padding + post_scale_bias in both PyTorch and Jax, D256 bprop in PyTorch, and SWA + dropout/post_scale_bias in Jax.
• Curated the L0 sweeps to keep CI time in check: deduplicated PyTorch tests, and tiered the newly enabled JAX tests across L0/L1/L2.

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 backend checks with cuDNN-frontend graph probes and introduces opaque configuration and parameter handles.

• Adds v2 backend-selection and forward/backward C APIs while retaining deprecated wrappers.
• Introduces process-wide, device-aware graph caches and optional cache diagnostics.
• Updates the PyTorch and JAX bindings and adjusts fused-attention tests and documentation.

Confidence Score: 4/5

The PR appears safe to merge, with only the non-blocking diagnostic-message lifetime concern remaining.

The previously blocking cache-key, deprecated-wrapper, attribute-setter, pointer-initialization, and probe-configuration defects are addressed at current HEAD; the remaining issue only limits how long callers may retain a diagnostic string pointer.

Files Needing Attention: transformer_engine/common/fused_attn/fused_attn.cpp

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds v2 graph-based backend selection and execution routing; previously reported wrapper, probe, and optional-pointer defects are fixed, while diagnostic storage retains its documented lifetime limitation.
transformer_engine/common/fused_attn/config_and_params.cpp	Implements opaque config/parameter derivation and attribute accessors, including reachable boolean setters for forward and backward parameters.
transformer_engine/common/fused_attn/config_and_params.h	Defines the unified attention configuration and a device-aware cache key used by all fused-attention graph caches.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Routes F16 support checks and execution through the shared configuration and process-wide device-aware graph cache.
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Routes FP8 support checks and execution through the shared configuration and fixes the previously reported diagnostic string construction errors.
transformer_engine/jax/cpp_extensions/attention.py	Updates JAX fused-attention availability checks to pass the expanded configuration and consume backend diagnostics.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Updates PyTorch backend-selection configuration and retains the corrected attention scale default.

Sequence Diagram

sequenceDiagram
  participant Frontend as PyTorch/JAX frontend
  participant Config as Opaque config/params
  participant Probe as Backend v2 probe
  participant Cache as Device-aware graph cache
  participant cuDNN as cuDNN frontend
  Frontend->>Config: Set attention attributes
  Frontend->>Probe: Query supported backend
  Probe->>cuDNN: Build/check matching graph
  cuDNN-->>Cache: Cache supported graph
  Probe-->>Frontend: Backend + diagnostic
  Frontend->>Config: Submit forward/backward params
  Config->>Cache: Lookup execution graph
  Cache->>cuDNN: Execute graph

Loading

_{Reviews (37): Last reviewed commit: "Merge branch 'main' into fe_check_suppor..." | Re-trigger Greptile}

[cyanguwa]

/te-ci L1

[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

[cyanguwa]

/te-ci L1

[cyanguwa]

/te-ci L1