Fix XQA reduction races in shared-cache GQA decoding by Copilot · Pull Request #33703 · microsoft/onnxruntime

Copilot · 2026-10-28T23:16:30Z

Description

Output ownership: Only the final reducing CTA publishes multi-block attention output; other CTAs cannot overwrite it with partial results.
Synchronization: All threads latch the final-CTA flag before its shared memory is reused. Four more cp.async drains / warp barriers on the same path (after GEMM0, before mergeAndSaveOutTile, at the start of loadBufAsync and after the waitGroup<1>() in the multi-block merge loop) were compiled only for XQA_PAGED_INT4; they are now enabled for every variant. Two more __syncwarp()s, after the K-tile wait in the GEMM0 loop and after the V-tile wait in syncVTileLoad, are ported from NVIDIA/TensorRT-LLM#18141: cp.async.wait_group only waits for the copies of the executing thread, but every lane reads tile data that other lanes of the warp copied.
Regression coverage: XqaMultiBlockTest (fp16 and bf16) runs 24 independent decode problems per launch on 4 concurrent streams, each with its own workspace and output, with K/V in mapped host memory. That is what makes the races observable; a single decode problem on one stream passes on the unpatched kernel. It covers head sizes 64, 128 and 256 (256 has a single GEMM1 warp group, so the CTA output is staged in a different shared-memory buffer), group sizes 1, 2 and 4, batch 1, 2 and 4, a sliding window with per-head attention sinks, 1 and 4 sub-sequences, direct launches and CUDA graph replay. Every variant captures a 4-stream fork/join graph and asserts that the multi-block grid is launched. The test is skipped below sm_80.

Motivation and Context

GQA decoding with aliased past/present KV caches can intermittently corrupt attention and change generated tokens, especially under CUDA graph capture. Eligible shared-cache decoding selects XQA independently of sdpa_kernel, so requesting Flash does not exclude the affected path.

Testing

The A100 section below was measured on the first revision of this PR (8 tests, without the two K / V __syncwarp()s and the head size 256 tests); the H200 section covers the current revision.

A100-SXM4-80GB (sm_80), first revision

CUDA 13.0, Release, built with CMAKE_CUDA_ARCHITECTURES=80, onnxruntime_BUILD_UNIT_TESTS=ON, onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON, targets onnxruntime_provider_test and onnxruntime_providers_cuda_ut:

cd build/<config>/Release
GTEST_FILTER='XqaMultiBlockTest*' ./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All'

The test detects the bug. The previous version of the test passed in 24/24 processes on the unpatched kernel. The current version, 24 independent processes per row with all 8 tests in each:

`mha_impl.cuh`	Result
`main` (both race-fix hunks reverted)	0/24 pass: all 8 tests fail in every run, typically within the first few launches per stream
the two race-fix hunks only	24/24 pass
this PR	24/24 pass (192 test runs, none skipped)

The four extra sync sites matter. compute-sanitizer --tool racecheck on the fp16 head size 64 test, one repeat per configuration:

	errors	racing instruction pairs	pairs involving `cp.async`
the two race-fix hunks only	634	347	328
this PR	168	19	0

The 19 remaining pairs (LDSM/STS, LDS/STS) are also present without this change. I believe they are the GEMM0 → GEMM1 tile hand-off, which the xBar / ctaRowMax mbarriers order and racecheck does not model, but I have not proven that.

Cost of the four sync sites. CUDA graph replay of 100 back-to-back launches, median of 3 runs on one A100, 12 shapes (head size 64/128, group size 1/4, batch 1/8, length 505–2503): +0.0% to +1.1%, for example 15.57 → 15.68 µs and 25.16 → 25.42 µs.

Regression suites (no failures; skipped cases need WebGPU, Flash Attention or a PagedAttention EP that this build does not enable, SM89+ for FP8, or SM90+ for cuDNN paged SDPA and H512 MEA):

