[Copilot]

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.

[azure-pipelines]

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

[Copilot]

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