[Akshay Sonawane (apsonawane)]

Summary

Allow the existing grouped INT8 DynamicSparseAttention kernel to handle three- and four-row speculative verification through ORT_DSA_INT8_GROUPED_MAX_ROWS=4.

• Preserve the default maximum of two rows and all other eligibility conditions.
• Extend grouped-decode and vector-staging reference shapes through four rows.
• Document the opt-in range and fresh-process test configuration.

Base and integration

Incremental PR against asonawane/main-ops-sparse-attention-indexer, not main: upstream main lacks the prerequisite operators. This is one component of the DFlash2 stack, alongside the INT2 MoE and hyperconnection tuning branches and GenAI asonawane/dflash2-device-pipeline. No model assets, binaries, benchmark/evaluation scripts, or Mobius changes are included.

Validation

• Changed CUDA implementation and C++ reference-test translation units compile with this branch's contrib headers against the existing integration core build.
• git diff --check passes; changed lines formatted with clang-format.
• The original combined experimental stack passed 24 numeric reference tests. Those results are not a full clean build/numeric qualification of this rebased operator branch. Keep this PR draft pending that qualification.

Combined H200 experiment (not an individual-kernel ablation)

INT2 target + FP8 DFlash2, greedy batch one; six-run means from the previously tested private integration:

Context	Decode tokens/s	TTFT ms
2K	333.964	580.497
8K	342.140	2421.365
16K	339.684	5030.850
32K	336.767	10520.556

Peak sampled memory: 59.350 GiB versus 64.322 GiB old-runtime control. These are combined-stack historical measurements, not measurements of this rebased PR in isolation.

Accuracy caveat: paired HumanEval DFlash/same-S4 perturbed-draft control scored 154/164 with exact output parity. Five-problem LiveCodeBench smoke scored 2/5 versus historical MTP 3/5, including a capacity-limited regression. This stack is not yet an accuracy-qualified default.

Companion draft PRs

• GenAI device pipeline and integration: Integrate the DFlash2 device pipeline with Qwen graph execution onnxruntime-genai#2771
• INT2 FC2/routing fusion: Fuse INT2 FC2 routing finalization for four-row verification #33547
• Hyperconnection S4 tuning: [CUDA] Optimize HyperConnectionLowRankMix for S<=8 speculative decode and deterministic prefill #33548

[Copilot]

🔵 Needs a closer look

Exact-branch CUDA numerical qualification remains pending, and automated tests do not enable the expanded grouped path.

1 open issue

• Add fresh-process tests for grouped and vector-staging dispatch · New

What changed in this PR

Extends the existing CUDA DynamicSparseAttention grouped INT8 path to support opt-in three- and four-row speculative verification, while preserving the default two-row limit.

Changes:

• Adds ORT_DSA_INT8_GROUPED_MAX_ROWS with a supported range of 2–4.
• Extends grouped-decode and vector-staging reference cases through four rows.
• Documents eligibility and fresh-process test configuration.

File	Description
onnxruntime/​test/​contrib_ops/​dynamic_sparse_attention_op_test.cc	Adds three- and four-row reference cases.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.cu	Makes the grouped INT8 row limit configurable.
docs/​contrib_ops/​cuda/​dynamic_sparse_attention.md	Documents the opt-in range and testing setup.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Squashed into #33472 (asonawane/main-ops-sparse-attention-indexer). Closing this stacked PR.