[Akshay Sonawane (apsonawane)]

Description

Extends com.microsoft::DynamicSparseAttention on the CUDA Execution Provider with symmetric signed INT8 per-channel KV cache support, per-batch committed_length cursor validation, and 12-warp $S \le 8$ grouped decode specialization, and extends com.microsoft::PackedSparseAttentionIndexer with mixed-precision T_KEY (BFloat16 activations with float past_key_state), fused $S \le 8$ QSA refresh/selection kernels, compact hierarchical TopK (ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1), and warp-specialized ragged QSA prefill (ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1).

Summary of Changes

com.microsoft::DynamicSparseAttention INT8 KV Cache, Cursor Validation & 12-Warp Decode

File	Change
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds optional inputs committed_length ([1] or [batch_size]), key_scale, and value_scale ([num_kv_heads, head_size] or [1, num_kv_heads, 1, head_size]) plus optional Output(3, "validation_status") (int32[1], codes 0..10), and extends T_CACHE to include tensor(int8)
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.{h,cc}, dynamic_sparse_attention_impl.{h,cu}, dynamic_sparse_attention_prefill.cuh	Implements symmetric signed INT8 KV cache append/decode/prefill with 8-byte vectorized staging (ORT_DSA_INT8_VECTOR_STAGING=1 default), shared-scale split merge for split_count <= 32 (ORT_DSA_SHARED_SCALE_MERGE=1 default), 12-warp (head_tiles=1) specialization in GroupedDecodeDynamicSparseAttentionKernel (query_heads_per_kv == 12, $S \le 8$), and device/strict-host validation including code 9 (kDynamicSparseAttentionInvalidCommittedLength)
onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc	Rejects CUDA-only INT8 KV cache and committed_length / validation_status extensions with clear status errors

com.microsoft::PackedSparseAttentionIndexer Mixed T_KEY, Compact Hierarchical TopK & Ragged Prefill

File	Change
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds T_KEY type constraint for past_key_state, present_key_state, and state_update so CUDA QSA can pair BFloat16 activations with float key state while preserving mainline output indices (output 6 = state_update, output 8 = state_status)
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.{h,cc}, packed_sparse_attention_indexer_impl.{h,cu}, qsa_topk.cuh	Adds fused $S \le 8$ QSA refresh/score/select kernels, compact hierarchical TopK (ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1 default: 2048-to-512 chunk compaction with QsaCompactTileTopKKernel and fused QsaFinalMergeAndEmitKernel for $\le 5$ chunks, falling back to QsaMergeTwoLevelsTopKKernel when =0), and warp-specialized ragged QSA prefill (ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1 default)
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc	Enforces matching activation and key-state types on WebGPU

Tests and Documentation

• onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc: Adds coverage for INT8 KV cache append/decode/prefill (with smooth_softmax=0 and head_sink cleared for prefill), 12-head grouped decode (FP16/BF16), scalar and per-batch committed_length, strict-validation cursor mismatch rejection without validation_status output, and fresh-process tests (Int8CacheGroupedDecodeAlternateSwitchesInFreshProcess_CUDA) for non-default ORT_DSA_INT8_VECTOR_STAGING / ORT_DSA_SHARED_SCALE_MERGE combinations.
• onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc, fresh_process_test_util.h: Adds coverage for BFloat16 activations with float past_key_state, compact/hierarchical TopK chunk boundaries, stable ties, partial merge groups, and fresh-process parity runs (QsaAlternateTopKModesInFreshProcess) with ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=0 and ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=0.
• docs/ContribOperators.md, docs/OperatorKernels.md, docs/contrib_ops/cuda/dynamic_sparse_attention.md, docs/contrib_ops/packed_sparse_attention_indexer.md: Documents schema updates, INT8 symmetric quantization, validation status codes (0..10), dispatch switches, and H200 benchmark deltas.

Performance

Measured on NVIDIA H200 (SM90, Qwen3.8-Flash-Next RC3 INT4 model, DFlash2 $S=8$ decode across 2K / 8K / 32K context):

• ORT_DSA_INT8_VECTOR_STAGING=1 (default): ~0.30 ms/cycle (~2%) faster than =0.
• ORT_DSA_SHARED_SCALE_MERGE=1 (default): ~0.05–0.10 ms/cycle (~0.3–0.5%) faster than =0.
• ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1 (default): ~0.20–0.27 ms/cycle (~1.5%) faster decode than =0.
• ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1 (default): reduces prefill TTFT by 3.7% at 8K and 14% at 32K compared to =0.

Testing

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

Checklist

• Tests added/updated
• No breaking changes (or documented in description)
• Documentation updated (if applicable)

[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.

[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}

[Copilot]

🟡 Changes recommended

The INT8 prefill test bypasses its intended kernel, and alternative TopK paths still need dispatch-forced reference-parity coverage.

2 open issues

• Add coverage for non-compact and radix TopK selection paths · New
• Disable smooth_softmax to reach INT8 prefill and match the reference · New

2 resolved since last review

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

_{Review effort: Balanced}

[David Fan (jiafatom)]

Please fix CI build.

[kunal-vaishnavi]

From Copilot:

Completed an implementation pass and a second edge-case/regression pass at cd4e2e4. One new P1 alignment finding is attached inline. Checked INT8 scaling and cache handling, grouped decode/prefill, compact Top-K padding and causal visibility, optional outputs, and non-CUDA rejection paths. Existing review threads were checked to avoid duplicate comments.

Validation: a standalone CUDA probe reproduced the float4 misalignment failure on H200. This was not a build or session-level run of the unmodified ORT operators. Full PR-specific ORT suites, WebGPU execution, cross-architecture qualification, and benchmarks were not run.

[Tianlei Wu (tianleiwu)]

The iOS build failure is fixed in c8c7b77 (cc David Fan (@jiafatom)). The only compile error on the previous head was error: 'system' is unavailable: not available on iOS, from fresh_process_test_util.h, which packed_sparse_attention_indexer_op_test.cc included unconditionally. The include and QsaAlternateTopKModesInFreshProcess are now under USE_CUDA, since the helper is only needed by the CUDA kernel tests. The other jobs on the previous head were green; iOS_CI_on_Mac and iphone_simulator are re-running on the new commit.