onnxruntime
1ad2a1f2 - webgpu: Bump FlashAttentionDecodeQKV workgroup to 128, tile_size_k_vec to 32 (#29586)

Commit
36 days ago
webgpu: Bump FlashAttentionDecodeQKV workgroup to 128, tile_size_k_vec to 32 (#29586) ## Summary - Bump `FlashAttentionDecodeQKVProgram` workgroup size from 64 to 128. - Bump `tile_size_k_vec` from 8 to 32. - `sub_tile_count` is derived from `WorkgroupSizeX() / tile_size_k_vec`, so it goes from 8 to 4 automatically. The shader template expresses all loops in terms of these parameters, so no shader-side changes are needed. ## Motivation `FlashAttentionDecodeQKVProgram` was dispatching with a smaller workgroup / tile_size_k_vec than `MatMulNBitsProgram` on the same class of hardware. Mirroring the MatMulNBits dispatch shape improves occupancy and reduces the per-call cost of the fused QK^T + softmax + V multiply step during token generation. Measured on Qwen3-1.7B-graph-prune (WebGPU EP, D3D12, decode): | Metric | Before | After | |---|---|---| | `GroupQueryAttention\|FlashAttentionDecodeQKV` per-call GPU time | 0.83 ms | 0.47 ms | | Token generation throughput | 257 tps | 282 tps | ## Test plan - [x] Built ORT with WebGPU EP (Release, D3D12) on Windows. - [x] `onnxruntime_provider_test.exe --gtest_filter="GroupQueryAttentionTest.*"` — 52 passed, 12 CUDA-only tests skipped. All WebGPU FlashAttention paths pass, including `BatchedRightPaddedRotaryPrefillFlashAttention_WebGPU`, `BatchedRightPaddedRotaryPrefillFlashAttentionLargeSpread_WebGPU`, and `WebGPU_SharedKV_*` variants. - [x] End-to-end correctness on Qwen3-1.7B-graph-prune (decode). - [x] Profiled Qwen3-1.7B-graph-prune with the WebGPU profiling pipeline; confirmed the per-op time improvement above.
Author
Parents
Loading