From aa995cee1df471afc0a6d6f0d8e2ad026538699f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 10 Aug 2026 16:18:37 -0700 Subject: [PATCH 1/3] Fix CUDA MHA shared-cache scratch lifetime Keep the sequence-length scratch allocation alive through asynchronous attention launches, populate it with correctly typed total lengths, and seed non-aliased shared-cache outputs before in-place append. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/multihead_attention.cc | 16 +++++++--------- .../contrib_ops/multihead_attention_op_test.cc | 7 +++++++ 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc index fa39435b0350b..afec927c61f2b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc @@ -226,12 +226,10 @@ Status MultiHeadAttention::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 in place. Copy the past cache when the runtime did not alias the outputs. auto* past_key_data = (past_key == nullptr) ? nullptr : past_key->Data(); auto* past_value_data = (past_value == nullptr) ? nullptr : past_value->Data(); auto* present_key_data = (present_key == nullptr) ? nullptr : present_key->MutableData(); @@ -549,11 +547,11 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) cons data.allow_debug_info = kernel_options_->AllowDebugInfo(); // For past-present buffer sharing. + IAllocatorUniquePtr seqlens_k_buffer; if (parameters.past_present_share_buffer) { - std::vector 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(seqlens_k_bytes, GetComputeStream(context)); + std::vector seqlens_k(parameters.batch_size, parameters.total_sequence_length); + const size_t seqlens_k_bytes = sizeof(seqlens_k[0]) * seqlens_k.size(); + seqlens_k_buffer = GetScratchBuffer(seqlens_k_bytes, GetComputeStream(context)); if (seqlens_k_buffer != nullptr) { data.seqlens_k_total = reinterpret_cast(seqlens_k_buffer.get()); CUDA_RETURN_IF_ERROR(cudaMemcpy(data.seqlens_k_total, seqlens_k.data(), seqlens_k_bytes, cudaMemcpyHostToDevice)); diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index a5fcbb25ca93f..75d6b4303f347 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -848,6 +848,13 @@ TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_UsingDMMHAInsi RunMultiHeadAttentionTests(data, DISABLE_CPU | DISABLE_WEBGPU | DISABLE_DML); } +TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_ConcatKVInPlace) { + AttentionTestData data; + GetSelfAttention_PastPresentBufferShare_UsingDMMHAInsideMHA(data); + data.cache_indir_data.clear(); + 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 From cf88473f80bff8dc97d819f272fe8cf42bb79169 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 15:19:44 -0700 Subject: [PATCH 2/3] Fill MHA sequence-length scratch on the compute stream Replace the synchronous host-to-device copy of the past-present sequence lengths with a stream-ordered device fill so the buffer is populated in order with the attention launches and remains valid during CUDA graph capture. Also guard the past-to-present cache seeding against missing tensors and expand the comment explaining why it is needed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/multihead_attention.cc | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc index afec927c61f2b..92200a06e73e6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc +++ b/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; @@ -229,19 +230,23 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) cons } if (parameters.past_present_share_buffer) { - // Buffer-sharing kernels append in place. Copy the past cache when the runtime did not alias the outputs. + // 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(); auto* past_value_data = (past_value == nullptr) ? nullptr : past_value->Data(); auto* present_key_data = (present_key == nullptr) ? nullptr : present_key->MutableData(); auto* present_value_data = (present_value == nullptr) ? nullptr : present_value->MutableData(); - 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,12 +554,14 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) cons // For past-present buffer sharing. IAllocatorUniquePtr seqlens_k_buffer; if (parameters.past_present_share_buffer) { - std::vector seqlens_k(parameters.batch_size, parameters.total_sequence_length); - const size_t seqlens_k_bytes = sizeof(seqlens_k[0]) * seqlens_k.size(); + const size_t seqlens_k_bytes = sizeof(int) * static_cast(parameters.batch_size); seqlens_k_buffer = GetScratchBuffer(seqlens_k_bytes, GetComputeStream(context)); if (seqlens_k_buffer != nullptr) { data.seqlens_k_total = reinterpret_cast(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(stream, data.seqlens_k_total, parameters.total_sequence_length, + parameters.batch_size); } } From f0e7fa0ea81ea8f526ba69104da8d01a08da875f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 21:31:46 -0700 Subject: [PATCH 3/3] Fix FP32 in-place KV cache append --- .../cuda/bert/attention_kv_cache.cu | 9 ++++-- .../multihead_attention_op_test.cc | 29 +++++++++++++++++-- 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_kv_cache.cu b/onnxruntime/contrib_ops/cuda/bert/attention_kv_cache.cu index f878f6794fa31..8670e78e48789 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_kv_cache.cu +++ b/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); diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index 75d6b4303f347..cc716acc3fc66 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -850,8 +850,33 @@ TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_UsingDMMHAInsi TEST(MultiHeadAttentionTest, SelfAttention_PastPresentBufferShare_ConcatKVInPlace) { AttentionTestData data; - GetSelfAttention_PastPresentBufferShare_UsingDMMHAInsideMHA(data); - data.cache_indir_data.clear(); + 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); }