[Akshay Sonawane (apsonawane)]

Extends com.microsoft::DynamicSparseAttention with INT8 per-channel KV cache support, cursor validation, and 12-warp $S \le 8$ grouped decode specialization, and extends com.microsoft::PackedSparseAttentionIndexer with FP32 past_key_state (T_KEY) support and fast $S \le 8$ QSA block scoring/TopK on CUDA.

Key Changes

Component	Files	Description
DynamicSparseAttention INT8 KV Cache & 12-Warp Decode	onnxruntime/core/graph/contrib_ops/bert_defs.cc, onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.{cc,h}, dynamic_sparse_attention_impl.{h,cu}	Adds optional committed_length, key_scale, value_scale inputs and validation_status output, INT8 KV cache append/attention, and 12-warp (head_tiles=1) specialization in GroupedDecodeDynamicSparseAttentionKernel (kTargetGroupedQuantizedDecodeBlocks=512) for query_heads_per_kv == 12 ($S \le 8$).
PackedSparseAttentionIndexer Mixed T_KEY & Fast QSA Decode	onnxruntime/core/graph/contrib_ops/bert_defs.cc, onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.{cc,h}, packed_sparse_attention_indexer_impl.{h,cu}, qsa_topk.cuh	Adds T_KEY type constraint (BFloat16 activations with float past_key_state), fused $S \le 8$ refresh/select kernels, and hierarchical/compact TopK selection while using mainline state_update (output 6) and state_status (output 8).
Tests & Docs	onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc, packed_sparse_attention_indexer_op_test.cc, docs/	Adds unit tests for INT8 KV cache, 12-head grouped decode, BFloat16+float QSA state, and hierarchical TopK boundaries.

Testing Notes

• Run unit tests: onnxruntime_provider_test --gtest_filter="DynamicSparseAttention*:PackedSparseAttentionIndexer*"
• Verified end-to-end on qwen_flash_next_int4_int8_rc2_contrib_dflash2 ($S=8$, state_update_capacity=7).

[Copilot]

🟡 Changes recommended

Runtime and inference defects remain unresolved, alongside missing grouped-path coverage and generated-documentation mismatches.

7 open issues

• Regenerate and commit synchronized operator documentation · New
• Allocate grouped decode workspace for single-split cases · New
• Add CUDA parity tests for floating-cache grouped paths · New
• Handle or reject reuse-mode outputs 8–10 · New
• Reject unsupported WebGPU validation options · New
• Conditionally include optional QSA position_ids binding · New
• Fix mixed BF16 activation and FP32 key-state type propagation · New

What changed in this PR

Extends ONNX Runtime’s sparse-attention operators with typed caches, optional state outputs, and optimized CUDA selection and attention paths.

Changes:

• Adds INT8 attention caches and BF16 QSA with FP32 key state.
• Adds snapshot/status outputs, optional QSA positions, and provider restrictions.
• Updates CUDA dispatch, regression tests, and operator documentation.

File	Description
onnxruntime/​test/​contrib_ops/​packed_sparse_attention_indexer_op_test.cc	Tests mixed-state types, positions, and tiled selection.
onnxruntime/​test/​contrib_ops/​dynamic_sparse_attention_op_test.cc	Tests INT8 caches, grouped decode, and validation.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Extends schemas and shape inference.
onnxruntime/​contrib_ops/​webgpu/​bert/​packed_sparse_attention_indexer.h	Adds optional-position shader configuration.
onnxruntime/​contrib_ops/​webgpu/​bert/​packed_sparse_attention_indexer.cc	Handles positions and unsupported modes.
onnxruntime/​contrib_ops/​webgpu/​bert/​dynamic_sparse_attention.cc	Adds cache-type and scale restrictions.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer.h	Separates activation and state template types.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer.cc	Wires typed state and optional outputs.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.h	Extends kernel interfaces.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.cu	Implements snapshots and optimized selection.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers mixed-state kernels.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention.cc	Integrates quantized caches and validation.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.h	Extends cache and validation interfaces.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.cu	Adds quantization and grouped attention kernels.
onnxruntime/​contrib_ops/​cpu/​sparse/​packed_sparse_attention_indexer_common.h	Defines additional output slots.
docs/​ContribOperators.md	Updates operator interface documentation.
docs/​contrib_ops/​packed_sparse_attention_indexer.md	Documents state extensions and experimental controls.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Strict cursor validation can allow cache updates after a mismatch, and new refresh outputs lack numerical regression coverage.

1 open issue

• Documented FLOAT16/FLOAT QSA type lacks kernel support · New

2 previously missed issues

In code that hasn't changed since last review

