[Tianlei Wu (tianleiwu)]

Description

Follow-up to #32493 that addresses the remaining review feedback after the original PR was merged.

• Resolve the default attention scale once for the cuDNN paged path and carry the exact value from the buildability probe through dispatch, keeping the byte-keyed graph cache stable during CUDA graph capture.
• Clear the stale FlashAttention selection when a late XQA-to-cuDNN retry succeeds.
• Remove the unused max_sequence_length_kv shape-gate argument, document the unbounded thread-local graph-cache growth, and include cuDNN in backend diagnostics.
• Exercise default-scale cache-key canonicalization with head_size=96 in the cuDNN paged tests.

Testing

• clang-format --dry-run --Werror on all changed files
• git diff --check
• Direct compilation of paged_attention.cc, paged_attention_impl.cu, and cudnn_flash_attention.cc in the CUDA 13.0 / cuDNN 9.23 H200 build
• Standalone compilation of paged_attention_op_test.cc with the CUDA provider flags

The available CUDA build was configured with onnxruntime_BUILD_UNIT_TESTS=OFF, so the focused PagedAttention.Cuda_CudnnPaged* runtime tests could not be executed locally.

[Copilot]

🟡 Changes recommended

The cache-key and late XQA fallback fixes lack regression tests that would fail if those behaviors regress.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Refines CUDA PagedAttention’s cuDNN paged-SDPA dispatch and cache behavior.

Changes:

• Reuses one canonical attention scale for probing and execution.
• Corrects XQA-to-cuDNN fallback state and diagnostics.
• Updates cuDNN tests to use head size 96.

File summaries

File	Description
attention_data.h	Carries the resolved cuDNN scale.
cudnn_flash_attention.cc	Simplifies eligibility and documents cache growth.
cudnn_flash_attention.h	Updates the paged-support API.
paged_attention.cc	Canonicalizes scale and corrects backend selection.
paged_attention_impl.cu	Uses the stored canonical scale.
paged_attention_op_test.cc	Changes cuDNN test shape to head size 96.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.