[All] Refactor nvte_get_fused_attn_backend with cudnn-frontend calls

[cyanguwa]

Description

This PR replaces the hand-maintained backend selection logic in nvte_get_fused_attn_backend with cudnn-frontend's production-grade support checks.

• It will build the same graph that runtime uses in execution, cache the graph if the build is successful, and provide a warmed-up cache for the next nvte_get_fused_attn_backend call.
• It provides a cleaner dispatch logic, avoids accidental regressions, and helps TE to stay in sync with cudnn-frontend's support surface.
• It also provides appropriate error messaging for when no backend is available, so users are notified and make config/architecture/cudnn version adjustments.

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

See Description.

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 fused-attention backend selection to use cudnn-frontend support checks. The main changes are:

• Adds an opaque fused-attention config handle and attribute API.
• Threads backend config fields through PyTorch and JAX availability checks.
• Replaces hand-written support logic with graph-based support probes.
• Updates fused-attention runtime paths to build richer probe configs.
• Adds tests for updated backend availability behavior.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Refactors fused-attention backend selection around the new config object and cudnn-frontend probes.
transformer_engine/common/fused_attn/config_and_params.cpp	Adds the backing implementation for fused-attention config creation, destruction, and attribute access.
transformer_engine/common/include/transformer_engine/fused_attn.h	Adds the public fused-attention config handle, attributes, and C++ wrapper helpers.
transformer_engine/jax/cpp_extensions/attention.py	Updates JAX fused-attention helper calls to pass the expanded backend query inputs.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Updates PyTorch backend availability checks to pass expanded fused-attention config values.

_{Reviews (12): Last reviewed commit: "replace with opaque handle" | Re-trigger Greptile}

[greptile-apps]

Thread-local buffer invalidated on next call on the same thread

*message is set to the .c_str() of the thread-local fused_attn_backend_message_buffer. Any subsequent call to nvte_get_fused_attn_backend on the same thread will clear and overwrite this buffer, invalidating the raw pointer before the caller has a chance to copy it. The internal probe calls currently pass nullptr for message, so the buffer won't be clobbered by them, but this invariant is fragile and easy to break in a future change.

[greptile-apps]

Breaking public API change: cudnnHandle_t added as a required parameter

nvte_get_fused_attn_backend is a public C API (no namespace, exported symbol). Adding cudnnHandle_t handle as a required parameter, and including <cudnn.h> in the public header, is a breaking change that forces every downstream consumer to hold and pass a cuDNN handle for what was previously a pure-metadata query.

[cyanguwa]

/te-ci L1

[greptile-apps]

Forward path re-probes backward with deterministic=false, wasting graph cache entries for deterministic training

nvte_fused_attn_fwd calls nvte_get_fused_attn_backend with is_training=is_training and hardcoded deterministic=false. When a user trains with deterministic=true, the forward invocation builds and caches a cuDNN graph for the non-deterministic backward. A few lines later, nvte_fused_attn_bwd calls nvte_get_fused_attn_backend again with deterministic=true, causing the backward graph to be rebuilt from scratch on every step. The deterministic flag should be threaded through to this site (as it already is in the nvte_fused_attn_bwd call), or the forward pass should skip the backward probe entirely and let nvte_fused_attn_bwd drive its own backend check.

[cyanguwa]

/te-ci L1

[cyanguwa]

/te-ci L1

[greptile-apps]

Deprecated wrapper silently discards is_training, accepting unsupported training configurations

The parameter is cast to (void) and then overridden to false, meaning any caller that passes is_training=true through this deprecated API now only runs the forward graph probe. Before this PR, the old hand-written logic applied training-specific restrictions (e.g. rejecting head_dim_qk > 128 for backward on Hopper with older cuDNN). Now those restrictions are bypassed — the forward probe passes, and the function returns a non-No_Backend value for configurations where the actual backward kernel would fail. Downstream code that calls the deprecated API with is_training=true to guard training-mode usage will silently get incorrect results.

[jberchtold-nvidia]

Should this be batch_dim = sequence_dim - 1 instead of the current 1 - sequence_dim?

For the (B, S, H, D) case they're equivalent, in both cases batch_dim = 0. But if we wanted to support an additional axis (_, B, S, H, D), the current formula would put the batch dim as batch_dim = 1 - seq_dim = 1 - 2 = -1, which is really the D axis.

[jberchtold-nvidia]

Please use the existing already exported JAXX_Scaling_Mode instead of exporting NVTEScalingMode

[pggPL]

There is a limitation for torch.compile dynamic shapes here. If b and s are passed to black-box function affecting the program flow, I will need to force recompilation on every b or s change.
If we have args like s_greater_or_eq_512 and we check s >= 512 in python then no recompilation would be needed for every change of s. It may be more problematic for b.

I don't think this is a blocker for this PR, but we need to keep it in mind.

[greptile-apps]

The probe uses cfg.is_training = false, so it builds and caches an inference-mode cuDNN forward graph (no softmax-stats output). When the actual kernel is called below with the real is_training value — which is true during training — the cuDNN graph cache key differs and there is a guaranteed cache miss. The training-mode graph is then rebuilt on the first forward pass of every training run, negating the prewarming benefit. Additionally, if a configuration happens to be inference-supported but training-forward-unsupported (cuDNN validates outputs differently in training mode), the backend check returns success but the actual kernel throws at execution time.

Suggested change

	cfg.is_training = false;
	cfg.return_max_logit = return_max_logit;
	cfg.deterministic = false;
	NVTE_Fused_Attn_Backend fused_attention_backend =
	nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr);
	cfg.is_training = is_training;
	cfg.return_max_logit = return_max_logit;
	cfg.deterministic = false;
	NVTE_Fused_Attn_Backend fused_attention_backend =
	nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr);