tune dynamic scheduler prefetch

Author: jayhshah

flash_attn/cute/flash_fwd_sm100.py

@@ -584,7 +584,6 @@ def __call__(
             mCuSeqlensQ is None and
             mSeqUsedQ is None
         )
-        print("Dynamic persistent is ", self.dynamic_persistent)
         if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
             TileScheduler = SingleTileVarlenScheduler
         elif const_expr(not self.dynamic_persistent):
@@ -717,7 +716,6 @@ class SharedStorage:
         if cutlass.const_expr(self.use_block_sparsity and mPageTable is not None):
             raise NotImplementedError("Block sparsity + paged KV not supported on SM100")
 
-        print("smem = ", self.shared_storage.size_in_bytes())
         # Launch the kernel synchronously
         self.kernel(
             mQ,
@@ -911,7 +909,7 @@ def kernel(
             assert tile_count_semaphore is not None
             work_info = storage.work_info.get_tensor((4, ))
             sched_pipeline_producer_group = pipeline.CooperativeGroup(
-                pipeline.Agent.Thread, 32
+                pipeline.Agent.Thread, 1
             )
             sched_pipeline_consumer_group = pipeline.CooperativeGroup(
                 pipeline.Agent.Thread, 32 * self.num_non_empty_warps

flash_attn/cute/tile_scheduler.py

@@ -418,23 +418,20 @@ def initial_work_tile_info(self, *, loc=None, ip=None):
     @cute.jit
     def prefetch_next_work(self, *, loc=None, ip=None):
         params = self.params
-        next_tile_idx = 0
-        if cute.arch.lane_idx() == 0:
-            next_tile_idx = cute.arch.grid_dim()[0] + utils.atomic_add_i32(
-                1, params.tile_count_semaphore,
-            )
-        next_tile_idx = cute.arch.shuffle_sync(next_tile_idx, 0)
-        work_info = self.get_current_work(next_tile_idx)
         self._scheduler_pipeline.producer_acquire(
             self._producer_pipeline_state,
             self._scheduler_pipeline.producer_try_acquire(self._producer_pipeline_state),
         )
-        block, head_idx, batch_idx, _ = work_info.tile_idx
         if cute.arch.lane_idx() == 0:
+            next_tile_idx = cute.arch.grid_dim()[0] + utils.atomic_add_i32(
+                1, params.tile_count_semaphore,
+            )
+            work_info = self.get_current_work(next_tile_idx)
+            block, head_idx, batch_idx, _ = work_info.tile_idx
             self._work_info[0] = block
             self._work_info[1] = head_idx
             self._work_info[2] = batch_idx
-        self._scheduler_pipeline.producer_commit(self._producer_pipeline_state)
+            self._scheduler_pipeline.producer_commit(self._producer_pipeline_state)
         self._producer_pipeline_state.advance()
 
     @cute.jit