webgpu: Fix GQA decode split-reduce head_size out-of-bounds race - #29593
Merged
tianleiwu merged 1 commit intoJul 7, 2026
Merged
Conversation
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
marked this pull request as ready for review
July 7, 2026 07:14
tianleiwu
approved these changes
Jul 7, 2026
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.
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.
Summary
flash_attention_decode_qkv.wgsl.template) that corrupts the upper half of each head whenv_head_size_vec < tile_size_k_vec.seqlens_kvalues in the WebGPUGroupQueryAttentionop tests (group-query-attention.jsonc).Motivation
The experimental "Run ort-web tests - WebGPU EP" CI stage had 9 failing
GroupQueryAttentionop 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](wheretile_outputis sizedv_head_size_vec) but is guarded only byif (local_idx < tile_size_k_vec)(tile_size_k_vec = 8). Whenv_head_size_vec < tile_size_k_vec(e.g.head_size = 8→v_head_size_vec = 2), threads withlocal_idxin2..7compute an out-of-bounds index. WGSL clamps OOB array indices to the last valid element, so those threads all race ontile_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_vecto the reduction guard in both the TurboQuant and non-TurboQuant branches.2. Stale test data. 16 cases in
group-query-attention.jsoncsetseqlens_k = total_sequence_length, but the canonical GQA convention isseqlens_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 andseqlens_k out of rangeerrors on the CPU/wasm backend. The stored expected outputs were already correct.Test plan
--build_wasm --use_webgpu --use_webnn) and deployed tojs/web/dist.npm test --webgpu-ep -- op group-query-attention -b=webgpu: all[webgpu]and[wasm]GQA cases pass (previously 9[webgpu]tensor mismatches + CPUseqlens_kerrors).[webgpu]failures 16 → 9, and the test-data fix cleared the remaining off-by-one cases on both backends.