[WebGPU] Implement SparseAttentionIndexer operator - #32528
Open
kunal-vaishnavi with Copilot wants to merge 17 commits into
Open
[WebGPU] Implement SparseAttentionIndexer operator#32528kunal-vaishnavi with Copilot wants to merge 17 commits into
kunal-vaishnavi with Copilot wants to merge 17 commits into
Conversation
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…on-indexer' into copilot/implement-sparse-attention-indexer-again Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Copilot created this pull request from a session on behalf of
kunal-vaishnavi
September 10, 2026 07:02
View session
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
kunal-vaishnavi
changed the base branch from
main
to
copilot/implement-sparse-attention-indexer
September 10, 2026 07:19
kunal-vaishnavi
marked this pull request as ready for review
September 10, 2026 07:19
kunal-vaishnavi
added this pull request to stack #32530
September 10, 2026 07:20
…on-indexer' into copilot/implement-sparse-attention-indexer-again # Conflicts: # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
WebGPU mishandles explicit zero scales and 64-bit positions, an opt-out test is ineffective, and the PR bundles major undocumented features.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Implements WebGPU SparseAttentionIndexer, while also bundling substantial CUDA, filesystem, and device-discovery changes.
Changes:
- Adds WebGPU QSA/CSA kernels, registration, tests, and documentation.
- Adds packed INT4 PagedAttention/XQA support.
- Adds CUDA device mapping and cross-platform random-access files.
File summaries
| File | Description |
|---|---|
onnxruntime/test/python/transformers/test_paged_attention.py |
Extends XQA scale and dispatch tests. |
onnxruntime/test/providers/cuda/plugin/cuda_plugin_arena_test.cc |
Tests runtime CUDA device discovery. |
onnxruntime/test/providers/cuda/plugin/cuda_device_mapping_test.cc |
Tests CUDA identity mapping. |
onnxruntime/test/platform/env_test.cc |
Tests random-access files. |
onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc |
Adds WebGPU QSA/CSA coverage. |
onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h |
Stores runtime-discovered devices. |
onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc |
Maps and creates CUDA devices. |
onnxruntime/core/providers/cuda/plugin/cuda_device_mapping.h |
Adds device-mapping helpers. |
onnxruntime/core/platform/windows/env.h |
Declares Windows random-access support. |
onnxruntime/core/platform/windows/env.cc |
Implements Windows positional reads. |
onnxruntime/core/platform/posix/env.cc |
Implements POSIX positional reads. |
onnxruntime/core/platform/env.h |
Defines the random-access file API. |
onnxruntime/core/graph/contrib_ops/bert_defs.cc |
Extends PagedAttention for packed INT4. |
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc |
Registers SparseAttentionIndexer. |
onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h |
Declares WebGPU programs and kernel. |
onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc |
Implements QSA/CSA WGSL execution. |
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc |
Registers INT4 PagedAttention kernels. |
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_spec_dec_fp16_int4_256.cu |
Instantiates speculative INT4 XQA. |
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h |
Documents INT4 XQA contracts. |
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu |
Dispatches INT4 XQA kernels. |
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int4_256.cu |
Instantiates decode INT4 XQA. |
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h |
Adds the INT4 quantization type. |
onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh |
Unpacks INT4 cache grains. |
onnxruntime/contrib_ops/cuda/bert/xqa/mha.h |
Defines packed cache heads. |
onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh |
Integrates INT4 scaling and synchronization. |
onnxruntime/contrib_ops/cuda/bert/xqa/int4_cache.cuh |
Adds INT4 dequantization. |
onnxruntime/contrib_ops/cuda/bert/paged_attention.h |
Adds per-channel XQA control. |
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc |
Routes INT4 and per-channel execution. |
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu |
Implements INT4 cache reads and writes. |
onnxruntime/contrib_ops/cuda/bert/attention_data.h |
Adds scale-normalizer scratch state. |
onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h |
Validates packed INT4 caches. |
docs/OperatorKernels.md |
Regenerates kernel type documentation. |
docs/ContribOperators.md |
Documents packed cache shapes and types. |
docs/contrib_ops/webgpu/sparse_attention_indexer.md |
Documents WebGPU support. |
docs/contrib_ops/cuda/sparse_attention_indexer.md |
References the WebGPU implementation. |
docs/contrib_ops/cuda/paged_attention.md |
Documents INT4 PagedAttention. |
docs/contrib_ops/cuda/gqa.md |
Updates the INT4 build default. |
cmake/CMakeLists.txt |
Enables INT4 KV cache by default. |
.github/workflows/android.yml |
Raises the Android size threshold. |
Review details
Suppressed comments (1)
onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc:467
- This second read also truncates
position_idsto its low 32 bits. The existingCsaInt64MaxPositioncontract expectsINT64_MAXto make every compressed entry eligible, but it becomes-1here and yields threshold zero, so WebGPU selects nothing. Compute the ceil-division threshold from the two-word int64 storage (also avoidingposition + 1overflow).
<< " let position = " << position_ids.GetByOffset("row") << ";\n"
<< " let threshold = select(0u, u32(position + 1) / uniforms.compress_ratio, position >= 0);\n"
- Files reviewed: 40/40 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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>
…_infer.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…on-indexer' into copilot/implement-sparse-attention-indexer-again # Conflicts: # docs/contrib_ops/cuda/sparse_attention_indexer.md # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc # onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py 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>
…on-indexer' into copilot/implement-sparse-attention-indexer-again Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…on-indexer' into copilot/implement-sparse-attention-indexer-again # Conflicts: # docs/contrib_ops/cuda/sparse_attention_indexer.md # onnxruntime/test/contrib_ops/sparse_attention_indexer_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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds WebGPU support for
com.microsoft.SparseAttentionIndexer, covering Qwen4-Exp QSA and DeepSeek V4 CSA policies using the existing provider-neutral schema.policy_mode="qsa"and"csa".Motivation and Context
Sparse attention models need policy-specific candidate selection before downstream attention execution. This change enables QSA and CSA index generation on WebGPU while preserving the shared operator ABI, deterministic selection semantics, and graph-visible state.