Skip to content

[WebGPU] Implement SparseAttentionIndexer operator - #32528

Open
kunal-vaishnavi with Copilot wants to merge 17 commits into
copilot/implement-sparse-attention-indexerfrom
copilot/implement-sparse-attention-indexer-again
Open

[WebGPU] Implement SparseAttentionIndexer operator#32528
kunal-vaishnavi with Copilot wants to merge 17 commits into
copilot/implement-sparse-attention-indexerfrom
copilot/implement-sparse-attention-indexer-again

Conversation

Copilot AI commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

Description

Adds WebGPU support for com.microsoft.SparseAttentionIndexer, covering Qwen4-Exp QSA and DeepSeek V4 CSA policies using the existing provider-neutral schema.

  • GPU execution
    • Implements WGSL programs for state updates, QSA grouping/scoring, CSA overlap compression, rotary embedding, and deterministic TopK selection.
    • Keeps persistent state and selected indices device-resident without host readback.
    • Uses FP32 accumulation for pooling, softmax, RMSNorm, and scoring.
  • Supported configurations
    • Supports policy_mode="qsa" and "csa".
    • Registers FP32 and FP16 batched execution.
    • Preserves explicit QSA key and CSA compressed-window state.
  • Integration
    • Registers the kernel with the WebGPU execution provider.
    • Adds reference-based QSA and CSA coverage, including empty compressed state.
    • Documents the supported subset and optimization follow-ups.

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.

Copilot AI and others added 4 commits September 10, 2026 02:32
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

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

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@kunal-vaishnavi
kunal-vaishnavi changed the base branch from main to copilot/implement-sparse-attention-indexer September 10, 2026 07:19
@kunal-vaishnavi
kunal-vaishnavi marked this pull request as ready for review September 10, 2026 07:19
@kunal-vaishnavi
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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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_ids to its low 32 bits. The existing CsaInt64MaxPosition contract expects INT64_MAX to make every compressed entry eligible, but it becomes -1 here and yields threshold zero, so WebGPU selects nothing. Compute the ceil-division threshold from the two-word int64 storage (also avoiding position + 1 overflow).
      << "  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.

Comment thread onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc Outdated
Comment thread onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc Outdated
Comment thread onnxruntime/test/python/transformers/test_paged_attention.py
Comment thread cmake/CMakeLists.txt
Copilot AI and others added 4 commits September 10, 2026 00:42
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>
Copilot AI and others added 2 commits September 10, 2026 07:49
…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>
@kunal-vaishnavi kunal-vaishnavi changed the title Implement SparseAttentionIndexer for WebGPU [WebGPU] Implement SparseAttentionIndexer operator Sep 10, 2026
Copilot AI and others added 2 commits September 10, 2026 08:46
…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>
Copilot AI and others added 2 commits September 10, 2026 21:13
…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>
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.

3 participants