Fix CUDA GQA graph replay test diagnostics - #32632
Draft
Justin Chu (justinchuby) with Copilot wants to merge 10 commits into
Draft
Justin Chu (justinchuby) with Copilot wants to merge 10 commits into
Justin Chu (justinchuby) with Copilot wants to merge 10 commits into
Conversation
## Summary Plan contiguous GQA FlashDecode split-KV launches from fixed KV-cache capacity during CUDA graph capture and replay, while retaining live-sequence-length planning for ordinary eager execution. ## Why CUDA graphs freeze launch geometry and workspace addresses at capture time. Planning NumSplits from the current live sequence length can become stale as the cache grows, preventing adaptive split-KV behavior from remaining valid across replay. Graph-enabled warmup now reserves capacity-sized workspace, capture uses the fixed-capacity plan, and eager decode avoids redundant capacity heuristic work. ## Behavior - Capture/replay uses fixed cache capacity for stable NumSplits and workspace sizing - Graph warmup reserves replay-sized workspace before capture - Eager execution continues to tune from the live sequence length - Active memset size remains limited to the launch plan - Debug output reports the resolved NumSplits ## Validation Focused host tests cover head sizes 64, 128, and 256; local-window and sequence-tail behavior; non-decode inputs; capture planning; and ordinary eager routing. For an SM108 configuration with live length 129 and capacity 4097, capture selected 17/17/22 splits for head sizes 64/128/256, while eager retained live-length plans. Independent review found one redundant eager heuristic computation, which is fixed in this commit. No CUDA kernel timing is claimed because this Windows host did not have nvcc. The change preserves eager routing and targets graph planning correctness and replay-stable adaptive split-KV behavior. Based on the mechanisms validated in justinchuby/onnx-genai#1340 and #1350. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
…uda-attention-capture-safe-split-kv
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Preserve both CUDA cache-aliasing coverage and the FlashDecode graph replay regression test. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Use head-major cache offsets for initialization, appended-value mirroring, and the output reference to match the bound BNSH tensor layout. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Fix failing GitHub Actions job for Linux CUDA plugin
Fix CUDA GQA graph replay test diagnostics
Sep 16, 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.
Description
CUDA GQA replay test
FlashDecodeSplitKvCudaGraphReplay.GetFlashAttentionSplitPlanusing the active CUDA device properties.Split-planning cleanup
defined(USE_FLASH_ATTENTION)guard style.Motivation and Context
The
Test Linux CUDA Plugin EP x64 Releasejob failed inGroupQueryAttentionTest.FlashDecodeSplitKvCudaGraphReplay. The test was treating attention debug output as a hard correctness signal even though CUDA attention options are initialized once per EP/test process, making that output unreliable in CI.The test now asserts the intended split behavior through the same planner used by the kernel path, while preserving the CUDA graph replay/value checks that cover the regression.