Suites	Passed	Skipped
`GroupQueryAttention{,GrowingCache,Position,ShapeInference}Test.*`	107	25
`PagedAttention.`, `SparsePagedAttention.`	45	53
`PagedAttentionXqa`, `SparsePagedAttentionInt8` (INT8 XQA)	12	2 (FP8)
internal `Xqa*` (8 multi-block, 2 H512)	10	0

H200 (sm_90a), current revision

CUDA 13.0, Release, built for sm_90a with the same options and targets; Flash Attention and memory-efficient attention are not compiled in. ORT runs the same mha_impl.cuh kernel on sm_90 (the Hopper gmma_impl.cuh is not built).

`mha_impl.cuh`	`XqaMultiBlockTest*` (10 tests)
`main`	0/10 pass: every test fails, including both head size 256 tests
this PR before the K / V `__syncwarp()`s	10/10 pass
this revision	10/10 pass in each of 7 parallel processes (one per GPU), none skipped

The two K / V __syncwarp()s are hardening, not a failure I can reproduce: the tests also pass without them on H200, and racecheck does not model intra-warp cp.async visibility. They follow from the cp.async.wait_group semantics and match NVIDIA/TensorRT-LLM#18141, which includes them in its fix for the same race. The head size 256 tests cover the original race fix for a kernel configuration that was not tested before.

Cost of the two K / V __syncwarp()s. CUDA graph replay of 50 back-to-back launches (min of 9 trials per run), median of 4 alternating A/B rounds, 10 shapes (head size 64/128/256, batch 1–32, length 1024–8192), fp16 and bf16: median change +0.02%, range −0.43% to +0.86%, within the round-to-round noise (about ±1–2%). For example 14.87 → 14.83 µs (head size 128, batch 1, length 8192), 12.01 → 12.06 µs (head size 64) and 20.30 → 20.47 µs (head size 256).

Regression suites on this revision:

Suites	Passed	Skipped	Failed
`GroupQueryAttention{,GrowingCache,Position,ShapeInference}Test.*`	108	24	0
`PagedAttention.`, `SparsePagedAttention.`	47	46	5
`PagedAttentionXqa`, `SparsePagedAttentionInt8` (INT8 and FP8 XQA)	14	0	0
internal `Xqa*` (10 multi-block, 2 H512)	12	0	0

The 5 PagedAttention failures (Cuda_SkipInputSanitizerCudaGraphReplayChangingLengths, Cuda_SkipInputSanitizerRaggedSmallBatchParity, Cuda_XqaNativeFp16CacheMultiTokenBoundFallsBack, Cuda_XqaSpecDecInt8CacheHeadSize128NonCausalFallsBack, Cuda_XqaSpecDecInt8CacheHeadSize128UnsupportedWidthFallsBack) are not caused by this PR: they expect a Flash Attention / memory-efficient fallback (PagedAttention requires FlashAttention ...), which this build does not compile in, and they fail the same way on main and on the PR before this revision.

Not covered: GPUs other than sm_80 and sm_90a, Windows (the _WIN32 environment-variable handling in the test was kept from the earlier revision but not compiled here), builds with Flash Attention enabled, and the Python GQA tests.

Comparison with TensorRT-LLM

Compared with NVIDIA/TensorRT-LLM#18141 and every later change to cpp/kernels/xqa in TensorRT-LLM (the numbers in the table are TensorRT-LLM pull requests):

