quack-kernels 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

quack/tile_scheduler.py

@@ -25,6 +25,13 @@ class RasterOrder(IntEnum):
     AlongN = 1
 
 
+class PersistenceMode(IntEnum):
+    NONE = 0
+    STATIC = 1
+    DYNAMIC = 2
+    CLC = 3
+
+
 @cute.jit
 def get_raster_order_from_option(
     raster_order_option: RasterOrderOption, problem_shape_ncluster_mn: cute.Shape, group_size: Int32
@@ -61,7 +68,7 @@ class TileSchedulerArguments(ArgumentsBase):
     cluster_shape_mnk: cutlass.Constexpr[cute.Shape]
     tile_count_semaphore: Optional[cute.Pointer] = None
     batch_idx_permute: Optional[cute.Tensor] = None
-    is_persistent: cutlass.Constexpr[bool] = False
+    persistence_mode: cutlass.Constexpr[PersistenceMode] = PersistenceMode.NONE
 
 
 class TileScheduler:
@@ -69,15 +76,15 @@ class TileScheduler:
     class Params(ParamsBase):
         problem_shape_ncluster_mnl: cute.Shape
         raster_order: RasterOrder
-        num_clusters_per_problem_divmod: FastDivmod
+        num_clusters_per_problem_fdd: FastDivmod
         num_groups_regular: Int32
-        group_size_divmod: FastDivmod
-        group_size_tail_divmod: FastDivmod
-        num_clusters_in_group_divmod: FastDivmod
+        group_size_fdd: FastDivmod
+        group_size_tail_fdd: FastDivmod
+        num_clusters_in_group_fdd: FastDivmod
         tile_count_semaphore: Optional[cute.Pointer]
         batch_idx_permute: Optional[cute.Tensor]
         cluster_shape_mn: cutlass.Constexpr[cute.Shape]
-        is_persistent: cutlass.Constexpr[bool]
+        persistence_mode: cutlass.Constexpr[PersistenceMode]
 
         @staticmethod
         @cute.jit
@@ -107,26 +114,30 @@ class TileScheduler:
             group_size_tail = ncluster_fast % group_size
             num_groups_regular = ncluster_fast // group_size
             num_clusters_in_group = group_size * ncluster_slow
+            if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC):
+                assert args.tile_count_semaphore is not None
             return TileScheduler.Params(
                 problem_shape_ncluster_mnl,
                 raster_order,
-                FastDivmod.create(num_clusters_per_problem),
+                FastDivmod(num_clusters_per_problem),
                 num_groups_regular,
-                FastDivmod.create(group_size),
+                FastDivmod(group_size),
                 # Don't divide by 0
-                FastDivmod.create(group_size_tail if group_size_tail > 0 else 1),
-                FastDivmod.create(num_clusters_in_group),
-                args.tile_count_semaphore if const_expr(args.is_persistent) else None,
+                FastDivmod(group_size_tail if group_size_tail > 0 else 1),
+                FastDivmod(num_clusters_in_group),
+                args.tile_count_semaphore
+                if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC)
+                else None,
                 args.batch_idx_permute,
                 cluster_shape_mn,
-                args.is_persistent,
+                args.persistence_mode,
             )
 
     def __init__(
         self,
         current_work_linear_idx: Int32,
         num_tiles_executed: Int32,
-        tile_count: Optional[cute.Tensor],
+        sched_smem: Optional[cute.Tensor],
         scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync],
         pipeline_state: PipelineStateWAdvance,
         params: Params,
@@ -136,7 +147,7 @@ class TileScheduler:
     ):
         self._current_work_linear_idx = current_work_linear_idx
         self.num_tiles_executed = num_tiles_executed
-        self._tile_count = tile_count
+        self._sched_smem = sched_smem
         self._scheduler_pipeline = scheduler_pipeline
         self._pipeline_state = pipeline_state
         self.params = params
@@ -151,16 +162,14 @@ class TileScheduler:
     @cute.jit
     def create(
         params: Params,
-        tile_count: Optional[cute.Tensor] = None,
+        sched_smem: Optional[cute.Tensor] = None,
         scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None,
-        is_scheduler_warp: bool | Boolean = False,
         *,
         loc=None,
         ip=None,
     ) -> "TileScheduler":
         """is_scheduler_warp should only be true for one warp in the whole cluster"""
-        stages = 0
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             cidx, cidy, _ = cute.arch.cluster_idx()
             cdimx, _, _ = cute.arch.cluster_dim()
             cluster_id = cidx + cidy * cdimx
@@ -168,16 +177,20 @@ class TileScheduler:
         else:
             _, _, bidz = cute.arch.block_idx()
             current_work_linear_idx = Int32(bidz)
-            if const_expr(params.tile_count_semaphore is not None):
-                assert tile_count is not None
-                assert scheduler_pipeline is not None
-                stages = const_expr(cute.size(tile_count))
+        stages = 0
+        if const_expr(
+            params.persistence_mode
+            in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC]
+        ):
+            assert sched_smem is not None
+            assert scheduler_pipeline is not None
+            stages = const_expr(cute.size(sched_smem, mode=[1]))
         return TileScheduler(
             current_work_linear_idx,
             Int32(0),  # num_tiles_executed
-            tile_count,
+            sched_smem,
             scheduler_pipeline,
-            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)),
+            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)),
             params,
             loc=loc,
             ip=ip,
@@ -195,7 +208,7 @@ class TileScheduler:
         num_ctas_mnl = tuple(
             x * y for x, y in zip(params.problem_shape_ncluster_mnl, params.cluster_shape_mn)
         ) + (params.problem_shape_ncluster_mnl[2],)
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             return num_ctas_mnl
         else:
             num_ctas_in_problem = cute.size(num_ctas_mnl, loc=loc, ip=ip)
