Skip to content

[WebGPU] Optimize WebGPU SparsePagedAttention softmax execution - #32613

Merged
kunal-vaishnavi merged 83 commits into
mainfrom
copilot/copilotcopilotadd-cuda-sparsepagedattention-s2-non
Oct 3, 2026
Merged

kunal-vaishnavi merged 83 commits into
mainfrom
copilot/copilotcopilotadd-cuda-sparsepagedattention-s2-non

Conversation

Copilot AI commented Sep 15, 2026 •

Copy link
Copy Markdown
Contributor

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

Copilot AI and others added 30 commits September 10, 2026 06:13
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ntion' into copilot/copilotadd-cuda-sparsepagedattention

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ntion' into copilot/copilotadd-cuda-sparsepagedattention

# Conflicts:
#	onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc
#	onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h
#	onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu
#	onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ntion' into copilot/copilotadd-cuda-sparsepagedattention

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ntion' into copilot/copilotadd-cuda-sparsepagedattention

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…ntion' into copilot/copilotadd-cuda-sparsepagedattention

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Copilot AI and others added 2 commits October 1, 2026 07:41
…gedattention' into copilot/copilotcopilotadd-cuda-sparsepagedattention-s2-non

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@kunal-vaishnavi kunal-vaishnavi changed the title Optimize single-source WebGPU SparsePagedAttention Optimize WebGPU SparsePagedAttention softmax execution Oct 1, 2026
@kunal-vaishnavi kunal-vaishnavi changed the title Optimize WebGPU SparsePagedAttention softmax execution [WebGPU] Optimize WebGPU SparsePagedAttention softmax execution Oct 1, 2026
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…gedattention' into copilot/copilotcopilotadd-cuda-sparsepagedattention-s2-non

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Base automatically changed from copilot/copilotadd-cuda-sparsepagedattention to main October 3, 2026 01:02
…tadd-cuda-sparsepagedattention-s2-non

# Conflicts:
#	docs/ContribOperators.md
#	docs/OperatorKernels.md
#	docs/contrib_ops/webgpu/sparse_paged_attention.md
#	onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc
#	onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu
#	onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention.cc
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention.h
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_auxiliary.wgsl.template
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_finalize.wgsl.template
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_main.wgsl.template
#	onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_token_meta.wgsl.template
#	onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc
#	onnxruntime/core/graph/contrib_ops/ms_opset.h
#	onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@kunal-vaishnavi
kunal-vaishnavi merged commit 9edd179 into main Oct 3, 2026
@kunal-vaishnavi
kunal-vaishnavi deleted the copilot/copilotcopilotadd-cuda-sparsepagedattention-s2-non branch October 3, 2026 18:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants