Merge branch 'Dao-AILab:main' into main

Author: sstamenk

csrc/flash_attn/flash_api.cpp

@@ -698,11 +698,14 @@ mha_varlen_fwd(at::Tensor &q,  // total_q x num_heads x head_size, total_q := \s
     params.page_block_size = page_block_size;
     // Keep references to these tensors to extend their lifetime
     at::Tensor softmax_lse_accum, out_accum;
-    if (paged_KV || seqlenq_ngroups_swapped) {
+    if (seqlenq_ngroups_swapped) {
         std::tie(softmax_lse_accum, out_accum) =
             set_params_splitkv(params, batch_size, num_heads, head_size,
                                max_seqlen_k, max_seqlen_q, head_size_rounded,
                                p_dropout, num_splits, get_num_sm(get_current_device()), opts);
+    } else if (paged_KV) {
+        TORCH_CHECK(num_splits <= 1, "num_splits > 1 is not supported for varlen paged KV");
+        params.num_splits = num_splits;
     }
 
     if (leftpad_k_.has_value()) {

flash_attn/cute/block_info.py

@@ -143,14 +143,14 @@ def get_n_block_max_for_m_block(
         self,
         seqlen_info: SeqlenInfoQK,
         m_block: Int32,
-        n_block_global_max: Int32,
     ) -> Int32:
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
         if const_expr(self.is_causal or self.window_size_right is not None):
             m_idx_max = (m_block + 1) * self.tile_m
             if const_expr(self.qhead_per_kvhead_packgqa > 1):
                 m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
             n_idx_right = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
             if const_expr(self.window_size_right is not None):
                 n_idx_right += self.window_size_right
-            return min(n_block_global_max, cute.ceil_div(n_idx_right, self.tile_n))
-        return n_block_global_max
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n))
+        return n_block_max