Fix CUDA MHA shared-cache scratch lifetime

onnxruntime/contrib_ops/cuda/bert/attention_kv_cache.cu

@@ -889,10 +889,13 @@ Status LaunchConcatKVInPlace(int batch_size,
                              const bool is_new_kv_bnsh_format,
                              cudaStream_t stream,
                              const int max_threads_per_block) {
-  // static_assert(sizeof(T) == 2);
-  assert(head_size % 4 == 0);
+  constexpr int elements_per_vector = sizeof(float2) / sizeof(T);
+  if (head_size % elements_per_vector != 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Head size must be divisible by ", elements_per_vector,
+                           " for vectorized kernel.");
+  }
 
-  const int H = head_size / 4;
+  const int H = head_size / elements_per_vector;
   if (H * kv_num_heads <= max_threads_per_block) {
     const dim3 grid(new_seq_len, batch_size, 1);
     const dim3 block(H, kv_num_heads, 1);

onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc

@@ -12,6 +12,7 @@
 #include "contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
 #include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
 #include "contrib_ops/cuda/bert/lean_attention/lean_api.h"
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
 
 using namespace onnxruntime::cuda;
 using namespace ::onnxruntime::common;
@@ -226,24 +227,26 @@ Status MultiHeadAttention<T, QK>::ComputeInternal(OpKernelContext* context) cons
   if (use_decoder_masked_multihead_attention) {
     // Kernel only works for token generation with beam search
     kernel_type = AttentionKernelType::AttentionKernel_DecoderAttention;
+  }
 
-    // No production use-case will incur this copy cost as the implementation of
-    // DecoderMaskedMultiHeadAttention is written in such a way that the past and present buffers
-    // must be shared to have parity in the outputs.
-    // This is just to circumvent the OpTester's limitation of not being able to bind a specific
-    // buffer to inputs/outputs.
+  if (parameters.past_present_share_buffer) {
+    // Buffer-sharing kernels append the new KV in place into present_key/present_value and never read
+    // past_key/past_value, so they assume the runtime aliased past onto present. Nothing else copies the
+    // past cache forward. When the buffers are not actually aliased (e.g. OpTester cannot bind a specific
+    // buffer to both an input and an output), the untouched history slots would be garbage, so seed them
+    // here. When they are aliased the pointers compare equal and this is a no-op.
     auto* past_key_data = (past_key == nullptr) ? nullptr : past_key->Data<T>();
     auto* past_value_data = (past_value == nullptr) ? nullptr : past_value->Data<T>();
     auto* present_key_data = (present_key == nullptr) ? nullptr : present_key->MutableData<T>();
     auto* present_value_data = (present_value == nullptr) ? nullptr : present_value->MutableData<T>();
 
-    if (present_key_data != past_key_data) {
-      DUMP_STRING("Copying past_key to present_key for OpTester");
+    if (past_key_data != nullptr && present_key_data != nullptr && present_key_data != past_key_data) {
+      DUMP_STRING("Copying past_key to present_key");
       CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(present_key_data, past_key_data, past_key->SizeInBytes(),
                                            cudaMemcpyDeviceToDevice, stream));
     }
-    if (present_value_data != past_value_data) {
-      DUMP_STRING("Copying past_value to present_value for OpTester");
+    if (past_value_data != nullptr && present_value_data != nullptr && present_value_data != past_value_data) {
+      DUMP_STRING("Copying past_value to present_value");
       CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(present_value_data, past_value_data, past_value->SizeInBytes(),
                                            cudaMemcpyDeviceToDevice, stream));
     }
@@ -549,14 +552,16 @@ Status MultiHeadAttention<T, QK>::ComputeInternal(OpKernelContext* context) cons
   data.allow_debug_info = kernel_options_->AllowDebugInfo();
 
   // For past-present buffer sharing.
+  IAllocatorUniquePtr<void> seqlens_k_buffer;
   if (parameters.past_present_share_buffer) {
-    std::vector<int64_t> seqlens_k(parameters.batch_size, parameters.total_sequence_length - 1);
-    size_t seqlens_k_bytes = 0;
-    seqlens_k_bytes = sizeof(int) * parameters.batch_size;
-    auto seqlens_k_buffer = GetScratchBuffer<void>(seqlens_k_bytes, GetComputeStream(context));
+    const size_t seqlens_k_bytes = sizeof(int) * static_cast<size_t>(parameters.batch_size);
+    seqlens_k_buffer = GetScratchBuffer<void>(seqlens_k_bytes, GetComputeStream(context));
     if (seqlens_k_buffer != nullptr) {
       data.seqlens_k_total = reinterpret_cast<int*>(seqlens_k_buffer.get());
-      CUDA_RETURN_IF_ERROR(cudaMemcpy(data.seqlens_k_total, seqlens_k.data(), seqlens_k_bytes, cudaMemcpyHostToDevice));
+      // Fill on the compute stream. A host-to-device copy would both race with the
+      // asynchronous launches below and be illegal during CUDA graph capture.
+      onnxruntime::cuda::Fill<int32_t>(stream, data.seqlens_k_total, parameters.total_sequence_length,
+                                       parameters.batch_size);
     }
   }
 

onnxruntime/test/contrib_ops/multihead_attention_op_test.cc

@@ -848,6 +848,38 @@ TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_UsingDMMHAInsi
   RunMultiHeadAttentionTests(data, DISABLE_CPU | DISABLE_WEBGPU | DISABLE_DML);
 }
 
+TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_ConcatKVInPlace) {
+  AttentionTestData data;
+  data.hidden_size = 4;
+  data.v_hidden_size = 4;
+  data.num_heads = 1;
+  data.batch_size = 1;
+  data.sequence_length = 1;
+  data.kv_sequence_length = 1;
+  data.mask_type = AttentionMaskType::MASK_NONE;
+  data.query_data = {1.0f, 0.0f, 0.0f, 0.0f};
+  data.key_data = {0.0f, 1.0f, 0.0f, 0.0f};
+  data.value_data = {5.0f, 6.0f, 7.0f, 8.0f};
+  data.past_key_data = {1.0f, 0.0f, 0.0f, 0.0f,
+                        0.0f, 0.0f, 0.0f, 0.0f,
+                        0.0f, 0.0f, 0.0f, 0.0f};
+  data.past_value_data = {1.0f, 2.0f, 3.0f, 4.0f,
+                          0.0f, 0.0f, 0.0f, 0.0f,
+                          0.0f, 0.0f, 0.0f, 0.0f};
+  data.past_seq_len_data = {1};
+  data.max_sequence_length = 3;
+  data.present_key_data = {1.0f, 0.0f, 0.0f, 0.0f,
+                           0.0f, 1.0f, 0.0f, 0.0f,
+                           0.0f, 0.0f, 0.0f, 0.0f};
+  data.present_value_data = {1.0f, 2.0f, 3.0f, 4.0f,
+                             5.0f, 6.0f, 7.0f, 8.0f,
+                             0.0f, 0.0f, 0.0f, 0.0f};
+  data.fp32_output_data = {2.5101626f, 3.5101626f, 4.5101624f, 5.5101624f};
+  data.is_static_kv = false;
+  data.buffer_share = true;
+  RunMultiHeadAttentionTests(data, DISABLE_CPU | DISABLE_WEBGPU | DISABLE_DML);
+}
+
 TEST(MultiHeadAttentionTest, CrossAttention_DiffSequenceLengths_UsingDMMHAInsideMHA) {
   // Whisper decoder cross attention with past_kv used directly as K and V, no mask, and bias
   // Used in decoder-with-past's cross-attention layers