From ebac9e560b87108b138e1ddfccaefc1b71f3e190 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 20:59:32 +0000 Subject: [PATCH] Enable grouped INT8 sparse attention for four-row verification Add an opt-in maximum-row gate while preserving the two-row default, and extend grouped/vector-staging reference shapes through four rows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/dynamic_sparse_attention.md | 14 ++++++++++++++ .../cuda/bert/dynamic_sparse_attention_impl.cu | 6 +++++- .../dynamic_sparse_attention_op_test.cc | 4 ++-- 3 files changed, 21 insertions(+), 3 deletions(-) create mode 100644 docs/contrib_ops/cuda/dynamic_sparse_attention.md diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md new file mode 100644 index 0000000000000..941798301094b --- /dev/null +++ b/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`. diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 4c8553950e4e7..b737873d1c92f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/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("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) && diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index cf678e02942b6..94f6081d8ad9e 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/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));