[Cute] Apply review feedback: rename to mCuTotalMBlocks, drop num_head, use shared utility

Addresses @reubenconducts's May 22 review comments on Dao-AILab#2520:

1. Rename mTileCumsum -> mCuTotalMBlocks across all 9 kernels + scheduler
   + interface for consistency with the convention introduced in Dao-AILab#2224
   (already in main; used by blocksparse, and by Dao-AILab#2559).

2. Drop num_head (and the related pack_gqa arch-conditional remap) from
   the host cumsum. Per-batch cumsum is now pure m_blocks; the scheduler
   handles num_head separately. Removes the SM80/SM120 vs SM90/100/110/MLA
   branching that previously mirrored the pack_gqa_layout reshape behavior.

3. Replace the inline binary search in _varlen_coord_map's cumsum-on
   branch with a call to utils.get_batch_from_cu_tensor (the shared
   utility from Dao-AILab#2556). The existing snap-to-group-boundary + warp-scan
   structure is preserved — the cumsum serves as a hint to skip ahead,
   and the warp-scan refines to the exact batch using _get_num_m_blocks
   (which already handles pack_gqa, q_stage, cluster, etc.). This matches
   the scheduler-side approach in Dao-AILab#2559.

The pack_gqa seqlen multiplier stays in _compute_cu_total_m_blocks so
that per-batch m_block counts match the kernel's _get_num_m_blocks
formula — the snap is forward-only, so under-estimating per-batch counts
is safe but over-estimating (which dropping the multiplier would cause
when pack_gqa is on) would land the snap past the correct batch and the
warp-scan couldn't recover.

Verified on SM100:
- 72 new tests (test_varlen_scheduler_binary_search_correctness{,_bwd}): pass
- existing test_varlen (B=20 slice, 576 cases): pass
- existing test_flash_attn_mla_absorbed_varlen (480 cases): pass

Author: horakka5

flash_attn/cute/flash_bwd.py

@@ -382,7 +382,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
         mdQ_semaphore: Optional[cute.Tensor] = None,
@@ -430,7 +430,7 @@ def __call__(
             qhead_per_kvhead_packgqa=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1,
             mCuSeqlensQ=mCuSeqlensK,
             mSeqUsedQ=mSeqUsedK,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_postprocess.py

@@ -215,7 +215,7 @@ def __call__(
         scale: cutlass.Float32,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
@@ -261,7 +261,7 @@ def __call__(
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_preprocess.py

@@ -136,7 +136,7 @@ def __call__(
         mCuSeqlensQ: Optional[cute.Tensor],  # (batch + 1,)
         mSeqUsedQ: Optional[cute.Tensor],  # (batch,)
         mdLSE: Optional[cute.Tensor],  # (batch, nheads, seqlen) or (nheads, total_q)
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
@@ -196,7 +196,7 @@ def __call__(
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_sm100.py

@@ -458,7 +458,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         window_size_left: Int32 | int | None = None,
@@ -737,7 +737,7 @@ def __call__(
             is_persistent=self.is_persistent,  # persistent mode not tested
             lpt=self.spt,
             head_swizzle=self.deterministic,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_sm90.py

@@ -350,7 +350,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         window_size_left: Int32 | int | None = None,
@@ -539,7 +539,7 @@ def _qkv_transpose(t):
             is_persistent=False,
             lpt=self.spt,
             head_swizzle=self.deterministic,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_fwd.py

@@ -631,7 +631,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         mPageTable: Optional[cute.Tensor] = None,
         window_size_left: Optional[Int32] = None,
         window_size_right: Optional[Int32] = None,
@@ -699,7 +699,7 @@ def __call__(
             qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)

flash_attn/cute/flash_fwd_mla_sm100.py

@@ -361,7 +361,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,  # (b + 1)
         mSeqUsedQ: Optional[cute.Tensor] = None,  # (b)
         mSeqUsedK: Optional[cute.Tensor] = None,  # (b)
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         mIndexTopk: Optional[
@@ -633,7 +633,7 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma):
             is_split_kv=False,
             cluster_shape_mn=self.cluster_shape_mn,
             use_cluster_idx=False,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(
             tile_sched_args, scheduling_mode=self.scheduling_mode

flash_attn/cute/flash_fwd_sm100.py

@@ -373,7 +373,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         mPageTable: Optional[cute.Tensor] = None,  # (b_k, max_num_pages_per_seq)
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
@@ -661,7 +661,7 @@ def __call__(
             is_split_kv=self.is_split_kv,
             cluster_shape_mn=self.cluster_shape_mn,
             use_cluster_idx=not self.is_persistent and self.cta_group_size > 1,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(
             tile_sched_args, scheduling_mode=self.scheduling_mode

flash_attn/cute/flash_fwd_sm90.py

@@ -166,7 +166,7 @@ def __call__(
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        mTileCumsum: Optional[
+        mCuTotalMBlocks: Optional[
             cute.Tensor
         ] = None,  # int32, (num_batch + 1,); see TileSchedulerArguments
         mPageTable: Optional[cute.Tensor] = None,  # (b_k, max_num_pages_per_seq)
@@ -344,7 +344,7 @@ def __call__(
             element_size=self.dtype.width // 8,
             is_persistent=False,
             lpt=self.is_causal or self.is_local,
-            mTileCumsum=mTileCumsum,
+            mCuTotalMBlocks=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)