[CUDA][WebGPU] Add PackedSparseAttentionIndexerMerge for MTP index sharing by kunal-vaishnavi · Pull Request #33224 · microsoft/onnxruntime

kunal-vaishnavi · 2026-10-09T01:42:26Z

Description

Introduce PackedSparseAttentionIndexerMerge and extend PackedSparseAttentionIndexer to support index sharing, configurable output capacity, and explicit state-sharing policies.

Add the merge operator for combining base and additional sparse-attention indices, with CUDA and WebGPU support and dedicated documentation.
Extend input/output slot definitions, operator schemas, and kernel registration for shared indices and merged indices, counts, and status, including in-place merged-tensor operations.
Update the CUDA indexer to handle indexshare_mode, validate interface consistency, and manage merged outputs.
Add max_output_entries to control output capacity consistently across output shapes and QSA/CSA kernel parameters.
Add indexshare_state_policy to control CSA state advancement and sharing, with validation and error handling for incompatible interface usage.

Motivation and Context

Advanced sparse-attention workflows need to reuse selected indices and combine them with additional selections rather than generate each index set independently. These changes provide explicit index-sharing and merging interfaces to support those workflows.

Configurable output capacity lets callers bound selection outputs, while an explicit CSA state policy distinguishes state advancement from sharing. Together, these controls make index reuse more flexible and state handling more predictable in large-scale transformer models.

Extend the indexer with refresh/reuse bundles, padded capacity, and QSA hold/CSA advance policies. Rename the merge operator to PackedSparseAttentionIndexerMerge without legacy aliases and cover validation, overflow, and CUDA/WebGPU behavior.

kunal-vaishnavi · 2026-10-09T22:27:05Z

Pushed a1017e9 to address the CI failures:

Made the shared capacity reader accept the actual kernel-info interface. This fixes the core/plugin adapter mismatch in CUDA, WebGPU plugin, and WASM compilation.
Renamed the QSA reuse norm local to avoid MSVC C4456 under warnings-as-errors, including the Windows documentation build.
Declared Type metadata dependencies for every WebGPU merge binding that uses value-type aliases. This addresses the seven Mac WebGPU merge tests failing shader validation.
Updated the two generated documentation files narrowly for the IndexShare interface and merge operator. The merge section matches repository-generator output; both CUDA kernel signatures match registered schemas.

The WASM Debug external-initializer failures came from the target merge rather than this PR head and are already fixed on main by #33226 (2297e77). Fresh target-merge CI should include that fix; no unrelated skips were added here.

Local validation passed: 22 capacity/policy cases with both kernel-info interfaces; actual CUDA wrappers in normal and plugin modes; source-local shadow warnings as errors; actual WebGPU merge wrapper in normal and plugin-adapter modes; generated-doc/signature checks; changed-range clang-format; and git diff --check.

No full ORT rebuild or replacement of the installed runtime was performed. Fresh provider/runtime tests, including Mac Metal execution, remain for new-head CI. No stale-head workflows were rerun.

Tianlei Wu (tianleiwu)

Review Summary

This PR adds com.microsoft.PackedSparseAttentionIndexerMerge (append_range on CUDA and WebGPU, and append_indices on CUDA) and extends com.microsoft.PackedSparseAttentionIndexer on CUDA with the IndexShare refresh/reuse interface (indexshare_mode, merged_indices, merged_counts, merged_status, max_output_entries, indexshare_state_policy, and indexshare_status).

Overall the design and test coverage are well structured. Please take a look at a few correctness and resource-sizing issues identified inline:

