onnxruntime
3188a2ee - [CUDA][WebGPU] Add packed sparse attention indexer for continuous batching (#32618)

Commit
19 days ago
[CUDA][WebGPU] Add packed sparse attention indexer for continuous batching (#32618) ### 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. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Author
Parents
Loading