Fix Flash and Lean attention split heuristics

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 <algorithm>
 #include <cutlass/numeric_types.h>
 #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<size_t>(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<float> efficiency;
   efficiency.reserve(max_splits);

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 <algorithm>
 #include <cutlass/numeric_types.h>
 
 #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<size_t, size_t, size_t, size_t, size_t, size_t, size_t, size_t> 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<size_t>(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<size_t, size_t, size_t, size_t, size_t, size_t, size_t, size_t> 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<size_t, size_t, size_t, size_t, size_t, size_t, size_t, size_t> 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<size_t>(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);
 

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 <cstddef>
+#include <tuple>
+
+#if defined(USE_FLASH_ATTENTION)
+namespace onnxruntime {
+namespace flash {
+std::tuple<size_t, size_t, size_t> 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<size_t, size_t, size_t, size_t, size_t, size_t, size_t, size_t>
+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