@@ -212,12 +225,12 @@ class TileScheduler:
     ) -> Tuple[Int32, Int32]:
         # CTA Swizzle to promote L2 data reuse
         params = self.params
-        group_id, id_in_group = params.num_clusters_in_group_divmod.divmod(cluster_id_in_problem)
+        group_id, id_in_group = divmod(cluster_id_in_problem, params.num_clusters_in_group_fdd)
         cid_fast_in_group, cid_slow = Int32(0), Int32(0)
         if group_id < params.num_groups_regular:
-            cid_slow, cid_fast_in_group = params.group_size_divmod.divmod(id_in_group)
+            cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_fdd)
         else:  # tail part
-            cid_slow, cid_fast_in_group = params.group_size_tail_divmod.divmod(id_in_group)
+            cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_tail_fdd)
         if group_id % 2 == 1:  # serpentine order
             ncluster_slow = (
                 params.problem_shape_ncluster_mnl[1]
@@ -225,56 +238,151 @@ class TileScheduler:
                 else params.problem_shape_ncluster_mnl[0]
             )
             cid_slow = ncluster_slow - 1 - cid_slow
-        cid_fast = group_id * params.group_size_divmod.divisor + cid_fast_in_group
+        cid_fast = group_id * params.group_size_fdd.divisor + cid_fast_in_group
         cid_m, cid_n = cid_fast, cid_slow
         if params.raster_order == RasterOrder.AlongN:
             cid_m, cid_n = cid_slow, cid_fast
         return cid_m, cid_n
 
     @cute.jit
-    def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo:
+    def _delinearize_work_idx(
+        self, *, block_zero_only: bool = False, loc=None, ip=None
+    ) -> cutlass.utils.WorkTileInfo:
         params = self.params
-        if const_expr(not params.is_persistent):
-            cluster_id_in_problem = self._current_work_linear_idx
-            _, _, bidz = cute.arch.block_idx()
-        else:
-            bidz, cluster_id_in_problem = params.num_clusters_per_problem_divmod.divmod(
-                self._current_work_linear_idx
-            )
-        cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip)
-        # Get the pid from cluster id
-        bidx_in_cluster = cute.arch.block_in_cluster_idx()
-        pid_m = cid_m * params.cluster_shape_mn[0] + bidx_in_cluster[0]
-        pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
-        batch_idx = (
-            bidz if const_expr(params.batch_idx_permute is None) else params.batch_idx_permute[bidz]
-        )
-        tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode == PersistenceMode.NONE):
             is_valid = self.num_tiles_executed == 0
         else:
             is_valid = self._current_work_linear_idx < cute.size(params.problem_shape_ncluster_mnl)
+        pid_m, pid_n, batch_idx = Int32(0), Int32(0), Int32(0)
+        if is_valid:
+            if const_expr(params.persistence_mode == PersistenceMode.NONE):
+                cluster_id_in_problem = self._current_work_linear_idx
+                _, _, bidz = cute.arch.block_idx()
+            else:
+                bidz, cluster_id_in_problem = divmod(
+                    self._current_work_linear_idx, params.num_clusters_per_problem_fdd
+                )
+            cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, loc=loc, ip=ip)
+            if const_expr(block_zero_only):
+                bidx_in_cluster = (Int32(0), Int32(0))
+            else:
+                # Get the pid from cluster id
+                bidx_in_cluster = cute.arch.block_in_cluster_idx()
+            pid_m = cid_m * params.cluster_shape_mn[0] + bidx_in_cluster[0]
+            pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
+            batch_idx = (
+                bidz
+                if const_expr(params.batch_idx_permute is None)
+                else params.batch_idx_permute[bidz]
+            )
+        tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
         return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid)
 
+    @cute.jit
+    def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo:
+        params = self.params
+        pid_m, pid_n, batch_idx, is_valid = Int32(0), Int32(0), Int32(0), Boolean(False)
+        if const_expr(params.persistence_mode == PersistenceMode.NONE):
+            pass
+        # elif const_expr(params.persistence_mode == PersistenceMode.STATIC):
+        #     return self._delinearize_work_idx(loc=loc, ip=ip)
+        else:
+            self._scheduler_pipeline.consumer_wait(self._pipeline_state)
+            pid_m, pid_n, batch_idx, is_valid_i32 = [
+                self._sched_smem[i, self._pipeline_state.index] for i in range(4)
+            ]
+            # Need this fence since the STAS from the producer is using the async proxy.
+            # Without this, we get race condition / deadlock.
+            if const_expr(cute.size(params.cluster_shape_mn) > 1):
+                cute.arch.fence_view_async_shared()
+            cute.arch.sync_warp()
+            with cute.arch.elect_one():
+                self._scheduler_pipeline.consumer_release(self._pipeline_state)
+            self._pipeline_state.advance()
+            is_valid = Boolean(is_valid_i32)
+        tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
+        return cutlass.utils.WorkTileInfo(tile_coord_mnkl, Boolean(is_valid))
+
+    # @cute.jit
     def initial_work_tile_info(self, *, loc=None, ip=None):
-        return self.get_current_work(loc=loc, ip=ip)
+        return self._delinearize_work_idx(loc=loc, ip=ip)
+        # if is_scheduler_warp:
+        # work_tile_info = self._delinearize_work_idx(block_zero_only=True, loc=loc, ip=ip)
+        # self.write_work_tile_to_smem(work_tile_info, loc=loc, ip=ip)
+        # self.write_work_tile_to_smem(self._delinearize_work_idx(block_zero_only=True, loc=loc, ip=ip), loc=loc, ip=ip)
 
     @cute.jit
