Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
3e309f4
Implement SparseAttentionIndexer CUDA operator
Copilot Sep 10, 2026
d0a8e7a
Harden SparseAttentionIndexer validation and tests
Copilot Sep 10, 2026
cb58b7e
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 10, 2026
9acfe23
Implement SparseAttentionIndexer for WebGPU
Copilot Sep 10, 2026
7adf139
Address SparseAttentionIndexer review feedback
Copilot Sep 10, 2026
f89ea1e
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 10, 2026
d96d08f
Implement SparseAttentionIndexer CUDA operator
Copilot Sep 10, 2026
b5e7985
Harden SparseAttentionIndexer validation and tests
Copilot Sep 10, 2026
213b731
Address SparseAttentionIndexer review feedback
Copilot Sep 10, 2026
f330287
Update onnxruntime/test/python/onnxruntime_test_python_symbolic_shape…
kunal-vaishnavi Sep 10, 2026
765eadc
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 10, 2026
f74c551
Fix WebGPU SparseAttentionIndexer review issues
Copilot Sep 10, 2026
c24495d
Fix SparseAttentionIndexer CI failures
Copilot Sep 10, 2026
38c341c
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 10, 2026
e3ea06a
Fix SparseAttentionIndexer CI failures
Copilot Sep 10, 2026
b028e65
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 10, 2026
59bc669
Gate FP8 XQA test on supported GPUs
Copilot Sep 10, 2026
4df139c
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 15, 2026
2652363
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 16, 2026
3d328b2
Adapt WebGPU indexer to shared cache ABI
Copilot Sep 16, 2026
6a13c37
Merge branch 'stack-32526' into stack-32528
kunal-vaishnavi Sep 18, 2026
fe9d059
Fuse query normalization in WebGPU sparse indexer
kunal-vaishnavi Sep 18, 2026
72d02b1
Merge branch 'stack-32526' into stack-32528
kunal-vaishnavi Sep 18, 2026
7d485c8
Merge branch 'stack-32526' into stack-32528
kunal-vaishnavi Sep 18, 2026
bc922c2
Use integer masks in WebGPU sparse indexer
kunal-vaishnavi Sep 18, 2026
75d6b8a
Merge branch 'stack-32526' into stack-32528
kunal-vaishnavi Sep 19, 2026
dbf7d34
Merge branch 'stack-32526' into stack-32528
kunal-vaishnavi Sep 20, 2026
01dc6e1
Support packed QK in WebGPU sparse attention indexer
kunal-vaishnavi Sep 20, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions docs/contrib_ops/cuda/sparse_attention_indexer.md
Original file line number Diff line number Diff line change
Expand Up @@ -115,9 +115,10 @@ follow from `Lb`, `S` and `r` alone (see [§5](#5-policy-csa)).
| `I` | `tensor(int64)` |
| `M` | `tensor(int32)` |

Only the CUDA execution provider registers a kernel. There is no CPU kernel; the header under
`contrib_ops/cpu/sparse/` only holds the CUDA-free constants that the schema, the kernel and the
tests must agree on.
CUDA registers all three `T` types. WebGPU registers `float` and `float16`; see the
[WebGPU implementation notes](../webgpu/sparse_attention_indexer.md). There is no CPU kernel; the
header under `contrib_ops/cpu/sparse/` only holds the provider-neutral constants that the schema,
kernels and tests must agree on.

### Output slot discipline

Expand Down
49 changes: 49 additions & 0 deletions docs/contrib_ops/webgpu/sparse_attention_indexer.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
# SparseAttentionIndexer on WebGPU

The WebGPU execution provider implements version 1 of
`com.microsoft.SparseAttentionIndexer` for the `qsa` and `csa` policies. It uses
the provider-neutral schema and state ABI described in the
[operator documentation](../cuda/sparse_attention_indexer.md).

## Supported subset

- batched inputs;
- `qsa` and `csa` policy modes;
- `float32` and `float16`;
- explicit graph-visible key and packed projection-buffer state;
- rank-2 INT64 QSA padding masks with internally derived causal visibility;
- deterministic score-descending, index-ascending TopK ties.

BF16, packed/variable-length inputs, and fixed-capacity caches using
`past_sequence_length` are not supported by the WebGPU kernel. Unknown policies
and policy-incompatible inputs or attributes are rejected.

## Execution

State concatenation, visible-token grouping, QSA pooling, CSA overlap
compression, query/key RMS normalization, rotary embedding, scoring, selection, and
output padding execute in WGSL. The implementation does not map GPU buffers,
read selected values back to the host, or retain state in the kernel object.
All reductions and softmax calculations accumulate in FP32, including for
FP16 inputs.

The rank-3 query projection is logically reshaped into heads inside the shader. Query and key
projections therefore both connect directly to the operator; their consecutive norm-weight inputs
are applied internally before rotary embedding.

The initial implementation prioritizes correctness and uses one independently
writable workgroup per query or completed CSA window. Candidate scoring during
selection is recomputed rather than materialized, avoiding candidate-count
limits and GPU-to-CPU synchronization at the cost of additional computation.

## Follow-up work

- packed/variable-length input;
- fixed-capacity cache updates;
- specialized large-candidate TopK;
- subgroup-optimized reductions;
- fused projection, pooling, and scoring;
- reduced recomputation and temporary-buffer use;
- selector/executor fusion;
- WebGPU `DynamicSparseAttention` and `SparsePagedAttention`;
- additional element types.
Loading
Loading