Hash-table capacity sizing and 100% load factor in PackedSparseAttentionIndexerMerge (packed_sparse_attention_indexer_merge_common.h): hash_capacity is sized from output_capacity (max_output_entries, up to 2^30) rather than the candidate insertion bound, and allows a 1.0 load factor with zero -1 sentinel slots when output_capacity is a power of two.
Signed 32-bit int overflow in MergeIndicesKernel (packed_sparse_attention_indexer_merge_impl.cu): Loop induction variables (column, slot, chunk) use 32-bit int with stride kThreads (256) even though total can reach INT32_MAX.
Host and device validation gaps in PackedSparseAttentionIndexer reuse mode (packed_sparse_attention_indexer.cc, packed_sparse_attention_indexer_impl.cu): ComputeReuse bypasses several policy-specific input/output checks from ComputeInternal, copies merged_* before ComputeCsa(context, true) validates CSA inputs, and ValidatePackedReuseKernel does not zero/fill rows when merged_status[row] != 0 or merged_counts[row] is out of range.
Shape-inference and schema doc updates (bert_defs.cc): PackedSparseAttentionIndexerTypeAndShapeInference and PackedSparseAttentionIndexerMergeSchema omit several attribute/input-shape checks, and the merge schema docstrings still describe additional_indices / additional_counts as reserved.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

🟡 Changes recommended

Accepted large dimensions can overflow CUDA traversal counters and cause out-of-bounds accesses.

2 open findings

What changed in this PR

Adds sparse-attention index merging and CUDA index reuse for multi-token prediction workflows.

Changes:

Adds CUDA and WebGPU merge kernels.
Extends the CUDA indexer with reuse modes, padded output capacity, and state-sharing policies.
Adds schemas, registrations, tests, and documentation.

File	Description
onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc	Tests refresh, reuse, and state advancement.
onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc	Tests merging and error handling.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the merge schema.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines merge and index-sharing interfaces.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers the WebGPU merge kernel.
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc	Rejects unsupported index sharing.
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h	Declares WebGPU merge classes.
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc	Implements WebGPU range merging.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h	Declares reuse support.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc	Implements reuse and state policies.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc	Validates and dispatches CUDA merges.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h	Declares merge launch functions.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu	Implements CUDA merge kernels.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h	Extends indexer parameters.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu	Supports padded outputs and reuse validation.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers the CUDA merge kernel.
onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h	Shares merge validation helpers.
onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h	Defines index-sharing slots.
docs/OperatorKernels.md	Updates kernel interfaces.
docs/ContribOperators.md	Documents schema additions.
docs/contrib_ops/PackedSparseAttentionIndexerMerge.md	Explains merge semantics and provider support.

🧠 Review effort: Balanced

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

Bound hash scratch by candidate width at a safe load factor and use wide hash/iteration arithmetic. Validate policy-specific reuse interfaces before copying outputs, sanitize invalid selections and metadata, and prevent rejected CSA rows from advancing state. Complete merge and IndexShare shape checks, document append_indices, and add regressions for validation and duplicate-heavy chunks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…ilures Replace capacity-specialized single-invocation hash/scan shaders with parallel first-occurrence checks and stable chunk compaction. Preserve validation precedence and failure output clearing. Cover shared-row chunk boundaries, late invalid prefixes, and large output capacity, and update implementation documentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kunal-vaishnavi · 2026-10-16T00:36:41Z

Pushed CI repair in 77897a1 (normal append-only push).

The Windows WebGPU failure starts in PackedSparseAttentionIndexerMerge: StableOverlapAndSharedRows reports DXC E_FAIL, and QwenCapacity then loses the device; later WebGPU tests cascade after that device loss. The old implementation ran one invocation per query and switched to a quadratic serial output-prefix scan above the workgroup-hash limit.

The merge now uses a 256-invocation cooperative workgroup with parallel first-occurrence checks and ordered chunk compaction, removing the capacity-specialized serial hash/scan shaders. Stable order, metadata-error precedence, range overlap, padding, and failure clearing are preserved. Added shared-row/chunk-tail and late-invalid-prefix regressions and updated the operator implementation documentation.

Validation:

