[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, taking the opaque handles above in place of flat argument lists.
• Graph construction split from execution (common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu, fused_attn_fp8.cu, graph_cache.h): graph building moved out of fused_attn_*_{fwd,bwd}_impl into create_graph_{f16,fp8}_{fwd,bwd}(const FusedAttnConfig &); the *_impl bodies now only fetch the graph, finish its plans, and bind the variant pack. The probe (support_verdict_*) and execution reach that builder through the same get_graph<Backend, Pass, create_graph_*> and cache, so they can't diverge; build_plans() is deferred so a probe-only miss doesn't pay for it.
• 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=<level>[:<ranks>]: 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; [:<ranks>] specifies select ranks for diagnostics; 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.
• Added test_fused_attn_graph_cache to test graph cache's behavior with NVTE_FUSED_ATTN_CACHE_DEBUG=2 on, and test_fused_attn_backend_message to test the surfacing of TE-specific or cuDNN-related error messages to users.

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]

Confidence Score: 5/5

The PR appears safe to merge; no new actionable issue was introduced after the previous review.

Summary

• Adds v2 fused-attention APIs while retaining deprecated compatibility shims.
• Unifies support probing and execution graph construction through a process-wide, device-aware graph cache.
• Updates PyTorch and JAX bindings, including packed-THD metadata and context-parallel capability checks.
• Adds backend diagnostics, cache instrumentation, documentation, and expanded tests.
• Changes since the previous review only reformat a diagnostic string and introduce no behavioral changes.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[PyTorch or JAX attention request] --> B[Opaque fused-attention config]
    B --> C[Backend v2 support query]
    C --> D[TE capability guards]
    D --> E[cuDNN frontend graph construction]
    E --> F[Process-wide device-aware graph cache]
    F --> G[Forward or backward execution]
    C -->|Unsupported| H[Diagnostic and compatible fallback]

Loading

[cyanguwa]

/te-ci L1

[cyanguwa]

/te-ci L0 L1 L2 L3

[cyanguwa]

/te-ci L0 L1 L2 L3

[cyanguwa]

/te-ci L0 L1 L2 L3

[KshitijLakhani]

nit: So the fact that we use an adjusted window size for CP in one of our primitives which is different than the global window size is not something this PR introduces, however, I'm wondering ff the caching becomes less efficient because of this ?

workspace query with global window
      → build/cache global-window graph in the thread-local cache

  runtime with adjusted CP window
      → different descriptor
      → cache miss
      → build/cache adjusted-window graph

it's a minor penalty I suppose, but I do not think we need to worry about it - thoughts ?

[cyanguwa]

nit: So the fact that we use an adjusted window size for CP in one of our primitives which is different than the global window size is not something this PR introduces, however, I'm wondering ff the caching becomes less efficient because of this ?

workspace query with global window
      → build/cache global-window graph in the thread-local cache

  runtime with adjusted CP window
      → different descriptor
      → cache miss
      → build/cache adjusted-window graph

it's a minor penalty I suppose, but I do not think we need to worry about it - thoughts ?

Done. I added a helper function to get the effective window size for CP + Ring + Striped + THD + SWA. Looking through other CP variants, it seems that they are already passing the exact config that the execution will run, during the probing/support query stage. Thanks.

[cyanguwa]

/te-ci L0 L1 L2 L3

[cyanguwa]

/te-ci L0 L1 L2 L3

[cyanguwa]

/te-ci L0 L1 L2 L3

[cyanguwa]

L1 and L2 jax distributed tests have some hangs but fused_attn related tests have run and passed.