TensorRT-LLM change	Status here
TRT-LLM#18141: only the final reducing CTA writes the output, `isLastCta` is latched by all threads before shared memory is reused, GEMM0 `cp.async` drain	In this PR
TRT-LLM#18141: `__syncwarp()`s inside the multi-block merge loop	Covered by the barriers at the start of `loadBufAsync` and after `waitGroup<1>()` that this PR enables for every variant
TRT-LLM#18141: `__syncwarp()` after the K-tile and V-tile `cp.async` waits	Added in this revision
TRT-LLM#18141: `gemm0ToPostProcBarrier` mbarrier	Not needed: the `__syncthreads()` at the end of GEMM0 pairs with the first `__syncthreads()` of the GEMM1 post-processing (`needMergeGlobal` or `mergeAndSaveOutTile`), which orders the same hand-off; warps that already exited are not waited for
TRT-LLM#18141 (sink handled by CTA 0 only) and TRT-LLM#19220 (sink in the initial row max, zero-denominator guard)	Not applicable: the sink is still added once in the final reduction, with the merged row max. The sliding-window tests check per-head sinks of 4.0–6.0 against the analytic reference with 1 and 4 sub-sequences, and a sink applied twice would move the result by at least 13%, far above the tolerance
TRT-LLM#14459 (paged sliding window with unmapped pages)	Not needed: `LaunchSanitizeBlockTable` maps page ids outside `[-1, num_blocks)` to -1, `HeadPtr::operator+` turns them into `nullptr`, and `copyPartialHeadsAsync` zero-fills them
TRT-LLM#15739 (speculative-decoding sinks)	Already equivalent here
TRT-LLM#17976 (SM120 performance)	`writesOutput` is part of this PR. The V-page cursor rewrite is not needed: ORT only instantiates 128-token pages, and for every head size and `cacheVTileSeqLen`, `nbVItersPerXIter > 1` only occurs with `xIterSeqStride <= tokensPerPage`, where the old and the new cursor stepping are identical
TRT-LLM#12170 (`__CUDA_ARCH__ == 1210` in the XQA per-architecture tables)	Not needed for correctness: ORT's unknown-architecture `#else` default already gives sm_121 the sm_120 values. Ported as explicit parity in #33762
TRT-LLM#10335 (performance), TRT-LLM#11066 and the Hopper part of TRT-LLM#14459 (`gmma_impl.cuh`, not built by ORT), TRT-LLM#15312 (host runner)	Not applicable

Fixes [CUDA EP] GroupQueryAttention with past/present bound to the same buffer gives nondeterministic, sometimes wrong results across runs (worse under CUDA graph capture) #33562

azure-pipelines · 2026-10-28T23:16:39Z

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

Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>

Tianlei Wu (tianleiwu)

Verdict: COMMENT — both production hunks are correct and verified (each fixes a race I reproduced on main), but the new test passes on the unpatched kernel. I'd strengthen it, and show that it fails on main, before this leaves draft. Everything else is non-blocking. The two hunks correspond to the race fixes in NVIDIA/TensorRT-LLM#18141 (merged), which this kernel is derived from.

What I verified (A100-SXM4-80GB, sm_80, nvcc 13.0, -O3 -DNDEBUG). I did not run the gtest binary itself (it needs a full CUDA build with onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON). A standalone harness compiles mha_impl.cuh from main and from this PR head and drives xq::launchMHA with the test's setup: head size 64, capacity 2503, XQA_NB_SUB_SEQ=4, zero query/key, lengths 505/1025/2503, the analytic mean and the test's tolerances. "Wrong" means outside the tolerance.

Configuration (fp16 unless noted)	`main` kernel	PR kernel
The test's configuration (1 head, 1 stream, device-memory K/V), 400,000 launches per length, fp16 and bf16	0 wrong	0 wrong
Same, with a 100 µs `__nanosleep` injected (A) in non-final CTAs before their output write, or (B) in warps other than warp 0 before they read the last-CTA flag; 200 launches per length	200/200 launches wrong at every length, for A and for B (max error up to 1.3)	0 wrong
24 heads × 4 streams, K/V in mapped host memory, 80,000 launches per length	6,108 / 12,502 / 12,105 wrong heads (length 505 / 1025 / 2503; max error 2.8)	0
Same, only 100 launches per length, repeated 20 times	wrong output in 20/20 runs	0/20 runs
96 heads × 4 streams, device-memory K/V, 80,000 launches per length	length 2503: 179 wrong heads; 505 and 1025: 0	0

