[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

This PR refactors the fused-attention APIs by replacing the hand-written static support table with live cuDNN-frontend probe calls, and introduces opaque NVTEFusedAttnConfig / NVTEFusedAttnFwdParams / NVTEFusedAttnBwdParams handles with attribute-based accessors for better API/ABI stability.

• Adds a process-wide GraphCache backed by std::map<FusedAttnConfig, shared_ptr<CacheEntry>> with double-checked locking and a build_plans once_flag; the deprecated nvte_get_fused_attn_backend wrapper is fixed to pass batch_size=1 (preventing log2(0) UB) and to correctly forward is_training.
• Introduces v2 execution-path APIs (nvte_fused_attn_fwd_v2, nvte_fused_attn_bwd_v2) that accept the new opaque handles and fix an uninitialized input_Bias/input_SoftmaxOffset pointer bug from the old paths.
• Removes deprecated cuDNN-v8 utility functions, removes FusedAttnBackend.__eq__/__hash__ to unblock torch.compile constant-folding, and upgrades the enum sync-check from assert to RuntimeError.

Confidence Score: 4/5

• This PR is safe to merge with one caveat: the single global frontend_build_mutex serializes all four build sites, which is functionally correct but a known scalability constraint worth documenting.
• The PR fixes several pre-existing correctness bugs (uninitialized pointers, log2(0) UB, multi-GPU cache aliasing, batch_size=0 in deprecated wrapper), replaces the fragile hand-written support table with live cuDNN-frontend probing, and introduces a well-designed cache with proper double-checked locking and once_flag-guarded plan builds. The remaining P1 finding (unbounded cache growth with no eviction) is a known constraint acceptable at this stage. The JAX FP8 o_dtype limitation was already acknowledged in Previous Threads. No security vulnerabilities introduced.
• transformer_engine/common/fused_attn/graph_cache.h (unbounded cache growth, single global build mutex) and transformer_engine/common/fused_attn/fused_attn.cpp (thread-local message buffer lifetime).

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/config_and_params.cpp	New 1264-line file: implements FusedAttnConfig::derive(), make_cache_key(), and make_config() for fwd/bwd params. Core logic is sound. All attribute switch statements are complete with proper break statements. Minor concern: FusedAttnFwdParams::make_config() hardcodes check_for_backward_support=false (intentional per design) and deterministic=false (correct for fwd).
transformer_engine/common/fused_attn/fused_attn.cpp	Replaces the ~120-line hand-maintained support table in nvte_get_fused_attn_backend with cuDNN-frontend probe calls. Deprecated wrapper now correctly sets batch_size=1, o_format, do_format, dqkv_layout, and is_training. nvte_fused_attn_fwd_v2/bwd_v2 properly initialize all pointers before use. A few TE-specific pre-filters (CUDA-graph, bottom-right SWA) are applied before dispatching to support_verdict_f16/fp8.
transformer_engine/common/fused_attn/graph_cache.h	New process-wide cache using a single frontend_build_mutex for all four build sites (F16/FP8 × fwd/bwd). Transient errors are NOT permanently cached — exceptions from query_support propagate without inserting the entry, so failed configs are retried on the next call. build_plans uses call_once with correct std::call_once exception-reset semantics. The global serialization of all builds (including across different backends) is intentional for cuDNN-frontend thread-safety.
transformer_engine/common/fused_attn/graph_cache_debug.h	Production debug facility gated behind NVTE_FUSED_ATTN_CACHE_DEBUG env var. Intentional memory leaks (thread registry, mutex) are justified to avoid static-destructor ordering races with the atexit summary. Overhead when disabled is just a few static arrays (~512 bytes). Well-structured and clearly documented.
transformer_engine/common/fused_attn/utils.cu	Added explicit guards for batch_size==0 and num_tokens==0 before calling log2 in get_max_batch_size and get_max_tokens, preventing UB from log2(0)=-inf cast to size_t. Removed dead utility functions (tensor_create, pw_desc_create, etc.) that were only used by the now-removed max512 and fp8_impl_v0 backends.
transformer_engine/common/include/transformer_engine/fused_attn.h	New opaque handle types (NVTEFusedAttnConfig, NVTEFusedAttnFwdParams, NVTEFusedAttnBwdParams) with create/destroy/get/set attribute accessors. v2 entry points added alongside deprecated shims. The cudnnHandle_t parameter added to nvte_get_fused_attn_backend_v2 (but not to the deprecated v1) is unavoidable given the probe semantics.
transformer_engine/jax/cpp_extensions/attention.py	FusedAttnHelper now passes batch_size, bottom_right_diagonal, and bias shape dimensions to the backend probe. NVTE_DEBUG/NVTE_DEBUG_LEVEL logging matches the PyTorch parity goal. Non-FP8 JAX attention is correctly handled; FP8 probe still passes o_dtype=q_type (pre-existing limitation, already flagged in prior reviews).
transformer_engine/pytorch/attention/dot_product_attention/utils.py	AttentionParams.softmax_scale now correctly defaults to 1.0 (matching the docstring). core_attention_bias_shape changed from a string to an Optional[Tuple[int,int,int,int]] for precise bias shape control. New FusedAttentionParams dataclass mirrors the C++ FusedAttnConfig struct layout for the v2 probe API.
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Removed custom eq/ne/hash overrides from FusedAttnBackend that were blocking torch.compile constant-folding. Sync-check converted from assert (skipped with -O) to RuntimeError. These are clean correctness improvements.

Sequence Diagram

sequenceDiagram
    participant User as User/Python
    participant Backend as nvte_get_fused_attn_backend_v2
    participant Cache as GraphCache process-wide
    participant FE as cudnn-frontend

    User->>Backend: "FusedAttnConfig check_fwd=T check_bwd=T"
    Backend->>Backend: derive convenience fields
    Backend->>Cache: support_verdict_f16 Fwd handle
    Cache->>FE: build validate build_op_graph check_support
    FE-->>Cache: OK graph inserted into cache
    Cache-->>Backend: supported
    Backend->>Cache: support_verdict_f16 Bwd handle
    FE-->>Cache: OK
    Cache-->>Backend: supported
    Backend-->>User: NVTE_F16_arbitrary_seqlen

    User->>Backend: "nvte_fused_attn_fwd_v2 make_config derive check_fwd=T check_bwd=F"
    Backend->>Cache: get_graph F16 Fwd cfg handle
    Cache-->>Backend: HIT graph from probe
    Backend->>FE: build_plans once per entry
    FE-->>Backend: compiled kernels
    Backend->>FE: graph execute handle device ptrs
    FE-->>User: output O aux tensors

    User->>Backend: "nvte_fused_attn_bwd_v2 make_config derive check_fwd=F check_bwd=T"
    Backend->>Cache: get_graph F16 Bwd cfg handle
    Cache-->>Backend: HIT graph from probe
    Backend->>FE: build_plans and execute
    FE-->>User: dQ dK dV

Loading

_{Reviews (48): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[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

[cyanguwa]

/te-ci L0 L1 L2 L3