-    def fetch_next_work(self, is_scheduler_warp: bool | Boolean = False, *, loc=None, ip=None):
+    def _fetch_next_work_idx(self, *, loc=None, ip=None) -> Int32:
         """is_scheduler_warp should only be true for one warp in the whole cluster"""
         params = self.params
-        if const_expr(params.is_persistent and params.tile_count_semaphore is not None):
-            current_work_linear_idx = self._current_work_linear_idx
-            if is_scheduler_warp:
-                if cute.arch.lane_idx() == 0:
-                    num_persistent_clusters = cute.arch.grid_dim()[2]
-                    current_work_linear_idx = num_persistent_clusters + utils.atomic_inc_i32(
+        num_persistent_clusters = Int32(cute.arch.grid_dim()[2])
+        if const_expr(params.persistence_mode == PersistenceMode.STATIC):
+            return self._current_work_linear_idx + num_persistent_clusters
+        elif const_expr(params.persistence_mode == PersistenceMode.DYNAMIC):
+            next_work_linear_idx = Int32(0)
+            if cute.arch.lane_idx() == 0:
+                # If varlen_m, problem_shape_ncluster_mnl[0] is None, so we use atomic_add
+                # instead of atomic_inc, and at the end of the kernel must reset the semaphore to 0.
+                #                 # cute.printf("before atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
+                if const_expr(params.problem_shape_ncluster_mnl[0] is not None):
+                    next_work_linear_idx = num_persistent_clusters + utils.atomic_inc_i32(
                         cute.size(params.problem_shape_ncluster_mnl) - 1,
                         params.tile_count_semaphore,
                     )
-                # lane 0 already has the right tile_idx, just need to broadcast
-                current_work_linear_idx = cute.arch.shuffle_sync(current_work_linear_idx, 0)
-            self._current_work_linear_idx = current_work_linear_idx
+                else:  # varlen_m
+                    next_work_linear_idx = num_persistent_clusters + utils.atomic_add_i32(
+                        1, params.tile_count_semaphore
+                    )
+                # cute.printf("after atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
+            return cute.arch.shuffle_sync(next_work_linear_idx, 0)
+        else:
+            return Int32(0)
+
+    @cute.jit
+    def write_work_tile_to_smem(
+        self, work_tile_info: cutlass.utils.WorkTileInfo, *, loc=None, ip=None
+    ):
+        params = self.params
+        if const_expr(self._sched_smem is not None):
+            # producer phase is always consumer_phase ^ 1
+            pipeline_state_producer = PipelineStateWAdvance(
+                self._pipeline_state.stages,
+                self._pipeline_state.count,
+                self._pipeline_state.index,
+                self._pipeline_state.phase ^ 1,
+            )
+            self._scheduler_pipeline.producer_acquire(pipeline_state_producer)
+            sched_data = [
+                work_tile_info.tile_idx[0],
+                work_tile_info.tile_idx[1],
+                work_tile_info.tile_idx[3],
+                Int32(work_tile_info.is_valid_tile),
+            ]
+            lane_idx = cute.arch.lane_idx()
+            if lane_idx < cute.size(params.cluster_shape_mn):
+                # cute.printf("Producer pid_m = {}, pid_n = {}, batch_idx = {}, is_valid = {}, after empty wait, idx = {}", sched_data[0], sched_data[1], sched_data[2], sched_data[3], self._current_work_linear_idx)
+                pipeline_idx = self._pipeline_state.index
+                if const_expr(cute.size(params.cluster_shape_mn) == 1):
+                    for i in cutlass.range_constexpr(4):
+                        self._sched_smem[i, pipeline_idx] = sched_data[i]
+                    self._scheduler_pipeline.producer_commit(self._pipeline_state)
+                else:
+                    peer_cta_rank_in_cluster = lane_idx
+                    # Here we assume that the block idx in cluster is linearized such that
+                    # x is the fastest moving direction.
+                    bidx_in_cluster = peer_cta_rank_in_cluster % params.cluster_shape_mn[0]
+                    bidy_in_cluster = peer_cta_rank_in_cluster // params.cluster_shape_mn[0]
+                    mbar_ptr = self._scheduler_pipeline.producer_get_barrier(self._pipeline_state)
+                    cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, 16, peer_cta_rank_in_cluster)
+                    utils.store_shared_remote_x4(
+                        sched_data[0] + bidx_in_cluster,
+                        sched_data[1] + bidy_in_cluster,
+                        sched_data[2],
+                        sched_data[3],
+                        smem_ptr=self._sched_smem[None, pipeline_idx].iterator,
+                        mbar_ptr=mbar_ptr,
+                        peer_cta_rank_in_cluster=peer_cta_rank_in_cluster,
+                    )
 
     @cute.jit
     def advance_to_next_work(
@@ -285,73 +393,37 @@ class TileScheduler:
         loc=None,
         ip=None,
     ):
-        tidx = cute.arch.thread_idx()[0]
-        bidx = cute.arch.block_idx()[0]
-        bidz = cute.arch.block_idx()[2]
+        """is_scheduler_warp should only be true for one warp in the whole cluster.
+        Moreover, we assume that only block zero in the cluster is calling this function.
+        If calling with is_scheduler_warp = True, advance_count must be 1.
+        """
         params = self.params
