Enable grouped INT8 sparse attention for four-row verification

docs/contrib_ops/cuda/dynamic_sparse_attention.md

@@ -0,0 +1,14 @@
+# CUDA DynamicSparseAttention grouped decode
+
+`ORT_DSA_INT8_GROUPED_MAX_ROWS` controls the maximum query rows eligible for
+the existing grouped INT8 decode kernel. The default is 2, preserving existing
+dispatch. Set it to 4 before starting the process to allow three- and four-row
+speculative verification; the supported range is 2 through 4.
+
+All other grouped-kernel eligibility conditions still apply. Ineligible cache
+formats, shapes, and devices retain their existing dispatch. This does not
+change the operator schema, cache layout, or sparse-index selection semantics.
+
+The INT8 grouped-decode and vector-staging reference tests cover one through
+four query rows. To exercise the opt-in path, run the INT8 grouped-decode tests
+in a fresh process with `ORT_DSA_INT8_GROUPED_MAX_ROWS=4`.

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu

@@ -1732,8 +1732,12 @@ bool UseGroupedDecodeAttention(const DynamicSparseAttentionParameters& parameter
                                size_t element_size,
                                size_t max_shared_memory_per_block,
                                bool quantized_cache = false) {
+  static const int max_decode_rows =
+      ParseEnvironmentVariableWithDefault<int>("ORT_DSA_INT8_GROUPED_MAX_ROWS", 2);
+  ORT_ENFORCE(max_decode_rows >= 2 && max_decode_rows <= 4,
+              "ORT_DSA_INT8_GROUPED_MAX_ROWS must be in [2, 4].");
   return (parameters.sequence_length == 1 ||
-          (quantized_cache && parameters.sequence_length == 2)) &&
+          (quantized_cache && parameters.sequence_length <= max_decode_rows)) &&
          parameters.num_heads / parameters.kv_num_heads >= 4 &&
          GetCandidateCapacity(parameters) > kSplitCandidateSize &&
          SupportsGroupedAttention(parameters, element_size, max_shared_memory_per_block) &&

onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc

@@ -583,7 +583,7 @@ TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA)
 }
 
 TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) {
-  for (int sequence_length : {1, 2}) {
+  for (int sequence_length : {1, 2, 3, 4}) {
     for (bool auxiliary : {false, true}) {
       SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) +
                    ", auxiliary=" + std::to_string(auxiliary));
@@ -603,7 +603,7 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) {
 }
 
 TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeVectorStagingChannels_CUDA) {
-  for (int sequence_length : {1, 2}) {
+  for (int sequence_length : {1, 2, 3, 4}) {
     for (bool auxiliary : {false, true}) {
       SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) +
                    ", auxiliary=" + std::to_string(auxiliary));