214 actual generated-shader oracle cases passed through Dawn on software Vulkan, covering Qwen/4097-entry cases, duplicates, overflow, invalid metadata, shared rows, chunk boundaries, multi-chunk ranges, and INT32 range endpoints.
Generated HLSL using the exact Tint revision pinned by ORT; compilation passes DXC 1.8 and 1.9 in HLSL 2018/2021 modes.
Both actual changed C++ translation units pass warnings-as-errors syntax checks; pinned clang-format and git diff --check pass.
Qwen-sized dispatch/readback on software Vulkan improved from approximately 76 ms to 8 ms. This is a local software-adapter comparison, not a Windows GPU performance claim.

Fresh CI is running for the new head. Windows D3D12 execution and the full ORT provider suite remain for CI; the original Windows-only DXC failure was not reproduced on Linux, so I am not claiming those jobs are already green.

## Description Stacked on #33224 (base `kvaishnavi/index-share`; review only the top commit). Adds an optional `state_status` output to `PackedSparseAttentionIndexer` that reports, per request, whether this call's state update was applied or rejected, and why. Today a rejected step (inconsistent cursor or `state_capacity` overflow) is a safe no-op on state with an empty selection, but the caller cannot tell it happened: the rejection flags live in a scratch buffer and `indexshare_status` is only produced on the IndexShare interface. A CUDA-graph-replayed decode/verify loop cannot synchronize inside the operator, so it needs a device tensor it can read on the host after replay. ## Summary of Changes ### Schema and docs | File | Change | |------|--------| | `bert_defs.cc` | New optional output 8 `state_status` (int32, shape `(batch_size)`); declared output range is now 6-9; type/shape inference; one doc sentence | | `packed_sparse_attention_indexer_common.h` | `kStateStatus = 8`, `kOutputCount = 9`, `StateStatusCode` {0 ok, 1 invalid device metadata, 2 capacity overflow} | | `docs/ContribOperators.md`, `docs/contrib_ops/packed_sparse_attention_indexer.md` | Output and status contract | ### CUDA / WebGPU | File | Change | |------|--------| | `packed_sparse_attention_indexer_impl.cu/.h` | The QSA and CSA update kernels write the code (1 or 2) instead of the boolean `1` into the existing `overflow_flags`; all readers still test `!= 0`, so selection behaviour is unchanged | | `packed_sparse_attention_indexer.cc` | When output 8 is requested it is used as the `overflow_flags` buffer (no extra kernel, no copy, no host sync). On the IndexShare reuse path that holds state (QSA, CSA `hold`) it is zero-filled with `cudaMemsetAsync` | | `webgpu/.../packed_sparse_attention_indexer.cc` | Fails if `state_status` is requested (CUDA only, like IndexShare) | ### Semantics - A non-zero request leaves its state unchanged and emits an empty selection (`selected_counts == 0`, indices `-1`), exactly as before. - Codes follow `PackedSparseAttentionIndexerMerge` (1 invalid device metadata, 2 capacity overflow). - `state_status` is independent of `indexshare_status` (per-query selection status, forwarded on reuse and zero on refresh). On CSA reuse with `advance`, both are set when a row is rejected. - Default behaviour is unchanged: the output is optional and absent in existing models. ## Testing Built with CUDA 13.0, sm_90, in-tree CUDA EP, and run on one H200: ``` CUDA_VISIBLE_DEVICES=5 ./onnxruntime_provider_test \ --gtest_filter='*PackedSparseAttentionIndexer*:*DynamicSparseAttention*:*SparseAttentionIndexer*:*SparsePagedAttention*' ``` 177 passed (17 WebGPU tests skipped, no WebGPU in this build); `PackedSparseAttentionIndexerTest.*` repeated 20 times, no flake. New tests in `packed_sparse_attention_indexer_op_test.cc` (`Qsa/CsaStateStatus*`, `IndexShareCsaAdvanceReportsStateStatus`, `QsaTwoRowStepEqualsChainedSingleRowSteps`): - status values for a healthy batch, capacity overflow, inconsistent cursor, malformed state lengths, and mixed batches `{ok, invalid, overflow}`, for QSA and CSA, checked against the host reference (which now also computes the codes); - requesting `state_status` leaves selections and state bit-identical to not requesting it; - IndexShare CSA reuse with state overflow reports `indexshare_status = 1` and `state_status = 2`; - a two-row QSA step (one request, two tokens, `compress_ratio` 2 and 4, all three cursor phases) is bit-identical in selections, block state and pending entries to two chained one-row steps. This passes without any kernel change, i.e. the multi-row verify step is already supported by `main`; the test pins that contract. Mutation checks: with the output plumbing removed 7 new tests fail; with code 2 reported as 1, 4 fail; existing tests pass in both. Not covered: WebGPU compile (not in this build), plugin-EP build, an end-to-end CUDA-graph capture/replay (the kernels are launch-only and the reuse-path zero fill is `cudaMemsetAsync`, so nothing in the path synchronizes). ## Motivation and Context Needed by decode/verify loops that capture the indexer in a CUDA graph and drive it with device-resident lengths (as in the multi-token verify path of #33328). Reuses the status convention introduced in #33224 so callers see one scheme. Touches different hunks from #33328 (scoring kernels); a trivial merge is expected for the test file. ## Checklist - [x] Tests added/updated - [x] No breaking changes (new optional output; the declared output range grows from 6-8 to 6-9) - [x] Documentation updated --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kunal-vaishnavi added 2 commits October 7, 2026 10:43