Strict mode ignores cursor validation errors

onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention.cc:264

Strict validation reads the error flag back inside ValidateDynamicSparseAttentionOnDevice, before this cursor kernel runs. With batch_size=1, sequence_length=1, cache_capacity=2, seqlens_k={0}, and committed_length={1}, general validation succeeds, then the cursor writes 9 without any further host check. Run still succeeds and appends to the cache; omitting validation_status discards the error entirely. Check the cursor result before launching attention in strict mode, and add a regression without the status output.

Add QSA refresh snapshot and validation status coverage

onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.cu:205

The new QSA snapshot/status outputs lack refresh-mode numerical coverage: RunQsaPackedTest never requests outputs 8–10, and the only tests connecting them expect IndexShare reuse rejection. Add CUDA refresh tests comparing the first-token snapshot with a separate one-token transition, including a compression boundary and an empty request. Also verify validation_status values 0, 9, and 0x99 for valid requests and rejected overflow or inconsistent metadata. Existing state_update capture tests exercise output 6, not these outputs.

7 resolved since last review

• Fix mixed BF16 activation and FP32 key-state type propagation
• Conditionally include optional QSA position_ids binding
• Reject unsupported WebGPU validation options
• Handle or reject reuse-mode outputs 8–10
• Add CUDA parity tests for floating-cache grouped paths
• Allocate grouped decode workspace for single-split cases
• Regenerate and commit synchronized operator documentation

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed 9356555 with the review clarification and WebGPU pipeline fix.

• Corrected mixed-state schema/generated docs to promise only CUDA QSA BFLOAT16 activations with FLOAT key state, not unsupported FLOAT16/FLOAT. WebGPU matching-type requirement is explicit.
• Changed the binding-limit test to pass limit.c_str() to ConfigOptions::AddConfigEntry. Windows, Linux, macOS and WebGPU plugin logs all report this same std::string-to-const-char* compilation error; the static aggregate failure is downstream of compilation.
• Android failed before compiling this PR while sdkmanager unpacked the emulator download: "Error on ZipFile unknown archive". This is an SDK download/infrastructure failure, not addressed by changing operator code. Fresh CI after this push will make another attempt; no unrelated workflow edits were made.

Validation: full-PR lintrunner -a and check-only lint, source-to-official-renderer documentation checks, and git diff --check passed. No local builds, compilation, native tests, or model execution were run as requested. Fresh CI is pending, not claimed green.

[Akshay Sonawane (apsonawane)]

Added the scoped QSA and attention decode experiments on top of the current PR head, preserving its workspace-allocation, state-validation, and WebGPU review fixes.

Published commit: d21946e6cb37621b9a2e5f623ea56e635b3e7bdf.

• Opt-in ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1: compact 2048-key chunks to 512 keys with device-side 512/1024/2048 sort probes, a fused merge/emit for at most five chunks, and existing pairwise merging above that. Preserves stable keys, causal tails, capacity stride/padding, and overflow handling. Rejects simultaneous radix and compact modes.
• Opt-in ORT_DSA_INT8_VECTOR_STAGING=1: eight-byte grouped INT8 staging with independent per-channel key/value scales and the existing shape/resource guards.
• Opt-in ORT_DSA_SHARED_SCALE_MERGE=1: share exponential factors across channels for grouped merges of at most 32 splits. Preserve nonzero-partial-sum activity and underflow/nonfinite accumulation semantics.
• Added QSA chunk/fallback boundaries and distinct-channel signed INT8/auxiliary-cache/merge-boundary tests. Updated operator documentation.
• No multi-snapshot state changes or dense-prefill normalization experiments were included.

Validation of the publishing sources: all changed CUDA translation units compile. Process-isolated native suites pass with all new switches both off and on: 48 attention tests and 36 packed-indexer tests per mode. This uses retained build dependencies, including build-compatible core headers and upstream test-utility definitions; it is not a clean full upstream build.

The earlier experimental tree had two malformed-state-length QSA/CSA failures with compact mode off and on. Those historical results must not be represented as a clean full suite. The latest publishing-version focused suites above pass, but the broader combined GDN run encountered an empty-request failure reproducible with the unmodified GDN PR-head kernel; clean-build CI remains required.

Historical Nsight measurements from experimental packages showed QSA summed kernel time 0.651 -> 0.409 ms/cycle, grouped attention compute 45.30 -> 31.02 us, and merge 9.28 -> 4.22 us. These are not new benchmark results for this transplanted PR tip. Full model accuracy remains deferred; there is no sustained-250-TPS-at-8K claim.

