Repository navigation
[WebGPU] Unroll bounded NVIDIA FlashAttention reductions in WGSL - #33395
Yang Gu (gyagp) wants to merge 3 commits into
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
fe63446 to
2648fb0
Compare
|
/run bots |
|
The request was received. Approve workflow runs. Thanks. |
2648fb0 to
f041c21
Compare
|
/run bots |
|
Jiajia Qin (@qjia7) Please review |
|
The request was received. Approve workflow runs. Thanks. |
| // The NVIDIA shared-memory path benefits from constant private-array indices. | ||
| // Bound source expansion to at most 32 vec4 terms per reduction. Both inputs | ||
| // to this choice already participate in the dense-prefill program cache hint. | ||
| const bool unroll_shm_inner = is_nvidia_ && qkv_head_size_ / 4 <= 32; |
There was a problem hiding this comment.
| const bool unroll_shm_inner = is_nvidia_ && qkv_head_size_ / 4 <= 32; | |
| const bool unroll_shm_inner = is_nvidia_ && use_shm_path_; |
Please also check whether we can remove qkv_head_size_ / 4 <= 32 for unroll_shm_inner.
There was a problem hiding this comment.
Added the explicit use_shm_path_ guard and checked removal of the head-size limit on current main f42c839f56.
I retained the limit because expanding head size 512 causes substantial regressions on RTX 5080/NVIDIA 617.42. At 1024 tokens (24 Q heads/8 KV heads, causal FP16), isolated GPU time changes as follows; positive percentages mean slower:
| Target | Head size | Rolled GPU time (ms) | Expanded GPU time (ms) | GPU time change |
|---|---|---|---|---|
| D3D12 / SM 6.10 | 128 | 0.850 | 0.748 | -11.99% |
| D3D12 / SM 6.10 | 192 | 1.384 | 1.372 | -0.86% |
| D3D12 / SM 6.10 | 256 | 1.943 | 1.944 | +0.05% |
| D3D12 / SM 6.10 | 512 | 7.812 | 17.273 | +121.10% |
| Vulkan | 128 | 0.712 | 0.712 | -0.12% |
| Vulkan | 192 | 1.088 | 1.085 | -0.25% |
| Vulkan | 256 | 2.285 | 2.288 | +0.16% |
| Vulkan | 512 | 9.521 | 25.782 | +170.79% |
| D3D12 / SM 6.8 | 128 | 1.309 | 1.375 | +5.05% |
| D3D12 / SM 6.8 | 192 | 3.332 | 3.333 | +0.01% |
| D3D12 / SM 6.8 | 256 | 2.521 | 2.522 | +0.01% |
| D3D12 / SM 6.8 | 512 | 18.507 | 18.493 | -0.08% |
| D3D12 / SM 6.9 | 128 | 0.850 | 0.748 | -12.02% |
| D3D12 / SM 6.9 | 192 | 1.384 | 1.371 | -0.93% |
| D3D12 / SM 6.9 | 256 | 1.940 | 1.943 | +0.14% |
| D3D12 / SM 6.9 | 512 | 7.815 | 17.300 | +121.37% |
Three alternating process rounds per measurement, identical inputs. D3D12 uses native DXIL GPU-timestamp replay; Vulkan uses native Dawn GPU-timestamp replay linked to the same Dawn build. Their sampling harnesses differ, so the comparisons above are within each API. This is kernel timing rather than full-model throughput. The 512/2048-token checks also show a head size 512 regression.
The final selection is is_nvidia_ && use_shm_path_ && qkv_head_size_ / 4 <= 32. I expanded the 65-token high-magnitude test to heads 64/128/192/256/512 and verified actual shaders:64/128 expand,192/256/512 retain loops. All three native tests, 89 template tests and 48 shader compilations pass. All 64 shader numerical cases pass. Current-main Phi-4 graph capture also passes at 1024 with Subgroup Matrix on both APIs, including exact A/B/A replay and finite/reference-checked logits.
I also refreshed the PR description with the dated October 29 full-model A/B, including Subgroup Matrix. The known SM 6.8 tradeoff and limited architecture coverage remain disclosed.
f041c21 to
a33792f
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
Maintainer judgment is needed on NVIDIA-wide unrolling because reported SM 6.8 regressions and limited architecture coverage leave the cross-target performance tradeoff unresolved.
No open issues 🎉
What changed in this PR
Optimizes WebGPU FlashAttention prefill on NVIDIA by unrolling shared-memory reductions for head sizes up to 128, while preserving accumulation precision and order. Larger heads and other attention paths retain their existing behavior.
Changes:
- Generates bounded, unrolled QK and PV reduction helpers for eligible NVIDIA shaders.
- Selects unrolled or loop-based reductions in the WGSL template.
- Adds high-magnitude FP16 coverage for partial tiles and head sizes 64–512.
| File | Description |
|---|---|
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc |
Adds head-size boundary coverage using 65-token high-magnitude FP16 inputs. |
onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template |
Conditionally invokes unrolled reduction helpers. |
onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc |
Generates helpers and gates unrolling to NVIDIA shared-memory paths with heads up to 128. |
Review effort: Balanced
4ac803e to
c95092f
Compare
Aion-1.1 median prefill on RTX 5080 / D3D12 SM 6.10, with graph capture, 1024 input tokens and 32 generated tokens: - DP4A: 6,506.29 -> 7,688.99 tokens/s (+18.18%). - Subgroup Matrix: 10,120.60 -> 13,293.90 tokens/s (+31.35%).
c95092f to
0d7e0cd
Compare
|
/run bots |
|
Jiajia Qin (@qjia7) please take another look |
|
The request was received. Approve workflow runs. Thanks. |
|
/run bots |
|
The request was received. Approve workflow runs. Thanks. |
Change
Unroll the QK/PV reductions in NVIDIA's dense shared-memory FlashAttention prefill path for head sizes up to 128. Preserve FP32 score accumulation, accumulation order and cache specialization.
Selection:
is_nvidia_ && use_shm_path_ && qkv_head_size_ / 4 <= 32. Larger heads retain rolled reductions; removing the cap regressed the isolated head-size-512 kernel substantially. The implementation has no shader-model gate.Aion performance
Aion-1.1-Instruct, RTX 5080 / NVIDIA 617.42, verified D3D12 SM 6.10, Agility SDK 721 and compatible DXC
dc3e6c48. ORT main1e4bfba8eeda255d97ffeade88616cd5f06ed777plus this change; GenAI743cc24c739cd19aa953e9d1ea4bc2cd8495b07e. The model uses 28 GroupQueryAttention layers, head size 128 and a 2048-token local attention window. No named adapter was loaded.Median prefill tokens/s over three alternating process rounds, one warmup and five measured repetitions per process. Batch 1, 1024 default input tokens, 32 generated tokens, maximum length 8192, graph capture and generator reuse. Binary-identical off/on variants use a test-only unroll switch and identical fixed-seed input IDs. Timing excludes initialization, profiling and shader logging.
DP4A/Subgroup Matrix describe the model GEMMs; the optimized attention kernel uses shared memory in both modes. Subgroup Matrix support is not required for this optimization. SM denotes the applied DirectX shader target, not CUDA compute capability. Aion was tested on SM 6.10 only; SM 6.8/6.9 and other GPUs were not measured. Vulkan shows no consistent improvement.
Both variants include the same existing Dawn/DXC compatibility prerequisites: the documented HLSL workgroup-bitcast and SPIR-V initialization fixes, compatible DXC and Agility integration. No additional ORT precision patches were included in these timings.
Validation
A separate repeated-token stress fixture exposes a pre-existing Aion accuracy gap (DP4A vs Subgroup Matrix cosine 0.89272; CPU vs Subgroup Matrix 0.98395 and CPU vs DP4A 0.93533). All 36 additional stress output/token pairs are byte-identical with unroll off/on. This gap remains unresolved; diagnostic precision patches were excluded from accepted timings. The validation above establishes preservation and agreement on the measured workload, not arbitrary-prompt model quality.