-        if const_expr(params.is_persistent):
-            num_persistent_clusters = cute.arch.grid_dim()[2]
-            if const_expr(params.tile_count_semaphore is None):  # Static persistent
-                self._current_work_linear_idx += advance_count * Int32(num_persistent_clusters)
-            else:  # Dynamic persistent
-                if const_expr(advance_count > 1):
-                    self._pipeline_state.advance_iters(advance_count - 1)
-                current_work_linear_idx = self._current_work_linear_idx
-                if is_scheduler_warp:
-                    self._scheduler_pipeline.producer_acquire(self._pipeline_state)
-                    lane_idx = cute.arch.lane_idx()
-                    if lane_idx < cute.size(params.cluster_shape_mn):
-                        # cute.printf("Producer bidx = {}, bidz = {}, tidx = {}, after empty wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
-                        if const_expr(cute.size(params.cluster_shape_mn) == 1):
-                            self._tile_count[self._pipeline_state.index] = current_work_linear_idx
-                            self._scheduler_pipeline.producer_commit(self._pipeline_state)
-                        else:
-                            peer_cta_rank_in_cluster = lane_idx
-                            mbar_ptr = self._scheduler_pipeline.producer_get_barrier(
-                                self._pipeline_state
-                            )
-                            cute.arch.mbarrier_arrive_and_expect_tx(
-                                mbar_ptr, 4, peer_cta_rank_in_cluster
-                            )
-                            utils.store_shared_remote(
-                                val=current_work_linear_idx,
-                                smem_ptr=self._tile_count.iterator + self._pipeline_state.index,
-                                mbar_ptr=mbar_ptr,
-                                peer_cta_rank_in_cluster=peer_cta_rank_in_cluster,
-                            )
-                        # cute.printf("Producer bidx = {}, bidz = {}, tidx = {}, after full arrive", bidx, bidz, tidx)
-                else:
-                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, before full wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
-                    self._scheduler_pipeline.consumer_wait(self._pipeline_state)
-                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, after full wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
-                    current_work_linear_idx = self._tile_count[self._pipeline_state.index]
-                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, after smem read, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
-                    # Need this fence since the STAS from the producer is using the async proxy.
-                    # Without this, we get race condition / deadlock.
-                    if const_expr(cute.size(params.cluster_shape_mn) > 1):
-                        cute.arch.fence_proxy(
-                            cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta
-                        )
-                    cute.arch.sync_warp()
-                    with cute.arch.elect_one():
-                        # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, before empty arrive", bidx, bidz, tidx)
-                        self._scheduler_pipeline.consumer_release(self._pipeline_state)
-                        # if tidx == 320: cute.printf("bidx = {}, bidz = {}, tidx = {}, idx = {}, after empty arrive", bidx, bidz, tidx, current_work_linear_idx)
-                    # if tidx == 320: cute.printf("bidx = {}, bidz = {}, tidx = {}, idx = {}, after empty arrive", bidx, bidz, tidx, current_work_linear_idx)
-                self._current_work_linear_idx = current_work_linear_idx
-                self._pipeline_state.advance()
         self.num_tiles_executed += Int32(advance_count)
+        if const_expr(self._pipeline_state is not None and advance_count > 1):
+            self._pipeline_state.advance_iters(advance_count - 1)
+        if const_expr(params.persistence_mode in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC]):
+            # We assume here that advance_count is 1 for scheduler_warp
+            if is_scheduler_warp:
+                self._current_work_linear_idx = self._fetch_next_work_idx(loc=loc, ip=ip)
+                work_tile_info = self._delinearize_work_idx(block_zero_only=True, loc=loc, ip=ip)
+                self.write_work_tile_to_smem(work_tile_info, loc=loc, ip=ip)
 
     def producer_tail(self):
-        if const_expr(self.params.is_persistent and self.params.tile_count_semaphore is not None):
-            self._scheduler_pipeline.producer_tail(self._pipeline_state)
+        if const_expr(self._scheduler_pipeline is not None):
+            pipeline_state_producer = PipelineStateWAdvance(
+                self._pipeline_state.stages,
+                self._pipeline_state.count,
+                self._pipeline_state.index,
+                self._pipeline_state.phase ^ 1,
+            )
+            self._scheduler_pipeline.producer_tail(pipeline_state_producer)
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
         for obj in [
             self._current_work_linear_idx,
             self.num_tiles_executed,
-            self._tile_count,
+            self._sched_smem,
             self._scheduler_pipeline,
             self._pipeline_state,
             self.params,
@@ -367,7 +439,7 @@ class TileScheduler:
             [
                 self._current_work_linear_idx,
                 self.num_tiles_executed,
-                self._tile_count,
+                self._sched_smem,
                 self._scheduler_pipeline,
                 self._pipeline_state,
                 self.params,
@@ -396,16 +468,16 @@ class TriangularTileScheduler(TileScheduler):
     @dataclass
     class Params(ParamsBase):
         problem_shape_ncluster_mnl: cute.Shape
-        num_clusters_per_problem_divmod: FastDivmod
+        num_clusters_per_problem_fdd: FastDivmod
         group_size_inv_f32: Float32
         num_groups_regular: Int32
-        group_size_divmod: FastDivmod
-        group_size_tail_divmod: FastDivmod
-        group_size_mul_group_size_divmod: FastDivmod
-        group_size_tail_mul_group_size_divmod: FastDivmod
+        group_size_fdd: FastDivmod
+        group_size_tail_fdd: FastDivmod
+        group_size_mul_group_size_fdd: FastDivmod
+        group_size_tail_mul_group_size_fdd: FastDivmod
         tile_count_semaphore: Optional[cute.Pointer]
         cluster_shape_mn: cutlass.Constexpr[cute.Shape]
-        is_persistent: cutlass.Constexpr[bool]
+        persistence_mode: cutlass.Constexpr[PersistenceMode]
 
         @staticmethod
         @cute.jit
@@ -425,19 +497,23 @@ class TriangularTileScheduler(TileScheduler):
             group_size = min(args.group_size, cluster_m)
             group_size_tail = cluster_m % group_size
             num_groups_regular = cluster_m // group_size
+            if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC):
+                assert args.tile_count_semaphore is not None
             return TriangularTileScheduler.Params(
                 problem_shape_ncluster_mnl,
-                FastDivmod.create(num_clusters_per_problem),
+                FastDivmod(num_clusters_per_problem),
                 Float32(1.0 / group_size),
                 num_groups_regular,
-                FastDivmod.create(group_size),
+                FastDivmod(group_size),
                 # Don't divide by 0
-                FastDivmod.create(group_size_tail if group_size_tail > 0 else 1),
-                FastDivmod.create(group_size * group_size),
-                FastDivmod.create((group_size_tail if group_size_tail > 0 else 1) * group_size),
-                args.tile_count_semaphore if const_expr(args.is_persistent) else None,
+                FastDivmod(group_size_tail if group_size_tail > 0 else 1),
+                FastDivmod(group_size * group_size),
+                FastDivmod((group_size_tail if group_size_tail > 0 else 1) * group_size),
+                args.tile_count_semaphore
+                if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC)
+                else None,
                 cluster_shape_mn,
-                args.is_persistent,
+                args.persistence_mode,
             )
 
     @staticmethod
