Skip to content

webgpu: Fix GQA decode split-reduce head_size out-of-bounds race - #29593

Merged
tianleiwu merged 1 commit into
microsoft:mainfrom
qjia7:webgpu-fix-gqa-decode-head-size-race
Jul 7, 2026
Merged

webgpu: Fix GQA decode split-reduce head_size out-of-bounds race#29593
tianleiwu merged 1 commit into
microsoft:mainfrom
qjia7:webgpu-fix-gqa-decode-head-size-race

Conversation

@qjia7

@qjia7 qjia7 commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Fix an out-of-bounds/data-race in the WebGPU GQA decode split-reduce shader (flash_attention_decode_qkv.wgsl.template) that corrupts the upper half of each head when v_head_size_vec < tile_size_k_vec.
  • Correct 16 stale seqlens_k values in the WebGPU GroupQueryAttention op tests (group-query-attention.jsonc).

Motivation

The experimental "Run ort-web tests - WebGPU EP" CI stage had 9 failing GroupQueryAttention op tests (tensor mismatches). Root cause is two independent bugs:

1. Kernel out-of-bounds race. In the fused QKV decode shader, the V-multiply reduction writes tile_output[m][k + local_idx] (where tile_output is sized v_head_size_vec) but is guarded only by if (local_idx < tile_size_k_vec) (tile_size_k_vec = 8). When v_head_size_vec < tile_size_k_vec (e.g. head_size = 8v_head_size_vec = 2), threads with local_idx in 2..7 compute an out-of-bounds index. WGSL clamps OOB array indices to the last valid element, so those threads all race on tile_output[m][1] with a +=, corrupting the second head_size vec4. The manifestation is GPU-dependent (zeros on some GPUs, off-by-8 stale values on the CI GPU); the first vec4 stays correct.

Fix: add && k + local_idx < v_head_size_vec to the reduction guard in both the TurboQuant and non-TurboQuant branches.

2. Stale test data. 16 cases in group-query-attention.jsonc set seqlens_k = total_sequence_length, but the canonical GQA convention is seqlens_k = total_sequence_length - 1 (0-based past length), as used by the C++ GQA tests and enforced by the CPU kernel bounds check. This produced an off-by-one KV read on WebGPU and seqlens_k out of range errors on the CPU/wasm backend. The stored expected outputs were already correct.

Test plan

  • Built wasm + WebGPU EP (--build_wasm --use_webgpu --use_webnn) and deployed to js/web/dist.
  • Ran npm test --webgpu-ep -- op group-query-attention -b=webgpu: all [webgpu] and [wasm] GQA cases pass (previously 9 [webgpu] tensor mismatches + CPU seqlens_k errors).
  • Confirmed the kernel fix alone reduced [webgpu] failures 16 → 9, and the test-data fix cleared the remaining off-by-one cases on both backends.

The fused QKV decode shader (flash_attention_decode_qkv.wgsl.template)
reduces per-tile V contributions into tile_output, which is sized
v_head_size_vec. The reduction was guarded only by
`if (local_idx < tile_size_k_vec)` (tile_size_k_vec = 8), but it
writes `tile_output[m][k + local_idx]`. When v_head_size_vec is
smaller than tile_size_k_vec (e.g. head_size = 8 -> v_head_size_vec = 2),
threads with local_idx in 2..7 compute an out-of-bounds index. WGSL
clamps out-of-bounds array indices to the last valid element, so those
threads all race on `tile_output[m][1]` with a += update, corrupting
the second head_size vec4. The result is GPU-dependent garbage in the
upper half of each head (zeros on some GPUs, stale values on others),
while the first vec4 stays correct.

Fix: add `&& k + local_idx < v_head_size_vec` to the reduction guard
in both the TurboQuant and non-TurboQuant branches so only threads that
map to a valid head_size element write to tile_output.

Also corrects stale seqlens_k values in the WebGPU GroupQueryAttention
op tests (group-query-attention.jsonc). The canonical GQA convention is
seqlens_k = total_sequence_length - 1 (0-based past length), as used by
the C++ GQA tests and enforced by the CPU kernel's bounds check. 16
cases incorrectly set seqlens_k = total_sequence_length, which produced
an off-by-one KV read on WebGPU and 'seqlens_k out of range' errors on
the CPU/wasm backend. The stored expected outputs were already correct.

Verified: rebuilt wasm + WebGPU EP and ran the full
'op group-query-attention' suite. All [webgpu] and [wasm] GQA cases now
pass (previously 9 [webgpu] tensor mismatches + CPU seqlens_k errors).
@qjia7
qjia7 marked this pull request as ready for review July 7, 2026 07:14
@tianleiwu
tianleiwu merged commit ced4d83 into microsoft:main Jul 7, 2026
88 of 89 checks passed
tianleiwu pushed a commit that referenced this pull request Jul 7, 2026
)

## Summary

- Fix an out-of-bounds/data-race in the WebGPU GQA decode split-reduce
shader (`flash_attention_decode_qkv.wgsl.template`) that corrupts the
upper half of each head when `v_head_size_vec < tile_size_k_vec`.
- Correct 16 stale `seqlens_k` values in the WebGPU
`GroupQueryAttention` op tests (`group-query-attention.jsonc`).

## Motivation

The experimental "Run ort-web tests - WebGPU EP" CI stage had 9 failing
`GroupQueryAttention` op tests (tensor mismatches). Root cause is two
independent bugs:

**1. Kernel out-of-bounds race.** In the fused QKV decode shader, the
V-multiply reduction writes `tile_output[m][k + local_idx]` (where
`tile_output` is sized `v_head_size_vec`) but is guarded only by `if
(local_idx < tile_size_k_vec)` (`tile_size_k_vec = 8`). When
`v_head_size_vec < tile_size_k_vec` (e.g. `head_size = 8` →
`v_head_size_vec = 2`), threads with `local_idx` in `2..7` compute an
out-of-bounds index. WGSL clamps OOB array indices to the last valid
element, so those threads all race on `tile_output[m][1]` with a `+=`,
corrupting the second head_size vec4. The manifestation is GPU-dependent
(zeros on some GPUs, off-by-8 stale values on the CI GPU); the first
vec4 stays correct.

Fix: add `&& k + local_idx < v_head_size_vec` to the reduction guard in
both the TurboQuant and non-TurboQuant branches.

**2. Stale test data.** 16 cases in `group-query-attention.jsonc` set
`seqlens_k = total_sequence_length`, but the canonical GQA convention is
`seqlens_k = total_sequence_length - 1` (0-based past length), as used
by the C++ GQA tests and enforced by the CPU kernel bounds check. This
produced an off-by-one KV read on WebGPU and `seqlens_k out of range`
errors on the CPU/wasm backend. The stored expected outputs were already
correct.

## Test plan

- [x] Built wasm + WebGPU EP (`--build_wasm --use_webgpu --use_webnn`)
and deployed to `js/web/dist`.
- [x] Ran `npm test --webgpu-ep -- op group-query-attention -b=webgpu`:
all `[webgpu]` and `[wasm]` GQA cases pass (previously 9 `[webgpu]`
tensor mismatches + CPU `seqlens_k` errors).
- [x] Confirmed the kernel fix alone reduced `[webgpu]` failures 16 → 9,
and the test-data fix cleared the remaining off-by-one cases on both
backends.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants