[Dante (KafKafrnZ)]

Problem

On the CUDA EP, GroupQueryAttention reads total_sequence_length as 0 when a CUDA node computes it. Fixes #33867, which reports the first row below. The other two rows are the same bug.

Graph, with the total computed on CUDA	Before (onnxruntime-gpu 1.31.0)	After
With attention_bias	attention_bias dimension 3 must be equal to total_sequence_length, got 8	correct output
No bias, growing or no cache	Flash total sequence length must be positive and fit int32	correct output
No bias, static cache, cuDNN SDPA	no error, wrong output	correct output

For the last row, on an RTX 5070 Ti, where cuDNN SDPA is auto-selected: with Q and K zero, a first prompt returned all zeros where the running mean of V is expected (got 0, expected 0.125). Over a prefill and 3 decode steps with random inputs, the output differed from a fed total at every step (by up to 2.65), and the present cache differed from the first decode step on (by up to 3.54).

This isn't a recent regression. On 1.26.0, before bias support, a computed total ran with no error and returned a present of length 0, (2, 2, 0, 16).

Cause

Input 6 is declared OrtMemTypeCPUInput, so a value from a CUDA node arrives through MemcpyToHost, whose output is in host-accessible memory with device type GPU:

[33867] total_seqlen Device:[DeviceType:1 MemoryType:5 VendorId:4318 DeviceId:0 Alignment:0] value=8

CheckInputs only accepts device.Type() == OrtDevice::CPU, so it treats that readable value as unreadable and uses 0. The check came from #26384, for WebGPU values that really are on the device.

Fix

Use device.UsesCpuMemory(), as cuda/tensor/reshape.cc does.

• The read is ordered after the copy. The planner adds a host wait when a consumer reads a CPU-memory output on the same stream (allocation_planner.cc, the MemCpyToHost case). With the length alternating between runs, 2000 runs gave 0 errors and 0 wrong outputs. As a control, with that planner condition removed, 1950 of 2000 runs failed on a stale length.
• Graph capture: capture isn't allowed with Memcpy nodes, so this path doesn't run during capture. This comes from reading the code.
• WebGPU: a value on the device is DEFAULT memory and still skips validation. I couldn't build WebGPU here, so this also comes from reading the code.
• seqlens_k is left alone. It has the same check, but it isn't a CPU input on the CUDA kernel, so a computed value really is on the device.

Behaviour change

A computed total now gets the same validation and first-prompt detection as a fed one. For example, a later multi-token chunk with batch 2 used to run unvalidated and now gets batch_size must be 1 when sequence_length > 1 and past context is given., the same as a fed total.

Testing

RTX 5070 Ti (sm_120), Windows, CUDA 13.4, cuDNN 9.20, on main ee5b21dd99.

• 4 new CudaTotalSequenceLengthFromCudaNode* tests (with and without a bias, each with and without a static cache): all fail without the fix and pass with it.
• Computed vs fed total over prefill + decode and two prefill chunks + decode, with static and growing caches: outputs and caches are identical at every step.
• GroupQueryAttention*:PagedAttention* 183 passed / 37 skipped, MultiHeadAttentionTest* 33, AttentionTest* 89, CUDA_EP_Unittest.All 613 / 8 skipped.
• Plugin EP: GQA's object compiles with this change, checked on a tree from 11-03 rather than this base. It wasn't run.

Not traced: which step on the cuDNN path produces the zeros. The kernel's debug line shows EffectiveKvLengthBound=0 instead of 8, and the new tokens are written to the cache correctly either way.

🤖 Generated with Claude Code

[azure-pipelines]

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

[Copilot]

🟢 Approval recommended

No active review issues remain for the change to recognize host-accessible sequence-length tensors.

No open issues 🎉

What changed in this PR

Fixes #33867 by allowing GroupQueryAttention to read total_sequence_length computed by a CUDA node and copied into host-accessible memory, rather than incorrectly treating it as zero.

Changes:

• Uses UsesCpuMemory() to recognize readable host memory while preserving validation skipping for device-only values.
• Adds four CUDA regression tests covering bias presence and static-cache usage, checking output values and present-cache shapes.

File	Description
onnxruntime/​test/​contrib_ops/​group_query_attention_op_test.cc	Adds regression coverage for CUDA-computed sequence lengths.
onnxruntime/​contrib_ops/​cpu/​bert/​group_query_attention_helper.h	Recognizes host-accessible sequence-length tensors using UsesCpuMemory().

_{Review effort: Balanced}