[Copilot]

Description

Adds com.microsoft.PackedSparseAttentionIndexer version 1 for packed, variable-length continuous batching, with CUDA and WebGPU implementations.

The operator supports both QSA and CSA policies over flattened token inputs and per-request metadata while preserving fixed-capacity state across prefill and decode calls. It is designed to produce selection metadata compatible with SparsePagedAttention without changing the existing dense SparseAttentionIndexer ABI or behavior.

Key changes:

• Add the shared operator schema, fixed generic QSA/CSA state slots, strict policy-specific validation, and fixed output shape inference.
• Add CUDA and WebGPU kernels and provider registrations for FP32, FP16, and BF16 where supported.
• Add device-side packed metadata validation without host synchronization in normal provider paths.
• Add deterministic TopK selection and complete per-request overflow handling that preserves prior state and returns empty selections.
• Support safe state aliasing where provider semantics permit it while avoiding unsafe CSA buffer aliasing.
• Share policy constants and math helpers with the existing sparse attention indexer where practical.
• Add symbolic shape inference support and focused Python coverage.
• Add C++ shape, QSA/CSA numerical, prefill/decode, zero-token request, malformed metadata, deterministic ordering, and overflow tests.
• Add operator and WebGPU implementation documentation and regenerate contrib operator/kernel documentation.

Validation performed:

• onnxruntime_provider_test CPU build completed successfully.
• All 12 PackedSparseAttentionIndexerShapeInferenceTest tests passed.
• Symbolic shape-inference tests passed.
• Repository linting, formatting checks, and secret scanning passed.
• CUDA/WebGPU numerical tests are included but could not be executed locally because those providers and compatible hardware/toolchains are unavailable in the environment.

Motivation and Context

The existing SparseAttentionIndexer contract uses dense [B, S, ...] inputs, dense masking, and growing contiguous state. OgaEngine-compatible paged-attention models instead use flattened packed tokens, cumulative_sequence_lengths, per-request past lengths, dynamic request scheduling, and fixed-address state.

This separate operator bridges that contract mismatch for sparse paged attention while keeping the existing dense operator stable. QSA selections reference request-local positions in the main paged cache; CSA selections reference fixed-capacity auxiliary compressed state suitable for SparsePagedAttention local-plus-selected execution.

Built on the sparse attention indexer work from #32528.

[Akshay Sonawane (apsonawane)]

I found four issues at the current head:

1. CUDA CSA state can be corrupted when buffers alias. The compaction copy runs across threads, so one thread can overwrite a source slot before another reads it. With compress_ratio=64, a prior buffer length of 65, and 64 new tokens, this can corrupt the pending state consumed by the next decode. See packed_sparse_attention_indexer_impl.cu:593–609. Please make compaction overlap-safe.

2. WebGPU shaders exceed the portable storage-buffer limit. The QSA update uses nine storage bindings; the CSA update and selection shaders use ten. WebGPU guarantees only eight per shader stage, so these kernels cannot run on some valid adapters. See packed_sparse_attention_indexer.cc:872 and :1047. Please reduce the per-dispatch binding count or gate support on the adapter limit, accounting for tensors split across multiple bindings.

3. All-empty packed batches can fail WebGPU validation. The update runs before the empty-token return and binds a zero-sized key or key_gate buffer. A valid batch with cumulative_sequence_lengths=[0,0] can therefore fail rather than return empty selections with preserved state. See packed_sparse_attention_indexer.cc:868 and :1043. Please handle this case before binding zero-sized token tensors, without skipping required state handling.

4. The schema does not enforce policy-specific output slot 4. Shape inference checks for six declared outputs but does not check that QSA leaves present_gate_buffer unnamed and CSA names it. See bert_defs.cc:2747. This was raised in an earlier review thread, but the guard is absent at the current head. Please restore the check and its regression tests.

[Akshay Sonawane (apsonawane)]

Thanks for adding an explicit storage-binding limit check. It improves the error, but does not make the WebGPU kernel usable on adapters exposing the spec-guaranteed limit of eight bindings: QSA update needs nine, and CSA update/selection need ten.

Because the node is already assigned to WebGPU, returning NOT_IMPLEMENTED during execution is not an execution-provider fallback. Could we reduce the per-dispatch binding count to the portable limit, or prevent WebGPU assignment on adapters that cannot run this operator?

[Copilot]

I found four issues at the current head:

1. CUDA CSA state can be corrupted when buffers alias. The compaction copy runs across threads, so one thread can overwrite a source slot before another...

Addressed all four findings and synced with current main (a84958ac):

1. CUDA CSA compaction now processes slots in ascending order with a block barrier between slots, making overlapping past/present buffers forward-memmove safe. Added the reported large-ratio compaction case.
2. WebGPU dispatches now pack related tensors so QSA/CSA update and selection each use at most the portable eight logical storage bindings; physical binding counts still account for buffer segmentation.
3. All-empty QSA/CSA WebGPU batches copy/preserve state and return before binding zero-sized token tensors, with regression tests.
4. Shape inference again requires QSA to omit output slot 4 and CSA to name it, with rejection tests for both malformed forms.

Thanks for adding an explicit storage-binding limit check. It improves the error, but does not make the WebGPU kernel usable on adapters exposing the spec-guaranteed limit of eight bindings: QSA updat...

Fixed by reducing the dispatches rather than relying only on runtime rejection. Related tensors are now packed so QSA update/selection and CSA update/selection each require no more than eight logical storage buffers. CSA selection recomputes rejection from packed immutable metadata, eliminating its separate overflow binding. Added QSA and CSA WebGPU numerical tests with maxStorageBuffersPerShaderStage explicitly forced to 8; segmented physical bindings remain counted by the existing limit check.