Enable cuDNN SDPA for contrib Attention

[mastryukov1990]

Description

Enable the existing cuDNN SDPA runner for the CUDA implementation of the contrib Attention operator when:

• the input type is FP16 or BF16;
• the mask is absent or is a compact 1D key sequence-length mask;
• past/present KV cache inputs are not used; and
• cuDNN reports support for the concrete attention shape.

On Hopper and newer GPUs, cuDNN SDPA is auto-preferred consistently with MultiHeadAttention. Explicit kernel selection continues to be respected. Raw attention masks and cache paths remain on their existing kernels.

The new test forces cuDNN SDPA and uses a partial sequence-length mask, verifying that the masked key/value token is excluded from both output positions.

Motivation and Context

Transformer optimizer fusions can produce com.microsoft::Attention nodes with compact sequence-length masks. Although cuDNN SDPA was enabled for MultiHeadAttention and GroupQueryAttention in #28849, the legacy fused Attention path did not dispatch to the shared cuDNN runner. On an H100, all 24 attention nodes in a BGE encoder consequently selected the math kernel.

This change reuses the existing cuDNN implementation and eligibility checks rather than adding another attention kernel.

Performance

Exact A/B comparison using the same patched build of current main, with cuDNN auto-dispatch disabled for the baseline via ORT_ENABLE_CUDNN_FLASH_ATTENTION=0:

Batch × sequence	Baseline rows/s	cuDNN SDPA rows/s	Speedup
65 × 696	156.80	546.11	3.48×
12 × 1600	50.50	209.29	4.14×
12 × 3072	20.03	93.43	4.67×

Environment: NVIDIA H100 (SM90), CUDA 12.8, cuDNN frontend 1.24, FP16 BGE encoder with right-padded batches. Debug dispatch confirmed CUDNN_FLASH_ATTENTION for 24/24 fused Attention nodes.

Numerical comparison against the same build with cuDNN disabled, over sampled embeddings:

• maximum absolute difference: 0.001465;
• minimum cosine similarity: 0.999955.

Validation

• Full CUDA build completed successfully.
• ContribOpAttentionTest.CudnnFlashAttentionWithKeySequenceLengthMask: passed and reported SdpaKernel=CUDNN_FLASH_ATTENTION.
• ContribOpAttentionTest.*: 46/46 passed (one pre-existing disabled test).
• clang-format --dry-run --Werror on all changed files.
• git diff --check.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[mastryukov1990]

@tianleiwu Could you please take a look and approve the pending GitHub Actions workflows when you have a chance? This PR extends the cuDNN SDPA integration from #28849 to the fused contrib Attention operator with compact 1D sequence-length masks. I validated the dispatch on H100 (24/24 BGE attention nodes selected cuDNN), the full ContribOpAttentionTest.* suite passed (46/46), and the exact same-build A/B showed a 3.48–4.67× throughput improvement. Thanks!

[Copilot]

Pull request overview

This PR extends the CUDA implementation of the contrib com.microsoft::Attention operator to reuse the existing cuDNN SDPA (Flash Attention) runner when eligible (FP16/BF16, no raw mask, no KV cache, and cuDNN reports support), aligning its dispatch behavior more closely with existing MultiHeadAttention/GroupQueryAttention integration.

Changes:

• Add cuDNN SDPA enable/auto-enable flags to the contrib CUDA Attention kernel and compute cuDNN eligibility early in dispatch.
• Route eligible contrib Attention executions to AttentionKernel_CudnnFlashAttention, including temp-space allocator wiring for the cuDNN runner.
• Add a new contrib Attention test using a 1D key sequence-length mask intended to exercise the cuDNN SDPA path.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/test/contrib_ops/attention_op_test.cc	Adds a new CUDA-focused test case using a 1D key sequence-length mask and env-var kernel selection controls.
onnxruntime/contrib_ops/cuda/bert/attention.h	Adds member flags to track explicit and auto cuDNN SDPA enablement for contrib Attention.
onnxruntime/contrib_ops/cuda/bert/attention.cc	Implements cuDNN SDPA eligibility checks and dispatch integration for contrib Attention, and propagates the selection into workspace sizing and execution data.

[mastryukov1990]

@tianleiwu The follow-up commit 55624f9 addresses the review feedback by asserting the actual cuDNN kernel selection. I validated it on H100: the targeted test passed (1/1) and the full ContribOpAttentionTest.* suite passed (46/46). All 31 workflows are awaiting approval again after the push. Could you please approve/start them when you have a chance? Thanks!

[mastryukov1990]

The Linux TensorRT CI failure appears unrelated to this change. The Build Linux TensorRT x64 Release job failed while the Gradle wrapper was downloading gradle-8.7-bin.zip from services.gradle.org:

java.net.SocketTimeoutException: Connect timed out

The sibling TensorRT CUDA Minimal build completed successfully, so this looks like a transient network/infrastructure failure. A rerun of the failed job should be sufficient.

Failed job: https://github.com/microsoft/onnxruntime/actions/runs/29374987406/job/87230486564