#	Severity	Where	Issue
1	High	`xqa_multiblock_test.cc` (inline)	The test passes on the unpatched kernel, so it cannot catch a regression
2	Suggestion	`mha_impl.cuh` (inline)	Four more barriers / `cp.async` drains on the same path are still `XQA_PAGED_INT4`-only
3	Nitpick	`xqa_multiblock_test.cc` (inline)	The graph variants are single-stream; the direct variants don't assert the multi-block grid
4	Nitpick	`xqa_multiblock_test.cc`	One configuration only (head size 64, group size 1, batch 1, unquantized cache; no paged KV, sliding window or sink)
5	Nitpick	`mha_impl.cuh` L2399-L2401 (outside the diff)	Possible follow-up, unrelated to #33562: the title of the upstream change also mentions attention-sink numerical accuracy, and it changed where the sink term is added in multi-block mode; this copy still adds it in the final reduction. I did not test whether that produces wrong results here
6	Nitpick	PR description	No validation evidence, and the checklist items are unchecked

Not run: the gtest binary, GPUs other than A100, Windows, quantized / paged-INT4 / fp8 variants, and group sizes above 1.

… the multi-block test Address review feedback on the XQA multi-block reduction fix. Kernel: four synchronization points around cp.async shared-memory reuse were only compiled for the INT4 paged variants. Enable them for every variant: - land the K prefetches at the end of the gemm0 branch before the post-processing warps reuse the shared memory, - land the copies of all lanes before the output tile reuses the shared memory, - sync the warp before a merge-buffer load overwrites a tile that other lanes may still be reading, - sync the warp after each lane waited for its own copies, before the whole warp reads the tile. compute-sanitizer racecheck no longer reports any shared-memory hazard that involves a cp.async write in the fp16 head-size-64 test (328 of 347 racing instruction pairs before, none after). Launch time changes by at most about 1% in CUDA-graph benchmarks on A100. Test: the previous test also passed on the unpatched kernel. Rewrite it so that it fails there: - 24 or more independent decode problems per launch, run by 4 concurrent streams that each own their workspace and output buffers, - K/V in mapped host memory so that the CTAs of a sequence finish out of order, - every round is captured through a multi-stream fork/join, and the grid of every captured kernel is asserted for all variants, not only for graph replay, - head size 64 and 128, GQA group size 1, 2 and 4, batch 1, 2 and 4, sliding window with attention sinks, fp16 and bf16, 1 and 4 sub-sequences, - the analytic reference is kept and every stream and round must also match the first output bit for bit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Tianlei Wu (tianleiwu) · 2026-10-30T00:10:17Z

Update for the review summary items that have no thread, all in 7ae2282. The three inline threads are answered and resolved.

Incremental updates. #4, one configuration only. The test now covers head sizes 64 and 128, group sizes 1, 2 and 4, batch 1, 2 and 4, a sliding window with per-head sinks, fp16 and bf16, 1 and 4 sub-sequences, and lengths 505 / 1025 / 2503. It does not cover quantized (INT8 / FP8 / INT4) or paged caches. The four sync sites are shared with those variants, and the existing INT8 XQA paged tests and the GQA / PagedAttention suites pass; the numbers are in the description. FP8 was not run (needs sm_89+).
Add doxygen web for the project #5, attention sink. XqaMultiBlockTest.{Float16,BFloat16}SlidingWindowAttentionSinks check per-head sinks of 4.0–6.0 against the analytic reference sum(V) / (visible + exp(sink)) with 1 and 4 sub-sequences, and pass in 24/24 processes. A sink applied twice would move the result by at least 13%, far above the tolerance. So adding the sink in the final reduction gave no wrong result here. I did not port the move to CTA 0 and did not test much larger sink magnitudes, so any numerical-accuracy motivation behind the upstream change is not ruled out; I left that as a possible follow-up.
Add doxygen generated website for the project #6, description. Updated with the build and run commands and the results: the new test fails 0/24 on main and passes 24/24 with this PR, plus racecheck, timing, the regression suites and what is not covered.

