Repository navigation
Conversation
On SM>=90, GroupQueryAttention auto-prefers cuDNN SDPA. cuDNN is given the present KV capacity as its sequence length. When past and present are not shared, that capacity grows by one every decode step, so cuDNN builds a new execution plan for every generated token. With Ministral-3-3B (fp16 GQA, past_present_share_buffer=false) on H100 this made each GQA call ~12 ms and decode ~3 tok/s; Flash attention gives ~250 tok/s on the same model. Only auto-select cuDNN when the KV buffer is shared (fixed capacity). An explicit cuDNN request through sdpa_kernel / ORT_SDPA_KERNEL is still honored. The workspace estimator keeps treating cuDNN as reachable, which remains a conservative upper bound. Files changed: - onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 458edacb-429b-4169-99fb-c2241f5113d4
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Missing regression tests for SM90+ cuDNN selection with shared buffers, unshared buffers, and explicit selection leave guard removal and recurring per-token plan building undetected.
1 open issue
What changed in this PR
This PR limits automatic cuDNN SDPA selection for CUDA GroupQueryAttention to shared KV buffers, avoiding per-token execution-plan rebuilding during decode with growing, unshared caches.
Changes:
- Requires shared past/present KV buffers for automatic cuDNN selection on SM90+.
- Preserves explicit cuDNN requests through
sdpa_kernelorORT_SDPA_KERNEL. - Documents the performance rationale for the new selection guard.
| File | Description |
|---|---|
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc |
Adds the shared-buffer requirement to automatic cuDNN selection while preserving explicit requests. |
Review effort: Balanced
Copilot comment #4541338598: cover SM90 shared/unshared auto-selection and explicit selection with unshared buffers. The production call uses IsGQACudnnSdpaEnabled at group_query_attention.cc:803; its definition is in group_query_attention_eligibility.h:127. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f390d3f0-4e55-45ca-a858-bdc339efc3e3
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
On SM>=90,
GroupQueryAttentionauto-selects cuDNN SDPA. cuDNN receivesseqlen_present_kv_cache(the present KV capacity) as its KV sequence length. Whenpast_present_share_bufferis false, that capacity grows by one on every decode step, so the cuDNN graph/plan is rebuilt for every generated token.This change only auto-selects cuDNN when the past/present buffer is shared (fixed capacity). Explicit cuDNN requests via
sdpa_kernel/ORT_SDPA_KERNELare still honored. The GQA workspace estimator still treats cuDNN as reachable, which remains a conservative upper bound.Motivation and Context
Ministral-3-3B (fp16 GQA,
past_present_share_buffer=false) decoded at ~3 tok/s on H100. Session profiling showed GQA at ~315 ms of a ~320 ms decode step, constant regardless of KV length;ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFOshowedSdpaKernel=CUDNN_FLASH_ATTENTION.Validation (H100, CUDA 12.8, cuDNN 9.8,
CMAKE_CUDA_ARCHITECTURES=90)Single GQA node (32 Q heads, 8 KV heads, head size 128, fp16, past 512, unshared KV), median over 60 decode steps:
sdpa_kernel=8, old auto behavior)sdpa_kernel=1)Max abs diff, default vs cuDNN outputs: 2.4e-4.
End-to-end Ministral-3-3B on H100 with
sdpa_kernel=1(equivalent to this change): ~3 -> ~250 tok/s decode; warm multi-turn image TTFT ~490 ms -> 30-68 ms; identical greedy token counts.