[mastryukov1990]

Sorry for tagging you again, @tianleiwu. The failed Linux TensorRT CI job appears to be a transient Gradle download timeout. Could you please rerun the failed job when you have a chance? Thank you!

[tianleiwu]

Sorry for tagging you again, @tianleiwu. The failed Linux TensorRT CI job appears to be a transient Gradle download timeout. Could you please rerun the failed job when you have a chance? Thank you!

@mastryukov1990, CI pipeline has some issues right now (some jobs are queued forever). I could help trigger CIs later.

[mastryukov1990]

Hi @tianleiwu, sorry to follow up again. The PR currently has 20 successful workflows and 11 cancelled ones (mostly Windows workflows, plus ONNX Runtime CUDA Builds); none are queued or running now. Is the CI infrastructure healthy enough to rerun the cancelled workflows at this point, or should we wait a bit longer? If it is safe now, could you please trigger them? Thank you!

[tianleiwu]

@mastryukov1990, please merge latest main branch, which has a commit for CI pool change.

[mastryukov1990]

@tianleiwu Done — I merged the latest main (f05b21861c) into the PR branch in 9ec8479931. The merge was conflict-free, and the PR diff remains limited to the original three files. The new workflows are now awaiting approval. Thank you!

[tianleiwu]

Reviewed the cuDNN SDPA enablement for the fused contrib Attention kernel. The change reuses the shared cudnn_sdpa::is_stable/is_supported predicate and AttentionKernel_CudnnFlashAttention dispatch already used by MultiHeadAttention, so no new kernel is introduced.

Verified against the surrounding paths:

• Mutual exclusion preserved — cudnn_sdpa_supported correctly gates off flash, the fused TRT runner, and memory-efficient attention; the QkvToContext <= 1 fused-kernel assert still holds.
• QKV format is valid — with cuDNN selected, PrepareQkv_Attention falls into the unfused branch producing Q_K_V_BNSH, which CudnnFlashAttention/build_graph explicitly accept and stride correctly.
• Workspace sizing is correct — GetAttentionWorkspaceSize(..., cudnn_sdpa_supported, false) returns the full qkv_bytes transpose buffer (previously this arg was hard-coded false).
• Safe default — AttentionData::kernel_type defaults to AttentionKernel_Default, so non-eligible calls keep their existing dispatch.
• Test validates both kernel selection (SdpaKernel=CUDNN_FLASH_ATTENTION) and numerics (masked key/value token excluded from both outputs), and skips cleanly when cuDNN is unavailable.

No blocking issues. One maintainability suggestion left inline regarding the implicit attention_bias safety invariant.

[mastryukov1990]

Hi @tianleiwu, the latest head (1f2642b2a7) currently has 31 workflows awaiting maintainer approval. The requested attention_bias documentation and test coverage have been added. Could you please approve/start the workflows when you have a chance? Thank you!

[mastryukov1990]

@microsoft-github-policy-service agree

[mastryukov1990]

@tianleiwu I merged the latest main (0656e94735) into the PR branch in 3f6c54567e. This includes #29754, which fixes the Windows QNN Java 8 / Spotless 7.2.1 CI failure by setting up Temurin 17. The PR diff remains limited to the original three files. The new 31 workflows are awaiting maintainer approval; could you please approve/start them when convenient? Thank you!

[mastryukov1990]

@tianleiwu The latest push only merged the current main into the PR branch to pick up the upstream CI fix; it did not add any new PR changes, and the diff is still limited to the original three files. Could you please re-enable auto-merge (squash) when convenient? Thank you!

[mastryukov1990]

@tianleiwu The newly failed builds appear to be an upstream Windows runner/toolchain issue rather than a regression from this PR. All of the red Windows workflows fail during CMake's compiler sanity check, before any ONNX Runtime source from this PR is compiled:

MSB8040: Spectre-mitigated libraries are required for this project.

The builds enable /Qspectre, but the runners do not have the required MSVC Spectre-mitigated libraries installed. The Linux CUDA and CUDA Plugin workflows pass. This may be related to the BinSkim/Spectre configuration introduced in #29624.

Could you please advise whether the Windows runner image/workflows are being fixed upstream and re-enable or rerun the workflows once that is ready? A rerun on the current runner image will likely fail in the same way.

[tianleiwu]

@mastryukov1990, it needs image update in CI machine pool. Please be patient. I can help re-run the failed CI when the image is updated.

[mastryukov1990]

@tianleiwu Thank you for the clarification. Do you happen to have a rough ETA for when the CI machine-pool image is expected to be updated? No urgency — I just want to understand when it would make sense to check back and rerun the failed workflows.

[mastryukov1990]

@tianleiwu Thank you — the CI image update appears to have resolved the previous Spectre/MSB8040 failures. Only three workflows remain red: VitisAI and XNNPACK lost communication with their self-hosted runners, and the WebGPU vcpkg job hit DXGI_ERROR_DEVICE_REMOVED. These appear unrelated to this PR. Could you please rerun the failed jobs when convenient? Thank you!