Fix CUDA MHA shared-cache scratch lifetime - #31968
Merged
Akshay Sonawane (apsonawane) merged 3 commits intoAug 13, 2026
Merged
Conversation
Keep the sequence-length scratch allocation alive through asynchronous attention launches, populate it with correctly typed total lengths, and seed non-aliased shared-cache outputs before in-place append. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Akshay Sonawane (apsonawane)
requested review from
Tianlei Wu (tianleiwu)
and
a lite review from Copilot
August 10, 2026 23:20
Akshay Sonawane (apsonawane)
enabled auto-merge (squash)
August 10, 2026 23:21
Copilot started reviewing on behalf of
Akshay Sonawane (apsonawane)
August 10, 2026 23:21
View session
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes CUDA MultiHeadAttention behavior when past_present_share_buffer is enabled but outputs are not actually aliased (e.g., in tests), ensuring in-place KV append paths have valid inputs and that scratch buffers remain alive across asynchronous CUDA launches.
Changes:
- Seed
present_{key,value}withpast_{key,value}when buffer-sharing kernels append in place but runtime outputs are not aliased. - Keep the
seqlens_k_totalscratch allocation alive until afterQkvToContextis invoked, and populate it with correctly typed total sequence lengths. - Add a regression test that exercises the non-
cache_indirectionin-place KV concat path under past/present buffer sharing.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc | Fixes scratch-buffer lifetime for seqlens_k_total, corrects its element type/value, and ensures non-aliased present outputs are initialized before in-place append. |
| onnxruntime/test/contrib_ops/multihead_attention_op_test.cc | Adds a regression test covering past/present buffer sharing with in-place KV concat (no cache_indirection). |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Replace the synchronous host-to-device copy of the past-present sequence lengths with a stream-ordered device fill so the buffer is populated in order with the attention launches and remains valid during CUDA graph capture. Also guard the past-to-present cache seeding against missing tensors and expand the comment explaining why it is needed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Ti-Tai Wang (titaiwangms)
approved these changes
Aug 13, 2026
Akshay Sonawane (apsonawane)
deleted the
fix/cuda-mha-seqlens-buffer-lifetime
branch
August 13, 2026 18:40
This was referenced Sep 10, 2026
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.
Keep the sequence-length scratch allocation alive through asynchronous attention launches, populate it with correctly typed total lengths, and seed non-aliased shared-cache outputs before in-place append.