@@ -448,30 +524,32 @@ class TriangularTileScheduler(TileScheduler):
     @cute.jit
     def create(
         params: Params,
-        tile_count: Optional[cute.Tensor] = None,
+        sched_smem: Optional[cute.Tensor] = None,
         scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None,
-        is_scheduler_warp: bool | Boolean = False,
         *,
         loc=None,
         ip=None,
     ) -> "TriangularTileScheduler":
         stages = 0
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             cluster_id, _, _ = cute.arch.cluster_idx()
             current_work_linear_idx = Int32(cluster_id)
         else:
             _, _, bidz = cute.arch.block_idx()
             current_work_linear_idx = Int32(bidz)
-            if const_expr(params.tile_count_semaphore is not None):
-                assert tile_count is not None
-                assert scheduler_pipeline is not None
-                stages = const_expr(cute.size(tile_count))
+        if const_expr(
+            params.persistence_mode
+            in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC]
+        ):
+            assert sched_smem is not None
+            assert scheduler_pipeline is not None
+            stages = const_expr(cute.size(sched_smem))
         return TriangularTileScheduler(
             current_work_linear_idx,
             Int32(0),  # num_tiles_executed
-            tile_count,
+            sched_smem,
             scheduler_pipeline,
-            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)),
+            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)),
             params,
             loc=loc,
             ip=ip,
@@ -486,15 +564,11 @@ class TriangularTileScheduler(TileScheduler):
         loc=None,
         ip=None,
     ) -> Tuple[Int32, Int32, Int32]:
-        clusters = (
-            params.num_clusters_per_problem_divmod.divisor,
-            1,
-            params.problem_shape_ncluster_mnl[2],
-        )
+        clusters = (params.num_clusters_per_problem_fdd.divisor, 1)
         num_ctas_mnl = tuple(x * y for x, y in zip(clusters, params.cluster_shape_mn)) + (
             params.problem_shape_ncluster_mnl[2],
         )
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             return num_ctas_mnl
         else:
             num_ctas_in_problem = cute.size(num_ctas_mnl, loc=loc, ip=ip)
@@ -506,17 +580,19 @@ class TriangularTileScheduler(TileScheduler):
             return (*params.cluster_shape_mn, num_persistent_clusters)
 
     @cute.jit
-    def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo:
+    def _delinearize_work_idx(
+        self, *, block_zero_only: bool = False, loc=None, ip=None
+    ) -> cutlass.utils.WorkTileInfo:
         params = self.params
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             cluster_id_in_problem = self._current_work_linear_idx
             _, _, bidz = cute.arch.block_idx()
         else:
-            bidz, cluster_id_in_problem = params.num_clusters_per_problem_divmod.divmod(
-                self._current_work_linear_idx
+            bidz, cluster_id_in_problem = divmod(
+                self._current_work_linear_idx, params.num_clusters_per_problem_fdd
             )
         # CTA Swizzle to promote L2 data reuse
-        group_size = params.group_size_divmod.divisor
+        group_size = params.group_size_fdd.divisor
         group_id = (
             utils.ceil(
                 (utils.sqrt(2 * cluster_id_in_problem + 2.25) - 0.5) * params.group_size_inv_f32
@@ -528,40 +604,40 @@ class TriangularTileScheduler(TileScheduler):
         group_size_actual = (
             group_size
             if group_id < params.num_groups_regular
-            else params.group_size_tail_divmod.divisor
+            else params.group_size_tail_fdd.divisor
         )
         group_col, group_remainder = Int32(0), Int32(0)
         if group_id < params.num_groups_regular:
-            group_col, group_remainder = params.group_size_mul_group_size_divmod.divmod(id_in_group)
+            group_col, group_remainder = divmod(id_in_group, params.group_size_mul_group_size_fdd)
         else:  # tail part
-            group_col, group_remainder = params.group_size_tail_mul_group_size_divmod.divmod(
-                id_in_group
+            group_col, group_remainder = divmod(
+                id_in_group, params.group_size_tail_mul_group_size_fdd
             )
         cid_m_in_group, cid_n_in_group = Int32(0), Int32(0)
         if id_in_group >= group_size_actual * group_size * group_id:  # triangular tail
             cid_m_in_group, cid_n_in_group = triangular_idx_to_coord(group_remainder)
         else:
             if group_id < params.num_groups_regular:
-                cid_n_in_group, cid_m_in_group = params.group_size_divmod.divmod(group_remainder)
+                cid_n_in_group, cid_m_in_group = divmod(group_remainder, params.group_size_fdd)
             else:
-                cid_n_in_group, cid_m_in_group = params.group_size_tail_divmod.divmod(
-                    group_remainder
-                )
+                cid_n_in_group, cid_m_in_group = divmod(group_remainder, params.group_size_tail_fdd)
         cid_m = cid_m_start + cid_m_in_group
         cid_n = group_col * group_size + cid_n_in_group
 
-        # Get the pid from cluster id
-        bidx_in_cluster = cute.arch.block_in_cluster_idx()
+        if const_expr(block_zero_only):
+            bidx_in_cluster = (Int32(0), Int32(0))
+        else:
+            # Get the pid from cluster id
+            bidx_in_cluster = cute.arch.block_in_cluster_idx()
         pid_m = cid_m * params.cluster_shape_mn[0] + bidx_in_cluster[0]
         pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
         tile_coord_mnkl = (pid_m, pid_n, None, bidz)
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode == PersistenceMode.NONE):
             is_valid = self.num_tiles_executed == 0
         else:
             is_valid = (
                 self._current_work_linear_idx
-                < params.num_clusters_per_problem_divmod.divisor
-                * params.problem_shape_ncluster_mnl[2]
+                < params.num_clusters_per_problem_fdd.divisor * params.problem_shape_ncluster_mnl[2]
             )
         # bidx, bidy, bidz = cute.arch.block_idx()
         # tidx, _, _ = cute.arch.thread_idx()
@@ -581,7 +657,7 @@ class VarlenMTileSchedulerArguments(ParamsBase):
     tile_shape_mn: cutlass.Constexpr[cute.Shape]
     cluster_shape_mnk: cutlass.Constexpr[cute.Shape]
     tile_count_semaphore: Optional[cute.Pointer] = None
-    is_persistent: cutlass.Constexpr[bool] = False
+    persistence_mode: cutlass.Constexpr[PersistenceMode] = PersistenceMode.NONE
 
 
 class VarlenMTileScheduler(TileScheduler):
@@ -592,13 +668,13 @@ class VarlenMTileScheduler(TileScheduler):
         cu_seqlens_m: cute.Tensor
         raster_order: cutlass.Constexpr[RasterOrder]
         group_size: Int32
-        group_size_divmod: Optional[FastDivmod]
-        group_size_tail_divmod: Optional[FastDivmod]
-        num_clusters_in_group_divmod: FastDivmod
+        group_size_fdd: Optional[FastDivmod]
+        group_size_tail_fdd: Optional[FastDivmod]
+        num_clusters_in_group_fdd: FastDivmod
         tile_shape_mn: cutlass.Constexpr[cute.Shape]
         tile_count_semaphore: Optional[cute.Pointer]
         cluster_shape_mn: cutlass.Constexpr[cute.Shape]
-        is_persistent: cutlass.Constexpr[bool]
+        persistence_mode: cutlass.Constexpr[PersistenceMode]
 
         @staticmethod
         @cute.jit
@@ -621,43 +697,40 @@ class VarlenMTileScheduler(TileScheduler):
                 if args.raster_order == RasterOrderOption.AlongM
                 else RasterOrder.AlongN  # For Heuristic we also use AlongN
             )
-            ncluster_fast = (
-                problem_shape_ncluster_mn[0]
-                if raster_order == RasterOrder.AlongM
-                else problem_shape_ncluster_mn[1]
-            )
-            ncluster_slow = (
-                problem_shape_ncluster_mn[1]
-                if raster_order == RasterOrder.AlongM
-                else problem_shape_ncluster_mn[0]
-            )
+            ncluster_fast = problem_shape_ncluster_mn[
+                0 if raster_order == RasterOrder.AlongM else 1
+            ]
+            ncluster_slow = problem_shape_ncluster_mn[
+                1 if raster_order == RasterOrder.AlongM else 0
+            ]
             if const_expr(ncluster_fast is not None):
                 group_size = min(args.group_size, ncluster_fast)
                 group_size_tail = ncluster_fast % group_size
             else:
                 group_size, group_size_tail = args.group_size, None
+            num_clusters_in_group = None
             if const_expr(ncluster_slow is not None):
                 num_clusters_in_group = group_size * ncluster_slow
-            else:
-                num_clusters_in_group = None
+            if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC):
+                assert args.tile_count_semaphore is not None
             return VarlenMTileScheduler.Params(
                 problem_shape_ncluster_mnl,
                 args.total_m,
                 args.cu_seqlens_m,
                 raster_order,
                 group_size,
-                FastDivmod.create(group_size) if ncluster_fast is not None else None,
+                FastDivmod(group_size) if ncluster_fast is not None else None,
                 # Don't divide by 0
-                FastDivmod.create(group_size_tail if group_size_tail > 0 else 1)
+                FastDivmod(group_size_tail if group_size_tail > 0 else 1)
                 if group_size_tail is not None
                 else None,
-                FastDivmod.create(num_clusters_in_group)
-                if num_clusters_in_group is not None
-                else None,
+                FastDivmod(num_clusters_in_group) if num_clusters_in_group is not None else None,
                 args.tile_shape_mn,
-                args.tile_count_semaphore if const_expr(args.is_persistent) else None,
+                args.tile_count_semaphore
+                if const_expr(args.persistence_mode == PersistenceMode.DYNAMIC)
+                else None,
                 cluster_shape_mn,
-                args.is_persistent,
+                args.persistence_mode,
             )
 
     def __init__(
@@ -666,7 +739,7 @@ class VarlenMTileScheduler(TileScheduler):
         num_tiles_executed: Int32,
         current_batch_idx: Int32,
         num_work_idx_before_cur_batch: Int32,
-        tile_count: Optional[cute.Tensor],
+        sched_smem: Optional[cute.Tensor],
         scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync],
         pipeline_state: PipelineStateWAdvance,
         params: Params,
@@ -678,7 +751,7 @@ class VarlenMTileScheduler(TileScheduler):
         self.num_tiles_executed = num_tiles_executed
         self._current_batch_idx = current_batch_idx
         self._num_work_idx_before_cur_batch = num_work_idx_before_cur_batch
-        self._tile_count = tile_count
+        self._sched_smem = sched_smem
         self._scheduler_pipeline = scheduler_pipeline
         self._pipeline_state = pipeline_state
         self.params = params
@@ -693,9 +766,8 @@ class VarlenMTileScheduler(TileScheduler):
     @cute.jit
     def create(
         params: Params,
-        tile_count: Optional[cute.Tensor] = None,
+        sched_smem: Optional[cute.Tensor] = None,
         scheduler_pipeline: Optional[cutlass.pipeline.PipelineAsync] = None,
-        is_scheduler_warp: bool | Boolean = False,
         *,
         loc=None,
         ip=None,
@@ -703,18 +775,21 @@ class VarlenMTileScheduler(TileScheduler):
         stages = 0
         _, _, bidz = cute.arch.block_idx()
         current_work_linear_idx = Int32(bidz)
-        if const_expr(params.tile_count_semaphore is not None):
-            assert tile_count is not None
+        if const_expr(
+            params.persistence_mode
+            in [PersistenceMode.STATIC, PersistenceMode.DYNAMIC, PersistenceMode.CLC]
+        ):
+            assert sched_smem is not None
             assert scheduler_pipeline is not None
-            stages = const_expr(cute.size(tile_count))
+            stages = const_expr(cute.size(sched_smem, mode=[1]))
         return VarlenMTileScheduler(
             current_work_linear_idx,
             Int32(0),  # num_tiles_executed
             Int32(0),  # current_batch_idx
             Int32(0),  # num_work_idx_before_cur_batch
-            tile_count,
+            sched_smem,
             scheduler_pipeline,
-            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)),
+            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(0)),
             params,
             loc=loc,
             ip=ip,
