Sync upstream

csrc/flash_attn/flash_api.cpp

@@ -708,6 +708,9 @@ mha_varlen_fwd(at::Tensor &q,  // total_q x num_heads x head_size, total_q := \s
             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