onnxruntime
ced4d830 - webgpu: Fix GQA decode split-reduce head_size out-of-bounds race (#29593)

Commit
8 days ago
webgpu: Fix GQA decode split-reduce head_size out-of-bounds race (#29593) ## 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.
Author
Parents
Loading