[Copilot]

Description

Adds WebGPU execution for com.microsoft.DynamicSparseAttention, covering Qwen4-Exp selected-token attention and DeepSeek V4 local-plus-selected attention without host readback.

• Attention
  • Selected main-cache and local-plus-auxiliary modes
  • Joint FP32 stable softmax with optional sink
  • Tiled QK reductions and guarded int32 indices
• QKV and cache
  • Packed and unpacked QKV
  • GQA head mapping
  • Alias-safe contiguous cache updates
  • Shared auxiliary K/V storage
• Position processing
  • QK RMSNorm
  • Partial and interleaved RoPE
  • Rotary offsets and explicit position IDs
• Coverage
  • FP32 and FP16
  • Cache append, sparse/empty selections, padding, invalid-index guards, hybrid attention, sinks, QK norm, and RoPE

Motivation and Context

Enables model-neutral sparse-attention execution on WebGPU while keeping selection metadata processing and attention computation device-resident.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Dispatch padding can cause out-of-bounds access, and several loops defeat sparse-attention performance.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds WebGPU support for com.microsoft.DynamicSparseAttention, keeping sparse-attention processing device-resident.

Changes:

• Implements query preparation, cache updates, RoPE, normalization, and sparse attention shaders.
• Registers the WebGPU kernel.
• Adds FP16/FP32 coverage for key attention modes and features.

File summaries

File	Description
dynamic_sparse_attention.cc	Implements WebGPU execution.
dynamic_sparse_attention.h	Defines kernel and shader programs.
webgpu_contrib_kernels.cc	Registers the kernel.
dynamic_sparse_attention_op_test.cc	Adds WebGPU tests.

Review details

Suppressed comments (2)

onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc:388

• The loop always executes max_selected full QK reductions even when selected_count is much smaller; candidate_valid is checked only after the reduction. A common row with one selected token and width 4096 consequently does 4096 dot products. Clamp the device-provided count for safety and use it as the loop bound.

         << "  for (var i = 0u; i < uniforms.max_selected; i++) {\n"

onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc:240

• As in query preparation, every key-channel invocation repeats the identical full-head RMS reduction, producing O(head_size²) key loads per appended KV head. This is especially costly for token decode with few selected entries. Reduce once per (batch, sequence, KV head) workgroup and share the inverse RMS across channel writes.

         << "  for (var c = 0u; c < uniforms.head_size; c++) {\n"
         << "    let kv = f32("
         << (packed_qkv_ ? query->GetByOffset("key_base + c") : key->GetByOffset("key_base + c"))
         << ");\n"
         << "    k_sumsq += kv * kv;\n"

• Files reviewed: 4/4 changed files
• Comments generated: 4
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

WebGPU shader-validation failures, incorrect cache shape inference, and an aborting no-exception test block approval.

Review effort: Balanced
Findings: 2 · 1

Open (3)

• Broadcast WGSL return condition before barriers · New
• Make WGSL barrier loop bounds workgroup-uniform · New
• Validate full-width int64 position IDs · New

Resolved since last review (4)

• Oversized 1-D dispatches are normalized into padded 2-D/3-D grids, but this shader treats every…
• This iterates the configured maximum window and performs a complete QK reduction even after `index…
• Each channel invocation recomputes the same full-head Q RMS sum, making normalization O(head_size²)…
• The new alias-dependent path is not explicitly exercised by the WebGPU tests: the generic OpTester…

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

Review effort: Lite
Findings: 2 · 2

Open (4)

• When head_sink != nullptr, the generated shader path is identical regardless of… · New
• When head_sink != nullptr, the generated shader path is identical regardless of… · New
• kAttentionValuesPerInvocation assumes kMaxHeadSize is an exact multiple of… · New
• The same #if defined(USE_CUDA) || defined(USE_WEBGPU) guard is repeated for two adjacent include… · New

Resolved since last review (3)

• Make WGSL barrier loop bounds workgroup-uniform
• Broadcast WGSL return condition before barriers
• Validate full-width int64 position IDs