Copilot

🔵 Needs a closer look

CUDA maintainer review is needed to resolve 19 unexplained racecheck instruction pairs and assess cross-architecture synchronization safety across XQA variants beyond the SM80-only validation.

No open issues 🎉

What changed in this PR

Fixes XQA multi-block reduction races that can corrupt GroupQueryAttention output during shared-cache decoding, addressing #33562.

Changes:

Restricts output writes to the final reducing thread block.
Enables shared-memory synchronization across all XQA variants, rather than only paged INT4.
Adds fp16 and bf16 regression tests covering concurrent streams, CUDA graph replay, grouped queries, and sliding windows with attention sinks.

No actionable findings were identified.

File	Description
`onnxruntime/test/contrib_ops/cuda_kernels/xqa_multiblock_test.cc`	Adds multi-stream reduction stress tests with analytic output checks and launch-grid validation.
`onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh`	Fixes output ownership and synchronization around shared-memory reuse and reduction.

_{Review effort: Balanced}

… XQA Compare this PR with NVIDIA/TensorRT-LLM#18141 and the later XQA commits of TensorRT-LLM. Two synchronization points of #18141 were still missing here: - after the K-tile cp.async wait in the gemm0 loop, - after the V-tile cp.async wait in syncVTileLoad. cp.async.wait_group only waits for the copies issued by the executing thread, but each lane reads K/V tile data that other lanes of the warp copied (ldmatrix), so the warp has to synchronize before the tile is read. Both call sites are warp-uniform. This is formal memory-model hardening plus upstream parity: compute-sanitizer racecheck does not model these intra-warp hazards, and on H200 the tests pass without the barriers as well. The other changes of #18141 and the later XQA commits are either already part of this PR, are not needed because of design differences (the sink is added once in the final reduction, the gemm0 to post-processing hand-off is ordered by __syncthreads, invalid page ids are zero-filled through the block-table contract), or only touch code that ORT does not build (Hopper gmma, sm_121, host runner). Test: add a head-size-256 configuration to XqaMultiBlockTest. Head size 256 has a single GEMM1 warp group, so the CTA output is staged in a different shared-memory buffer than for head sizes 64 and 128. On H200 (sm_90a) both new tests fail on the kernel of main and pass with this PR. Validated on H200 (sm_90a, CUDA 13.0): XqaMultiBlockTest passes 10/10 before and after this change (7/7 parallel processes after it) and fails 10/10 on main. CUDA-graph benchmark over 10 shapes in fp16 and bf16, 4 alternating rounds: median change +0.02%, range -0.43% to +0.86%. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Akshay Sonawane (apsonawane) · 2026-10-30T19:50:04Z

Review at `dfe553ee348abab1d1f9c0a3ef3005dd364b8ec1`

I reviewed the complete two-file diff, traced the surrounding producer/consumer and multi-block reduction paths, checked the launch/workspace helpers and test registration, and read the existing resolved feedback. I did not find a new confirmed correctness defect in the current changes.

Production fixes and previous feedback

Final-output ownership is now gated by the latched isLastCta flag. Single-block operation retains true; non-final multi-block CTAs cannot overwrite the merged output with partial results. The block barrier after reading the flag prevents shared-memory reuse before all participating threads latch it.
The added K/V warp barriers follow per-thread cp.async.wait_group waits. The unconditional GEMM0 drain precedes its block-wide hand-off; the consumer's first block barrier is in needMergeGlobal or mergeAndSaveOutTile. The merge-loop barriers protect buffer reuse and cross-lane reads. I did not identify an unmatched/divergent barrier introduced by these hunks.
The current regression is substantially stronger than the earlier single-problem test: four concurrent streams with independent workspace/output, mapped K/V when available, distinct split means, NaN-prefilled outputs, analytic and bitwise checks, direct launches and fork/join graph replay, and explicit (splits, KV heads, batch) grid assertions.
FP16/BF16 cases cover head sizes 64/128/256, grouped queries and sliding-window sinks. The forced-SM-count workspace calculation agrees with GetXQAScratchSize's subsequence formula for the tested one/four splits. The source is registered through the internal-CUDA-test glob; those tests require onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS.