@@ -733,54 +808,33 @@ class VarlenMTileScheduler(TileScheduler):
         num_batch = params.problem_shape_ncluster_mnl[2]
         total_clusters_m_max = (params.total_m + num_batch * (block_size - 1)) // block_size
         total_clusters_max = total_clusters_m_max * params.problem_shape_ncluster_mnl[1]
-        if const_expr(not params.is_persistent):
+        if const_expr(params.persistence_mode in [PersistenceMode.NONE, PersistenceMode.CLC]):
             return (*params.cluster_shape_mn, total_clusters_max)
         else:
             num_persistent_clusters = cutlass.min(max_active_clusters, total_clusters_max)
             return (*params.cluster_shape_mn, num_persistent_clusters)
 
-    @cute.jit
-    def _get_num_m_blocks(
-        self, lane: Int32, bidb_start: Int32, block_size: cutlass.Constexpr[int]
-    ) -> Int32:
-        num_batch = self.params.problem_shape_ncluster_mnl[2]
-        batch_idx = lane + bidb_start
-        cur_cu_seqlen = Int32(0)
-        if batch_idx <= num_batch:
-            cur_cu_seqlen = self.params.cu_seqlens_m[batch_idx]
-        next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
-        seqlen = next_cu_seqlen - cur_cu_seqlen
-        return (
-            cute.ceil_div(seqlen, block_size)
-            if batch_idx < num_batch and lane < cute.arch.WARP_SIZE - 1
-            else Int32(0)
-        )
-
     @cute.jit
     def _swizzle_cta(
         self, cluster_id_in_problem: Int32, num_clusters_m: Int32, *, loc=None, ip=None
     ) -> Tuple[Int32, Int32]:
         params = self.params
         # CTA Swizzle to promote L2 data reuse
-        if const_expr(params.num_clusters_in_group_divmod is not None):
-            group_id, id_in_group = params.num_clusters_in_group_divmod.divmod(
-                cluster_id_in_problem
-            )
-            num_clusters_in_group = params.num_clusters_in_group_divmod.divisor
+        if const_expr(params.num_clusters_in_group_fdd is not None):
+            group_id, id_in_group = divmod(cluster_id_in_problem, params.num_clusters_in_group_fdd)
+            num_clusters_in_group = params.num_clusters_in_group_fdd.divisor
         else:
             assert params.raster_order == RasterOrder.AlongN
             num_clusters_in_group = params.group_size * num_clusters_m
             group_id = cluster_id_in_problem // num_clusters_in_group
             id_in_group = cluster_id_in_problem - group_id * num_clusters_in_group
         cid_fast_in_group, cid_slow = Int32(0), Int32(0)
-        if const_expr(
-            params.group_size_divmod is not None and params.group_size_tail_divmod is not None
-        ):
+        if const_expr(params.group_size_fdd is not None and params.group_size_tail_fdd is not None):
             num_clusters = num_clusters_m * params.problem_shape_ncluster_mnl[1]
             if (group_id + 1) * num_clusters_in_group <= num_clusters:
-                cid_slow, cid_fast_in_group = params.group_size_divmod.divmod(id_in_group)
+                cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_fdd)
             else:  # tail part
