[Sayan Shaw (sayanshaw24)]

Summary

Follow-up to #32528 (the initial WebGPU SparseAttentionIndexer implementation). This PR reduces the QSA recomputation called out in its follow-up list. The sparse-indexer prerequisite #32656 has merged into main; this PR now targets main directly. The optimization changes only the WebGPU QSA selection path.

• Cache visible token positions and block scores per query in bounded workgroup memory. Reuse query/key RMS normalization within each block score.
• Reuse cached positions on visibility overflow, scanning only the uncached tail; retain exact visibility and score recomputation fallbacks outside the cache bounds, preserving score-descending/index-ascending TopK semantics and non-prefix masks.
• Bound both caches plus the count to 15,364 bytes, below WebGPU's guaranteed 16 KiB workgroup-storage limit.
• Add FP32/FP16 non-prefix-mask parity tests and document the optimized/fallback paths.

Validation

• The original Qwen integration run passed 18 focused SparseAttentionIndexerShapeInferenceTest cases. On current main, a Linux Release WebGPU provider/provider-test build and all 23 focused QSA/CSA and packed-indexer cases passed on NVIDIA H200/Vulkan, including both bounded-cache fallback thresholds. Focused formatting and git diff --check passed.
• End-to-end NVIDIA H200 Vulkan testing used a non-paged, no-MTP Qwen3.8 raw INT4 RTN/INT8-last-matmul Generator export, with 16-token prefill chunks:

Context	Before	After
16-token prefill	43.27 s	1.51 s
Next decode step after 16 tokens	42.63 s	0.70 s
1K prefill	Did not complete within 181 s of prefill	68.33 s (14.99 tok/s)
1K decode	No valid rate	1.29 tok/s over eight tokens
2K prefill	No prior completed result	145.37 s (14.09 tok/s)
2K decode	About 30.8 s/step at the old cache boundary	1.03 tok/s over eight tokens

The eight generated pieces from the four-token smoke prompt were unchanged, as was the first token ID from the 16-token probe. Model-load time (~58–60 s) is excluded from prefill measurements. The H200 static WebGPU build now executes the WebGPU-specific parity tests; all 23 focused QSA/CSA and packed-indexer cases passed on current main.

The EngramGate and GenAI builder/runtime fixes in the working tree are not part of this PR.

The performance table above was measured on the earlier Qwen3.8 integration branch; GPU regression tests, not full-model timing, were rerun on current main.

Current main update (October 6)

Kunal's #32656 merged, and this PR is now based on main with a three-file QSA-only review diff. Merged current main again after that restack. The WebGPU provider-test build and all 23 focused H200 QSA/CSA plus packed-indexer tests passed. The October 6 Windows WebGPU job exposed a long-running visibility-overflow boundary test that lost the D3D12 device. Commit 9499988 avoids redundant prefix rescans; all 22 focused WebGPU QSA/CSA tests pass on H200/Vulkan, with that boundary test taking 241 ms. Android minimal-size CI passed on the preceding head; CI for the fix is pending.

[Akshay Sonawane (apsonawane)]

PR looks good, Sayan Shaw (@sayanshaw24) or kunal-vaishnavi if anyone could fix the merge conflict I can approve it again

[Sayan Shaw (sayanshaw24)]

Akshay Sonawane (@apsonawane) Thanks for flagging this. I rebased this one-commit QSA optimization onto the latest #32656 head (0d0d060), preserving its newer packed-indexer and serial CSA changes while keeping the QSA cache path parallel. Pushed b12610e; the static WebGPU provider rebuilt and 20 focused H200 QSA/CSA tests passed (two CUDA-only tests skipped). The branch is now directly based on the updated target; GitHub is recalculating mergeability.

[Copilot]

Copilot review overview

🟡 Changes recommended

The newly introduced out-of-cache fallback branches lack WebGPU test coverage.

Review effort: Balanced
Findings: 1

Open (1)

• Add WebGPU tests covering score and visibility cache fallbacks · New

What changed in this PR

Optimizes WebGPU QSA selection by caching visibility data and scores in workgroup memory while retaining fallback paths.

Changes:

• Adds bounded workgroup caches and shared RMS computations.
• Adds FP32/FP16 non-prefix-mask parity tests.
• Documents cache limits and fallback behavior.

File	Description
onnxruntime/​contrib_ops/​webgpu/​bert/​sparse_attention_indexer.cc	Implements cached parallel QSA scoring.
onnxruntime/​test/​contrib_ops/​sparse_attention_indexer_op_test.cc	Adds cached-path parity tests.
docs/​contrib_ops/​webgpu/​sparse_attention_indexer.md	Documents caching and fallbacks.

---

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