A standalone host check confirmed that the selected lengths' nonempty split-only means differ from the analytic full result by more than the BF16 tolerance, and that double-adding the tested sinks also exceeds that tolerance. This validates the oracle's discrimination, not GPU race reproduction. The A100/H200 before/after execution results in the description remain the author's measurements.

Remaining qualification, not a newly proved race

The description's first-revision A100 racecheck report still contains 19 unclassified instruction pairs. They should not be called proven false positives, nor assumed to prove a race in this current revision. Please attach the tool version and current-revision report, map each pair to source/shared-memory storage, and identify the exact xBar/row-stat barrier happens-before edge that orders it. Keep this distinct from the now-removed cp.async hazards. Quantized/paged/spec-dec and Windows coverage should likewise be reported explicitly rather than inferred solely from this nonquantized test.

Ensure the regression runs in the intended CUDA CI configuration without skipping, with at least one supported architecture's before/after evidence retained. The current test falls back to device K/V if host mapping fails; in that mode, its sensitivity to the original race should be qualified rather than assumed identical to the mapped-memory stress run.

WASM Debug is a separate test-loading failure

The failed job JUnit artifact reports the same three optimizer external-data tests failing with Module.MountedFiles is not available: DiskExternalConstantsUseOrtLoadingAndPropagateFailures, DiskExternalScalarAndSingletonConstantsPreserveExactTypeAndRank, and EmptyExternalConstantsValidateFilesAndOffsetsWithoutReadingWholeFiles. They are not XQA regression failures. Address/pick up the WASM harness/external-data-loading fix and rerun, without disabling the CUDA regression coverage.

Validation limits: complete changed-source review, helper/registration trace, host oracle checks and downloaded CI results. No local CUDA build, GPU execution, independent performance measurement or Compute Sanitizer run. No code changes, thread resolutions or approval were submitted in this review.

