work PR 2520 into interface and kernels

Author: reubenconducts

benchmarks/benchmark_varlen_sched.py

@@ -181,6 +181,13 @@ def build_ctx(
 
 
 def _make_meta(ctx):
+    tile_m = 128
+    qhead_per_kvhead = ctx["nheads"] // ctx["nheads_kv"]
+    arch = torch.cuda.get_device_capability()[0]
+    if arch == 10 and ctx["max_seqlen_q"] * qhead_per_kvhead > tile_m:
+        q_stage = 2
+    else:
+        q_stage = 1
     return get_scheduler_metadata(
         num_batch=ctx["batch"],
         max_seqlen_q=ctx["max_seqlen_q"],
@@ -189,12 +196,13 @@ def _make_meta(ctx):
         nheads_kv=ctx["nheads_kv"],
         headdim=ctx["headdim"],
         num_splits=ctx["num_splits"],
-        tile_m=128,
+        tile_m=tile_m,
         tile_n=128,
         causal=ctx["causal"],
         pack_gqa=ctx["pack_gqa"],
         cu_seqlens_q=ctx["cu_q"],
         cu_seqlens_k=ctx["cu_k"],
+        q_stage=q_stage,
     )
 
 
@@ -322,8 +330,12 @@ def parse_args():
     p.add_argument("--headdim", type=int, default=128)
     p.add_argument("--nheads", type=int, default=16)
     p.add_argument("--nheads-kv", type=int, default=2)
-    p.add_argument("--pack-gqa", action="store_true", default=True,
-                   help="Force pack_gqa=True (default). --no-pack-gqa to disable.")
+    p.add_argument(
+        "--pack-gqa",
+        action="store_true",
+        default=True,
+        help="Force pack_gqa=True (default). --no-pack-gqa to disable.",
+    )
     p.add_argument("--no-pack-gqa", dest="pack_gqa", action="store_false")
     p.add_argument("--seeds", type=int, default=3)
     p.add_argument("--warmup", type=int, default=2)
@@ -434,7 +446,9 @@ def main():
                     eff = sq * sk - sq * (sq - 1) // 2
                 else:
                     # clamp to non-negative for queries near 0
-                    first_visible_q = -shift  # smallest q with sk - sq + q + 1 > 0 is q = sq - sk
+                    first_visible_q = (
+                        -shift
+                    )  # smallest q with sk - sq + q + 1 > 0 is q = sq - sk
                     visible = sq - first_visible_q
                     eff = visible * sk - visible * (visible - 1) // 2
                 eff = max(0, eff)

flash_attn/cute/flash_bwd.py

@@ -389,6 +389,7 @@ def __call__(
         mdV_semaphore: Optional[cute.Tensor] = None,
         aux_tensors: Optional[list] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -429,6 +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,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_postprocess.py

@@ -215,6 +215,7 @@ def __call__(
         scale: cutlass.Float32,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -258,6 +259,7 @@ def __call__(
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_preprocess.py

@@ -136,6 +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)
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -193,6 +194,7 @@ def __call__(
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_bwd_sm100.py

@@ -466,6 +466,7 @@ def __call__(
         aux_tensors: Optional[list] = None,
         # Block-sparse tensors (Q direction - for iterating m_blocks per n_block):
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -732,6 +733,7 @@ def __call__(
             qhead_per_kvhead_packgqa=1,  # pack_gqa disabled for bwd
             element_size=self.k_dtype.width // 8,
             is_persistent=self.is_persistent,  # persistent mode not tested
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
             lpt=self.spt,
             head_swizzle=self.deterministic,
         )

flash_attn/cute/flash_bwd_sm90.py

@@ -357,6 +357,7 @@ def __call__(
         mdV_semaphore: Optional[cute.Tensor] = None,
         aux_tensors: Optional[list] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -536,6 +537,7 @@ def _qkv_transpose(t):
             is_persistent=False,
             lpt=self.spt,
             head_swizzle=self.deterministic,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
 
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)

flash_attn/cute/flash_fwd.py

@@ -631,6 +631,8 @@ def __call__(
         learnable_sink: Optional[cute.Tensor] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         aux_tensors=None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuTotalSplitsMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -692,6 +694,8 @@ def __call__(
             qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
+            cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)

flash_attn/cute/flash_fwd_sm100.py

@@ -398,8 +398,8 @@ def __call__(
         tile_count_semaphore: Optional[cute.Tensor] = None,
         virtual_batch_idx_ptr: Optional[cute.Tensor] = None,
         num_nheads_in_l2_ptr: Optional[cute.Tensor] = None,
-        cu_total_m_blocks_ptr: Optional[cute.Tensor] = None,
-        cu_total_splits_m_blocks_ptr: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuTotalSplitsMBlocks: Optional[cute.Tensor] = None,
         max_seqlen_q: Int32 | int | None = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
@@ -688,8 +688,8 @@ def __call__(
             num_splits_dynamic_ptr=num_splits_dynamic_ptr,
             virtual_batch_idx_ptr=virtual_batch_idx_ptr,
             num_nheads_in_l2_ptr=num_nheads_in_l2_ptr,
-            cu_total_m_blocks_ptr=cu_total_m_blocks_ptr,
-            cu_total_splits_m_blocks_ptr=cu_total_splits_m_blocks_ptr,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
+            cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks,
             tile_count_semaphore=tile_count_semaphore.iterator if tile_count_semaphore is not None else None,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(

flash_attn/cute/flash_fwd_sm90.py

@@ -172,6 +172,8 @@ def __call__(
         learnable_sink: Optional[cute.Tensor] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         aux_tensors: Optional[list] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuTotalSplitsMBlocks: Optional[cute.Tensor] = None,
         # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
@@ -312,6 +314,7 @@ def __call__(
                 (self.tile_m, self.tile_hdimv),  # No mcast
             )
         if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
+            # TODO: dispatch to DynamicPersistentVarlenScheduler when appropriate
             TileScheduler = SingleTileVarlenScheduler
         else:
             TileScheduler = (
@@ -341,6 +344,8 @@ def __call__(
             element_size=self.dtype.width // 8,
             is_persistent=False,
             lpt=self.is_causal or self.is_local,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
+            cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)