Add stable sparse selection merge for CUDA and WebGPU

22e70ae

kunal-vaishnavi changed the title ~~[CUDA] Add PackedSparseAttentionIndexerMerge for MTP index sharing~~ [CUDA][WebGPU] Add PackedSparseAttentionIndexerMerge for MTP index sharing Oct 9, 2026

Fix IndexShare plugin builds and WebGPU metadata validation

a1017e9

Tianlei Wu (tianleiwu) mentioned this pull request Oct 15, 2026

[CUDA] Add state_status output to PackedSparseAttentionIndexer #33383

Merged

3 tasks done

Tianlei Wu (tianleiwu) reviewed Oct 15, 2026

View reviewed changes

kunal-vaishnavi marked this pull request as ready for review October 15, 2026 22:39

Copilot AI balanced review requested due to automatic review settings October 15, 2026 22:39

Copilot started reviewing on behalf of kunal-vaishnavi October 15, 2026 22:40

View session

Merge current main into IndexShare PR

0568f6c

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI reviewed Oct 15, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu Outdated

Comment thread onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu Outdated

Copilot AI and others added 2 commits October 15, 2026 22:59

Merge latest main into IndexShare PR

9736482

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kunal-vaishnavi added this pull request to stack #33406 October 16, 2026 00:20

Tianlei Wu (tianleiwu) approved these changes Oct 16, 2026

View reviewed changes

Tianlei Wu (tianleiwu) merged commit f33c174 into main Oct 16, 2026
98 checks passed

Tianlei Wu (tianleiwu) deleted the kvaishnavi/index-share branch October 16, 2026 05:40

Navigation Menu

Repository navigation

[CUDA][WebGPU] Add PackedSparseAttentionIndexerMerge for MTP index sharing - #33224

[CUDA][WebGPU] Add PackedSparseAttentionIndexerMerge for MTP index sharing #33224
Tianlei Wu (tianleiwu) merged 7 commits into
mainfrom
kvaishnavi/index-share

kunal-vaishnavi commented Oct 9, 2026

Uh oh!

kunal-vaishnavi commented Oct 9, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

kunal-vaishnavi commented Oct 16, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

kunal-vaishnavi commented Oct 9, 2026

Description

Motivation and Context

Uh oh!

kunal-vaishnavi commented Oct 9, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Review Summary

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

kunal-vaishnavi commented Oct 16, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants