From 82915bd4864d37d001d18c2c490a0dfb7545861f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 24 Aug 2026 11:17:44 -0700 Subject: [PATCH 1/4] [CUDA] Make adaptive FlashDecode splits CUDA-graph safe ## Summary Plan contiguous GQA FlashDecode split-KV launches from fixed KV-cache capacity during CUDA graph capture and replay, while retaining live-sequence-length planning for ordinary eager execution. ## Why CUDA graphs freeze launch geometry and workspace addresses at capture time. Planning NumSplits from the current live sequence length can become stale as the cache grows, preventing adaptive split-KV behavior from remaining valid across replay. Graph-enabled warmup now reserves capacity-sized workspace, capture uses the fixed-capacity plan, and eager decode avoids redundant capacity heuristic work. ## Behavior - Capture/replay uses fixed cache capacity for stable NumSplits and workspace sizing - Graph warmup reserves replay-sized workspace before capture - Eager execution continues to tune from the live sequence length - Active memset size remains limited to the launch plan - Debug output reports the resolved NumSplits ## Validation Focused host tests cover head sizes 64, 128, and 256; local-window and sequence-tail behavior; non-decode inputs; capture planning; and ordinary eager routing. For an SM108 configuration with live length 129 and capacity 4097, capture selected 17/17/22 splits for head sizes 64/128/256, while eager retained live-length plans. Independent review found one redundant eager heuristic computation, which is fixed in this commit. No CUDA kernel timing is claimed because this Windows host did not have nvcc. The change preserves eager routing and targets graph planning correctness and replay-stable adaptive split-KV behavior. Based on the mechanisms validated in justinchuby/onnx-genai#1340 and #1350. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/group_query_attention.cc | 112 ++++++++++++++---- .../cuda/bert/group_query_attention.h | 1 + .../cuda/bert/group_query_attention_impl.h | 23 ++++ .../attention_split_heuristic_test.cc | 88 ++++++++++++++ 4 files changed, 202 insertions(+), 22 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index c6a64a75b7a5d..4caa11c50cf62 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -48,6 +48,68 @@ KVQuantizationType StringToKVQuantizationType(std::string s) { } } // namespace +#if USE_FLASH_ATTENTION +FlashAttentionSplitPlan GetFlashAttentionSplitPlan( + int batch_size, + int sequence_length, + int total_sequence_length, + int cache_capacity, + int num_heads, + int kv_num_heads, + int head_size, + int multi_processor_count, + int local_window_size, + bool use_fast_decode, + bool reserve_for_capture, + bool is_capturing) { + auto get_plan = [&](size_t kv_length, size_t heads_for_split) { + return onnxruntime::flash::get_num_splits_and_buffer_sizes( + static_cast(batch_size), static_cast(sequence_length), kv_length, + heads_for_split, static_cast(head_size), static_cast(multi_processor_count)); + }; + + auto effective_kv_length = [&](int kv_length) { + size_t length = static_cast(std::max(kv_length, 0)); + if (use_fast_decode && local_window_size > 0) { + length = std::min(length, static_cast(local_window_size)); + } + return length; + }; + + const size_t heads_for_split = static_cast(use_fast_decode ? kv_num_heads : num_heads); + const auto live_plan = get_plan(effective_kv_length(total_sequence_length), heads_for_split); + const size_t live_num_splits = std::get<0>(live_plan); + if (!use_fast_decode) { + return {live_num_splits, std::get<1>(live_plan), std::get<2>(live_plan)}; + } + + // Warm-up must reserve the replay-invariant workspace before capture starts, while eager launches + // continue using the live-length split count to avoid empty-split overhead at short context. + const bool needs_capture_workspace = reserve_for_capture || is_capturing; + size_t capture_num_splits = live_num_splits; + if (needs_capture_workspace) { + const auto capture_plan = get_plan(effective_kv_length(cache_capacity), heads_for_split); + capture_num_splits = std::get<0>(capture_plan); + } + + const size_t num_splits = is_capturing ? capture_num_splits : live_num_splits; + const size_t workspace_num_splits = + needs_capture_workspace ? std::max(live_num_splits, capture_num_splits) : live_num_splits; + if (workspace_num_splits <= 1) { + return {num_splits, 0, 0}; + } + + const size_t lse_bytes = onnxruntime::flash::get_softmax_lse_accum_size( + workspace_num_splits, static_cast(batch_size), static_cast(num_heads), + static_cast(sequence_length)); + const size_t rounded_head_size = (static_cast(head_size) + 31) / 32 * 32; + const size_t out_bytes = onnxruntime::flash::get_out_accum_size( + workspace_num_splits, static_cast(batch_size), static_cast(num_heads), + static_cast(sequence_length), rounded_head_size); + return {num_splits, lse_bytes, out_bytes}; +} +#endif + #define REGISTER_KERNEL_TYPED(T, U) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ GroupQueryAttention, \ @@ -94,6 +156,7 @@ constexpr int kHeadSinkInputIndex = 11; template GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) : CudaKernel(info) { + enable_cuda_graph_ = info.GetExecutionProvider()->IsGraphCaptureEnabled(); int64_t num_heads = 0; int64_t kv_num_heads = 0; ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); @@ -705,24 +768,25 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons // Allocate Flash specific buffers (Softmax LSE, Accum) size_t softmax_lse_bytes = onnxruntime::flash::get_softmax_lse_size(parameters.sequence_length, parameters.batch_size, parameters.num_heads); - int num_heads_for_split = data.use_flash_attention_fast_decode ? parameters.kv_num_heads : parameters.num_heads; - size_t sequence_length_for_split = static_cast(parameters.total_sequence_length); - if (data.use_flash_attention_fast_decode && parameters.local_window_size > 0) { - sequence_length_for_split = std::min(sequence_length_for_split, static_cast(parameters.local_window_size)); - } - - auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = onnxruntime::flash::get_num_splits_and_buffer_sizes( - parameters.batch_size, parameters.sequence_length, sequence_length_for_split, num_heads_for_split, - parameters.head_size, device_prop.multiProcessorCount); - - parameters.num_splits = static_cast(num_splits); - - if (data.use_flash_attention_fast_decode && num_splits > 1) { - // The heuristic used kv_num_heads to maximize occupancy for the GQA-aware kernel. - // However, the LSE and Accum buffers must store results for ALL num_heads. - softmax_lse_accum_bytes = onnxruntime::flash::get_softmax_lse_accum_size(num_splits, parameters.batch_size, parameters.num_heads, parameters.sequence_length); - auto round_multiple = [](size_t x, size_t m) { return (x + m - 1) / m * m; }; - out_accum_bytes = onnxruntime::flash::get_out_accum_size(num_splits, parameters.batch_size, parameters.num_heads, parameters.sequence_length, round_multiple(parameters.head_size, 32)); + const auto split_plan = GetFlashAttentionSplitPlan( + parameters.batch_size, parameters.sequence_length, parameters.total_sequence_length, + parameters.seqlen_present_kv_cache, parameters.num_heads, parameters.kv_num_heads, + parameters.head_size, device_prop.multiProcessorCount, parameters.local_window_size, + data.use_flash_attention_fast_decode, + enable_cuda_graph_, + onnxruntime::llm::common::isCapturing(Stream(context))); + parameters.num_splits = static_cast(split_plan.num_splits); + size_t softmax_lse_accum_bytes = split_plan.softmax_lse_accum_bytes; + size_t out_accum_bytes = split_plan.out_accum_bytes; + size_t active_softmax_lse_accum_bytes = 0; + size_t active_out_accum_bytes = 0; + if (parameters.num_splits > 1) { + active_softmax_lse_accum_bytes = onnxruntime::flash::get_softmax_lse_accum_size( + split_plan.num_splits, parameters.batch_size, parameters.num_heads, parameters.sequence_length); + const size_t rounded_head_size = (static_cast(parameters.head_size) + 31) / 32 * 32; + active_out_accum_bytes = onnxruntime::flash::get_out_accum_size( + split_plan.num_splits, parameters.batch_size, parameters.num_heads, + parameters.sequence_length, rounded_head_size); } softmax_lse_buffer = GetScratchBuffer(softmax_lse_bytes, GetComputeStream(context)); @@ -730,13 +794,14 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons out_accum_buffer = GetScratchBuffer(out_accum_bytes, GetComputeStream(context)); auto cuda_stream = Stream(context); - if (softmax_lse_accum_bytes > 0) { + if (active_softmax_lse_accum_bytes > 0) { // Initialize to 0 is fine because Flash kernel will write -inf to it if needed. // However, the standard Flash kernel often doesn't zero it globally. - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(softmax_lse_accum_buffer.get(), 0, softmax_lse_accum_bytes, cuda_stream)); + CUDA_RETURN_IF_ERROR( + cudaMemsetAsync(softmax_lse_accum_buffer.get(), 0, active_softmax_lse_accum_bytes, cuda_stream)); } - if (out_accum_bytes > 0) { - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(out_accum_buffer.get(), 0, out_accum_bytes, cuda_stream)); + if (active_out_accum_bytes > 0) { + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(out_accum_buffer.get(), 0, active_out_accum_bytes, cuda_stream)); } data.softmax_lse = reinterpret_cast(softmax_lse_buffer.get()); @@ -896,6 +961,9 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons debug_info.use_flash_attention = data.use_flash_attention; debug_info.use_efficient_attention = data.use_memory_efficient_attention; debug_info.use_cudnn_flash_attention = data.use_cudnn_sdpa; + if (data.use_flash_attention) { + debug_info.num_splits = std::max(1, parameters.num_splits); + } debug_info.Print("GroupQueryAttention", this->Node().Name(), diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index 8cdcedc3726f5..3212b3ad02a38 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -46,6 +46,7 @@ class GroupQueryAttention final : public CudaKernel { bool enable_xqa_; // True when ORT_ENABLE_XQA != 0 (default: on) and T is fp16/bf16. bool enable_cudnn_flash_attention_; // cuDNN SDPA explicitly enabled (env / sdpa_kernel) bool auto_enable_cudnn_flash_attention_; // auto-prefer cuDNN SDPA on SM>=90 when no explicit kernel pinned + bool enable_cuda_graph_{false}; KVQuantizationType k_quant_type_; KVQuantizationType v_quant_type_; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h index 59858c83dc171..f9c686c8fc181 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h @@ -2,6 +2,9 @@ // Licensed under the MIT License. #pragma once + +#include + #include "core/providers/cuda/shared_inc/cuda_utils.h" #include #include @@ -15,6 +18,26 @@ namespace onnxruntime { namespace contrib { namespace cuda { +struct FlashAttentionSplitPlan { + size_t num_splits = 0; + size_t softmax_lse_accum_bytes = 0; + size_t out_accum_bytes = 0; +}; + +FlashAttentionSplitPlan GetFlashAttentionSplitPlan( + int batch_size, + int sequence_length, + int total_sequence_length, + int cache_capacity, + int num_heads, + int kv_num_heads, + int head_size, + int multi_processor_count, + int local_window_size, + bool use_fast_decode, + bool reserve_for_capture, + bool is_capturing); + template Status QkvToContext( const cudaDeviceProp& device_prop, diff --git a/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc b/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc index a289010a957e3..4e22dfe8ea11a 100644 --- a/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc @@ -7,6 +7,8 @@ #include #if defined(USE_FLASH_ATTENTION) +#include "contrib_ops/cuda/bert/group_query_attention_impl.h" + namespace onnxruntime { namespace flash { std::tuple get_num_splits_and_buffer_sizes(size_t batch_size, size_t seqlen_q, @@ -106,6 +108,92 @@ TEST(FlashAttentionTest, GetNumSplitsAvoidsShortContextOverhead) { #endif } +TEST(FlashAttentionTest, GqaCaptureSplitPlanUsesCacheCapacity) { +#if defined(USE_FLASH_ATTENTION) + constexpr int batch_size = 1; + constexpr int sequence_length = 1; + constexpr int total_sequence_length = 129; + constexpr int cache_capacity = 4097; + constexpr int num_heads = 32; + constexpr int kv_num_heads = 4; + constexpr int multi_processor_count = 108; + + for (int head_size : {64, 128, 256}) { + const auto eager_plan = ::onnxruntime::contrib::cuda::GetFlashAttentionSplitPlan( + batch_size, sequence_length, total_sequence_length, cache_capacity, + num_heads, kv_num_heads, head_size, multi_processor_count, + -1, true, true, false); + const auto capture_plan = ::onnxruntime::contrib::cuda::GetFlashAttentionSplitPlan( + batch_size, sequence_length, total_sequence_length, cache_capacity, + num_heads, kv_num_heads, head_size, multi_processor_count, + -1, true, true, true); + const auto expected_eager_plan = flash::get_num_splits_and_buffer_sizes( + batch_size, sequence_length, total_sequence_length, kv_num_heads, + head_size, multi_processor_count); + const auto expected_capture_plan = flash::get_num_splits_and_buffer_sizes( + batch_size, sequence_length, cache_capacity, kv_num_heads, + head_size, multi_processor_count); + + EXPECT_EQ(eager_plan.num_splits, std::get<0>(expected_eager_plan)); + EXPECT_EQ(capture_plan.num_splits, std::get<0>(expected_capture_plan)); + EXPECT_GT(capture_plan.num_splits, eager_plan.num_splits); + EXPECT_EQ(eager_plan.softmax_lse_accum_bytes, capture_plan.softmax_lse_accum_bytes); + EXPECT_EQ(eager_plan.out_accum_bytes, capture_plan.out_accum_bytes); + EXPECT_GT(capture_plan.softmax_lse_accum_bytes, 0U); + EXPECT_GT(capture_plan.out_accum_bytes, 0U); + } +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(FlashAttentionTest, GqaCaptureSplitPlanHonorsLocalWindowAndSequenceTail) { +#if defined(USE_FLASH_ATTENTION) + constexpr int local_window_size = 257; + const auto capture_plan = ::onnxruntime::contrib::cuda::GetFlashAttentionSplitPlan( + 1, 1, 4097, 8193, 32, 4, 128, 108, + local_window_size, true, true, true); + const auto expected_plan = flash::get_num_splits_and_buffer_sizes( + 1, 1, local_window_size, 4, 128, 108); + + EXPECT_EQ(capture_plan.num_splits, std::get<0>(expected_plan)); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(FlashAttentionTest, NonDecodeSplitPlanKeepsLiveSequenceLength) { +#if defined(USE_FLASH_ATTENTION) + constexpr int total_sequence_length = 257; + const auto capture_plan = ::onnxruntime::contrib::cuda::GetFlashAttentionSplitPlan( + 1, 4, total_sequence_length, 8193, 32, 4, 64, 108, + -1, false, true, true); + const auto [expected_splits, expected_lse_bytes, expected_out_bytes] = + flash::get_num_splits_and_buffer_sizes( + 1, 4, total_sequence_length, 32, 64, 108); + + EXPECT_EQ(capture_plan.num_splits, expected_splits); + EXPECT_EQ(capture_plan.softmax_lse_accum_bytes, expected_lse_bytes); + EXPECT_EQ(capture_plan.out_accum_bytes, expected_out_bytes); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(FlashAttentionTest, GqaEagerSplitPlanDoesNotReserveCaptureWorkspace) { +#if defined(USE_FLASH_ATTENTION) + const auto eager_plan = ::onnxruntime::contrib::cuda::GetFlashAttentionSplitPlan( + 1, 1, 129, 4097, 32, 4, 64, 108, + -1, true, false, false); + + EXPECT_EQ(eager_plan.num_splits, 0U); + EXPECT_EQ(eager_plan.softmax_lse_accum_bytes, 0U); + EXPECT_EQ(eager_plan.out_accum_bytes, 0U); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + TEST(LeanAttentionTest, GetNumSplitsHandlesZeroSmCount) { #if defined(USE_LEAN_ATTENTION) const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes, sync_flag_bytes, From faf9be4ba34cd2c198f0c4b4aea75a1875f32f67 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 31 Aug 2026 17:03:43 -0700 Subject: [PATCH 2/4] [CUDA] Cover FlashDecode CUDA graph replay Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/attention_data.h | 1 + .../group_query_attention_op_test.cc | 226 +++++++++++++++++- 2 files changed, 225 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 9cb66f383bd36..c4224cebede9a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -5,6 +5,7 @@ #include #include +#include "core/framework/allocator.h" #include "contrib_ops/cpu/bert/attention_common.h" #include "contrib_ops/cpu/bert/attention_parameters.h" diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index eed2138578c7b..63a96eea80f02 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -5,26 +5,36 @@ #include #include #include +#include #include #include +#include #include #include +#include #include #include "gtest/gtest.h" +#include "contrib_ops/cpu/bert/attention_common.h" #include "core/platform/env.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" #include "test/util/include/scoped_env_vars.h" -#ifdef USE_WEBGPU +#if defined(USE_CUDA) || defined(USE_WEBGPU) #include "core/graph/model.h" -#include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/inference_session.h" #include "core/session/IOBinding.h" #include "test/test_environment.h" #include "test/unittest_util/framework_test_utils.h" #endif +#ifdef USE_CUDA +#include "core/providers/cuda/cuda_provider_options.h" +#include "test/common/cuda_op_test_utils.h" +#endif +#ifdef USE_WEBGPU +#include "core/providers/webgpu/webgpu_provider_options.h" +#endif namespace onnxruntime { namespace test { @@ -3127,6 +3137,218 @@ TEST(GroupQueryAttentionTest, CudaAttentionBiasParityVsCpu) { } } +#if defined(USE_CUDA) && defined(USE_FLASH_ATTENTION) +TEST(GroupQueryAttentionTest, FlashDecodeSplitKvCudaGraphReplay) { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableCudnnFlashAttention, "0"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}, + {"ORT_ENABLE_XQA", "0"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + constexpr int batch_size = 1; + constexpr int sequence_length = 1; + constexpr int num_heads = 32; + constexpr int kv_num_heads = 4; + constexpr int head_size = 64; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int cache_capacity = 4097; + + std::unordered_map domain_to_version = {{kMSDomain, 1}}; + auto model = std::make_unique( + "gqa_flash_decode_graph_replay", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector{}, + DefaultLoggingManager().DefaultLogger(), ModelOptions(true, true)); + Graph& graph = model->MainGraph(); + + ONNX_NAMESPACE::TypeProto fp16_type; + fp16_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + ONNX_NAMESPACE::TypeProto int32_type; + int32_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + auto& empty_arg = graph.GetOrCreateNodeArg("", nullptr); + std::vector input_defs = { + &graph.GetOrCreateNodeArg("query", &fp16_type), + &graph.GetOrCreateNodeArg("key", &fp16_type), + &graph.GetOrCreateNodeArg("value", &fp16_type), + &graph.GetOrCreateNodeArg("past_key", &fp16_type), + &graph.GetOrCreateNodeArg("past_value", &fp16_type), + &graph.GetOrCreateNodeArg("seqlens_k", &int32_type), + &graph.GetOrCreateNodeArg("total_sequence_length", &int32_type), + &empty_arg, // cos_cache + &empty_arg, // sin_cache + &empty_arg, // position_ids + &empty_arg, // attention_bias + &empty_arg, // head_sink + }; + std::vector output_defs = { + &graph.GetOrCreateNodeArg("output", &fp16_type), + &graph.GetOrCreateNodeArg("present_key", &fp16_type), + &graph.GetOrCreateNodeArg("present_value", &fp16_type), + }; + auto& node = graph.AddNode("gqa", "GroupQueryAttention", "GQA FlashDecode graph test", + input_defs, output_defs, nullptr, kMSDomain); + node.AddAttribute("num_heads", int64_t{num_heads}); + node.AddAttribute("kv_num_heads", int64_t{kv_num_heads}); + ASSERT_STATUS_OK(graph.Resolve()); + + std::string model_data; + ASSERT_TRUE(model->ToProto().SerializeToString(&model_data)); + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.do_copy_in_default_stream = true; + provider_options.use_tf32 = false; + provider_options.enable_cuda_graph = true; + provider_options.sdpa_kernel = static_cast(contrib::attention::AttentionBackend::FLASH_ATTENTION); + auto cuda_ep = CudaExecutionProviderWithOptions(&provider_options); + ASSERT_NE(cuda_ep, nullptr); + IExecutionProvider* ep = cuda_ep.get(); + + SessionOptions session_options; + InferenceSession session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(cuda_ep))); + std::istringstream model_stream(model_data); + ASSERT_STATUS_OK(session.Load(model_stream)); + ASSERT_STATUS_OK(session.Initialize()); + + const auto device_allocators = ep->CreatePreferredAllocators(); + const OrtMemoryInfo* device_memory_info = nullptr; + for (const auto& allocator : device_allocators) { + if (allocator->Info().device.Type() == OrtDevice::GPU && + allocator->Info().mem_type == OrtMemTypeDefault) { + device_memory_info = &allocator->Info(); + break; + } + } + ASSERT_NE(device_memory_info, nullptr); + auto gpu_allocator = session.GetAllocator(*device_memory_info); + ASSERT_NE(gpu_allocator, nullptr); + auto cpu_allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + auto make_gpu_value = [&](const auto& data, const TensorShape& shape) { + using T = typename std::decay_t::value_type; + Tensor cpu_tensor(DataTypeImpl::GetType(), shape, const_cast(data.data()), cpu_allocator->Info()); + Tensor gpu_tensor(DataTypeImpl::GetType(), shape, gpu_allocator); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue value; + Tensor::InitOrtValue(std::move(gpu_tensor), value); + return value; + }; + + std::vector query_data(batch_size * sequence_length * hidden_size, MLFloat16(0.0f)); + std::vector key_data(batch_size * sequence_length * kv_hidden_size, MLFloat16(0.0f)); + std::vector value_data(batch_size * sequence_length * kv_hidden_size, MLFloat16(0.75f)); + std::vector past_key_data(batch_size * kv_num_heads * cache_capacity * head_size, MLFloat16(0.0f)); + std::vector past_value_data(batch_size * kv_num_heads * cache_capacity * head_size); + for (int s = 0; s < cache_capacity; ++s) { + for (int h = 0; h < kv_num_heads; ++h) { + for (int d = 0; d < head_size; ++d) { + past_value_data[(s * kv_num_heads + h) * head_size + d] = + MLFloat16(0.01f + 0.0001f * static_cast((s + h + d) % cache_capacity)); + } + } + } + std::vector seqlens_data = {128}; + std::vector total_sequence_length_data = {cache_capacity}; + + auto query_value = make_gpu_value(query_data, TensorShape{batch_size, sequence_length, hidden_size}); + auto key_value = make_gpu_value(key_data, TensorShape{batch_size, sequence_length, kv_hidden_size}); + auto value_value = make_gpu_value(value_data, TensorShape{batch_size, sequence_length, kv_hidden_size}); + auto past_key_value = make_gpu_value( + past_key_data, TensorShape{batch_size, kv_num_heads, cache_capacity, head_size}); + auto past_value_value = make_gpu_value( + past_value_data, TensorShape{batch_size, kv_num_heads, cache_capacity, head_size}); + auto seqlens_value = make_gpu_value(seqlens_data, TensorShape{batch_size}); + + Tensor total_sequence_length_tensor( + DataTypeImpl::GetType(), TensorShape{1}, + total_sequence_length_data.data(), cpu_allocator->Info()); + OrtValue total_sequence_length_value; + Tensor::InitOrtValue(std::move(total_sequence_length_tensor), total_sequence_length_value); + Tensor output_tensor( + DataTypeImpl::GetType(), TensorShape{batch_size, sequence_length, hidden_size}, gpu_allocator); + OrtValue output_value; + Tensor::InitOrtValue(std::move(output_tensor), output_value); + + std::unique_ptr io_binding; + ASSERT_STATUS_OK(session.NewIOBinding(&io_binding)); + ASSERT_STATUS_OK(io_binding->BindInput("query", query_value)); + ASSERT_STATUS_OK(io_binding->BindInput("key", key_value)); + ASSERT_STATUS_OK(io_binding->BindInput("value", value_value)); + ASSERT_STATUS_OK(io_binding->BindInput("past_key", past_key_value)); + ASSERT_STATUS_OK(io_binding->BindInput("past_value", past_value_value)); + ASSERT_STATUS_OK(io_binding->BindInput("seqlens_k", seqlens_value)); + ASSERT_STATUS_OK(io_binding->BindInput("total_sequence_length", total_sequence_length_value)); + ASSERT_STATUS_OK(io_binding->BindOutput("output", output_value)); + ASSERT_STATUS_OK(io_binding->BindOutput("present_key", past_key_value)); + ASSERT_STATUS_OK(io_binding->BindOutput("present_value", past_value_value)); + ASSERT_STATUS_OK(io_binding->SynchronizeInputs()); + + auto update_seqlens = [&](int live_sequence_length) { + seqlens_data[0] = live_sequence_length - 1; + Tensor cpu_tensor(DataTypeImpl::GetType(), TensorShape{batch_size}, + seqlens_data.data(), cpu_allocator->Info()); + ORT_THROW_IF_ERROR( + ep->GetDataTransfer()->CopyTensor(cpu_tensor, *seqlens_value.GetMutable())); + }; + auto verify_output = [&](int live_sequence_length, size_t run_index) { + const int appended_position = live_sequence_length - 1; + for (int h = 0; h < kv_num_heads; ++h) { + for (int d = 0; d < head_size; ++d) { + past_value_data[(appended_position * kv_num_heads + h) * head_size + d] = + value_data[h * head_size + d]; + } + } + + Tensor cpu_output(DataTypeImpl::GetType(), + TensorShape{batch_size, sequence_length, hidden_size}, cpu_allocator); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor(output_value.Get(), cpu_output)); + for (int h = 0; h < num_heads; ++h) { + const int kv_head = h * kv_num_heads / num_heads; + for (int d = 0; d < head_size; ++d) { + float expected = 0.0f; + for (int s = 0; s < live_sequence_length; ++s) { + expected += past_value_data[(s * kv_num_heads + kv_head) * head_size + d].ToFloat(); + } + expected /= static_cast(live_sequence_length); + EXPECT_NEAR(cpu_output.Data()[h * head_size + d].ToFloat(), expected, 2e-3f) + << "run=" << run_index << ", head=" << h << ", dim=" << d; + } + } + }; + + RunOptions run_options; + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry("gpu_graph_id", "1")); + // The CUDA EP captures after two warm-up runs. Replay then grows the live KV length. + const std::vector live_sequence_lengths = {129, 129, 129, 2049, cache_capacity}; + testing::internal::CaptureStdout(); + for (size_t run_index = 0; run_index < live_sequence_lengths.size(); ++run_index) { + const int live_sequence_length = live_sequence_lengths[run_index]; + update_seqlens(live_sequence_length); + ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + verify_output(live_sequence_length, run_index); + if (run_index == 2) { + EXPECT_TRUE(ep->IsGraphCaptured(1)); + } + } + const std::string debug_output = testing::internal::GetCapturedStdout(); + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; + const size_t split_pos = debug_output.rfind("NumSplits="); + ASSERT_NE(split_pos, std::string::npos) << debug_output; + EXPECT_GT(std::stoi(debug_output.substr(split_pos + std::string("NumSplits=").size())), 1) << debug_output; +} +#endif + #ifdef USE_WEBGPU // WebGPU graph capture test for kv_empty (Gemma4 shared-KV) layers. // From 8e2d9dbcef7419534915bca84e820abbfd8bdeb1 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 14 Sep 2026 11:38:22 -0700 Subject: [PATCH 3/4] [CUDA] Fix BNSH indexing in FlashDecode graph replay test Use head-major cache offsets for initialization, appended-value mirroring, and the output reference to match the bound BNSH tensor layout. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/group_query_attention_op_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 3e0987cc9eb74..87ff7d04f5397 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -3957,7 +3957,7 @@ TEST(GroupQueryAttentionTest, FlashDecodeSplitKvCudaGraphReplay) { for (int s = 0; s < cache_capacity; ++s) { for (int h = 0; h < kv_num_heads; ++h) { for (int d = 0; d < head_size; ++d) { - past_value_data[(s * kv_num_heads + h) * head_size + d] = + past_value_data[(h * cache_capacity + s) * head_size + d] = MLFloat16(0.01f + 0.0001f * static_cast((s + h + d) % cache_capacity)); } } @@ -4009,7 +4009,7 @@ TEST(GroupQueryAttentionTest, FlashDecodeSplitKvCudaGraphReplay) { const int appended_position = live_sequence_length - 1; for (int h = 0; h < kv_num_heads; ++h) { for (int d = 0; d < head_size; ++d) { - past_value_data[(appended_position * kv_num_heads + h) * head_size + d] = + past_value_data[(h * cache_capacity + appended_position) * head_size + d] = value_data[h * head_size + d]; } } @@ -4022,7 +4022,7 @@ TEST(GroupQueryAttentionTest, FlashDecodeSplitKvCudaGraphReplay) { for (int d = 0; d < head_size; ++d) { float expected = 0.0f; for (int s = 0; s < live_sequence_length; ++s) { - expected += past_value_data[(s * kv_num_heads + kv_head) * head_size + d].ToFloat(); + expected += past_value_data[(kv_head * cache_capacity + s) * head_size + d].ToFloat(); } expected /= static_cast(live_sequence_length); EXPECT_NEAR(cpu_output.Data()[h * head_size + d].ToFloat(), expected, 2e-3f) From 3341e351d8badb9410aa7e028d5cf46045158494 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 14 Sep 2026 11:40:04 -0700 Subject: [PATCH 4/4] [CUDA] Preserve graph replay test includes after merge Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 87ff7d04f5397..62b7510d221a5 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -7,20 +7,24 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include "gtest/gtest.h" +#include "contrib_ops/cpu/bert/attention_common.h" #include "core/platform/env.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" #include "test/util/include/scoped_env_vars.h" #ifdef USE_CUDA +#include "core/providers/cuda/cuda_provider_options.h" #include "test/common/cuda_op_test_utils.h" #endif #if defined(USE_CUDA) || defined(USE_WEBGPU)