[WebGPU] Add DynamicSparseAttention

[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.