From e91df8f3bb74b352b985532ddba4f5b4f30f4d17 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 5 Jul 2026 19:17:35 +0000 Subject: [PATCH 1/2] fix github issue 29550 --- .../contrib_ops/cuda/bert/flash_attention/flash_api.cc | 5 +++++ .../contrib_ops/cuda/bert/lean_attention/lean_api.cc | 10 +++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc index 0d994d4060e6c..efda3f48b9cfc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc @@ -5,6 +5,7 @@ #if USE_FLASH_ATTENTION #include "contrib_ops/cuda/bert/flash_attention/flash_api.h" +#include #include #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/bert/flash_attention/flash.h" @@ -198,6 +199,7 @@ void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream, bool force_split size_t num_splits_heuristic(size_t batch_size, size_t seqlen_q, size_t seqlen_k, size_t num_heads, size_t head_size, size_t num_SMs, size_t max_splits) { // This needs to match with run_mha_fwd_splitkv_dispatch + num_SMs = std::max(num_SMs, 1); const size_t block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); const size_t num_n_blocks = (seqlen_k + block_n - 1) / block_n; // Technically kBlockM = 64 only for the splitKV kernels, not the standard kernel. @@ -209,6 +211,9 @@ size_t num_splits_heuristic(size_t batch_size, size_t seqlen_q, size_t seqlen_k, return 1; } max_splits = std::min({max_splits, num_SMs, num_n_blocks}); + if (max_splits <= 1) { + return 1; + } float max_efficiency = 0.f; std::vector efficiency; efficiency.reserve(max_splits); diff --git a/onnxruntime/contrib_ops/cuda/bert/lean_attention/lean_api.cc b/onnxruntime/contrib_ops/cuda/bert/lean_attention/lean_api.cc index 81301ebc7ba64..1f6e9ac9ab1b2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/lean_attention/lean_api.cc +++ b/onnxruntime/contrib_ops/cuda/bert/lean_attention/lean_api.cc @@ -9,6 +9,7 @@ #if USE_LEAN_ATTENTION #include "contrib_ops/cuda/bert/lean_attention/lean_api.h" +#include #include #include "contrib_ops/cuda/bert/lean_attention/flash.h" @@ -173,6 +174,7 @@ void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream) { std::tuple get_num_splits_and_buffer_sizes(size_t batch_size, size_t max_seqlen_q, size_t max_seqlen_k, size_t num_heads, size_t num_heads_k, size_t head_size, size_t num_SMs, bool is_causal) { // This needs to match with run_mha_fwd_splitkv_dispatch + num_SMs = std::max(num_SMs, 1); const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); const int block_m = head_size <= 64 ? 64 : (head_size <= 128 ? 64 : 64); const int num_m_blocks = (max_seqlen_q + block_m - 1) / block_m; @@ -207,6 +209,9 @@ std::tuple get_n tiles_per_head = num_m_blocks * num_n_blocks; } size_t total_tiles = tiles_per_head * batch_size * num_heads_k; + if (total_tiles == 0 || num_n_blocks == 0) { + return {0, 0, 0, 0, 1, 1, 0, tiles_per_head}; + } // StreamK Lean has as many threadblocks as SMs // This should be a function of tile size and number of scratchpad space @@ -222,13 +227,16 @@ std::tuple get_n // to account for ceil lean_griddimz = std::min(2 * num_SMs, 32 * num_heads_k * batch_size * num_m_blocks); } + lean_griddimz = std::max(1, std::min(lean_griddimz, total_tiles)); size_t max_tiles_per_tb = (total_tiles + lean_griddimz - 1) / lean_griddimz; // Find max number of splits size_t num_splits = 0; if (total_tiles % lean_griddimz == 0) { num_splits = 1 + ((num_n_blocks + max_tiles_per_tb - 2) / (max_tiles_per_tb)); - } else { + } else if (max_tiles_per_tb > 1) { num_splits = 1 + ((num_n_blocks + max_tiles_per_tb - 3) / (max_tiles_per_tb - 1)); + } else { + num_splits = 1; } size_t high_load_tbs = total_tiles - ((max_tiles_per_tb - 1) * lean_griddimz); From 74e216557e0e21556137bc336b4fa5b41fe961a1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 5 Jul 2026 21:53:18 +0000 Subject: [PATCH 2/2] add test --- .../attention_split_heuristic_test.cc | 127 ++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc 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 new file mode 100644 index 0000000000000..78293c3d2bbad --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc @@ -0,0 +1,127 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include +#include + +#if defined(USE_FLASH_ATTENTION) +namespace onnxruntime { +namespace flash { +std::tuple get_num_splits_and_buffer_sizes(size_t batch_size, size_t seqlen_q, + size_t seqlen_k, size_t num_heads, + size_t head_size, size_t num_SMs); +} // namespace flash +} // namespace onnxruntime +#endif + +#if defined(USE_LEAN_ATTENTION) +namespace onnxruntime { +namespace lean { +std::tuple +get_num_splits_and_buffer_sizes(size_t batch_size, size_t seqlen_q, size_t seqlen_k, size_t num_heads, + size_t num_heads_k, size_t head_size, size_t num_SMs, bool is_causal); +} // namespace lean +} // namespace onnxruntime +#endif + +namespace onnxruntime { +namespace cuda { +namespace test { + +TEST(FlashAttentionTest, GetNumSplitsHandlesZeroSmCount) { +#if defined(USE_FLASH_ATTENTION) + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = + flash::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 384, // seqlen_k: 3 N-blocks when head_size is 128 + 1, // num_heads + 128, // head_size + 0); // num_SMs: regression coverage for divide-by-zero in PR #29550 + + EXPECT_EQ(num_splits, 0U); + EXPECT_EQ(softmax_lse_accum_bytes, 0U); + EXPECT_EQ(out_accum_bytes, 0U); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(FlashAttentionTest, GetNumSplitsHandlesZeroKeyTiles) { +#if defined(USE_FLASH_ATTENTION) + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = + flash::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 0, // seqlen_k: no N-blocks + 1, // num_heads + 128, // head_size + 2); // num_SMs + + EXPECT_EQ(num_splits, 0U); + EXPECT_EQ(softmax_lse_accum_bytes, 0U); + EXPECT_EQ(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, + grid_dim_z, max_tiles_per_tb, high_load_tbs, tiles_per_head] = + lean::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 384, // seqlen_k: 3 N-blocks when head_size is 128 + 1, // num_heads + 1, // num_heads_k + 128, // head_size + 0, // num_SMs: regression coverage for divide-by-zero in PR #29550 + true); + + EXPECT_EQ(num_splits, 3U); + EXPECT_EQ(softmax_lse_accum_bytes, 12U); + EXPECT_EQ(out_accum_bytes, 1536U); + EXPECT_EQ(sync_flag_bytes, 4U); + EXPECT_EQ(grid_dim_z, 2U); + EXPECT_EQ(max_tiles_per_tb, 2U); + EXPECT_EQ(high_load_tbs, 1U); + EXPECT_EQ(tiles_per_head, 3U); +#else + GTEST_SKIP() << "Lean Attention is not enabled in this build."; +#endif +} + +TEST(LeanAttentionTest, GetNumSplitsHandlesZeroKeyTiles) { +#if defined(USE_LEAN_ATTENTION) + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes, sync_flag_bytes, + grid_dim_z, max_tiles_per_tb, high_load_tbs, tiles_per_head] = + lean::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 0, // seqlen_k: no N-blocks + 1, // num_heads + 1, // num_heads_k + 128, // head_size + 2, // num_SMs + true); + + EXPECT_EQ(num_splits, 0U); + EXPECT_EQ(softmax_lse_accum_bytes, 0U); + EXPECT_EQ(out_accum_bytes, 0U); + EXPECT_EQ(sync_flag_bytes, 0U); + EXPECT_EQ(grid_dim_z, 1U); + EXPECT_EQ(max_tiles_per_tb, 1U); + EXPECT_EQ(high_load_tbs, 0U); + EXPECT_EQ(tiles_per_head, 0U); +#else + GTEST_SKIP() << "Lean Attention is not enabled in this build."; +#endif +} + +} // namespace test +} // namespace cuda +} // namespace onnxruntime