Size Flash split-KV from valid KV length in ONNX Attention opset-24 external-cache path (CUDA)

onnxruntime/core/providers/cuda/llm/attention.cc

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 
+#include "core/common/inlined_containers.h"
 #include "core/common/safeint.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cpu/llm/attention.h"
@@ -254,10 +255,42 @@ Status Attention<T>::RunFlashAttention(
   size_t softmax_lse_bytes = onnxruntime::flash::get_softmax_lse_size(
       parameters.q_sequence_length, parameters.batch_size, parameters.q_num_heads);
 
+  // Size the Flash split-KV launch (num_splits) from the VALID KV length, not the
+  // KV cache buffer length. In the external-cache decode path (nonpad_kv_seqlen,
+  // opset 24), K/V are the full pre-allocated cache, so parameters.total_sequence_length
+  // is the buffer length. get_num_splits_and_buffer_sizes uses that length to derive
+  // num_n_blocks and choose num_splits; the split-KV kernel then partitions the buffer
+  // range into num_splits and reduces over all of them in the combine pass. Sizing this
+  // from the buffer over-partitions the launch and does wasted work over the padding
+  // region — for a fixed valid length, enlarging the cache buffer makes the same kernel
+  // dramatically slower (see issue #29686).
+  //
+  // The valid KV length lives device-side in nonpad_kv_seqlen, so read the per-batch max
+  // onto the host (a small batch_size-element copy) to size num_splits from the real work.
+  // This is a perf-only knob: the kernel's seqlen_k stays = total_sequence_length so the
+  // KV cache strides remain correct, and the device-side seqlens_k masking still bounds the
+  // math per batch, so correctness is unaffected.
+  int split_kv_seqlen = parameters.total_sequence_length;
+  if (nonpad_kv_seqlen != nullptr && parameters.batch_size > 0) {
+    InlinedVector<int64_t> host_nonpad_kv_seqlen(parameters.batch_size);
+    CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(
+        host_nonpad_kv_seqlen.data(), nonpad_kv_seqlen->Data<int64_t>(),
+        sizeof(int64_t) * parameters.batch_size, cudaMemcpyDeviceToHost, cuda_stream));
+    CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
+    int64_t max_valid_kv_seqlen = 0;
+    for (int b = 0; b < parameters.batch_size; ++b) {
+      max_valid_kv_seqlen = std::max(max_valid_kv_seqlen, host_nonpad_kv_seqlen[b]);
+    }
+    // Clamp to the buffer length (defensive; valid length can't exceed the cache buffer)
+    // and keep at least 1 so the heuristic is well-defined for an all-empty cache.
+    max_valid_kv_seqlen = std::min<int64_t>(max_valid_kv_seqlen, parameters.total_sequence_length);
+    split_kv_seqlen = std::max<int>(1, static_cast<int>(max_valid_kv_seqlen));
+  }
+
   auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] =
       onnxruntime::flash::get_num_splits_and_buffer_sizes(
           parameters.batch_size, parameters.q_sequence_length,
-          parameters.total_sequence_length, parameters.q_num_heads,
+          split_kv_seqlen, parameters.q_num_heads,
           parameters.head_size, device_prop.multiProcessorCount);
 
   auto softmax_lse_buffer = GetScratchBuffer<void>(softmax_lse_bytes, GetComputeStream(context));

onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py

@@ -567,6 +567,44 @@ def test_tensorscatter_attention_cuda_fp16(
         numpy.testing.assert_allclose(present_k, ref_present_k, rtol=rtol["fp16"], atol=atol["fp16"])
         numpy.testing.assert_allclose(present_v, ref_present_v, rtol=rtol["fp16"], atol=atol["fp16"])
 
+    @parameterized.expand(
+        [
+            # (name, batch, q_heads, kv_heads, total_kv (buffer), scatter_pos, nonpad_seqlens)
+            # q_seq_len is fixed at 1 here (decode), so for each batch b the new token is
+            # written at scatter_pos[b] and the resulting valid KV length is
+            # nonpad_seqlens[b] == scatter_pos[b] + q_seq_len (i.e. scatter_pos[b] + 1).
+            # A large pre-allocated cache buffer (total_kv) with a small valid KV length is the
+            # long-max-context decode scenario. The Flash split-KV launch must be sized
+            # from the valid length (max of nonpad_kv_seqlen), not the buffer length
+            # (issue #29686). These cases exercise that host-side sizing branch with a
+            # large num_n_blocks while keeping the valid range small; output must still
+            # match the reference regardless of the chosen num_splits.
+            ("mha_bigbuf_smallvalid", 1, 4, 4, 2048, [63], [64]),
+            ("gqa_bigbuf_diff_lens", 2, 8, 2, 2048, [63, 127], [64, 128]),
+            ("gqa_bigbuf_one_empty", 2, 8, 2, 2048, [0, 95], [1, 96]),
+        ]
+    )
+    def test_tensorscatter_attention_cuda_fp16_large_buffer(
+        self, name, batch, q_heads, kv_heads, total_kv, scatter_pos, seqlens
+    ):
+        output, ref_output, present_k, present_v, ref_present_k, ref_present_v = run_tensorscatter_attention(
+            batch_size=batch,
+            total_kv_seq_len=total_kv,
+            q_seq_len=1,
+            q_num_heads=q_heads,
+            kv_num_heads=kv_heads,
+            head_size=_HEAD_SIZE,
+            nonpad_seqlens=seqlens,
+            scatter_positions=scatter_pos,
+            ep="CUDAExecutionProvider",
+            torch_type=torch.float16,
+            ort_type=TensorProto.FLOAT16,
+            is_causal=0,
+        )
+        numpy.testing.assert_allclose(output, ref_output, rtol=rtol["fp16"], atol=atol["fp16"])
+        numpy.testing.assert_allclose(present_k, ref_present_k, rtol=rtol["fp16"], atol=atol["fp16"])
+        numpy.testing.assert_allclose(present_v, ref_present_v, rtol=rtol["fp16"], atol=atol["fp16"])
+
 
 @unittest.skipIf(not has_cuda_device(53), "CUDA device not available, skipping tests.")
 class TestTensorScatterAttentionCUDAFP32(unittest.TestCase):