## Description Ports the architecture-table part of NVIDIA/TensorRT-LLM#12170 ("Compile XQA with sm_120f") to ORT's XQA decode kernels: `__CUDA_ARCH__ == 1210` (sm_121 — GB10 / DGX Spark and other Blackwell Arm devices) is now listed next to sm_120 in the two compile-time arch tables. sm_121 has the same 99 KB opt-in shared-memory limit as sm_120 and uses the same K/V tile configuration. This is explicit-parity hardening, not a user-visible bug fix: sm_121 already built and ran through the "unknown architecture" `#else` default, which has the same values as sm_120. See Testing for the evidence that the generated code is identical before and after. ## Summary of Changes ### XQA arch tables | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh` | Add `__CUDA_ARCH__ == 1210` to the `860 / 890 / 1200` branch that selects `preferedKHeadPartBytes = 64` and `cacheVTileSeqLen = 32` (previously reached only via the `#else` default). Mention sm_121 in the related comment and in the shared-memory-limit hint text. | | `onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh` | Add `__CUDA_ARCH__ == 1210` to the 99 KB `kMAX_SMEM_SIZE` branch. Before, 1210 fell through to the 48 KB default. `kMAX_SMEM_SIZE` is currently only used by a `#if 0` `static_assert`, so this has no runtime effect today; it just keeps the table correct. | ### Not ported from TRT-LLM #12170 (not applicable) The rest of #12170 changes TensorRT-LLM's NVRTC JIT wrapper (`-arch=sm_120f`, dropping the compute_89 PTX to sm_120 two-stage compile through nvPTXCompiler, and the `nvptxcompiler_static` link). ORT compiles XQA ahead of time with nvcc, so there is nothing to port. In TRT-LLM the default arch branch was an `#error` and `kMAX_SMEM_SIZE` was undefined for 1210, which is why sm_121 failed to compile there; ORT's `#else` fallbacks mean it never did. ## Do the shipped libraries include sm_121? Yes, so no packaging or CMake change is needed: - Linux aarch64 CUDA 13 wheel: `87-real;89-real;90-real;100-real;103-real;110-real;120-real;121-real` (`py-cuda13-packaging-pipeline.yml`). - Unified Linux aarch64 wheel (`py-cuda13-unified-packaging-pipeline.yml`): passes no `-a`, so it takes the `build_linux_python_package.sh` default for CUDA 13.0 on aarch64, which is the same list. - CUDA plugin (`plugin-cuda13-pipeline.yml` / `plugin-cuda-packaging-stage.yml`): Linux aarch64 uses the same list; Windows ARM64 uses `120-real;121-real`. - x86_64 lists only carry `120-real` (sm_121 parts are Arm). - CMake normalizes `121-real` to `121a-real`, and XQA sources are compiled for every selected arch >= 80, so XQA gets sm_121a SASS (no PTX). The CUDA plugin's XQA target (`onnxruntime_providers_cuda_plugin_xqa`) uses the same `MIN_SM 80` filter, with no virtual replacement (that option is only used for the plugin's separate LLM target), so it gets the same SASS on Linux and Windows. ## Testing Built the XQA object library of the CUDA provider (`onnxruntime_providers_cuda_xqa`, CUDA 13.0) with `CMAKE_CUDA_ARCHITECTURES='120-real;121-real'`. The configure log shows `CMAKE_CUDA_ARCHITECTURES: 120a-real;121a-real`. - All 44 XQA objects contain both `sm_120a` and `sm_121a` cubins (`cuobjdump -lelf`), with 158 `kernel_mha` kernels each. Per-kernel register / shared / local / constant usage is identical between sm_120a and sm_121a. - A/B against `main` (without this change): SASS is identical for all 44 objects on sm_120a and sm_121a (600,232 instructions compared; all 88 cubins byte-identical), and also identical on sm_90a for 8 representative translation units (loader, paged, spec-dec, fp8/int8/int4, head sizes 64-512). So existing archs are unaffected and sm_121 behaves exactly as before. - The XQA kernels are present in the final `libonnxruntime_providers_cuda.so` (checked on an sm_90a build: 158 `kernel_mha` kernels). - `lintrunner` and `clang-format` are clean. - Not run on real sm_121 hardware (none available), so a smoke test on a DGX Spark would be welcome. To reproduce the cubin check: ``` python tools/ci_build/build.py --build_dir build/xqa_sm121 --config Release --use_cuda \ --cuda_home <cuda13> --cudnn_home <cudnn> --build_shared_lib --skip_tests --update \ --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=120-real;121-real' ninja -C build/xqa_sm121/Release onnxruntime_providers_cuda_xqa cuobjdump -lelf <obj dir>/xqa_loader_fp16_128.cu.o # lists sm_120a and sm_121a cubins ``` ## Motivation and Context Keeps ORT's XQA arch tables in sync with TensorRT-LLM (NVIDIA/TensorRT-LLM#12170) so sm_121 is handled explicitly rather than by the unknown-arch fallback, and so the `kMAX_SMEM_SIZE` table is right if that check is re-enabled. Independent of the XQA race fixes in #33703. Overlap with #33373: its commit "xqa: treat SM121 like SM120 in the per-architecture constants" makes the same two `|| __CUDA_ARCH__ == 1210` edits with identical spelling. This PR carries the XQA part on its own so it does not have to wait for that larger CMake change. The two PRs merge cleanly in either order (checked with `git merge-tree` against the head of #33373), so whichever lands second just has a no-op hunk. The extra edits here (the sm_121 mention in the comment and the shared-memory-limit hint text) are in different hunks. ## Checklist - [ ] Tests added/updated (N/A: arch-table-only change with identical generated code; verified by cubin comparison above) - [x] No breaking changes - [x] Documentation updated (N/A: only an error-message hint and comments changed) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Initial plan

