onnxruntime

9dde2edb - [CUDA] [WebGPU] Optimize sparse attention indexers and add packed QSA state updates (#32656)

[CUDA] [WebGPU] Optimize sparse attention indexers and add packed QSA state updates (#32656) ### Description - Optimize dense CUDA QSA decode with one-time query normalization/rotation, causal-prefix mask handling, in-place raw-cache updates, Qwen-specialized scoring, and deterministic exact TopK paths spanning small, bounded, and distributed long-context selections. - Optimize packed CUDA QSA decode and prefill by hoisting query preparation, specializing block scoring, and applying deterministic bounded and hierarchical selection while preserving the correctness fallback. - Extend dense and packed indexers with optional packed QK input, INT64 padding masks with causal filtering, and shared rank-2 or batched rank-3 rotary caches; align CUDA, WebGPU, schema validation, and symbolic shape inference with the updated contracts while retaining legacy CSA compatibility. - Add optional transactional QSA state-update capture only to `PackedSparseAttentionIndexer`. Capture is fused into existing packed cache-update kernels so speculative decoding can replay an accepted prefix after restoring a snapshot. `SparseAttentionIndexer` retains its original two-output QSA contract and raw-key cache representation and does not expose transactional state-update outputs. - Preserve score-descending/index-ascending tie ordering, exact handling of irregular mask holes and ragged requests, fixed-capacity state semantics, and CUDA graph compatibility. - Add generated operator documentation, CUDA and WebGPU numerical/schema tests, symbolic shape tests, 8K–256K long-context parity, TopK boundary and malformed-state coverage, CUDA graph structural checks and latency metrics, and a Python decode benchmark sweep. ### Motivation and Context A 65K-context CUDA profile showed QSA selection dominating decode time, with additional context-scaled overhead from visibility compaction, repeated query preparation, cache copying, and intermediate traffic. This PR removes repeated full-row selection work, bounds or hierarchically partitions TopK scratch traffic, avoids redundant work on common causal-prefix and in-place decode paths, and adds model-shape specializations without changing exact ordering semantics. Exact query-specific scoring and padding-mask inspection remain context-linear because every query can change. The optimized selection, query preparation, and cache maintenance reduce the avoidable work around that constraint. Packed transactional capture additionally lets speculative runtimes commit accepted packed-indexer cache transitions without rerunning the target model or retaining rejected state, while the dense indexer remains free of that tracking contract. --------- 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> Co-authored-by: Akshay Sonawane <asonawane@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>