-                cid_slow, cid_fast_in_group = params.group_size_tail_divmod.divmod(id_in_group)
+                cid_slow, cid_fast_in_group = divmod(id_in_group, params.group_size_tail_fdd)
         else:
             assert params.raster_order == RasterOrder.AlongM
             group_size_actual = cutlass.min(
@@ -802,7 +856,26 @@ class VarlenMTileScheduler(TileScheduler):
         return cid_m, cid_n
 
     @cute.jit
-    def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo:
+    def _get_num_m_blocks(
+        self, lane: Int32, bidb_start: Int32, block_size: cutlass.Constexpr[int]
+    ) -> Int32:
+        num_batch = self.params.problem_shape_ncluster_mnl[2]
+        batch_idx = lane + bidb_start
+        cur_cu_seqlen = Int32(0)
+        if batch_idx <= num_batch:
+            cur_cu_seqlen = self.params.cu_seqlens_m[batch_idx]
+        next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
+        seqlen = next_cu_seqlen - cur_cu_seqlen
+        return (
+            cute.ceil_div(seqlen, block_size)
+            if batch_idx < num_batch and lane < cute.arch.WARP_SIZE - 1
+            else Int32(0)
+        )
+
+    @cute.jit
+    def _delinearize_work_idx(
+        self, *, block_zero_only: bool = False, loc=None, ip=None
+    ) -> cutlass.utils.WorkTileInfo:
         params = self.params
         lane_idx = cute.arch.lane_idx()
         num_batch = self.params.problem_shape_ncluster_mnl[2]
@@ -819,7 +892,6 @@ class VarlenMTileScheduler(TileScheduler):
         )
         problems_end_tile = self._num_work_idx_before_cur_batch + clusters_in_problems
         # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, problems_end_tile = %d, num_clusters_m=%d, num_clusters_cumulative = %d, problems_end_tile = %d", self._tile_idx, problems_end_tile, num_clusters_m, num_clusters_cumulative, problems_end_tile)
-        cid_m, cid_n = Int32(0), Int32(0)
         next_tile_idx = self._current_work_linear_idx
         while problems_end_tile <= next_tile_idx:
             batch_idx += cute.arch.WARP_SIZE - 1
@@ -836,11 +908,14 @@ class VarlenMTileScheduler(TileScheduler):
                     num_clusters_cumulative, cute.arch.WARP_SIZE - 1
                 )
                 problems_end_tile += clusters_in_problems
+        if const_expr(params.persistence_mode == PersistenceMode.NONE):
+            is_valid = self.num_tiles_executed == 0 and batch_idx < num_batch
+        else:
+            is_valid = batch_idx < num_batch
         # Just a placeholer value in case batch_idx >= num_batch
         num_work_idx_before_cur_batch = problems_end_tile - clusters_in_problems
-        if batch_idx >= num_batch:
-            cid_m, cid_n, batch_idx = Int32(0), Int32(0), Int32(num_batch)
-        else:
+        cid_m, cid_n = Int32(0), Int32(0)
+        if is_valid:
             problems_start_tile = problems_end_tile - clusters_in_problems
             # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, problems_end_tile = %d, num_clusters_m=%d, batch_idx = %d", self._tile_idx, problems_end_tile, num_clusters_m, batch_idx)
             # The next problem to process is the first one that does not have ending tile position
@@ -859,42 +934,21 @@ class VarlenMTileScheduler(TileScheduler):
             num_clusters_m = cute.arch.shuffle_sync(num_clusters_m, batch_idx_in_problems)
             num_work_idx_before_cur_batch = problems_start_tile + num_clusters_prev_lane
             cluster_id_in_problem = next_tile_idx - num_work_idx_before_cur_batch
-            # cid_n = cluster_id_in_problem // num_clusters_m
-            # cid_m = cluster_id_in_problem - cid_n * num_clusters_m
             # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, cid_n=%d, cid_m=%d, is_valid = %d", self._tile_idx, batch_idx, cid_n, cid_m, is_valid)
             cid_m, cid_n = self._swizzle_cta(cluster_id_in_problem, num_clusters_m, loc=loc, ip=ip)
         self._current_batch_idx = batch_idx
         self._num_work_idx_before_cur_batch = num_work_idx_before_cur_batch
 
-        # Get the pid from cluster id
-        bidx_in_cluster = cute.arch.block_in_cluster_idx()
+        if const_expr(block_zero_only):
+            bidx_in_cluster = (Int32(0), Int32(0))
+        else:
+            # Get the pid from cluster id
+            bidx_in_cluster = cute.arch.block_in_cluster_idx()
         pid_m = cid_m * params.cluster_shape_mn[0] + bidx_in_cluster[0]
         pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
         tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
-        if const_expr(not params.is_persistent):
-            is_valid = self.num_tiles_executed == 0 and batch_idx < num_batch
-        else:
-            is_valid = batch_idx < num_batch
         return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid)
 
-    @cute.jit
-    def fetch_next_work(self, is_scheduler_warp: bool | Boolean = False, *, loc=None, ip=None):
-        """is_scheduler_warp should only be true for one warp in the whole cluster"""
-        if const_expr(self.params.tile_count_semaphore is not None):
-            params = self.params
-            current_work_linear_idx = self._current_work_linear_idx
-            if is_scheduler_warp:
-                if cute.arch.lane_idx() == 0:
-                    # cute.printf("before atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
-                    num_persistent_clusters = cute.arch.grid_dim()[2]
-                    current_work_linear_idx = num_persistent_clusters + utils.atomic_add_i32(
-                        1, params.tile_count_semaphore
-                    )
-                    # cute.printf("after atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
-                # lane 0 already has the right tile_idx, just need to broadcast
-                current_work_linear_idx = cute.arch.shuffle_sync(current_work_linear_idx, 0)
-            self._current_work_linear_idx = current_work_linear_idx
-
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
         for obj in [
@@ -902,7 +956,7 @@ class VarlenMTileScheduler(TileScheduler):
             self.num_tiles_executed,
             self._current_batch_idx,
             self._num_work_idx_before_cur_batch,
-            self._tile_count,
+            self._sched_smem,
             self._scheduler_pipeline,
             self._pipeline_state,
             self.params,
@@ -920,7 +974,7 @@ class VarlenMTileScheduler(TileScheduler):
                 self.num_tiles_executed,
                 self._current_batch_idx,
                 self._num_work_idx_before_cur_batch,
-                self._tile_count,
+                self._sched_smem,
                 self._scheduler_pipeline,
                 self._pipeline_state,
                 self.params,