cddcceb

Copilot AI assigned Copilot and Tianlei Wu (tianleiwu) Oct 28, 2026

Copilot started work on behalf of Tianlei Wu (tianleiwu) October 28, 2026 23:16

View session

Copilot AI linked an issue Oct 28, 2026 that may be closed by this pull request

[CUDA EP] GroupQueryAttention with past/present bound to the same buffer gives nondeterministic, sometimes wrong results across runs (worse under CUDA graph capture) #33562

Closed

Copilot AI and others added 3 commits October 28, 2026 23:30

Fix XQA multi-block output race in shared-cache GQA decoding

7870c67

Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>

Ensure XQA regression split overrides work on Windows

db4c856

Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>

Avoid deprecated CRT getenv in Windows XQA tests

0322476

Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>

Copilot AI changed the title ~~[WIP] Fix nondeterministic results in GroupQueryAttention with bound buffers~~ Fix XQA reduction races in shared-cache GQA decoding Oct 28, 2026

Copilot finished work on behalf of Tianlei Wu (tianleiwu) October 28, 2026 23:35

Copilot AI requested a review from Tianlei Wu (tianleiwu) October 28, 2026 23:35

Tianlei Wu (tianleiwu) reviewed Oct 29, 2026

View reviewed changes

Comment thread onnxruntime/test/contrib_ops/cuda_kernels/xqa_multiblock_test.cc Resolved

Comment thread onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh Resolved

Comment thread onnxruntime/test/contrib_ops/cuda_kernels/xqa_multiblock_test.cc Resolved

Tianlei Wu (tianleiwu) marked this pull request as ready for review October 30, 2026 04:38

Copilot AI balanced review requested due to automatic review settings October 30, 2026 04:38

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) October 30, 2026 04:39

View session

Copilot AI reviewed Oct 30, 2026

View reviewed changes

Tianlei Wu (tianleiwu) mentioned this pull request Oct 30, 2026

[CUDA] Add sm_121 to XQA arch-specific branches #33762

Merged

2 of 3 tasks

Akshay Sonawane (apsonawane) approved these changes Oct 30, 2026

View reviewed changes

Tianlei Wu (tianleiwu) merged commit e7cc17f into main Oct 30, 2026
94 checks passed

Tianlei Wu (tianleiwu) deleted the copilot/fix-group-query-attention-issue branch October 30, 2026 21:11

Navigation Menu

Repository navigation

Fix XQA reduction races in shared-cache GQA decoding - #33703

Fix XQA reduction races in shared-cache GQA decoding#33703
Tianlei Wu (tianleiwu) merged 6 commits into
mainfrom
copilot/fix-group-query-attention-issue

Copilot AI commented Oct 28, 2026 •

edited by tianleiwu

Loading

Uh oh!

azure-pipelines Bot commented Oct 28, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Tianlei Wu (tianleiwu) commented Oct 30, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Akshay Sonawane (apsonawane) commented Oct 30, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

Copilot AI commented Oct 28, 2026 • edited by tianleiwu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Testing

A100-SXM4-80GB (sm_80), first revision

H200 (sm_90a), current revision

Comparison with TensorRT-LLM

Uh oh!

azure-pipelines Bot commented Oct 28, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Tianlei Wu (tianleiwu) commented Oct 30, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Akshay Sonawane (apsonawane) commented Oct 30, 2026

Review at dfe553ee348abab1d1f9c0a3ef3005dd364b8ec1

Production fixes and previous feedback

Remaining qualification, not a newly proved race

WASM Debug is a separate test-loading failure

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Copilot AI commented Oct 28, 2026 •

edited by tianleiwu

Loading

Review at `dfe553ee348abab1d1f9c0a3ef3005dd364b8ec1`