Companion updates: GDN/convolution #33471, MoE/FP8 #33473, and GenAI packaging microsoft/onnxruntime-genai#2759.

[Akshay Sonawane (apsonawane)]

Fresh-build integration validation update (no source changes pushed to this ORT PR in this update):

• Rebuilt ORT from scratch for H200 (sm90a), CUDA provider as a plugin, with FP4 QMoE, FP8/INT4 KV cache, GQA value layout, and FPA-INTB GEMM enabled. onnxruntime_QUICK_BUILD=OFF; no retained build objects were reused.
• Rebuilt GenAI against the fresh ORT headers/libraries, and rebuilt the fold custom library.
• Focused convolution/GDN/fused-tensor/hyper-connection/INT2 tests: 193 passed.
• Sparse/indexer/FP8 suite plus the isolated QSA capacity test: 140 additional passes.
• Known limitations: QSA and CSA MalformedStateLengthsAreClamped tests still fail, including with specialized top-k disabled. These are not being claimed as passing. This is selected-suite validation, not a clean full-suite qualification.
• Selected GenAI Qwen/executor/weight-pool tests: 69 passed, including the opt-in real-GPU mapped-weight test with proper CUDA plugin registration. One additional test is disabled.

The latest fused INT2 model completed one warmup plus three measured requests at each of 2K/8K/16K/32K, generating 512 tokens each. All 16 requests matched the previous-runtime token IDs exactly, with zero generation failures or standard fallbacks. Runtime-map inspection confirmed only the fresh ORT/GenAI libraries were used.

Prompt	Pooled decode TPS	Mean TTFT
2048	269.90	0.584 s
8192	250.17	2.448 s
16384	266.92	5.065 s
32768	264.12	10.578 s

Evidence scope: this is the full dirty experimental source stack containing this PR's optimizations plus other local work, not a standalone build/benchmark of this exact PR head. It is not an isolated speedup attribution, a fresh paired reference comparison, or proof of cross-process sustained 250 TPS. Full model accuracy qualification remains deferred.

The rebuild exposed local expert-cache function-pointer and INT2 direct-link test wiring issues, which were fixed locally. Those expert-cache extensions are not in the published PR and have deliberately not been mixed into this update.

[Akshay Sonawane (apsonawane)]

Resolved conflicts with main (1f4b965) in merge commit b42576e.

• Preserved main’s established output 8 = state_status contract, including per-request invalid-metadata/capacity status semantics. Experimental snapshot outputs now occupy 9 = mid_kv_buffer, 10 = mid_state_lengths, 11 = validation_status. Earlier experimental graphs/exporters using slots 8–10 must insert an empty output at slot 8 before these three outputs. Mainline graph layouts are unchanged. Schema, kernel enum, inference, provider rejection paths, tests, and docs were updated consistently.
• Kept mixed-type key-state support and main’s BF16 optimized prefill dispatch, explicitly excluding quantized caches from that BF16-only cache path.
• Corrected raw-pointer/unique-pointer automatic-merge mismatches and preserved both branches’ tests. Added a regression requesting snapshot outputs both with and without state_status, including BF16 activations/FP32 state.
• Compiled merged sources/tests and linked an isolated CUDA provider/test executable against existing integration-build dependencies. Combined suite for these two merges on H200 GPU 1: 119 passed, 2 architecture/image-specific skips, 0 failures.

This is merged-source validation using existing integration-build dependencies, not a full clean standalone-branch build or WebGPU qualification. GitHub currently reports MERGEABLE (no merge conflicts); required checks/reviews remain separate. No active model package or frozen GPU 3 LiveCodeBench snapshot was changed.

[github-actions]

You can commit the suggested changes from lintrunner.

[Akshay Sonawane (apsonawane)]

Follow-up validation: ran lintrunner -a on an isolated checkout of this PR against origin/main, reviewed and applied the formatter fixes, then reran both lintrunner -a and lintrunner successfully with no lint issues. Pushed formatting-only commit b71498b2660189a4d3638de2095cf1d74e691ce7 (3 files). Verified identical C++/CUDA token sequences before and after formatting, including preserved string/comment contents. No runtime behavior changes; active integration worktree and evaluation snapshot untouched.

[github-actions]

You can commit the suggested changes from lintrunner.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

🟡 Changes recommended

The FP16 floating-cache 12-head decode path reads uninitialized attention scores, and grouped-prefill coverage remains incomplete.

2 open issues

• FP16 ratio-12 MMA path reads unwritten attention scores · New
• Add a long-sequence test for grouped-prefill FP16 MMA parity · New

1 resolved since last review

• Documented FLOAT16/FLOAT QSA type lacks kernel support

_{Review effort: Balanced}