onnxruntime
9edd1796 - [WebGPU] Optimize WebGPU SparsePagedAttention softmax execution (#32613)

Commit
36 days ago
[WebGPU] Optimize WebGPU SparsePagedAttention softmax execution (#32613) ### Description Optimize the WebGPU `SparsePagedAttention` execution pipeline: - write directly to the activation output for single-source attention without a head sink, avoiding FP32 partial-state allocation and the finalize dispatch - compute online-softmax scalar state once per workgroup and share it across lanes - run finalization with one workgroup per `(token, head)` and share merge coefficients across output dimensions - keep staged joint-softmax finalization for main-plus-auxiliary attention and head-sink configurations - account for the updated workgroup-memory requirements and preserve shader-cache specialization - add focused auxiliary direct-output coverage and update the WebGPU operator documentation This follows #32533, which is now merged into `main`. ### Motivation and Context The original WebGPU implementation always materialized FP32 partial softmax state and launched a separate finalize pass. It also repeated per-head scalar calculations across output lanes. These changes reduce scratch-buffer traffic and dispatch overhead for the common single-source path while reducing redundant work in staged configurations without changing attention semantics. Validation: - WGSL template tests: 89 passed - `clang-format --dry-run --Werror` on the changed C++ files - `git diff --check` - focused code review found no correctness, synchronization, bounds, or numerical issues --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Co-authored-by: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Author
Parents
Loading