# ********************************************************************************
# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
# ********************************************************************************

# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import enum
import math
import operator
from functools import partial
from typing import Callable, Optional, Tuple, Type, Union

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.utils as utils
import cutlass.utils.hopper_helpers as sm90_utils
import torch
from cutlass import Float32, Int32, const_expr
from cutlass._mlir.dialects import llvm, vector
from cutlass.cute.nvgpu import cpasync, warp, warpgroup
from cutlass.cute.runtime import from_dlpack
from cutlass.cutlass_dsl import T, dsl_user_op
from quack.copy_utils import sm90_get_smem_load_op
from quack.cute_dsl_utils import ParamsBase
from quack.layout_utils import make_acc_tensor_mn_view

# return PipelineStateWAdvance instead of PipelineState
from quack.pipeline import PipelineTmaCpAsync, make_pipeline_state
from quack.sm90_utils import partition_for_epilogue
from quack.tensormap_manager import TensorMapManagerSm90
from quack.tile_scheduler import RasterOrderOption, TileSchedulerArguments, VarlenMTileSchedulerArguments

from .tile_scheduler import SonicMoETileScheduler, SonicMoEVarlenMTileScheduler


class NamedBarrierGemm(enum.IntEnum):
    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
    EpilogueLoad = enum.auto()
    MmaWG0 = enum.auto()
    MmaWG1 = enum.auto()
    EpiWG0 = enum.auto()
    EpiWG1 = enum.auto()
    Prolog = enum.auto()


class HopperWgmma_MoE_kernel:
    def __init__(
        self,
        E: int,
        acc_dtype: Type[cutlass.Numeric],
        tile_shape_mnk: Tuple[int, int, int],
        cluster_shape_mnk: Tuple[int, int, int],
        pingpong: bool = False,
        is_persistent: bool = True,
        compute_dz_and_partial_ds_and_y1s: bool = False,
        compute_weight_gradient: bool = False,
        compute_relu: bool = False,
        compute_silu: bool = False,
        compute_gelu: bool = False,
        compute_relu_sq: bool = False,
        compute_swiglu: bool = False,
        compute_reglu: bool = False,
        compute_geglu: bool = False,
        is_normal_act: bool = False,
        is_glu: bool = False,
        is_A_gather: bool = False,
        is_scatter_idx_prefetched: bool = False,
        epi_tile_size: int = 32,
        initial_d_epi_stage: int = 4,
        index_dtype: Type[cutlass.Numeric] = cutlass.Int32,
        prefetch_idx_store_to_smem: int = 2048,
        inference_mode: bool = False,
        L2_group_size: int = 8,
        raster_order: RasterOrderOption = RasterOrderOption.Heuristic,
    ):
        self.epi_tile_size = epi_tile_size
        self.initial_d_epi_stage = initial_d_epi_stage

        self.is_A_gather = is_A_gather
        self.is_scatter_idx_prefetched = is_scatter_idx_prefetched

        self.compute_swiglu = compute_swiglu
        self.compute_geglu = compute_geglu
        self.compute_reglu = compute_reglu

        self.compute_relu = compute_relu
        self.compute_silu = compute_silu
        self.compute_gelu = compute_gelu
        self.compute_relu_sq = compute_relu_sq

        self.is_glu = is_glu or (compute_swiglu or compute_geglu or compute_reglu)
        self.is_normal_act = is_normal_act or (compute_gelu or compute_relu_sq or compute_relu or compute_silu)

        self.compute_dz_and_partial_ds_and_y1s = compute_dz_and_partial_ds_and_y1s
        self.compute_weight_gradient = compute_weight_gradient

        self.need_adhoc_epilogue_store = self.is_glu or self.is_normal_act or compute_dz_and_partial_ds_and_y1s
        self.need_epilogue_load = compute_dz_and_partial_ds_and_y1s

        self.L2_group_size = L2_group_size
        self.raster_order = raster_order

        self.E = E
        self.acc_dtype = acc_dtype
        assert self.acc_dtype == cutlass.Float32
        self.pingpong = pingpong
        self.is_persistent = is_persistent
        if self.pingpong:
            assert self.is_persistent, "Pingpong gemm requires persistent scheduler"

        self.cluster_shape_mnk = cluster_shape_mnk
        self.tile_shape_mnk = tuple(tile_shape_mnk)
        tile_M, tile_N = tile_shape_mnk[0], tile_shape_mnk[1]
        # check the cta tile shape
        if not self.pingpong:
            if tile_M not in [64, 128, 192, 256, 320]:
                raise ValueError("CTA tile shape M must be 64/128/192/256/320")
            if tile_M in [192, 320]:  # special case
                tile_N_max = 256 if tile_M == 192 else 160
                if not (tile_N % 32 == 0 and tile_N <= tile_N_max):
                    raise ValueError(
                        f"If tile_m == {tile_M}, CTA tile shape N must be divisible by 32 and <= {tile_N_max}"
                    )
            else:
                if not ((tile_N % 16 == 0 and tile_N <= 256) or (tile_N % 32 == 0 and tile_N <= 512)):
                    raise ValueError(
                        "CTA tile shape N must be divisible by 16 and <= 256, or divisible by 32 and <= 512"
                    )
        else:
            if tile_M not in [64, 128, 192]:
                raise ValueError("CTA tile shape M must be 64/128/192 if pingpong")
            tile_N_max = 256 if tile_M == 64 else (208 if tile_M == 128 else 128)
            if not (tile_N % 16 == 0 and tile_N <= tile_N_max):
                raise ValueError(f"CTA tile shape N must be divisible by 16 and <= {tile_N_max}")
        if not self.tile_shape_mnk[2] % 16 == 0:
            raise ValueError("CTA tile shape K must be divisible by 16")

        self.tile_M, self.tile_N, self.tile_K = tile_shape_mnk

        if not self.pingpong:
            if tile_M == 320:  # tile_M / 64 is not even so we have to split along N
                atom_layout_m, atom_layout_n = 1, 2
            elif tile_M == 192:
                if tile_N <= 128:
                    atom_layout_m, atom_layout_n = 3, 1
                else:
                    atom_layout_m, atom_layout_n = 1, 2
            else:
                atom_layout_m = tile_shape_mnk[0] // 64 if tile_shape_mnk[0] < 256 else 2
                atom_layout_n = 1
            assert atom_layout_m in [1, 2, 3] and atom_layout_n in [1, 2]
        else:
            atom_layout_m, atom_layout_n = 1, 1
        self.atom_layout_mnk = (atom_layout_m, atom_layout_n, 1)

        if is_A_gather:
            assert self.cluster_shape_mnk[1] == 1
            self.num_mcast_ctas_a = None
            self.is_a_mcast = False
        else:
            self.num_mcast_ctas_a = self.cluster_shape_mnk[1]
            self.is_a_mcast = self.num_mcast_ctas_a > 1
        self.num_mcast_ctas_b = self.cluster_shape_mnk[0]
        self.is_b_mcast = self.num_mcast_ctas_b > 1

        self.occupancy = 1
        self.mma_warp_groups = math.prod(self.atom_layout_mnk) * (1 if not self.pingpong else 2)
        if self.pingpong:
            assert self.mma_warp_groups == 2
        self.num_threads_per_warp_group = 128
        self.threads_per_cta = (self.mma_warp_groups + 1) * self.num_threads_per_warp_group
        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_90")
        self.num_mma_threads = (self.mma_warp_groups if not self.pingpong else 1) * self.num_threads_per_warp_group
        self.num_epi_threads = (self.mma_warp_groups if not self.pingpong else 1) * self.num_threads_per_warp_group
        self.tma_warp_id = self.mma_warp_groups * 4
        self.universal_copy_bits = 128
        # assumed BF16 now
        self.num_load_A_threads = (
            min(self.tile_M * self.tile_K // 8, self.threads_per_cta - self.tma_warp_id * cute.arch.WARP_SIZE)
            if is_A_gather
            else 0
        )
        if self.compute_weight_gradient and self.is_A_gather:
            if tile_M == 192:  # contiguous dimension
                self.num_load_A_threads = 3 * 32
            assert tile_M in [64, 128, 192, 256]

        self.num_epi_load_threads = 0
        if self.need_epilogue_load:
            # 3 warps to load A, 1 warp to load C, (and 1 warp to load S)
            self.num_load_A_threads = 4 * cute.arch.WARP_SIZE
            self.num_epi_load_threads = self.num_epi_threads

        regs_per_thread = math.prod(self.tile_shape_mnk[:2]) // self.num_mma_threads
        heavy_register_pressure = regs_per_thread >= 208

        if not is_A_gather:
            if self.mma_warp_groups == 3:
                self.num_regs_load, self.num_regs_mma = 32, 160
            else:
                heavy_register_pressure = regs_per_thread >= 208
                self.num_regs_load, self.num_regs_mma = (40, 232) if not heavy_register_pressure else (24, 240)
        else:
            if self.mma_warp_groups == 3:
                self.num_regs_load, self.num_regs_mma = 56, 152
            else:
                self.num_regs_load, self.num_regs_mma = (56, 224)

        self.ab_stage = None
        self.c_epi_stage = None
        self.d_epi_stage = None

        self.a_smem_layout_staged = None
        self.b_smem_layout_staged = None
        self.d_epi_smem_layout_staged = None
        self.d_epi_tile = None

        self.shared_storage = None
        self.buffer_align_bytes = 1024

        self.tensormap_update_mode = cutlass.utils.TensorMapUpdateMode.SMEM
        self.bytes_per_tensormap = 128
        self.tensor_memory_management_bytes = 12

        self.inference_mode = inference_mode

        if is_A_gather:
            if self.need_adhoc_epilogue_store:
                if self.inference_mode:
                    self.num_tensormaps = self.mma_warp_groups if self.pingpong else 1
                else:
                    self.num_tensormaps = 2 * self.mma_warp_groups if self.pingpong else 2
            else:
                self.num_tensormaps = 1 * self.mma_warp_groups if self.pingpong else 1
        else:
            if self.need_adhoc_epilogue_store:
                if self.inference_mode:
                    self.num_tensormaps = 2 * self.mma_warp_groups + 1 if self.pingpong else 3
                else:
                    self.num_tensormaps = self.mma_warp_groups + 1 if self.pingpong else 2
            else:
                self.num_tensormaps = 1 * self.mma_warp_groups + 1 if self.pingpong else 2

        if self.need_epilogue_load:
            self.num_tensormaps += 2 * self.mma_warp_groups if self.pingpong else 1

        if self.compute_weight_gradient:
            if self.is_A_gather:
                self.num_tensormaps = 1
                self.prefetch_token_idx_size = prefetch_idx_store_to_smem
                self.index_dtype = index_dtype

                assert (
                    self.prefetch_token_idx_size % self.tile_K == 0
                    and self.prefetch_token_idx_size >= self.tile_K
                    and self.prefetch_token_idx_size % self.num_load_A_threads == 0
                )
            else:
                self.num_tensormaps = 2
                self.prefetch_token_idx_size = 0
                self.index_dtype = None
        else:
            self.prefetch_token_idx_size = 0
            self.index_dtype = None

        self.tensormap_bytes_total = self.num_tensormaps * self.bytes_per_tensormap

    def _setup_attributes(self):
        self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk)

        self.d_epi_tile = self._sm90_compute_tile_shape_or_override(
            self.tile_shape_mnk,
            self.atom_layout_mnk,
            self.d_dtype,
        )
        self.c_epi_tile = self.d_epi_tile
        if const_expr(self.compute_dz_and_partial_ds_and_y1s):
            self.y_epi_tile = self.d_epi_tile
        elif const_expr(self.is_glu):
            self.y_epi_tile = (self.d_epi_tile[0], self.d_epi_tile[1] // 2)
        elif const_expr(self.is_normal_act):
            self.y_epi_tile = self.d_epi_tile
        else:
            self.y_epi_tile = None

        if const_expr(self.use_bias):
            self.bias_epi_tile = self.d_epi_tile
            self.initial_d_epi_stage -= 1  # for safety
        else:
            self.bias_epi_tile = None

        # Compute stage before compute smem layout
        self.ab_stage, self.c_epi_stage, self.d_epi_stage, self.y_epi_stage = self._compute_stages(
            self.tile_shape_mnk,
            self.initial_d_epi_stage,
            # epi_smem will reuse smem ab if not persistent.
            self.d_epi_tile,
            self.c_epi_tile,
            self.y_epi_tile,
            self.a_dtype,
            self.b_dtype,
            self.d_dtype,
            self.c_dtype,
            self.y_dtype,
            self.smem_capacity,
            self.occupancy,
            # epi_smem will reuse smem ab if not persistent.
            overlap_sD_sA=not self.is_persistent,
        )

        if const_expr((not self.inference_mode) and self.need_adhoc_epilogue_store):
            assert self.d_epi_stage == self.y_epi_stage

        self.sched_stage = 2 if self.pingpong else 1

        (
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.c_epi_smem_layout_staged,
            self.bias_epi_smem_layout_staged,
            self.d_epi_smem_layout_staged,
            self.y_epi_smem_layout_staged,
            self.s_epi_smem_layout_staged,
            self.prefetch_AIdx_smem_layout_staged,
        ) = self._make_smem_layouts(
            self.tile_shape_mnk,
            self.c_epi_tile,
            self.bias_epi_tile,
            self.d_epi_tile,
            self.y_epi_tile,
            self.a_dtype,
            self.a_layout,
            self.b_dtype,
            self.b_layout,
            self.prefetch_token_idx_size,
            self.ab_stage,
            self.c_dtype,
            self.c_layout,
            self.bias_dtype,
            self.bias_layout,
            self.d_dtype,
            self.d_layout,
            self.y_dtype,
            self.y_layout,
            self.s_dtype,
            self.c_epi_stage,
            self.d_epi_stage,
            self.y_epi_stage,
        )

    @dsl_user_op
    def tanh(self, a: float | Float32, *, loc=None, ip=None) -> Float32:
        return Float32(
            llvm.inline_asm(
                T.f32(),
                [Float32(a).ir_value(loc=loc, ip=ip)],
                "tanh.approx.f32 $0, $1;",
                "=f,f",
                has_side_effects=False,
                is_align_stack=False,
                asm_dialect=llvm.AsmDialect.AD_ATT,
            )
        )

    @dsl_user_op
    def fma(self, a: float | Float32, b: float | Float32, c: float | Float32, *, loc=None, ip=None) -> Float32:
        return Float32(
            llvm.inline_asm(
                T.f32(),
                [
                    Float32(a).ir_value(loc=loc, ip=ip),
                    Float32(b).ir_value(loc=loc, ip=ip),
                    Float32(c).ir_value(loc=loc, ip=ip),
                ],
                "fma.rn.f32 $0, $1, $2, $3;",
                "=f,f,f,f",
                has_side_effects=False,
                is_align_stack=False,
                asm_dialect=llvm.AsmDialect.AD_ATT,
            )
        )

    @dsl_user_op
    def silu(self, a: float | Float32, *, loc=None, ip=None) -> Float32:
        """
        silu(a) = a * sigmoid(a) = a * (1 + tanh(a / 2)) / 2 = (0.5 * a) * tanh(0.5 * a) + (0.5 * a)
        This compiles down to 3 SASS instructions: FMUL to get 0.5 * a, MUFU.TANH, and FFMA.
        """
        # return a / (1.0 + cute.arch.exp2(-a * math.log2(math.e)))
        a_half = 0.5 * a
        # return a_half * self.tanh(a_half) + a_half
        return self.fma(a_half, self.tanh(a_half), a_half)

    @dsl_user_op
    def relu(self, a: float | Float32, *, loc=None, ip=None) -> Float32:
        return cute.arch.fmax(a, 0.0)

    @dsl_user_op
    def relu_sq(self, a: float | Float32, *, loc=None, ip=None) -> Float32:
        return a * cute.arch.fmax(a, 0.0)

    @dsl_user_op
    def gelu(self, a: Float32, *, loc=None, ip=None) -> Float32:
        # gelu(x) ≈ 0.5*x*(1 + tanh(√(2/π)*(x + 0.044715*x^3)))
        c0 = const_expr(math.sqrt(2 / math.pi))  # √(2/π)
        c1 = 0.044715
        a2 = a * a
        # inner = √(2/π) * (x + 0.044715*x^3)
        inner = c0 * self.fma(c1, a2 * a, a)
        return 0.5 * a * self.fma(1.0, self.tanh(inner), 1.0)

    @dsl_user_op
    def elem_pointer(self, x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
        return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)

    @dsl_user_op
    def min_i32(self, a: int | Int32, b: int | Int32, *, loc=None, ip=None) -> Int32:
        return Int32(
            llvm.inline_asm(
                T.i32(),  # return type
                [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
                "min.s32 $0, $1, $2;",
                "=r,r,r",  # output, input constraints
                has_side_effects=False,
                is_align_stack=False,
                asm_dialect=llvm.AsmDialect.AD_ATT,
            )
        )

    @cute.jit
    def prefetch_gather_idx_for_A_when_vary_M(
        self, mAIdx: cute.Tensor, M_offset: int, M_boundary: int, copy_elems_per_thr_load: int  # m n k l
    ) -> cute.Tensor:
        assert const_expr(not self.compute_weight_gradient)
        M, K = self.tile_M, self.tile_K

        tidx, _, _ = cute.arch.thread_idx()
        tidx = tidx - self.tma_warp_id * cute.arch.WARP_SIZE

        stride_1_tile, other_tile = K, M

        threads_per_stride_1_dim = const_expr(stride_1_tile // copy_elems_per_thr_load)
        num_other_dim_per_load = const_expr(self.num_load_A_threads // threads_per_stride_1_dim)

        num_other_dim_per_thread = const_expr(other_tile // num_other_dim_per_load)
        tmAIdx = cute.make_rmem_tensor((num_other_dim_per_load,), dtype=mAIdx.element_type)

        for i in cutlass.range_constexpr(num_other_dim_per_thread):
            other_dim_offset = const_expr(i * num_other_dim_per_load) + tidx // threads_per_stride_1_dim

            if other_dim_offset < M_boundary:
                M_i = M_offset + other_dim_offset
                tmAIdx[i] = mAIdx[M_i]

        return tmAIdx

    @cute.jit
    def prefetch_scatter_idx_for_D_when_vary_M(
        self,
        mD: cute.Tensor,  # unused, kept for symmetry
        mDIdx: cute.Tensor,
        D_r2g_thr_copy,
        tcDgcD_flat_partition: cute.Tensor,
        epi_tile_layout: cute.Layout,
        epi_tile_num: int,
        copy_elems_per_thr_load: int,  # unused here, but fine to keep
        tile_coord_mnkl: Tuple[int, int, None, int],  # (block_M, block_N, _, batch)
        MIdx_cur_group: int,
        MIdx_next_group: int,
    ) -> cute.Tensor:
        # Same base M offset as store_D_scatter
        block_M, block_N = tile_coord_mnkl[0], tile_coord_mnkl[1]
        M_offset = block_M * const_expr(self.tile_M) + MIdx_cur_group

        tDcD0 = D_r2g_thr_copy.partition_D(tcDgcD_flat_partition[None, None, *epi_tile_layout.get_hier_coord(0)])
        num_load_per_thread = const_expr(cute.size(tDcD0, mode=[1]))

        tmDIdx = cute.make_rmem_tensor((epi_tile_num * num_load_per_thread,), dtype=mDIdx.element_type)
        tmDIdx = cute.make_rmem_tensor((epi_tile_num * num_load_per_thread,), dtype=mDIdx.element_type)

        for epi_idx in cutlass.range_constexpr(epi_tile_num):
            tDcD_slice = D_r2g_thr_copy.partition_D(
                tcDgcD_flat_partition[None, None, *epi_tile_layout.get_hier_coord(epi_idx)]
            )

            for i in cutlass.range_constexpr(num_load_per_thread):
                # Same coordinate source as in store_D_scatter
                MIdx_in_tile, _ = tDcD_slice[0, i, 0]
                MIdx = M_offset + MIdx_in_tile

                if MIdx < MIdx_next_group:
                    tmDIdx[epi_idx * num_load_per_thread + i] = mDIdx[MIdx]

        return tmDIdx

    @cute.jit
    def prefetch_gather_idx_for_A_when_vary_K(
        self,
        mAIdx: cute.Tensor,
        sAIdx: cute.Tensor,
        token_group_size: int,
        K_offset: int,
    ) -> cute.Tensor:
        assert const_expr(self.compute_weight_gradient and self.is_A_gather)

        tidx, _, _ = cute.arch.thread_idx()
        tidx = tidx - self.tma_warp_id * cute.arch.WARP_SIZE

        # !!! cannot be removed for correctness !!!
        cute.arch.barrier(barrier_id=NamedBarrierGemm.Prolog, number_of_threads=self.num_load_A_threads)

        for i in cutlass.range_constexpr(cute.ceil_div(self.prefetch_token_idx_size, self.num_load_A_threads)):
            offset = const_expr(i * self.num_load_A_threads) + tidx
            kidx = K_offset + offset

            if kidx < token_group_size:
                sAIdx[offset] = mAIdx[kidx]

        # !!! cannot be removed for correctness !!!
        cute.arch.barrier(barrier_id=NamedBarrierGemm.Prolog, number_of_threads=self.num_load_A_threads)

    @cute.jit
    def load_A_gather(
        self,
        mA: cute.Tensor,
        tmAIdx: Optional[cute.Tensor],
        sAIdx_prefetch: cute.Tensor,
        M_offset: cutlass.Int32,
        tAsA: cute.Tensor,
        tApA: cute.Tensor,
        A_g2s_thr_copy,
        K_offset: cutlass.Int32,
        token_group_size: cutlass.Int32,
        copy_elems_per_thr_load: cutlass.Int32,
    ):
        M, K = self.tile_M, self.tile_K

        tidx, _, _ = cute.arch.thread_idx()
        tidx = tidx - self.tma_warp_id * cute.arch.WARP_SIZE

        if const_expr(self.compute_weight_gradient):
            stride_1_tile, other_tile = M, K
        else:
            stride_1_tile, other_tile = K, M

        threads_per_stride_1_dim = const_expr(stride_1_tile // copy_elems_per_thr_load)
        num_other_dim_per_load = const_expr(self.num_load_A_threads // threads_per_stride_1_dim)

        K_offset_mod_smem_load = K_offset % const_expr(self.prefetch_token_idx_size)
        for i in cutlass.range_constexpr(cute.ceil_div(other_tile, num_other_dim_per_load)):
            stride_1_dim_offset = (tidx % threads_per_stride_1_dim) * copy_elems_per_thr_load
            other_dim_offset = const_expr(i * num_other_dim_per_load) + tidx // threads_per_stride_1_dim

            if const_expr(self.compute_weight_gradient):
                MIdx = M_offset + stride_1_dim_offset
                KIdx_global = K_offset + other_dim_offset

                if KIdx_global < token_group_size and MIdx < mA.shape[0]:
                    KIdx = sAIdx_prefetch[K_offset_mod_smem_load + other_dim_offset]
                    # KIdx = mAIdx_mk[K_offset + other_dim_offset]
                    tPrAptr = self.elem_pointer(mA, (MIdx, KIdx)).align(
                        self.universal_copy_bits // copy_elems_per_thr_load
                    )
                    mA_cur_copy = cute.make_tensor(tPrAptr, ((copy_elems_per_thr_load, 1), 1))

                    cute.copy(A_g2s_thr_copy, mA_cur_copy, tAsA[None, None, i])

            else:
                MIdx = tmAIdx[i]
                KIdx = K_offset + stride_1_dim_offset

                tPrAptr = self.elem_pointer(mA, (MIdx, KIdx)).align(
                    self.universal_copy_bits // copy_elems_per_thr_load
                )
                mA_cur_copy = cute.make_tensor(tPrAptr, ((copy_elems_per_thr_load, 1), 1))
                cute.copy(A_g2s_thr_copy, mA_cur_copy, tAsA[None, i, None], pred=tApA[None, i, None])

    @cute.jit
    def store_D_scatter(
        self,
        mD: cute.Tensor,  # m, n, k, l
        mDIdx: cute.Tensor,
        tmDIdx: cute.Tensor,  # assume to have same size as mD
        tDrD: cute.Tensor,
        tDcD_slice: cute.Tensor,  # ((8, 1), 16, 1)
        D_r2g_thr_copy,
        epi_idx: cutlass.Int32,
        copy_elems_per_thr_load: cutlass.Int32,
        tile_coord_mnkl: Tuple[int, int, None, int],  # m n k l
        MIdx_cur_group: int,
        MIdx_next_group: int,
    ):
        block_M, block_N = tile_coord_mnkl[0], tile_coord_mnkl[1]

        M_offset = block_M * const_expr(self.tile_M) + MIdx_cur_group
        N_offset = block_N * const_expr(self.tile_N)

        num_load_per_thread = const_expr(cute.size(tDcD_slice, mode=[1]))
        for i in cutlass.range_constexpr(num_load_per_thread):
            MIdx_in_epi_tile, NIdx_in_epi_tile = tDcD_slice[0, i, 0]

            MIdx = M_offset + MIdx_in_epi_tile
            NIdx = N_offset + NIdx_in_epi_tile

            if MIdx < MIdx_next_group and NIdx < mD.shape[1]:
                if const_expr(self.is_scatter_idx_prefetched):
                    SIdx = tmDIdx[i + epi_idx * num_load_per_thread]
                else:
                    SIdx = mDIdx[MIdx]  # equivalent
                tPDptr = self.elem_pointer(mD, (SIdx, NIdx)).align(self.universal_copy_bits // copy_elems_per_thr_load)

                mD_cur_copy = cute.make_tensor(tPDptr, ((copy_elems_per_thr_load, 1), 1))

                cute.copy(
                    D_r2g_thr_copy,
                    tDrD[None, i, None],
                    mD_cur_copy,
                )

    @cute.jit
    def fetch_scattered_S(
        self,
        tidx: int,
        mS: cute.Tensor,
        mS_scatter_idx: cute.Tensor,
        sS_staged: cute.Tensor,
        tile_coord_mnkl: Tuple[int, int, None, int],  # m n k l
        MIdx_cur_group: int,
        MIdx_next_group: int,
    ):
        block_M = tile_coord_mnkl[0]
        M = self.tile_M

        M_s = block_M * M + MIdx_cur_group

        for i in cutlass.range_constexpr(cute.ceil_div(M, self.num_epi_threads)):
            sS_offset = const_expr(i * self.num_epi_threads) + tidx
            M_i = M_s + sS_offset

            if M_i < MIdx_next_group and sS_offset < M:
                sIdx = mS_scatter_idx[M_i]
                sS_staged[sS_offset] = self.s_dtype(mS[sIdx])

    @dsl_user_op
    def prmt(self, a: int | Int32, b: int | Int32, c: int | Int32, *, loc=None, ip=None) -> Int32:
        return Int32(
            llvm.inline_asm(
                T.i32(),
                [
                    Int32(a).ir_value(loc=loc, ip=ip),
                    Int32(b).ir_value(loc=loc, ip=ip),
                    Int32(c).ir_value(loc=loc, ip=ip),
                ],
                "prmt.b32 $0, $1, $2, $3;",
                "=r,r,r,r",
                has_side_effects=False,
                is_align_stack=False,
                asm_dialect=llvm.AsmDialect.AD_ATT,
            )
        )

    @dsl_user_op
    def pack2x16_as_f32(
        self,
        a: Union[cutlass.BFloat16, cutlass.Float16],
        b: Union[cutlass.BFloat16, cutlass.Float16],
        *,
        loc=None,
        ip=None,
    ) -> cutlass.Float32:
        vec_src_type = T.bf16() if a.dtype == cutlass.BFloat16 else T.f16()

        vec_f16x2 = vector.from_elements(T.vector(2, vec_src_type), (a.ir_value(), b.ir_value()), loc=loc, ip=ip)
        vec_f32x1 = vector.bitcast(T.vector(1, T.f32()), vec_f16x2)
        return cutlass.Float32(vector.extract(vec_f32x1, dynamic_position=[], static_position=[0], loc=loc, ip=ip))

    @dsl_user_op
    def unpack2x16_as_2xf32(
        self, a: Float32, dtype: cutlass.Numeric, *, loc=None, ip=None
    ) -> Tuple[cutlass.Float32, cutlass.Float32]:

        vec_dst_type = T.bf16() if dtype == cutlass.BFloat16 else T.f16()

        vec_f32x1 = vector.from_elements(T.vector(1, T.f32()), (a.ir_value(),), loc=loc, ip=ip)
        vec_f16x2 = vector.bitcast(T.vector(2, vec_dst_type), vec_f32x1)
        res0 = Float32(vector.extract(vec_f16x2, dynamic_position=[], static_position=[0], loc=loc, ip=ip))
        res1 = Float32(vector.extract(vec_f16x2, dynamic_position=[], static_position=[1], loc=loc, ip=ip))
        return res0, res1

    @cute.jit
    def permute_gated_Cregs_b16(self, t: cute.Tensor) -> None:
        assert t.element_type.width == 16
        assert cute.size(t.shape) % 4 == 0, "Tensor size must be a multiple of 4 for b16 permutation"
        t_u32 = cute.recast_tensor(t, Int32)

        quad_idx = cute.arch.lane_idx() % 4
        lane_03 = quad_idx == 0 or quad_idx == 3
        selector_upper = Int32(0x5410) if lane_03 else Int32(0x1054)
        selector_lower = Int32(0x7632) if lane_03 else Int32(0x3276)
        # upper_map = [0, 3, 1, 2]
        # lower_map = [1, 2, 0, 3]
        # upper_idx = upper_map[quad_idx]
        # indexing isn't supported so we have to do arithmetic
        upper_idx = quad_idx // 2 if quad_idx % 2 == 0 else 3 - quad_idx // 2
        lower_idx = upper_idx ^ 1

        # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000
        width = 4
        mask = cute.arch.WARP_SIZE - width
        clamp = cute.arch.WARP_SIZE - 1
        mask_and_clamp = const_expr(mask << 8 | clamp)

        for i in cutlass.range_constexpr(cute.size(t_u32.shape) // 2):
            upper, lower = t_u32[i * 2 + 0], t_u32[i * 2 + 1]
            upper0 = upper if lane_03 else lower
            lower0 = lower if lane_03 else upper
            upper0 = cute.arch.shuffle_sync(upper0, offset=upper_idx, mask_and_clamp=mask_and_clamp)
            lower0 = cute.arch.shuffle_sync(lower0, offset=lower_idx, mask_and_clamp=mask_and_clamp)
            t_u32[i * 2 + 0] = self.prmt(upper0, lower0, selector_upper)
            t_u32[i * 2 + 1] = self.prmt(upper0, lower0, selector_lower)

    @cute.jit
    def __call__(
        self,
        mA: cute.Tensor,
        mB: cute.Tensor,
        mC: Optional[cute.Tensor],
        mBias: Optional[cute.Tensor],
        mD: cute.Tensor,
        mY: Optional[cute.Tensor],
        mS: Optional[cute.Tensor],
        mDS_partial: Optional[cute.Tensor],
        mMoffset: cute.Tensor,
        mAIdx: Optional[cute.Tensor],
        mDIdx: Optional[cute.Tensor],
        mS_scatter_idx: Optional[cute.Tensor],
        mA_tensormap: Optional[cute.Tensor],
        mB_tensormap: Optional[cute.Tensor],
        mC_tensormap: Optional[cute.Tensor],
        mD_tensormap: Optional[cute.Tensor],
        mY_tensormap: Optional[cute.Tensor],
        mTileCount_semaphore: Optional[cute.Pointer],
        mBatchIdx_schedule_order: Optional[cute.Tensor],
        max_active_clusters: Int32,
        stream: cuda.CUstream,
    ):
        # setup static attributes before smem/grid/tma computation
        self.a_dtype = mA.element_type
        self.b_dtype = mB.element_type
        self.c_dtype = mC.element_type if mC is not None else None
        self.d_dtype = mD.element_type
        self.s_dtype = cutlass.Float32

        self.a_layout = utils.LayoutEnum.from_tensor(mA)
        self.b_layout = utils.LayoutEnum.from_tensor(mB)
        self.c_layout = cutlass.utils.LayoutEnum.from_tensor(mC) if mC is not None else None
        self.d_layout = utils.LayoutEnum.from_tensor(mD)

        self.use_bias = const_expr(mBias is not None)
        if const_expr(self.use_bias):
            assert not self.compute_weight_gradient, "Bias addition is not supported when computing weight gradients"
            self.bias_dtype = mBias.element_type
            self.bias_layout = utils.LayoutEnum.from_tensor(mBias)
        else:
            self.bias_dtype = None
            self.bias_layout = None

        if const_expr(self.need_adhoc_epilogue_store):
            self.y_dtype = mY.element_type
            self.y_layout = utils.LayoutEnum.from_tensor(mY)
        else:
            self.y_layout = self.y_dtype = None

        if const_expr(mC is not None):
            assert self.acc_dtype == cutlass.Float32
            assert self.need_epilogue_load, "Set need_epilogue_load = True or set mC = None"

        if const_expr(mS is not None):
            assert self.compute_dz_and_partial_ds_and_y1s, "Set compute_dz_and_partial_ds = True or set mS = None"
            assert mDS_partial is not None
            assert mY is not None

        if const_expr(self.a_dtype.width == 16 and self.a_dtype != self.b_dtype):
            raise TypeError(f"Type mismatch: {self.a_dtype} != {self.b_dtype}")
        if const_expr(self.a_dtype.width != self.b_dtype.width):
            raise TypeError(f"Type width mismatch: {self.a_dtype.width} != {self.b_dtype.width}")
        if const_expr(self.a_dtype.width != 16 and self.a_dtype.width != 8):
            raise TypeError("a_dtype should be float16 or float8")

        if const_expr(mBatchIdx_schedule_order is not None):
            assert (
                mTileCount_semaphore is None
            ), "we only define a static scheduling order for static persistent tile scheduler"

        self.tensormap_management_bytes = (
            self.tensormap_bytes_total
            if const_expr(self.tensormap_update_mode == cutlass.utils.TensorMapUpdateMode.SMEM)
            else 0
        ) + self.tensor_memory_management_bytes

        self._setup_attributes()

        tiled_mma = sm90_utils.make_trivial_tiled_mma(
            self.a_dtype,
            self.b_dtype,
            self.a_layout.sm90_mma_major_mode(),
            self.b_layout.sm90_mma_major_mode(),
            self.acc_dtype,
            self.atom_layout_mnk,
            tiler_mn=(64, self.tile_shape_mnk[1] // self.atom_layout_mnk[1]),
        )
        if const_expr(self.atom_layout_mnk[1] > 1):
            # If N dimension is split among 2 WGs, we need to permute the N dimension so
            # that in the epilogue, WG0 and WG1 can write to epi smem of size e.g. (64, 32)
            # containing accumulators that are next to each other in the N dimension.
            # Without permutation WG0 would write to epi smem of size (64, 16) and
            # WG1 would write to a separate epi smem of size (64, 16) that's far away.
            atom_n = self.atom_layout_mnk[1]
            permutation_n = cute.make_ordered_layout(
                (8, self.tile_shape_mnk[1] // atom_n // 8, atom_n), order=(0, 2, 1)
            )
            tiled_mma = cute.make_tiled_mma(
                cute.make_mma_atom(tiled_mma.op),
                self.atom_layout_mnk,
                permutation_mnk=(None, permutation_n, None),
            )

        if const_expr(self.is_A_gather):
            A_tiled_copy = self._make_tiled_copy_2D(
                mA,
                self.tile_M,
                self.tile_K,
                self.a_layout == cutlass.utils.LayoutEnum.ROW_MAJOR,
                self.num_load_A_threads,
                self.universal_copy_bits,
                is_g2s=True,
            )
            tma_atom_a = tma_tensor_a = None
        else:
            A_tiled_copy = None
            tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors(
                mA,
                self.a_smem_layout_staged,
                (self.tile_shape_mnk[0], self.tile_shape_mnk[2]),
                self.cluster_shape_mnk[1],
            )

        tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors(
            mB,
            self.b_smem_layout_staged,
            (self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
            self.cluster_shape_mnk[0],
        )

        if const_expr(self.need_epilogue_load):
            tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors(
                mC, self.c_epi_smem_layout_staged, self.c_epi_tile, store_or_load="load"
            )
        else:
            tma_atom_c, tma_tensor_c = None, None

        atom_bias = None
        if const_expr(self.use_bias):
            atom_bias = cute.make_copy_atom(
                cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.cpasync.LoadCacheMode.ALWAYS),
                mBias.element_type,
                num_bits_per_copy=self.universal_copy_bits,
            )

            thread_per_row = self.tile_shape_mnk[1] // (self.universal_copy_bits // mBias.element_type.width)
            thread_layout = cute.make_ordered_layout((1, thread_per_row), order=(1, 0))
            value_layout = cute.make_layout((1, self.universal_copy_bits // mBias.element_type.width))
            atom_bias = cute.make_tiled_copy_tv(atom_bias, thread_layout, value_layout)

        if const_expr(self.d_epi_smem_layout_staged is not None):
            tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors(
                mD, self.d_epi_smem_layout_staged, self.d_epi_tile, store_or_load="store"
            )
        else:
            tma_atom_d, tma_tensor_d = None, None

        if const_expr(mDIdx is not None):
            copy_elems = self.universal_copy_bits // mD.element_type.width
            assert self.num_epi_threads % (self.d_epi_tile[1] // copy_elems) == 0

            D_tiled_copy = self._make_tiled_copy_2D(
                mD,
                self.d_epi_tile[0],
                self.d_epi_tile[1],
                self.d_layout.is_n_major_c(),
                self.num_epi_threads,
                self.universal_copy_bits,
                is_g2s=False,
            )
        else:
            D_tiled_copy = None

        if const_expr(self.need_adhoc_epilogue_store):
            tma_atom_y, tma_tensor_y = self._make_tma_epi_atoms_and_tensors(
                mY, self.y_epi_smem_layout_staged, self.y_epi_tile, store_or_load="store"
            )
        else:
            tma_atom_y, tma_tensor_y = None, None

        if const_expr(self.compute_weight_gradient):
            assert const_expr(
                not self.compute_dz_and_partial_ds_and_y1s
            ), "weight grad computation conflicts with activation grad computation"

            problem_shape_ntile_mnl = cute.ceil_div(mD.shape[:2], self.tile_shape_mnk[:2]) + (mD.shape[2],)
            TileScheduler = SonicMoETileScheduler
            tile_sched_args = TileSchedulerArguments(
                problem_shape_ntile_mnl=problem_shape_ntile_mnl,
                raster_order=self.raster_order,
                group_size=self.L2_group_size,
                cluster_shape_mnk=self.cluster_shape_mnk,
                is_persistent=self.is_persistent,
                tile_count_semaphore=mTileCount_semaphore,
                batch_idx_permute=mBatchIdx_schedule_order,
            )
        else:
            problem_shape_ntile_mnl = (
                None,
                cute.ceil_div(mD.shape[1], self.tile_shape_mnk[1]),
                mMoffset.shape[0] - 1,
            )
            TileScheduler = SonicMoEVarlenMTileScheduler
            tile_sched_args = VarlenMTileSchedulerArguments(
                problem_shape_ntile_mnl=problem_shape_ntile_mnl,
                total_m=mD.shape[0],
                cu_seqlens_m=mMoffset,
                raster_order=self.raster_order,
                group_size=self.L2_group_size,
                tile_shape_mn=self.tile_shape_mnk[:2],
                cluster_shape_mnk=self.cluster_shape_mnk,
                is_persistent=self.is_persistent,
                tile_count_semaphore=mTileCount_semaphore,
            )

        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
        grid = TileScheduler.get_grid_shape(tile_sched_params, max_active_clusters)

        c_epi_smem_size = cute.cosize(self.c_epi_smem_layout_staged) if const_expr(self.need_epilogue_load) else 0
        bias_epi_smem_size = cute.cosize(self.bias_epi_smem_layout_staged) if const_expr(self.use_bias) else 0
        d_epi_smem_size = (
            cute.cosize(self.d_epi_smem_layout_staged)
            if const_expr(self.is_persistent and (self.d_epi_stage > 0))
            else 0
        )
        y_epi_smem_size = (
            cute.cosize(self.y_epi_smem_layout_staged)
            if const_expr(self.need_adhoc_epilogue_store) and self.is_persistent
            else 0
        )
        s_epi_smem_size = (
            cute.cosize(self.s_epi_smem_layout_staged) if const_expr(self.compute_dz_and_partial_ds_and_y1s) else 0
        )

        @cute.struct
        class SharedStorage:
            mainloop_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2]
            tensormap_buffer: cute.struct.Align[cute.struct.MemRange[cutlass.Int64, self.num_tensormaps], 64]
            sD: cute.struct.Align[
                cute.struct.MemRange[self.d_dtype, d_epi_smem_size],
                self.buffer_align_bytes,
            ]
            sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2]
            tile_count: cute.struct.MemRange[cutlass.Int32, self.sched_stage]
            if const_expr(self.need_epilogue_load):
                sC: cute.struct.Align[
                    cute.struct.MemRange[self.c_dtype, c_epi_smem_size],
                    self.buffer_align_bytes,
                ]
                epi_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.c_epi_stage * 2]

            if const_expr(self.use_bias):
                sBias: cute.struct.Align[
                    cute.struct.MemRange[self.bias_dtype, bias_epi_smem_size],
                    self.buffer_align_bytes,
                ]

            if const_expr(self.compute_dz_and_partial_ds_and_y1s):
                sS: cute.struct.Align[
                    cute.struct.MemRange[self.s_dtype, s_epi_smem_size],
                    self.buffer_align_bytes,
                ]

            if const_expr(self.need_adhoc_epilogue_store):
                sY: cute.struct.Align[
                    cute.struct.MemRange[self.y_dtype, y_epi_smem_size],
                    self.buffer_align_bytes,
                ]
            sA: cute.struct.Align[
                cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged)],
                self.buffer_align_bytes,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged)],
                self.buffer_align_bytes,
            ]
            if const_expr(self.compute_weight_gradient and self.is_A_gather):
                sAIdx_prefetch: cute.struct.Align[
                    cute.struct.MemRange[self.index_dtype, self.prefetch_token_idx_size],
                    self.buffer_align_bytes,
                ]

        self.shared_storage = SharedStorage
        allocated_smem_size = self.shared_storage.size_in_bytes() + self.tensormap_management_bytes
        # Launch the kernel synchronously
        self.kernel(
            A_tiled_copy,
            mA,
            tma_atom_a,
            tma_tensor_a,
            mB,
            tma_atom_b,
            tma_tensor_b,
            tma_atom_c,
            tma_tensor_c,
            mC,
            mBias,
            atom_bias,
            D_tiled_copy,
            tma_atom_d,
            tma_tensor_d,
            mD,
            tma_atom_y,
            tma_tensor_y,
            mY,
            mS,
            mDS_partial,
            mMoffset,
            mAIdx,
            mDIdx,
            mS_scatter_idx,
            mA_tensormap,
            mB_tensormap,
            mC_tensormap,
            mD_tensormap,
            mY_tensormap,
            tiled_mma,
            self.cta_layout_mnk,
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.prefetch_AIdx_smem_layout_staged,
            self.c_epi_smem_layout_staged,
            self.bias_epi_smem_layout_staged,
            self.d_epi_smem_layout_staged,
            self.y_epi_smem_layout_staged,
            self.s_epi_smem_layout_staged,
            tile_sched_params,
            TileScheduler,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=self.cluster_shape_mnk,
            smem=allocated_smem_size,
            stream=stream,
            min_blocks_per_mp=1,
        )

    @cute.jit
    def update_tma_desc_ptr(
        self,
        mTensor: cute.Tensor,
        tma_atom: cute.CopyAtom,
        tensormap_manager: TensorMapManagerSm90,
        tensormap_ptr: cute.Pointer,
        token_start: Int32,
        token_group_size: Int32,
        is_tma_warp: bool,
        tensormap_smem_ptr: Optional[cute.Pointer] = None,
        address_space: cute.AddressSpace = cute.AddressSpace.generic,
    ) -> cute.Pointer:
        if const_expr(self.compute_weight_gradient):
            tensor_shape = (mTensor.shape[0], token_group_size)
            start_ptr = (mTensor.iterator + token_start * mTensor.stride[1]).toint()
        else:
            tensor_shape = (token_group_size, mTensor.shape[1])
            start_ptr = (mTensor.iterator + token_start * mTensor.stride[0]).toint()

        tensor_gmem_ptr = cute.make_ptr(
            mTensor.element_type,
            start_ptr,
            cute.AddressSpace.gmem,
            assumed_align=16,
        )
        real_tensor = cute.make_tensor(
            tensor_gmem_ptr,
            cute.make_layout(tensor_shape, stride=mTensor.stride),
        )
        if const_expr(self.tensormap_update_mode == cutlass.utils.TensorMapUpdateMode.GMEM):
            tensormap_manager.update_tensormap(
                (real_tensor,),
                (tma_atom,),
                tensormap_gmem_ptr=(tensormap_ptr,),
                is_manager_warp=is_tma_warp,
                tensormap_smem_ptr=None,
            )
        else:
            assert tensormap_smem_ptr is not None
            tensormap_manager.update_tensormap(
                (real_tensor,),
                (tma_atom,),
                tensormap_gmem_ptr=(tensormap_ptr,),
                is_manager_warp=is_tma_warp,
                tensormap_smem_ptr=(tensormap_smem_ptr,),
            )

        tensormap_manager.fence_tensormap_update(tensormap_ptr)

    @cute.jit
    def align_tensormap_smem_ptr(self, base_ptr: cute.Pointer):
        return cute.make_ptr(
            cutlass.Int64,
            base_ptr.toint(),
            cute.AddressSpace.smem,
            assumed_align=128,
        )

    @cute.jit
    def allocate_new_tensormap_smem_ptr(self, tensormap_smem_ptr: cute.Pointer):
        return self.align_tensormap_smem_ptr(tensormap_smem_ptr + self.bytes_per_tensormap // 8)

    @cute.jit
    def swiglu_derivative(self, g: Float32, u: Float32, dy1: Float32) -> Tuple[Float32, Float32, Float32]:
        half_g = 0.5 * g
        tanh_half_g = self.tanh(half_g)

        sig_g = self.fma(0.5, tanh_half_g, 0.5)
        sig_n_g = 1 - sig_g

        silu_g = self.fma(half_g, tanh_half_g, half_g)

        dg = dy1 * (u * self.fma(silu_g, sig_n_g, sig_g))
        du = dy1 * silu_g

        swiglu_output = silu_g * u
        return dg, du, swiglu_output

    @cute.jit
    def reglu_derivative(self, g: Float32, u: Float32, dy1: Float32) -> Tuple[Float32, Float32, Float32]:
        relu_g = cute.arch.fmax(0.0, g)

        relu_prime_g = 1.0
        if g < Float32(0.0):
            relu_prime_g = 0.0  # derivative of ReLU

        dg = dy1 * u * relu_prime_g
        du = dy1 * relu_g

        reglu_output = u * relu_g
        return dg, du, reglu_output

    @cute.jit
    def geglu_derivative(self, g: Float32, u: Float32, dy1: Float32) -> Tuple[Float32, Float32, Float32]:
        # gelu(g) = 0.5 * g * (1 + tanh(sqrt(2/pi) * (g + 0.044715*g^3)))
        sqrt_2_over_pi = const_expr(math.sqrt(2 / math.pi))
        c = 0.044715

        g2 = g * g
        g3 = g2 * g

        # t = sqrt(2/pi) * (g + c*g^3)
        t = sqrt_2_over_pi * self.fma(c, g3, g)

        tanh_t = self.tanh(t)
        one_plus_th = 1.0 + tanh_t  # 1 + tanh(t)
        gelu_g = 0.5 * g * one_plus_th  # gelu(g)

        # d th / d g = (1 - tanh(t)^2) * sqrt(2/pi) * (1 + 3c g^2)
        sech2 = self.fma(-tanh_t, tanh_t, 1.0)
        dt_dg = sech2 * sqrt_2_over_pi * self.fma(3.0 * c, g2, 1.0)

        # d gelu / d g = 0.5*(1 + tanh(t)) + 0.5*g*dt_dg
        gelu_prime = 0.5 * self.fma(g, dt_dg, one_plus_th)

        # Chain rule for y = gelu(g) * u
        dg = dy1 * u * gelu_prime
        du = dy1 * gelu_g

        geglu_output = u * gelu_g
        return dg, du, geglu_output

    @cute.jit
    def silu_derivative(self, x: Float32, dy1: Float32) -> Tuple[Float32, Float32]:
        half_x = 0.5 * x
        tanh_half_x = self.tanh(half_x)

        sig_x = self.fma(0.5, tanh_half_x, 0.5)
        sig_n_x = 1 - sig_x

        silu_x = self.fma(half_x, tanh_half_x, half_x)
        dx = dy1 * self.fma(silu_x, sig_n_x, sig_x)

        return dx, silu_x

    @cute.jit
    def relu_derivative(self, x: Float32, dy1: Float32) -> Tuple[Float32, Float32]:
        relu_x = cute.arch.fmax(0.0, x)

        relu_prime_x = 1.0
        if x < Float32(0.0):
            relu_prime_x = 0.0  # derivative of ReLU

        dx = dy1 * relu_prime_x
        return dx, relu_x

    @cute.jit
    def relu_sq_derivative(self, x: Float32, dy1: Float32) -> Tuple[Float32, Float32]:
        relu_x = cute.arch.fmax(x, 0.0)
        relu_sq_output = relu_x * x
        dx = dy1 * (2.0 * relu_x)
        return dx, relu_sq_output

    @cute.jit
    def gelu_derivative(self, x: Float32, dy1: Float32) -> Tuple[Float32, Float32]:
        # gelu(g) = 0.5 * g * (1 + tanh(sqrt(2/pi) * (g + 0.044715*g^3)))
        sqrt_2_over_pi = const_expr(math.sqrt(2 / math.pi))
        c = 0.044715

        x2 = x * x
        x3 = x2 * x

        # t = sqrt(2/pi) * (g + c*g^3)
        t = sqrt_2_over_pi * self.fma(c, x3, x)

        tanh_t = self.tanh(t)
        one_plus_tanh_t = 1.0 + tanh_t  # 1 + tanh(t)
        gelu_x = 0.5 * x * one_plus_tanh_t

        # d th / d g = (1 - tanh(t)^2) * sqrt(2/pi) * (1 + 3c g^2)
        sech2 = self.fma(-tanh_t, tanh_t, 1.0)
        dt_dg = sech2 * sqrt_2_over_pi * self.fma(3.0 * c, x2, 1.0)

        # d gelu / d g = 0.5*(1 + tanh(t)) + 0.5*g*dt_dg
        gelu_prime = 0.5 * self.fma(x, dt_dg, one_plus_tanh_t)

        # Chain rule for y = gelu(g) * u
        dx = dy1 * gelu_prime

        return dx, gelu_x

    @cute.jit
    def compute_activation(self, tRS_rD, tRS_rY):
        if const_expr(self.is_glu):
            # tRS_sY: (((2, 4), 1), 1, 1, (1, 4))
            # (((2, 4), 1), 1, 1)
            if const_expr(self.compute_swiglu):
                act_func = self.silu
            elif const_expr(self.compute_reglu):
                act_func = self.relu
            elif const_expr(self.compute_geglu):
                act_func = self.gelu
            else:
                raise NotImplementedError()

            for i in cutlass.range_constexpr(cute.size(tRS_rD) // 2):
                tRS_rY[i] = (act_func(tRS_rD[const_expr(2 * i)]) * tRS_rD[const_expr(2 * i + 1)]).to(self.y_dtype)

            self.permute_gated_Cregs_b16(tRS_rY)

        elif const_expr(self.is_normal_act):
            assert cute.size(tRS_rD) == cute.size(tRS_rY)
            if const_expr(self.compute_relu_sq):
                act_func = self.relu_sq
            elif const_expr(self.compute_relu):
                act_func = self.relu
            elif const_expr(self.compute_silu):
                act_func = self.silu
            elif const_expr(self.compute_gelu):
                act_func = self.gelu
            else:
                raise NotImplementedError()

            for i in cutlass.range_constexpr(cute.size(tRS_rD)):
                tRS_rY[i] = act_func(tRS_rD[i]).to(self.y_dtype)

        else:
            raise NotImplementedError()

    @cute.jit
    def compute_backward_activation(self, tRS_rAcc, sS, tRS_rcD, tRS_rC, tRS_rD, tRS_rD_out, tRS_rY, epi_idx: Int32):
        if const_expr(self.is_glu):
            # if we compute glu activation,
            #   we will assume the incoming C dtype as FP32, and we will output final result in FP32 (decompress to BF16 in caller side)

            if const_expr(self.compute_swiglu):
                bwd_act_func = self.swiglu_derivative
            elif const_expr(self.compute_reglu):
                bwd_act_func = self.reglu_derivative
            elif const_expr(self.compute_geglu):
                bwd_act_func = self.geglu_derivative
            else:
                raise NotImplementedError()

            for i in cutlass.range_constexpr(cute.size(tRS_rD)):
                g, u = self.unpack2x16_as_2xf32(tRS_rC[i], self.a_dtype)
                dy = tRS_rD[i]
                dg, du, fwd_output = bwd_act_func(g, u, dy)
                tRS_rAcc[const_expr(epi_idx * cute.size(tRS_rD) + i)] = dy * fwd_output
                s = sS[tRS_rcD[i]]
                tRS_rD_out[i] = self.pack2x16_as_f32(self.a_dtype(dg * s), self.a_dtype(du * s))
                tRS_rY[i] = self.y_dtype(fwd_output * s)

        elif const_expr(self.is_normal_act):
            if const_expr(self.compute_relu_sq):
                bwd_act_func = self.relu_sq_derivative
            elif const_expr(self.compute_relu):
                bwd_act_func = self.relu_derivative
            elif const_expr(self.compute_gelu):
                bwd_act_func = self.gelu_derivative
            elif const_expr(self.compute_silu):
                bwd_act_func = self.silu_derivative
            else:
                raise NotImplementedError()

            for i in cutlass.range_constexpr(cute.size(tRS_rD)):
                z = tRS_rC[i]
                dy = tRS_rD[i]
                dz, fwd_output = bwd_act_func(z, dy)
                tRS_rAcc[const_expr(epi_idx * cute.size(tRS_rD) + i)] = dy * fwd_output
                s = sS[tRS_rcD[i]]
                tRS_rD_out[i] = self.a_dtype(dz * s)
                tRS_rY[i] = self.y_dtype(fwd_output * s)

        else:
            raise NotImplementedError()

    #  GPU device kernel
    @cute.kernel
    def kernel(
        self,
        A_tiled_copy: Optional[cute.TiledCopy],
        mA_mkl: cute.Tensor,
        tma_atom_a: Optional[cute.CopyAtom],
        mA_mkl_tma: Optional[cute.Tensor],
        mB_nkl: cute.Tensor,
        tma_atom_b: cute.CopyAtom,
        mB_nkl_tma: cute.Tensor,
        tma_atom_c: Optional[cute.CopyAtom],
        mC_mnl_tma: Optional[cute.Tensor],
        mC_mnl: cute.Tensor,
        mBias_nl: Optional[cute.Tensor],
        cpasync_atom_bias: Optional[cute.CopyAtom],
        D_tiled_copy: Optional[cute.TiledCopy],
        tma_atom_d: Optional[cute.CopyAtom],
        mD_mnl_tma: cute.Tensor,
        mD_mnl: cute.Tensor,
        tma_atom_y: Optional[cute.CopyAtom],
        mY_mnl_tma: Optional[cute.Tensor],
        mY_mnl: Optional[cute.Tensor],
        mS_ml: Optional[cute.Tensor],
        mDS_partial: Optional[cute.Tensor],
        mTokenoffset: cute.Tensor,
        mAIdx_mkl: cute.Tensor,
        mDIdx_mnl: Optional[cute.Tensor],
        mS_scatter_idx: Optional[cute.Tensor],
        mA_tensormap: Optional[cute.Tensor],
        mB_tensormap: Optional[cute.Tensor],
        mC_tensormap: Optional[cute.Tensor],
        mD_tensormap: cute.Tensor,
        mY_tensormap: Optional[cute.Tensor],
        tiled_mma: cute.TiledMma,
        cta_layout_mnk: cute.Layout,
        a_smem_layout_staged: cute.ComposedLayout,
        b_smem_layout_staged: cute.ComposedLayout,
        prefetch_AIdx_smem_layout_staged: Optional[cute.Layout],
        c_epi_smem_layout_staged: Optional[cute.ComposedLayout],
        bias_epi_smem_layout_staged: Optional[cute.Layout],
        d_epi_smem_layout_staged: cute.ComposedLayout,
        y_epi_smem_layout_staged: Optional[cute.ComposedLayout],
        s_epi_smem_layout_staged: Optional[cute.Layout],
        tile_sched_params: ParamsBase,
        TileScheduler: cutlass.Constexpr[Callable],
    ):
        tidx, _, _ = cute.arch.thread_idx()
        # Assume: M: 2048, N: 512, K: 1024, L: 4
        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())

        if warp_idx == self.tma_warp_id:
            if const_expr(not self.is_A_gather):
                cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b)
            if const_expr(not self.inference_mode):
                cpasync.prefetch_descriptor(tma_atom_d)
            if const_expr(self.need_adhoc_epilogue_store):
                cpasync.prefetch_descriptor(tma_atom_y)
            if const_expr(tma_atom_c is not None):
                cpasync.prefetch_descriptor(tma_atom_c)

        A_thr_copy_elems = self.universal_copy_bits // mA_mkl.element_type.width

        a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, 0))
        b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0))
        if const_expr(self.is_A_gather):
            tma_copy_bytes = cute.size_in_bytes(self.b_dtype, b_smem_layout)
        else:
            tma_copy_bytes = cute.size_in_bytes(self.a_dtype, a_smem_layout) + cute.size_in_bytes(
                self.b_dtype, b_smem_layout
            )

        smem = cutlass.utils.SmemAllocator()
        shared_storage = smem.allocate(self.shared_storage)

        # Threads/warps participating in this pipeline
        if const_expr(self.is_A_gather):
            mainloop_pipeline_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread, 1 + self.num_load_A_threads
            )
            # Each warp will constribute to the arrive count with the number of mcast size
            mcast_size = self.num_mcast_ctas_b
            pipeline_class = PipelineTmaCpAsync
        else:
            # Threads/warps participating in this pipeline
            mainloop_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
            # Each warp will constribute to the arrive count with the number of mcast size
            mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
            pipeline_class = pipeline.PipelineTmaAsync

        consumer_arrive_cnt = mcast_size * (self.num_mma_threads // cute.arch.WARP_SIZE)
        mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt)
        cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape))
        mainloop_pipeline = pipeline_class.create(
            barrier_storage=shared_storage.mainloop_pipeline_array_ptr.data_ptr(),
            num_stages=self.ab_stage,
            producer_group=mainloop_pipeline_producer_group,
            consumer_group=mainloop_pipeline_consumer_group,
            tx_count=tma_copy_bytes,
            cta_layout_vmnk=cta_layout_vmnk,
        )

        if const_expr(self.need_epilogue_load):
            # Threads/warps participating in this pipeline
            epi_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
            # Each warp will contribute 1 to the arrive count
            consumer_arrive_cnt = self.num_epi_threads // cute.arch.WARP_SIZE
            epi_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt)
            c_smem_layout = cute.slice_(c_epi_smem_layout_staged, (None, None, 0))
            tma_copy_c_bytes = cute.size_in_bytes(self.c_dtype, c_smem_layout)
            epi_pipeline = pipeline.PipelineTmaAsync.create(
                barrier_storage=shared_storage.epi_pipeline_array_ptr.data_ptr(),
                num_stages=self.c_epi_stage,
                producer_group=epi_pipeline_producer_group,
                consumer_group=epi_pipeline_consumer_group,
                tx_count=tma_copy_c_bytes,
            )
        else:
            epi_pipeline = None

        sA = shared_storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
        sB = shared_storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)

        if const_expr(not self.is_persistent):
            sD_ptr = cute.recast_ptr(sA.iterator, d_epi_smem_layout_staged.inner, dtype=self.d_dtype)
            sD = cute.make_tensor(sD_ptr, d_epi_smem_layout_staged.outer)

            if const_expr(self.inference_mode and self.need_adhoc_epilogue_store):
                next_ptr = sD_ptr
            else:
                next_ptr = sD_ptr + cute.cosize(d_epi_smem_layout_staged)

            if const_expr(self.need_adhoc_epilogue_store):
                sY_ptr = cute.recast_ptr(next_ptr, y_epi_smem_layout_staged.inner, dtype=self.y_dtype)
                sY = cute.make_tensor(sY_ptr, y_epi_smem_layout_staged.outer)
                next_ptr = sY_ptr + cute.cosize(y_epi_smem_layout_staged)
            else:
                sY = None

        else:
            if const_expr(self.need_adhoc_epilogue_store):
                sY = shared_storage.sY.get_tensor(
                    y_epi_smem_layout_staged.outer, swizzle=y_epi_smem_layout_staged.inner
                )
                if const_expr(self.inference_mode):
                    sD = cute.make_tensor(
                        cute.recast_ptr(sY.iterator, d_epi_smem_layout_staged.inner, dtype=self.d_dtype),
                        d_epi_smem_layout_staged.outer,
                    )
                else:
                    sD = shared_storage.sD.get_tensor(
                        d_epi_smem_layout_staged.outer, swizzle=d_epi_smem_layout_staged.inner
                    )
            else:
                sY = None
                sD = shared_storage.sD.get_tensor(
                    d_epi_smem_layout_staged.outer, swizzle=d_epi_smem_layout_staged.inner
                )

        if const_expr(self.compute_weight_gradient and self.is_A_gather):
            sAIdx_prefetch = shared_storage.sAIdx_prefetch.get_tensor(prefetch_AIdx_smem_layout_staged)
        else:
            sAIdx_prefetch = None

        if const_expr(self.compute_dz_and_partial_ds_and_y1s):
            sS = shared_storage.sS.get_tensor(s_epi_smem_layout_staged, dtype=self.s_dtype)
        else:
            sS = None

        if const_expr(self.need_epilogue_load):
            sC = shared_storage.sC.get_tensor(c_epi_smem_layout_staged.outer, swizzle=c_epi_smem_layout_staged.inner)
        else:
            sC = None

        if const_expr(self.use_bias):
            sBias = shared_storage.sBias.get_tensor(bias_epi_smem_layout_staged)
        else:
            sBias = None

        sched_pipeline = None
        tile_count = None
        if const_expr(tile_sched_params.tile_count_semaphore is not None):
            sched_pipeline = self.make_sched_pipeline(
                cta_layout_mnk,
                sched_pipeline_mbar_ptr=shared_storage.sched_pipeline_array_ptr.data_ptr(),
            )
            tile_count = shared_storage.tile_count.get_tensor((self.sched_stage,))

        a_tensormap_smem_ptr = b_tensormap_smem_ptr = c_tensormap_smem_ptr = d_tensormap_smem_ptr = (
            y_tensormap_smem_ptr
        ) = None
        if cutlass.const_expr(self.tensormap_update_mode == utils.TensorMapUpdateMode.SMEM):
            tensormap_smem_ptr = shared_storage.tensormap_buffer.data_ptr()
            tensormap_smem_ptr = self.align_tensormap_smem_ptr(tensormap_smem_ptr)

            if const_expr(self.compute_weight_gradient):
                if const_expr(not self.is_A_gather):
                    tensormap_smem_ptr = a_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                        tensormap_smem_ptr
                    )
                tensormap_smem_ptr = b_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(tensormap_smem_ptr)
            else:
                if const_expr(self.pingpong):
                    if const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                        tensormap_smem_ptr = d0_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        tensormap_smem_ptr = d1_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        d_tensormap_smem_ptr = d0_tensormap_smem_ptr if warp_idx // 4 == 0 else d1_tensormap_smem_ptr

                    if const_expr(self.need_adhoc_epilogue_store):
                        tensormap_smem_ptr = y0_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        tensormap_smem_ptr = y1_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        y_tensormap_smem_ptr = y0_tensormap_smem_ptr if warp_idx // 4 == 0 else y1_tensormap_smem_ptr

                    if const_expr(self.need_epilogue_load):
                        tensormap_smem_ptr = c0_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        tensormap_smem_ptr = c1_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                        c_tensormap_smem_ptr = c0_tensormap_smem_ptr if warp_idx // 4 == 0 else c1_tensormap_smem_ptr

                else:
                    if const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                        tensormap_smem_ptr = d_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )
                    if const_expr(self.need_adhoc_epilogue_store):
                        tensormap_smem_ptr = y_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )

                    if const_expr(self.need_epilogue_load):
                        tensormap_smem_ptr = c_tensormap_smem_ptr = self.allocate_new_tensormap_smem_ptr(
                            tensormap_smem_ptr
                        )

        grid_dim = cute.arch.grid_dim()
        bid = cute.arch.block_idx()
        tensormap_workspace_idx = bid[2] * grid_dim[1] * grid_dim[0] + bid[1] * grid_dim[0] + bid[0]
        tensormap_manager = TensorMapManagerSm90(self.tensormap_update_mode, self.bytes_per_tensormap)

        if const_expr(self.compute_weight_gradient):
            if const_expr(not self.is_A_gather and (mA_tensormap is not None)):
                a_tensormap_ptr = tensormap_manager.get_tensormap_ptr(
                    mA_tensormap[tensormap_workspace_idx, None].iterator
                )
            else:
                a_tensormap_ptr = None

            if const_expr(mB_tensormap is not None):
                b_tensormap_ptr = tensormap_manager.get_tensormap_ptr(
                    mB_tensormap[tensormap_workspace_idx, None].iterator
                )
            else:
                b_tensormap_ptr = None

            if cutlass.const_expr(self.tensormap_update_mode == utils.TensorMapUpdateMode.SMEM):
                if const_expr(not self.is_A_gather):
                    tensormap_a_init_ptr = a_tensormap_smem_ptr
                tensormap_b_init_ptr = b_tensormap_smem_ptr
            else:
                if const_expr(not self.is_A_gather):
                    tensormap_a_init_ptr = b_tensormap_ptr
                tensormap_b_init_ptr = b_tensormap_ptr

        else:
            if const_expr(self.pingpong):
                tensormap_workspace_idx = tensormap_workspace_idx * 2 + warp_idx // 4

            if const_expr(
                (mD_tensormap is not None) and (not (self.inference_mode and self.need_adhoc_epilogue_store))
            ):
                d_tensormap_ptr = tensormap_manager.get_tensormap_ptr(
                    mD_tensormap[tensormap_workspace_idx, None].iterator
                )
            else:
                d_tensormap_ptr = None

            if const_expr(self.need_adhoc_epilogue_store):
                assert mY_tensormap is not None
                y_tensormap_ptr = tensormap_manager.get_tensormap_ptr(
                    mY_tensormap[tensormap_workspace_idx, None].iterator
                )
            else:
                y_tensormap_ptr = None

            if const_expr(self.need_epilogue_load):
                assert mC_tensormap is not None
                c_tensormap_ptr = tensormap_manager.get_tensormap_ptr(
                    mC_tensormap[tensormap_workspace_idx, None].iterator
                )
            else:
                c_tensormap_ptr = None

            if cutlass.const_expr(self.tensormap_update_mode == utils.TensorMapUpdateMode.SMEM):
                tensormap_d_init_ptr = d_tensormap_smem_ptr
                tensormap_y_init_ptr = y_tensormap_smem_ptr
                tensormap_c_init_ptr = c_tensormap_smem_ptr
            else:
                tensormap_d_init_ptr = d_tensormap_ptr
                tensormap_y_init_ptr = y_tensormap_ptr
                tensormap_c_init_ptr = c_tensormap_ptr

        TileSchedulerCls = partial(TileScheduler.create, tile_sched_params, tile_count, sched_pipeline)

        k_tile_cnt = cute.ceil_div(cute.size(mA_mkl.shape[1]), self.tile_shape_mnk[2])
        c_tile_cnt = (
            cute.size(cute.ceil_div(self.tile_shape_mnk[:2], self.c_epi_tile))
            if const_expr(self.need_epilogue_load)
            else Int32(0)
        )

        if warp_idx >= self.tma_warp_id:
            cute.arch.setmaxregister_decrease(self.num_regs_load)
            cute.arch.setmaxregister_decrease(self.num_regs_load)

            prolog_loading_warp_ids = (
                [const_expr(self.tma_warp_id + i) for i in range(self.num_load_A_threads // cute.arch.WARP_SIZE)]
                if const_expr(self.is_A_gather)
                else [const_expr(self.tma_warp_id)]
            )

            if warp_idx in prolog_loading_warp_ids:
                is_tma_warp = cutlass.Boolean(warp_idx == self.tma_warp_id)
                cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
                cluster_coord_mnk = cta_layout_mnk.get_flat_coord(cta_rank_in_cluster)

                a_mcast_mask = cute.make_layout_image_mask(cta_layout_mnk, cluster_coord_mnk, mode=1)
                a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0
                b_mcast_mask = cute.make_layout_image_mask(cta_layout_mnk, cluster_coord_mnk, mode=0)
                b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0

                mainloop_producer_state = make_pipeline_state(pipeline.PipelineUserType.Producer, self.ab_stage)
                is_scheduler_warp = warp_idx == self.tma_warp_id
                if const_expr(cute.size(cta_layout_mnk) > 1):
                    is_scheduler_warp = is_scheduler_warp and cute.arch.block_idx_in_cluster() == 0

                tile_scheduler = TileSchedulerCls(is_scheduler_warp=is_scheduler_warp)
                work_tile = tile_scheduler.initial_work_tile_info()

                last_batch_idx = cutlass.Int32(-1)
                token_group_size = cutlass.Int32(0)

                mcA_mkl = cute.make_identity_tensor((mA_mkl.shape[0], mA_mkl.shape[1]))

                TIdx_cur_group = TIdx_next_group = cutlass.Int32(0)
                mAIdx_mk = cute.domain_offset((0,), mAIdx_mkl)

                gA_mk = None
                A_g2s_thr_copy = None
                if const_expr(self.is_A_gather):
                    A_g2s_thr_copy = A_tiled_copy.get_slice(tidx - self.tma_warp_id * cute.arch.WARP_SIZE)
                    gA_mk = cute.local_tile(mA_mkl, (self.tile_M, self.tile_K), (0, None))
                    tAgA = A_g2s_thr_copy.partition_S(gA_mk)

                if const_expr(self.compute_weight_gradient):
                    if const_expr(not self.is_A_gather):
                        tensormap_manager.init_tensormap_from_atom(
                            tma_atom_a,
                            tensormap_a_init_ptr,
                            is_manager_warp=is_tma_warp,
                        )
                    tensormap_manager.init_tensormap_from_atom(
                        tma_atom_b,
                        tensormap_b_init_ptr,
                        is_manager_warp=is_tma_warp,
                    )
                    tensormap_manager.fence_tensormap_initialization()

                while work_tile.is_valid_tile:
                    tile_coord_mnkl = work_tile.tile_idx
                    batch_idx = tile_coord_mnkl[3]
                    # (bM, bK, RestK)
                    if batch_idx != last_batch_idx:
                        TIdx_cur_group, TIdx_next_group = cute.arch.make_warp_uniform(
                            mTokenoffset[batch_idx]
                        ), cute.arch.make_warp_uniform(mTokenoffset[batch_idx + 1])
                        token_group_size = TIdx_next_group - TIdx_cur_group

                        if const_expr(self.is_A_gather):
                            if const_expr(self.compute_weight_gradient):
                                mcA_mkl = cute.make_identity_tensor((mA_mkl.shape[0], token_group_size))
                            else:
                                mcA_mkl = cute.make_identity_tensor((token_group_size, mA_mkl.shape[1]))

                            mAIdx_mk = cute.domain_offset((TIdx_cur_group,), mAIdx_mkl)

                        if const_expr(self.compute_weight_gradient):
                            if const_expr(not self.is_A_gather):
                                assert a_tensormap_ptr is not None
                                self.update_tma_desc_ptr(
                                    mA_mkl,
                                    tma_atom_a,
                                    tensormap_manager,
                                    a_tensormap_ptr,
                                    TIdx_cur_group,
                                    token_group_size,
                                    is_tma_warp,
                                    tensormap_smem_ptr=a_tensormap_smem_ptr,
                                )
                            if const_expr(b_tensormap_ptr is not None):
                                self.update_tma_desc_ptr(
                                    mB_nkl,
                                    tma_atom_b,
                                    tensormap_manager,
                                    b_tensormap_ptr,
                                    TIdx_cur_group,
                                    token_group_size,
                                    is_tma_warp,
                                    tensormap_smem_ptr=b_tensormap_smem_ptr,
                                    # cute.AddressSpace.generic
                                )
                            k_tile_cnt = cute.ceil_div(token_group_size, self.tile_shape_mnk[2])

                        last_batch_idx = batch_idx

                    if const_expr(self.is_A_gather):
                        cA = cute.local_tile(mcA_mkl, (self.tile_M, self.tile_K), (tile_coord_mnkl[0], None))

                        tAsA = A_g2s_thr_copy.partition_D(sA)
                        tAcA = A_g2s_thr_copy.partition_D(cA)

                        tApA = cute.make_rmem_tensor(
                            cute.make_layout(
                                (
                                    tAgA.shape[0][1],
                                    cute.size(tAgA, mode=[1]),
                                    cute.size(tAgA, mode=[2]),
                                ),
                                stride=(cute.size(tAgA, mode=[1]), 1, 0),
                            ),
                            cutlass.Boolean,
                        )

                        for rest_v in cutlass.range_constexpr(tApA.shape[0]):
                            for m in cutlass.range_constexpr(tApA.shape[1]):
                                if const_expr(self.compute_weight_gradient):
                                    tApA[rest_v, m, 0] = cute.elem_less(tAcA[(0, rest_v), m, 0, 0][0], mA_mkl.shape[0])
                                else:
                                    tApA[rest_v, m, 0] = cute.elem_less(
                                        tAcA[(0, rest_v), m, 0, 0][0], token_group_size
                                    )
                    else:
                        if const_expr(self.compute_weight_gradient):
                            # update TMA map instead
                            mA_mk = cute.domain_offset((0, 0), mA_mkl_tma)
                        else:
                            mA_mk = cute.domain_offset((TIdx_cur_group, 0), mA_mkl_tma)

                        gA_mk_cur = cute.local_tile(mA_mk, (self.tile_M, self.tile_K), (tile_coord_mnkl[0], None))

                        a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape)
                        a_cta_crd = cluster_coord_mnk[1]

                        tAsA, tAgA_mkl = cpasync.tma_partition(
                            tma_atom_a,
                            a_cta_crd,
                            a_cta_layout,
                            cute.group_modes(sA, 0, 2),
                            cute.group_modes(gA_mk_cur, 0, 2),
                        )

                    if const_expr(self.compute_weight_gradient):
                        gB_nk = cute.local_tile(mB_nkl_tma, (self.tile_N, self.tile_K), (tile_coord_mnkl[1], None))
                    else:
                        gB_nk = cute.local_tile(mB_nkl_tma, self.tile_shape_mnk, tile_coord_mnkl, proj=(None, 1, 1))

                    b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape)
                    b_cta_crd = cluster_coord_mnk[0]
                    tBsB, tBgB_nkl = cpasync.tma_partition(
                        tma_atom_b,
                        b_cta_crd,
                        b_cta_layout,
                        cute.group_modes(sB, 0, 2),
                        cute.group_modes(gB_nk, 0, 2),
                    )

                    peek_ab_empty_status = cutlass.Boolean(True)
                    if 0 < k_tile_cnt:
                        peek_ab_empty_status = mainloop_pipeline.producer_try_acquire(mainloop_producer_state)

                    if const_expr(self.is_A_gather):
                        M_offset = cute.arch.make_warp_uniform(tile_coord_mnkl[0] * const_expr(self.tile_M))
                        if const_expr(self.compute_weight_gradient):
                            tmAIdx = None
                            M_boundary = mA_mkl.shape[0]
                        else:
                            M_boundary = cute.arch.make_warp_uniform(
                                self.min_i32(const_expr(self.tile_M), token_group_size - M_offset)
                            )
                            tmAIdx = self.prefetch_gather_idx_for_A_when_vary_M(
                                mAIdx_mk, M_offset, M_boundary, A_thr_copy_elems
                            )

                    if const_expr(self.compute_weight_gradient):
                        if const_expr(self.is_A_gather):
                            a_tma_desc_ptr = None
                        else:
                            a_tma_desc_ptr = tensormap_manager.get_tensormap_ptr(
                                a_tensormap_ptr, cute.AddressSpace.generic
                            )

                        b_tma_desc_ptr = tensormap_manager.get_tensormap_ptr(
                            b_tensormap_ptr, cute.AddressSpace.generic
                        )
                    else:
                        a_tma_desc_ptr = None
                        b_tma_desc_ptr = None

                    for k_tile in cutlass.range(k_tile_cnt, unroll=1):
                        if const_expr(self.is_A_gather):
                            mainloop_pipeline.producer_acquire(
                                mainloop_producer_state, peek_ab_empty_status, is_tma_warp=is_tma_warp
                            )
                        else:
                            mainloop_pipeline.producer_acquire(mainloop_producer_state, peek_ab_empty_status)

                        if is_tma_warp:
                            cute.copy(
                                tma_atom_b,
                                tBgB_nkl[None, k_tile],
                                tBsB[None, mainloop_producer_state.index],
                                tma_bar_ptr=mainloop_pipeline.producer_get_barrier(mainloop_producer_state),
                                mcast_mask=b_mcast_mask,
                                tma_desc_ptr=b_tma_desc_ptr,
                            )

                        K_offset = k_tile * const_expr(self.tile_K)

                        if const_expr(self.compute_weight_gradient and self.is_A_gather):
                            if K_offset % const_expr(self.prefetch_token_idx_size) == 0:
                                self.prefetch_gather_idx_for_A_when_vary_K(
                                    mAIdx_mk, sAIdx_prefetch, token_group_size, K_offset
                                )

                        if const_expr(self.is_A_gather):
                            self.load_A_gather(
                                mA_mkl,
                                tmAIdx,
                                sAIdx_prefetch,
                                M_offset,
                                tAsA[None, None, None, mainloop_producer_state.index],
                                tApA,
                                A_g2s_thr_copy,
                                K_offset,
                                token_group_size,
                                A_thr_copy_elems,
                            )
                        else:
                            cute.copy(
                                tma_atom_a,
                                tAgA_mkl[None, k_tile],
                                tAsA[None, mainloop_producer_state.index],
                                tma_bar_ptr=mainloop_pipeline.producer_get_barrier(mainloop_producer_state),
                                mcast_mask=a_mcast_mask,
                                tma_desc_ptr=a_tma_desc_ptr,
                            )

                        if const_expr(not self.is_A_gather):
                            # Mainloop pipeline's producer commit is a NOP
                            mainloop_pipeline.producer_commit(mainloop_producer_state)
                        else:
                            mainloop_pipeline.producer_cpasync_commit(mainloop_producer_state)
                        mainloop_producer_state.advance()

                        peek_ab_empty_status = cutlass.Boolean(True)
                        if k_tile + 1 < k_tile_cnt:
                            peek_ab_empty_status = mainloop_pipeline.producer_try_acquire(mainloop_producer_state)

                    tile_scheduler.fetch_next_work(is_scheduler_warp=is_scheduler_warp)
                    tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp)
                    work_tile = tile_scheduler.get_current_work()

                if const_expr(self.pingpong):
                    # Need to write the tile_idx to smem for the next WG in the pingpong mode
                    tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp)
                    # End of persistent scheduler loop
                mainloop_pipeline.producer_tail(mainloop_producer_state)
                if is_scheduler_warp:
                    tile_scheduler.producer_tail()

        if warp_idx < self.tma_warp_id:
            cute.arch.setmaxregister_increase(self.num_regs_mma)
            cute.arch.setmaxregister_increase(self.num_regs_mma)
            is_tma_warp = cutlass.Boolean(
                (not self.pingpong and warp_idx == 0) or (self.pingpong and (warp_idx == 0 or warp_idx == 4))
            )
            if const_expr(not self.compute_weight_gradient):
                if const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                    tensormap_manager.init_tensormap_from_atom(
                        tma_atom_d,
                        tensormap_d_init_ptr,
                        is_manager_warp=is_tma_warp,
                    )
                if const_expr(self.need_adhoc_epilogue_store):
                    tensormap_manager.init_tensormap_from_atom(
                        tma_atom_y,
                        tensormap_y_init_ptr,
                        is_manager_warp=is_tma_warp,
                    )
                if const_expr(self.need_epilogue_load):
                    tensormap_manager.init_tensormap_from_atom(
                        tma_atom_c,
                        tensormap_c_init_ptr,
                        is_manager_warp=is_tma_warp,
                    )

            tidx, _, _ = cute.arch.thread_idx()
            warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group)
            if const_expr(self.pingpong):
                tidx = tidx % self.num_threads_per_warp_group
            warp_group_thread_layout = cute.make_layout(
                self.mma_warp_groups if not self.pingpong else 1,
                stride=self.num_threads_per_warp_group,
            )
            thr_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx if not self.pingpong else 0))

            tCrA = tiled_mma.make_fragment_A(thr_mma.partition_A(sA))
            tCrB = tiled_mma.make_fragment_B(thr_mma.partition_B(sB))

            acc_shape = tiled_mma.partition_shape_C(cute.select(self.tile_shape_mnk, mode=[0, 1]))
            acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
            acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype)

            if const_expr(self.pingpong):
                if warp_group_idx == 0:
                    # WG0 needs a start signal at the very beginning
                    self.pingpong_barrier_arrive(warp_group_idx=0, stage="mma")
                    self.pingpong_barrier_arrive(warp_group_idx=0, stage="epi")

            mainloop_consumer_read_state = make_pipeline_state(pipeline.PipelineUserType.Consumer, self.ab_stage)
            epi_read_state = make_pipeline_state(pipeline.PipelineUserType.Consumer, self.c_epi_stage)
            epi_producer_state = make_pipeline_state(pipeline.PipelineUserType.Producer, self.c_epi_stage)

            if const_expr(not self.compute_weight_gradient):
                tensormap_manager.fence_tensormap_initialization()

            tile_scheduler = TileSchedulerCls()
            if const_expr(self.pingpong):
                if warp_idx >= 4:
                    # Advance 2nd Math WG pipeline states to the end of 1st Math WG
                    if const_expr(self.compute_weight_gradient):
                        wg0_batch_idx = tile_scheduler.initial_work_tile_info().tile_idx[-1]
                        wg0_token_group_size = cute.arch.make_warp_uniform(
                            mTokenoffset[wg0_batch_idx + 1]
                        ) - cute.arch.make_warp_uniform(mTokenoffset[wg0_batch_idx])
                        tile_scheduler.advance_to_next_work()
                        mainloop_consumer_read_state.advance_iters(
                            cute.ceil_div(wg0_token_group_size, self.tile_shape_mnk[2])
                        )
                    else:
                        tile_scheduler.advance_to_next_work()
                        mainloop_consumer_read_state.advance_iters(k_tile_cnt)

                    # mainloop_consumer_read_state.advance_iters(k_tile_cnt)
                    if const_expr(self.need_epilogue_load):
                        epi_read_state.advance_iters(c_tile_cnt)
                        epi_producer_state.advance_iters(c_tile_cnt)

            work_tile = tile_scheduler.initial_work_tile_info()
            last_batch_idx = cutlass.Int32(-1)
            token_group_size = cutlass.Int32(0)

            TIdx_cur_group = TIdx_next_group = cutlass.Int32(0)
            while work_tile.is_valid_tile:
                tile_coord_mnkl = work_tile.tile_idx
                batch_idx = tile_coord_mnkl[3]
                is_group_changed = batch_idx != last_batch_idx
                if is_group_changed:
                    # construct tensor D based on real address, shape and stride information
                    TIdx_cur_group, TIdx_next_group = cute.arch.make_warp_uniform(
                        mTokenoffset[batch_idx]
                    ), cute.arch.make_warp_uniform(mTokenoffset[batch_idx + 1])
                    token_group_size = cute.arch.make_warp_uniform(TIdx_next_group - TIdx_cur_group)
                    if const_expr(self.compute_weight_gradient):
                        k_tile_cnt = cute.arch.make_warp_uniform(
                            cute.ceil_div(token_group_size, self.tile_shape_mnk[2])
                        )
                    else:
                        if const_expr((not self.inference_mode) or (not self.need_adhoc_epilogue_store)):
                            assert d_tensormap_smem_ptr is not None and d_tensormap_ptr is not None
                            self.update_tma_desc_ptr(
                                mD_mnl,
                                tma_atom_d,
                                tensormap_manager,
                                d_tensormap_ptr,
                                TIdx_cur_group,
                                token_group_size,
                                is_tma_warp,
                                tensormap_smem_ptr=d_tensormap_smem_ptr,
                                # cute.AddressSpace.generic
                            )
                        if const_expr(self.need_adhoc_epilogue_store):
                            assert y_tensormap_smem_ptr is not None and y_tensormap_ptr is not None
                            self.update_tma_desc_ptr(
                                mY_mnl,
                                tma_atom_y,
                                tensormap_manager,
                                y_tensormap_ptr,
                                TIdx_cur_group,
                                token_group_size,
                                is_tma_warp,
                                tensormap_smem_ptr=y_tensormap_smem_ptr,
                                # cute.AddressSpace.generic
                            )
                        if const_expr(self.need_epilogue_load):
                            assert c_tensormap_smem_ptr is not None and c_tensormap_ptr is not None
                            self.update_tma_desc_ptr(
                                mC_mnl,
                                tma_atom_c,
                                tensormap_manager,
                                c_tensormap_ptr,
                                TIdx_cur_group,
                                token_group_size,
                                is_tma_warp,
                                tensormap_smem_ptr=c_tensormap_smem_ptr,
                                # cute.AddressSpace.generic
                            )
                        last_batch_idx = batch_idx

                k_pipe_mmas = 1
                mainloop_consumer_release_state = mainloop_consumer_read_state.clone()
                num_prologue_mma = min(k_pipe_mmas, k_tile_cnt)
                if const_expr(self.pingpong):
                    self.pingpong_barrier_sync(warp_group_idx, stage="mma")

                peek_ab_full_status = cutlass.Boolean(True)

                if const_expr(self.compute_weight_gradient):
                    if k_tile_cnt == 0:
                        acc.fill(0.0)

                if k_tile_cnt > 0:
                    peek_ab_full_status = mainloop_pipeline.consumer_try_wait(mainloop_consumer_read_state)

                tiled_mma.set(warpgroup.Field.ACCUMULATE, False)
                num_k_blocks = cute.size(tCrA, mode=[2])

                for k_tile in cutlass.range(num_prologue_mma):
                    # Wait for A/B buffer to be ready
                    mainloop_pipeline.consumer_wait(mainloop_consumer_read_state, peek_ab_full_status)
                    warpgroup.fence()
                    for k_blk_idx in cutlass.range(num_k_blocks, unroll_full=True):
                        k_blk_coord = (None, None, k_blk_idx, mainloop_consumer_read_state.index)
                        cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
                        tiled_mma.set(warpgroup.Field.ACCUMULATE, True)
                    warpgroup.commit_group()
                    mainloop_consumer_read_state.advance()
                    peek_ab_full_status = cutlass.Boolean(1)
                    if k_tile + 1 < k_tile_cnt:
                        peek_ab_full_status = mainloop_pipeline.consumer_try_wait(mainloop_consumer_read_state)

                for k_tile in cutlass.range(num_prologue_mma, k_tile_cnt, unroll=1):
                    # Wait for TMA copies to complete
                    mainloop_pipeline.consumer_wait(mainloop_consumer_read_state, peek_ab_full_status)
                    # WGMMA
                    warpgroup.fence()
                    for k_blk_idx in cutlass.range(num_k_blocks, unroll_full=True):
                        k_blk_coord = (None, None, k_blk_idx, mainloop_consumer_read_state.index)
                        cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
                    warpgroup.commit_group()
                    # Wait on the wgmma barrier for previous k_pipe_mmas wgmmas to complete
                    warpgroup.wait_group(k_pipe_mmas)
                    mainloop_pipeline.consumer_release(mainloop_consumer_release_state)
                    mainloop_consumer_read_state.advance()
                    mainloop_consumer_release_state.advance()
                    peek_ab_full_status = cutlass.Boolean(1)
                    if k_tile + 1 < k_tile_cnt:
                        peek_ab_full_status = mainloop_pipeline.consumer_try_wait(mainloop_consumer_read_state)
                if const_expr(self.pingpong):
                    # Cue for next WG's MMA to start
                    self.pingpong_barrier_arrive(1 - warp_group_idx, stage="mma")
                warpgroup.wait_group(0)
                for k_tile in cutlass.range(num_prologue_mma, unroll=1):
                    mainloop_pipeline.consumer_release(mainloop_consumer_release_state)
                    mainloop_consumer_release_state.advance()

                if const_expr(self.pingpong):
                    if const_expr(self.compute_weight_gradient):
                        other_batch_idx = tile_scheduler.prefetch_next_work().tile_idx[-1]
                        other_token_group_size = cute.arch.make_warp_uniform(
                            mTokenoffset[other_batch_idx + 1]
                        ) - cute.arch.make_warp_uniform(mTokenoffset[other_batch_idx])
                        mainloop_consumer_read_state.advance_iters(
                            cute.ceil_div(other_token_group_size, self.tile_shape_mnk[2])
                        )
                    else:
                        mainloop_consumer_read_state.advance_iters(k_tile_cnt)

                    # Update starting mainloop pipeline state for the next tile

                if const_expr(self.pingpong):
                    self.pingpong_barrier_sync(warp_group_idx, "epi")

                epilogue_barrier = pipeline.NamedBarrier(
                    barrier_id=int(NamedBarrierGemm.Epilogue), num_threads=self.num_epi_threads
                )

                # Wait for all warp groups in the thread block to finish, because smem for tensor
                # A in the mainloop is reused in the epilogue if not persistent.
                if const_expr(not self.is_persistent):
                    epilogue_barrier.arrive_and_wait()

                copy_atom_D_r2s = sm90_utils.sm90_get_smem_store_op(
                    self.d_layout,
                    elem_ty_d=self.d_dtype,
                    elem_ty_acc=self.acc_dtype,
                )
                copy_atom_D = cute.make_copy_atom(
                    warp.StMatrix8x8x16bOp(self.d_layout.is_m_major_c(), 4),
                    self.d_dtype,
                )
                tiled_copy_D_atom = cute.make_tiled_copy_C_atom(copy_atom_D, tiled_mma)
                tiled_copy_D_r2s = cute.make_tiled_copy_S(copy_atom_D_r2s, tiled_copy_D_atom)
                # (R2S, R2S_M, R2S_N, PIPE_D)
                tRS_sD = tiled_copy_D_r2s.get_slice(tidx).partition_D(sD)
                tRS_rD_layout = cute.make_layout(tiled_copy_D_r2s.get_slice(tidx).partition_S(sD).shape[:3])
                tRS_rD = cute.make_rmem_tensor(tRS_rD_layout, self.acc_dtype)

                if const_expr(self.need_epilogue_load):
                    copy_atom_C = cute.make_copy_atom(
                        warp.StMatrix8x8x16bOp(
                            self.c_layout.is_m_major_c(),
                            num_matrices=(4 if self.c_epi_tile[1] % 16 == 0 else 2),
                        ),
                        cutlass.Float16,  # this is just to get the right source layout
                    )
                    tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
                    copy_atom_C_s2r = sm90_get_smem_load_op(self.c_layout, self.c_dtype)
                    tiled_copy_C_s2r = cute.make_tiled_copy_S(copy_atom_C_s2r, tiled_copy_C_atom)
                    thr_copy_C_s2r = tiled_copy_C_s2r.get_slice(tidx)
                    tSR_sC = thr_copy_C_s2r.partition_S(sC)
                    tRS_rC = cute.make_rmem_tensor(tRS_rD_layout, self.c_dtype)
                    tRS_rC = cute.make_rmem_tensor(tRS_rD_layout, self.c_dtype)
                    tSR_rC = thr_copy_C_s2r.retile(tRS_rC)
                else:
                    thr_copy_C_s2r, tSR_sC, tRS_rC, tSR_rC = None, None, None, None

                if const_expr(self.need_adhoc_epilogue_store):
                    copy_atom_Y_r2s = sm90_utils.sm90_get_smem_store_op(
                        self.y_layout,
                        elem_ty_d=self.y_dtype,
                        elem_ty_acc=self.acc_dtype,
                    )
                    copy_atom_Y = cute.make_copy_atom(
                        warp.StMatrix8x8x16bOp(self.y_layout.is_m_major_c(), 4),
                        self.y_dtype,
                    )
                    tiled_copy_Y_atom = cute.make_tiled_copy_C_atom(copy_atom_Y, tiled_mma)
                    tiled_copy_Y_r2s = cute.make_tiled_copy_S(copy_atom_Y_r2s, tiled_copy_Y_atom)
                    tRS_sY = tiled_copy_Y_r2s.get_slice(tidx).partition_D(sY)

                # (R2S, R2S_M, R2S_N)
                tRS_rAcc = tiled_copy_D_r2s.retile(acc)
                # tRS_rAcc: tensor<ptr<f32, rmem, align<32>> o ((8,8),3,1):((1,8),64,0)>

                # (bM, bN)
                batch_idx = tile_coord_mnkl[3]
                if const_expr(self.compute_weight_gradient):
                    gD_mn = cute.local_tile(
                        mD_mnl_tma[None, None, batch_idx], (self.tile_M, self.tile_N), tile_coord_mnkl[:2]
                    )
                else:
                    gD_mn = cute.local_tile(mD_mnl_tma, (self.tile_M, self.tile_N), tile_coord_mnkl[:2])

                copy_elems_D = self.universal_copy_bits // mD_mnl.element_type.width
                tdgd_for_tma_partition = cute.zipped_divide(gD_mn, self.d_epi_tile)

                if const_expr(self.need_adhoc_epilogue_store):
                    y_tile_size = (self.tile_M, self.tile_N)
                    if const_expr(self.is_glu and not self.compute_dz_and_partial_ds_and_y1s):
                        y_tile_size = (self.tile_M, self.tile_N // 2)

                    gY_mn = cute.local_tile(mY_mnl_tma, y_tile_size, tile_coord_mnkl[:2])

                    tygy_for_tma_partition = cute.zipped_divide(gY_mn, self.y_epi_tile)
                # bSG_sD: tensor<ptr<bf16, smem, align<1024>, S<2,4,3>> o ((2048,1),(1,4)):((1,0),(0,2048))>
                # bSG_gD: tensor<(?{div=128},?{div=192},?) o (((32,64),1),(3,4)):(((1@0,1@1),0),(64@1,32@0))>
                if const_expr(self.inference_mode and self.need_adhoc_epilogue_store):
                    bSG_sD = bSG_gD = None
                else:
                    bSG_sD, bSG_gD = cpasync.tma_partition(
                        tma_atom_d,
                        0,
                        cute.make_layout(1),
                        cute.group_modes(sD, 0, 2),
                        tdgd_for_tma_partition,
                    )

                if const_expr(self.need_adhoc_epilogue_store):
                    bSG_sY, bSG_gY = cpasync.tma_partition(
                        tma_atom_y,
                        0,
                        cute.make_layout(1),
                        cute.group_modes(sY, 0, 2),
                        tygy_for_tma_partition,
                    )
                    assert const_expr(cute.size(tdgd_for_tma_partition, mode=[1])) == const_expr(
                        cute.size(tygy_for_tma_partition, mode=[1])
                    )

                if const_expr(self.use_bias):
                    expert_idx = tile_coord_mnkl[-1]
                    expert_elem_load = const_expr(self.universal_copy_bits // mBias_nl.element_type.width)
                    gBias = cute.local_tile(mBias_nl, (1, self.tile_shape_mnk[1]), (expert_idx, tile_coord_mnkl[1]))
                    cBias = cute.local_tile(
                        cute.make_identity_tensor((1, mBias_nl.shape[1])),
                        (1, self.tile_shape_mnk[1]),
                        (0, tile_coord_mnkl[1]),
                    )

                    thr_copy_bias = cpasync_atom_bias.get_slice(tidx)
                    tBiasgBias = thr_copy_bias.partition_S(gBias)
                    tBiassBias = thr_copy_bias.partition_D(sBias)
                    tBiascBias = thr_copy_bias.partition_S(cBias)

                    thread_per_row = const_expr(self.tile_shape_mnk[1] // expert_elem_load)
                    if tidx < thread_per_row:
                        if tBiascBias[0][1] < mBias_nl.shape[1]:
                            cute.copy(thr_copy_bias, tBiasgBias, tBiassBias)
                        else:
                            tBiassBias.fill(0.0)

                    cute.arch.cp_async_commit_group()
                    cute.arch.cp_async_wait_group(0)
                    # cannot be removed for correctness!
                    epilogue_barrier.arrive_and_wait()

                    partition_for_epi_fn = partial(
                        partition_for_epilogue,
                        epi_tile=self.d_epi_tile,
                        tiled_copy=tiled_copy_D_r2s,
                        tidx=tidx,
                        reference_src=True,
                    )
                    sBias_retiled = partition_for_epi_fn(
                        cute.make_tensor(sBias.iterator, cute.make_layout((self.tile_M, self.tile_N), stride=(0, 1)))
                    )

                epi_tile_num = const_expr(cute.size(tdgd_for_tma_partition, mode=[1]))

                epi_tile_shape = tdgd_for_tma_partition.shape[1]
                num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num
                epi_tile_layout = cute.make_layout(epi_tile_shape, stride=(epi_tile_shape[1], 1))

                if const_expr(mDIdx_mnl is not None):
                    mcD = cute.make_identity_tensor((self.tile_M, self.tile_N))

                    tcDgcD_flat_partition = cute.flat_divide(mcD, self.d_epi_tile)
                    D_r2g_thr_copy = D_tiled_copy.get_slice(tidx)

                    TIdx_cur_group, TIdx_next_group = mTokenoffset[batch_idx], mTokenoffset[batch_idx + 1]
                    if const_expr(self.is_scatter_idx_prefetched):
                        tmDIdx = self.prefetch_scatter_idx_for_D_when_vary_M(
                            mD_mnl,
                            mDIdx_mnl,
                            D_r2g_thr_copy,
                            tcDgcD_flat_partition,
                            epi_tile_layout,
                            epi_tile_num,
                            const_expr(self.universal_copy_bits // mD_mnl.element_type.width),
                            tile_coord_mnkl,
                            TIdx_cur_group,
                            TIdx_next_group,
                        )
                    else:
                        tmDIdx = None
                else:
                    mcD = tcDgcD_flat_partition = D_r2g_thr_copy = tmDIdx = None

                if const_expr(not self.compute_weight_gradient):
                    if const_expr(self.inference_mode and self.need_adhoc_epilogue_store):
                        d_tma_desc_ptr = None
                    else:
                        d_tma_desc_ptr = tensormap_manager.get_tensormap_ptr(
                            d_tensormap_ptr,
                            cute.AddressSpace.generic,
                        )
                    if const_expr(self.need_adhoc_epilogue_store):
                        y_tma_desc_ptr = tensormap_manager.get_tensormap_ptr(
                            y_tensormap_ptr,
                            cute.AddressSpace.generic,
                        )
                    if const_expr(self.need_epilogue_load):
                        c_tma_desc_ptr = tensormap_manager.get_tensormap_ptr(
                            c_tensormap_ptr,
                            cute.AddressSpace.generic,
                        )
                else:
                    d_tma_desc_ptr = y_tma_desc_ptr = c_tma_desc_ptr = None

                if const_expr(self.compute_dz_and_partial_ds_and_y1s):
                    TIdx_cur_group, TIdx_next_group = cute.arch.make_warp_uniform(
                        mTokenoffset[batch_idx]
                    ), cute.arch.make_warp_uniform(mTokenoffset[batch_idx + 1])
                    self.fetch_scattered_S(
                        tidx,
                        mS_ml,
                        mS_scatter_idx,
                        sS,
                        tile_coord_mnkl,
                        TIdx_cur_group,
                        TIdx_next_group,
                    )
                    epilogue_barrier.arrive_and_wait()

                    cD = cute.make_identity_tensor((self.tile_M, self.tile_N))
                    tDcD = tiled_mma.get_slice(tidx).partition_C(cD)
                    tRS_rcD_retiled = tiled_copy_D_r2s.retile(tDcD)
                    tRS_rcD = cute.make_rmem_tensor_like(tRS_rD, dtype=mS_scatter_idx.element_type)

                if const_expr(self.need_epilogue_load):
                    # mC_mn = cute.domain_offset((mTokenoffset[batch_idx], 0), mC_mnl_tma)
                    gC = cute.local_tile(mC_mnl_tma, (self.tile_M, self.tile_N), tile_coord_mnkl[:2])
                    tCgC_for_tma_partition = cute.zipped_divide(gC, self.c_epi_tile)
                    bGS_sC, bGS_gC = cpasync.tma_partition(
                        tma_atom_c,
                        0,
                        cute.make_layout(1),
                        cute.group_modes(sC, 0, 2),
                        tCgC_for_tma_partition,
                    )

                    for epi_idx in cutlass.range(min(epi_tile_num, self.c_epi_stage), unroll=1):
                        if is_tma_warp:
                            epi_pipeline.producer_acquire(epi_producer_state)
                            # Get the global memory coordinate for the current epi tile
                            gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
                            cute.copy(
                                tma_atom_c,
                                bGS_gC[None, gmem_coord],
                                bGS_sC[None, epi_producer_state.index],
                                tma_bar_ptr=epi_pipeline.producer_get_barrier(epi_producer_state),
                                tma_desc_ptr=c_tma_desc_ptr,
                            )
                            # Epi pipeline's producer commit is a NOP
                            epi_pipeline.producer_commit(epi_producer_state)
                        epi_producer_state.advance()

                for epi_idx in cutlass.range_constexpr(epi_tile_num):
                    # Copy from acc to D registers
                    # tRS_sD: (((2, 4), 1), 1, 2, (1, 4))
                    # tRS_rD = cute.make_fragment_like(tRS_sD[None, None, None, 0], self.acc_dtype) # (((2, 4), 1), 1, 2)

                    # tRS_rD: tensor<ptr<f32, rmem, align<16>> o (((2,4),1),1,2):(((1,2),0),0,8)>
                    for epi_v in cutlass.range_constexpr(cute.size(tRS_rD)):  # cute.size(tRS_rD): 16
                        tRS_rD[epi_v] = tRS_rAcc[const_expr(epi_idx * cute.size(tRS_rD) + epi_v)]
                        if const_expr(self.compute_dz_and_partial_ds_and_y1s):
                            tRS_rcD[epi_v] = tRS_rcD_retiled[const_expr(epi_idx * cute.size(tRS_rD) + epi_v)][0]

                    if const_expr(self.need_epilogue_load):
                        epi_pipeline.consumer_wait(epi_read_state)
                        cute.copy(thr_copy_C_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC)
                        # Fence to make sure shared memory read is visible to TMA load
                        cute.arch.fence_proxy(cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta)
                        cute.arch.sync_warp()
                        with cute.arch.elect_one():
                            epi_pipeline.consumer_release(epi_read_state)
                        epi_read_state.advance()
                        if const_expr(epi_idx + self.c_epi_stage < epi_tile_num):
                            if is_tma_warp:
                                epi_pipeline.producer_acquire(epi_producer_state)
                                # Get the global memory coordinate for the current epi tile
                                gmem_coord = epi_tile_layout.get_hier_coord(epi_idx + self.c_epi_stage)
                                cute.copy(
                                    tma_atom_c,
                                    bGS_gC[None, gmem_coord],
                                    bGS_sC[None, epi_producer_state.index],
                                    tma_bar_ptr=epi_pipeline.producer_get_barrier(epi_producer_state),
                                    tma_desc_ptr=c_tma_desc_ptr,
                                )
                                # Epi pipeline's producer commit is a NOP
                                epi_pipeline.producer_commit(epi_producer_state)
                            epi_producer_state.advance()

                    if const_expr(self.use_bias):
                        sBias_retiled_and_grouped = cute.group_modes(sBias_retiled, 3, cute.rank(sBias_retiled))
                        sBias_retiled_and_grouped_epi = sBias_retiled_and_grouped[
                            None, None, None, epi_tile_layout.get_hier_coord(epi_idx)
                        ]
                        rBias_retiled_epi_r = cute.make_rmem_tensor(
                            sBias_retiled_and_grouped_epi.layout, dtype=mBias_nl.element_type
                        )
                        cute.autovec_copy(
                            cute.filter_zeros(sBias_retiled_and_grouped_epi), cute.filter_zeros(rBias_retiled_epi_r)
                        )

                        for epi_v in cutlass.range_constexpr(cute.size(tRS_rD)):
                            tRS_rD[epi_v] = tRS_rD[epi_v] + self.acc_dtype(rBias_retiled_epi_r[epi_v])

                    if const_expr(self.compute_dz_and_partial_ds_and_y1s):
                        tRS_rD_out = cute.make_rmem_tensor_like(
                            tRS_rD, (cutlass.Float32 if const_expr(self.is_glu) else self.d_dtype)
                        )
                        tRS_rY = cute.make_rmem_tensor_like(tRS_sY[None, None, None, 0], self.y_dtype)
                        self.compute_backward_activation(
                            tRS_rAcc, sS, tRS_rcD, tRS_rC, tRS_rD, tRS_rD_out, tRS_rY, epi_idx
                        )

                    elif const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                        tRS_rD_out = cute.make_rmem_tensor_like(tRS_rD, self.d_dtype)
                        tRS_rD_out.store(tRS_rD.load().to(self.d_dtype))

                    if const_expr((self.is_glu or self.is_normal_act) and not self.compute_dz_and_partial_ds_and_y1s):
                        tRS_rY = cute.make_rmem_tensor_like(tRS_sY[None, None, None, 0], self.y_dtype)
                        self.compute_activation(tRS_rD, tRS_rY)

                    # Copy from D registers to shared memory
                    if const_expr(self.inference_mode and self.need_adhoc_epilogue_store):
                        epi_buffer = (num_prev_subtiles + epi_idx) % cute.size(tRS_sY, mode=[3])
                    else:
                        epi_buffer = (num_prev_subtiles + epi_idx) % cute.size(tRS_sD, mode=[3])

                    if const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                        cute.copy(tiled_copy_D_r2s, tRS_rD_out, tRS_sD[(None, None, None, epi_buffer)])
                    if const_expr(self.need_adhoc_epilogue_store):
                        cute.copy(tiled_copy_Y_r2s, tRS_rY, tRS_sY[(None, None, None, epi_buffer)])

                    if const_expr(mDIdx_mnl is not None):
                        epilogue_barrier.arrive_and_wait()
                        tDsD = D_r2g_thr_copy.partition_S(sD[None, None, epi_buffer])
                        tDrD = cute.make_rmem_tensor_like(tDsD)
                        tDrD = cute.make_rmem_tensor_like(tDsD)
                        cute.autovec_copy(tDsD, tDrD)

                        tDcD_slice = D_r2g_thr_copy.partition_D(
                            tcDgcD_flat_partition[None, None, *epi_tile_layout.get_hier_coord(epi_idx)]
                        )
                        self.store_D_scatter(
                            mD_mnl,
                            mDIdx_mnl,
                            tmDIdx,
                            tDrD,
                            tDcD_slice,
                            D_r2g_thr_copy,
                            epi_idx,
                            copy_elems_D,
                            tile_coord_mnkl,
                            TIdx_cur_group,
                            TIdx_next_group,
                        )
                        epilogue_barrier.arrive_and_wait()

                    else:
                        # Fence and barrier to make sure shared memory store is visible to TMA store
                        cute.arch.fence_proxy(cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta)
                        epilogue_barrier.arrive_and_wait()
                        # Get the global memory coordinate for the current epi tile.
                        gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
                        # Copy from shared memory to global memory
                        if is_tma_warp:
                            if const_expr(not (self.inference_mode and self.need_adhoc_epilogue_store)):
                                cute.copy(
                                    tma_atom_d,
                                    bSG_sD[None, epi_buffer],
                                    bSG_gD[None, gmem_coord],
                                    tma_desc_ptr=d_tma_desc_ptr,
                                )
                            if const_expr(self.need_adhoc_epilogue_store):
                                cute.copy(
                                    tma_atom_y,
                                    bSG_sY[None, epi_buffer],
                                    bSG_gY[None, gmem_coord],
                                    tma_desc_ptr=y_tma_desc_ptr,
                                )
                            cute.arch.cp_async_bulk_commit_group()
                            if const_expr(self.inference_mode and self.need_adhoc_epilogue_store):
                                cute.arch.cp_async_bulk_wait_group(const_expr(self.y_epi_stage - 1), read=True)
                            else:
                                cute.arch.cp_async_bulk_wait_group(const_expr(self.d_epi_stage - 1), read=True)

                        epilogue_barrier.arrive_and_wait()

                if const_expr(self.compute_dz_and_partial_ds_and_y1s):
                    y1 = make_acc_tensor_mn_view(acc)
                    cD = cute.make_identity_tensor((self.tile_M, self.tile_N))
                    tDcD = tiled_mma.get_slice(tidx).partition_C(cD)
                    tDcD_mn = make_acc_tensor_mn_view(tDcD)

                    tile_M_offset = cute.arch.make_warp_uniform(TIdx_cur_group + tile_coord_mnkl[0] * self.tile_M)

                    mDS_partial_M, mDS_partial_N = mDS_partial.shape
                    mDS_partial_flatten_view = cute.make_tensor(mDS_partial.iterator, (mDS_partial_M * mDS_partial_N,))
                    for r in cutlass.range_constexpr(cute.size(y1, mode=[0])):
                        col_sum = cutlass.Float32(0.0)

                        M_tile_idx = tDcD_mn[r, 0][0]
                        for c in cutlass.range_constexpr(cute.size(y1, mode=[1])):
                            col_sum = col_sum + y1[r, c]

                        col_sum = cute.arch.warp_reduction(col_sum, operator.add, threads_in_group=4)

                        M_idx_raw = tile_M_offset + M_tile_idx
                        if tidx % 4 == 0 and M_idx_raw < TIdx_next_group:
                            M_idx = mS_scatter_idx[M_idx_raw]
                            N_idx = tile_coord_mnkl[1]
                            mDS_partial_flatten_view[M_idx * mDS_partial_N + N_idx] = col_sum.to(
                                mDS_partial.element_type
                            )

                if const_expr(self.pingpong):
                    # With pingpong, 2 WGs write two different output tiles to the same smem,
                    # so we have to make sure the smem content is done reading before signalling
                    # the next WG's epilogue.
                    if const_expr(self.need_epilogue_load):
                        epi_read_state.advance_iters(c_tile_cnt)
                        epi_producer_state.advance_iters(c_tile_cnt)
                    if warp_idx == 0 or warp_idx == 4:
                        cute.arch.cp_async_bulk_wait_group(0, read=True)
                    self.pingpong_barrier_arrive(1 - warp_group_idx, stage="epi")

                tile_scheduler.advance_to_next_work(advance_count=1 if not self.pingpong else self.mma_warp_groups)
                work_tile = tile_scheduler.get_current_work()
                # End of persistent scheduler loop

            if const_expr(not self.pingpong):
                if warp_idx == 0:
                    cute.arch.cp_async_bulk_wait_group(0, read=True)

    def generate_tensormap(self, m, n, l):
        if not self.is_persistent:
            total_m = m * l
            block_size_m = self.tile_M * self.cluster_shape_mnk[0]
            block_size_n = self.tile_N * self.cluster_shape_mnk[1]
            total_clusters_m_max = (total_m + l * (block_size_m - 1)) // block_size_m
            total_clusters_max = total_clusters_m_max * ((n + block_size_n - 1) // block_size_n)
            total_ctas = total_clusters_max * self.cluster_shape_mnk[0] * self.cluster_shape_mnk[1]
        else:
            total_ctas = cutlass.utils.HardwareInfo().get_device_multiprocessor_count()
        if self.pingpong:
            total_ctas *= 2
        # 128 bytes per tensormap
        tensormaps_torch = torch.empty(total_ctas, 128 // 8, dtype=torch.int64, device="cuda")
        tensormaps_tensor = from_dlpack(tensormaps_torch, assumed_align=128).mark_compact_shape_dynamic(
            mode=0, stride_order=(0, 1)
        )
        return tensormaps_tensor

    def pingpong_barrier_sync(self, warp_group_idx: Int32, stage: str):
        assert stage in ["mma", "epi"]
        barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0
        cute.arch.barrier(
            barrier_id=int(barrier) + warp_group_idx,
            number_of_threads=2 * self.num_threads_per_warp_group,
        )

    def pingpong_barrier_arrive(self, warp_group_idx: Int32, stage: str):
        assert stage in ["mma", "epi"]
        barrier = NamedBarrierGemm.MmaWG0 if stage == "mma" else NamedBarrierGemm.EpiWG0
        cute.arch.barrier_arrive(
            barrier_id=int(barrier) + warp_group_idx,
            number_of_threads=2 * self.num_threads_per_warp_group,
        )

    def make_sched_pipeline(self, cluster_layout_mnk: cute.Layout, sched_pipeline_mbar_ptr: cute.Pointer):
        # Threads/warps participating in this pipeline
        sched_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        cluster_size = cute.size(cluster_layout_mnk)
        # Each warp that are not the scheduler warp will contribute 1 to the arrive count
        consumer_arrive_cnt = (
            (self.mma_warp_groups if not self.pingpong else 1) * 4
            + max(self.num_load_A_threads // cute.arch.WARP_SIZE, 1)
        ) * cluster_size - 1
        sched_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_arrive_cnt)
        return pipeline.PipelineAsync.create(
            barrier_storage=sched_pipeline_mbar_ptr,
            num_stages=self.sched_stage,
            producer_group=sched_pipeline_producer_group,
            consumer_group=sched_pipeline_consumer_group,
            # If there's cluster, the consumers must arrive at the mbar of CTA 0 in the cluster.
            consumer_mask=None if const_expr(cluster_size == 1) else 0,
        )

    def _compute_stages(
        self,
        tile_shape_mnk: Tuple[int, int, int],
        initial_d_epi_stage: int,
        d_epi_tile: Optional[Tuple[int, int]],
        c_epi_tile: Optional[Tuple[int, int]],
        y_epi_tile: Optional[Tuple[int, int]],
        a_dtype: Type[cutlass.Numeric],
        b_dtype: Type[cutlass.Numeric],
        d_dtype: Type[cutlass.Numeric],
        c_dtype: Optional[Type[cutlass.Numeric]],
        y_dtype: Optional[Type[cutlass.Numeric]],
        smem_capacity: int,
        occupancy: int,
        overlap_sD_sA: bool,
    ) -> Tuple[int, int, int]:
        d_epi_stage = initial_d_epi_stage if const_expr(not self.need_epilogue_load) else initial_d_epi_stage // 2
        y_epi_stage = d_epi_stage

        if self.inference_mode and self.need_adhoc_epilogue_store:
            d_epi_stage = 0

        if overlap_sD_sA:
            epi_bytes = 0
        else:
            d_bytes_per_stage = cute.size(d_epi_tile) * d_dtype.width // 8
            epi_bytes = d_bytes_per_stage * d_epi_stage

            if y_dtype is not None or const_expr(self.need_adhoc_epilogue_store):
                y_bytes_per_stage = cute.size(y_epi_tile) * y_dtype.width // 8
                epi_bytes += y_bytes_per_stage * y_epi_stage
            else:
                y_bytes_per_stage = 0

        c_epi_stage = 0 if (c_dtype is None or const_expr(not self.need_epilogue_load)) else d_epi_stage
        if c_dtype is not None and const_expr(self.need_epilogue_load):
            c_bytes_per_stage = cute.size(c_epi_tile) * c_dtype.width // 8 * c_epi_stage
            epi_bytes += c_bytes_per_stage * c_epi_stage
            d_epi_stage = c_epi_stage
        else:
            c_bytes_per_stage = 0

        a_shape = cute.slice_(tile_shape_mnk, (None, 0, None))
        b_shape = cute.slice_(tile_shape_mnk, (0, None, None))
        ab_bytes_per_stage = cute.size(a_shape) * a_dtype.width // 8 + cute.size(b_shape) * b_dtype.width // 8
        mbar_helpers_bytes = 1024

        remaining_bytes = (
            (smem_capacity - occupancy * 1024) // occupancy
            - mbar_helpers_bytes
            - epi_bytes
            - self.prefetch_token_idx_size * 4
            - (self.tile_shape_mnk[1] * (self.bias_dtype.width // 8) if self.use_bias else 0)
            - 1024  # aligned  self.tensormap_management_bytes
        )
        ab_stage = remaining_bytes // ab_bytes_per_stage

        # Refine epilogue stages:
        # Calculate remaining smem after allocating for A/B stages and reserved bytes
        # Add remaining unused smem to epilogue
        if not overlap_sD_sA:
            if self.inference_mode and self.need_adhoc_epilogue_store:
                epi_stage_delta = (remaining_bytes - ab_bytes_per_stage * ab_stage) // (
                    y_bytes_per_stage + c_bytes_per_stage
                )
                y_epi_stage += epi_stage_delta
            else:
                epi_stage_delta = (remaining_bytes - ab_bytes_per_stage * ab_stage) // (
                    d_bytes_per_stage + y_bytes_per_stage + c_bytes_per_stage
                )
                d_epi_stage += epi_stage_delta
                y_epi_stage += epi_stage_delta

            if c_epi_stage > 0:
                c_epi_stage += epi_stage_delta

        if not self.need_adhoc_epilogue_store:
            y_epi_stage = 0

        return ab_stage, c_epi_stage, d_epi_stage, y_epi_stage

    def _sm90_compute_tile_shape_or_override(
        self,
        tile_shape_mnk: Tuple[int, int, int],
        atom_layout_mnk: Tuple[int, int, int],
        element_type: Type[cutlass.Numeric],
        epi_tile_override: Tuple[int, int] | None = None,
    ) -> Tuple[int, int]:
        """Compute the epilogue tile shape or use override if provided.

        :param tile_shape_mnk: CTA tile shape (M,N,K)
        :type tile_shape_mnk: Tuple[int, int, int]
        :param element_type: Data type of elements
        :type element_type: type[cutlass.Numeric]
        :param is_cooperative: Whether to use cooperative approach
        :type is_cooperative: bool
        :param epi_tile_override: Optional override for epilogue tile shape
        :type epi_tile_override: Tuple[int, int] or None

        :return: Computed epilogue tile shape
        :rtype: Tuple[int, int]
        """
        if epi_tile_override is not None:
            return epi_tile_override
        if tile_shape_mnk[0] % 128 == 0 and atom_layout_mnk[0] > 1:
            tile_m = math.gcd(128, cute.size(tile_shape_mnk, mode=[0]))
            tile_n = math.gcd(self.epi_tile_size, cute.size(tile_shape_mnk, mode=[1]))
            return (tile_m, tile_n)
        elif tile_shape_mnk[0] % 192 == 0 and atom_layout_mnk[0] > 1:
            tile_m = math.gcd(192, cute.size(tile_shape_mnk, mode=[0]))
            tile_n = math.gcd(self.epi_tile_size, cute.size(tile_shape_mnk, mode=[1]))
            return (tile_m, tile_n)
        else:
            # In the case of tile shape 128 x N but atom_layout 1 x 2, we need to set
            # epi_tile_m = 64. If epi_tile_m = 128, the epilogue would iterate along the
            # M dimension first, then move to the N dimension. But the accumulator in registers
            # iterate along the N dimension first, then move to the M dimension.
            # We could change the epilogue to accommodate this,
            # but it's easier to just set epi_tile_m = 64.
            n_perf = 64 if element_type.width == 8 else min(self.epi_tile_size, tile_shape_mnk[1])
            tile_m = math.gcd(64, cute.size(tile_shape_mnk, mode=[0]))
            tile_n = math.gcd(n_perf, cute.size(tile_shape_mnk, mode=[1]))
            return (tile_m, tile_n)

    @staticmethod
    def _make_smem_layouts(
        tile_shape_mnk: Tuple[int, int, int],
        c_epi_tile: Tuple[int, int],
        bias_epi_tile: Tuple[int, int],
        d_epi_tile: Tuple[int, int],
        y_epi_tile: Optional[Tuple[int, int]],
        a_dtype: Type[cutlass.Numeric],
        a_layout: utils.LayoutEnum,
        b_dtype: Type[cutlass.Numeric],
        b_layout: utils.LayoutEnum,
        prefetch_idx_size: Optional[int],
        ab_stage: int,
        c_dtype: Optional[Type[cutlass.Numeric]],
        c_layout: Optional[cutlass.utils.LayoutEnum],
        bias_dtype: Optional[Type[cutlass.Numeric]],
        bias_layout: Optional[cutlass.utils.LayoutEnum],
        d_dtype: Type[cutlass.Numeric],
        d_layout: utils.LayoutEnum,
        y_dtype: Optional[Type[cutlass.Numeric]],
        y_layout: Optional[utils.LayoutEnum],
        s_dtype: Optional[Type[cutlass.Numeric]],
        c_epi_stage: int,
        d_epi_stage: int,
        y_epi_stage: int,
    ) -> Tuple[
        cute.ComposedLayout,
        cute.ComposedLayout,
        Optional[cute.ComposedLayout],
        cute.ComposedLayout,
        Optional[cute.ComposedLayout],
    ]:
        a_smem_shape = cute.slice_(tile_shape_mnk, (None, 0, None))

        a_is_k_major = a_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K
        b_is_k_major = b_layout.sm90_mma_major_mode() == warpgroup.OperandMajorMode.K

        a_major_mode_size = tile_shape_mnk[2 if a_is_k_major else 0]
        a_smem_layout_atom = warpgroup.make_smem_layout_atom(
            sm90_utils.get_smem_layout_atom(
                a_layout,
                a_dtype,
                a_major_mode_size,
            ),
            a_dtype,
        )
        a_smem_layout_staged = cute.tile_to_shape(
            a_smem_layout_atom,
            cute.append(a_smem_shape, ab_stage),
            order=(0, 1, 2) if a_is_k_major else (1, 0, 2),
        )

        b_smem_shape = cute.slice_(tile_shape_mnk, (0, None, None))

        b_major_mode_size = tile_shape_mnk[2 if b_is_k_major else 1]
        b_smem_layout_atom = warpgroup.make_smem_layout_atom(
            sm90_utils.get_smem_layout_atom(
                b_layout,
                b_dtype,
                b_major_mode_size,
            ),
            b_dtype,
        )
        b_smem_layout_staged = cute.tile_to_shape(
            b_smem_layout_atom,
            cute.append(b_smem_shape, ab_stage),
            order=(0, 1, 2) if b_is_k_major else (1, 0, 2),
        )

        d_smem_shape = d_epi_tile
        d_major_mode_size = d_epi_tile[1] if d_layout.is_n_major_c() else d_epi_tile[0]
        d_smem_layout_atom = warpgroup.make_smem_layout_atom(
            sm90_utils.get_smem_layout_atom(
                d_layout,
                d_dtype,
                d_major_mode_size,
            ),
            d_dtype,
        )
        if d_epi_stage > 0:
            d_epi_smem_layout_staged = cute.tile_to_shape(
                d_smem_layout_atom,
                cute.append(d_smem_shape, d_epi_stage),
                order=(1, 0, 2) if d_layout.is_m_major_c() else (0, 1, 2),
            )
        else:
            # calculating the layout
            d_epi_smem_layout_staged = cute.tile_to_shape(
                d_smem_layout_atom,
                cute.append(d_smem_shape, 1),
                order=(1, 0, 2) if d_layout.is_m_major_c() else (0, 1, 2),
            )

        if y_epi_tile is not None:
            y_smem_shape = y_epi_tile
            # we force `y` to have same major mode as `z`. Otherwise the epilogue write is tricky
            y_major_mode_size = y_epi_tile[1] if y_layout.is_n_major_c() else y_epi_tile[0]
            y_smem_layout_atom = warpgroup.make_smem_layout_atom(
                sm90_utils.get_smem_layout_atom(
                    y_layout,
                    y_dtype,
                    y_major_mode_size,
                ),
                y_dtype,
            )
            y_epi_smem_layout_staged = cute.tile_to_shape(
                y_smem_layout_atom,
                cute.append(y_smem_shape, y_epi_stage),
                order=(1, 0, 2) if y_layout.is_m_major_c() else (0, 1, 2),
            )
        else:
            y_epi_smem_layout_staged = None

        if c_dtype is not None:
            assert c_layout is not None
            c_smem_shape = c_epi_tile
            c_major_mode_size = c_epi_tile[1] if c_layout.is_n_major_c() else c_epi_tile[0]
            c_smem_layout_atom = warpgroup.make_smem_layout_atom(
                sm90_utils.get_smem_layout_atom(c_layout, c_dtype, c_major_mode_size),
                c_dtype,
            )
            c_epi_smem_layout_staged = cute.tile_to_shape(
                c_smem_layout_atom,
                cute.append(c_smem_shape, c_epi_stage),
                order=(1, 0, 2) if c_layout.is_m_major_c() else (0, 1, 2),
            )
        else:
            c_epi_smem_layout_staged = None

        if bias_dtype is not None and bias_layout is not None and bias_epi_tile is not None:
            bias_epi_smem_layout_staged = cute.make_layout((1, tile_shape_mnk[1]))
        else:
            bias_epi_smem_layout_staged = None

        if s_dtype is not None:
            s_epi_smem_layout_staged = cute.make_layout((tile_shape_mnk[0],))
        else:
            s_epi_smem_layout_staged = None

        if prefetch_idx_size > 0:
            prefetched_token_idx_smem_layout = cute.make_layout((prefetch_idx_size,))
        else:
            prefetched_token_idx_smem_layout = None

        return (
            a_smem_layout_staged,
            b_smem_layout_staged,
            c_epi_smem_layout_staged,
            bias_epi_smem_layout_staged,
            d_epi_smem_layout_staged,
            y_epi_smem_layout_staged,
            s_epi_smem_layout_staged,
            prefetched_token_idx_smem_layout,
        )

    @staticmethod
    def _make_tma_epi_atoms_and_tensors(
        tensor_d: cute.Tensor,
        epi_smem_layout_staged: cute.ComposedLayout,
        epi_tile: Tuple[int, int],
        store_or_load: str,
    ) -> Tuple[cute.CopyAtom, cute.Tensor]:
        """Create TMA atoms and tensors for storing D or loading C.

        :param tensor_d: Output tensor D
        :type tensor_d: cute.Tensor
        :param epi_smem_layout_staged: Shared memory layout for epilogue
        :type epi_smem_layout_staged: cute.ComposedLayout
        :param epi_tile: Epilogue tile shape
        :type epi_tile: Tuple[int, int]

        :return: TMA atom and tensor for C
        :rtype: Tuple[cute.CopyAtom, cute.Tensor]
        """
        assert store_or_load in ["load", "store"]
        epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
        d_cta_v_layout = cute.composition(cute.make_identity_layout(tensor_d.shape), epi_tile)
        op = cpasync.CopyBulkTensorTileG2SOp() if store_or_load == "load" else cpasync.CopyBulkTensorTileS2GOp()
        tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom(op, tensor_d, epi_smem_layout, d_cta_v_layout)
        return tma_atom_d, tma_tensor_d

    @staticmethod
    def _make_tma_atoms_and_tensors(
        tensor: cute.Tensor,
        smem_layout_staged: cute.ComposedLayout,
        smem_tile: Tuple[int, int],
        mcast_dim: int,
    ) -> Tuple[cute.CopyAtom, cute.Tensor]:
        """Create TMA atoms and tensors for input tensors.

        :param tensor: Input tensor (A or B)
        :type tensor: cute.Tensor
        :param smem_layout_staged: Shared memory layout for the tensor
        :type smem_layout_staged: cute.ComposedLayout
        :param smem_tile: Shared memory tile shape
        :type smem_tile: Tuple[int, int]
        :param mcast_dim: Multicast dimension
        :type mcast_dim: int

        :return: TMA atom and tensor
        :rtype: Tuple[cute.CopyAtom, cute.Tensor]
        """
        op = cpasync.CopyBulkTensorTileG2SOp() if mcast_dim == 1 else cpasync.CopyBulkTensorTileG2SMulticastOp()

        smem_layout = cute.slice_(smem_layout_staged, (None, None, 0))
        tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
            op,
            tensor,
            smem_layout,
            smem_tile,
            num_multicast=mcast_dim,
        )
        return tma_atom, tma_tensor

    def _make_tiled_copy_2D(
        self,
        tensor: cute.Tensor,
        tile_shape_0: cute.Int32,
        tile_shape_1: cute.Int32,
        is_row_major: bool,
        threads_for_copy: Union[cutlass.Int32, int],
        universal_copy_bits: cutlass.Int32,
        is_g2s: Optional[bool] = True,
    ) -> cute.TiledCopy:
        copy_atom = cute.make_copy_atom(
            (
                cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL)
                if const_expr(is_g2s)
                else cute.nvgpu.CopyUniversalOp()
            ),
            tensor.element_type,
            num_bits_per_copy=universal_copy_bits,
        )
        copy_elems = universal_copy_bits // tensor.element_type.width
        shape_dim_1 = cute.size(tile_shape_1) // copy_elems
        # thread layout for copy
        thread_layout = cute.make_layout((threads_for_copy // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1))
        if not is_row_major:
            shape_dim_0 = cute.size(tile_shape_0) // copy_elems
            thread_layout = cute.make_layout((shape_dim_0, threads_for_copy // shape_dim_0), stride=(1, shape_dim_0))
        # Value layout for copy
        value_layout = cute.make_layout((1, copy_elems)) if is_row_major else cute.make_layout((copy_elems, 1))
        return cute.make_tiled_copy_tv(copy_atom, thread_layout, value_layout)

    @staticmethod
    def is_valid_dtypes(
        a_dtype: Type[cutlass.Numeric],
        b_dtype: Type[cutlass.Numeric],
        acc_dtype: Type[cutlass.Numeric],
        out_dtype: Type[cutlass.Numeric],
        a_major: str,
        b_major: str,
    ) -> bool:
        """
        Check if the dtypes are valid

        :param a_dtype: The data type of tensor A
        :type a_dtype: Type[cutlass.Numeric]
        :param b_dtype: The data type of tensor B
        :type b_dtype: Type[cutlass.Numeric]
        :param acc_dtype: The data type of the accumulator
        :type acc_dtype: Type[cutlass.Numeric]
        :param d_dtype: The data type of the output tensor
        :type d_dtype: Type[cutlass.Numeric]
        :param a_major: major mode of tensor A
        :type a_major: str
        :param b_major: major mode of tensor B
        :type b_major: str

        :return: True if the dtypes are valid, False otherwise
        :rtype: bool
        """
        is_valid = True
        # tested a_dtype
        if a_dtype not in {
            cutlass.Float16,
            cutlass.BFloat16,
        }:
            is_valid = False
        # tested b_dtype
        if b_dtype not in {
            cutlass.Float16,
            cutlass.BFloat16,
        }:
            is_valid = False
        # tested acc_dtype
        if acc_dtype not in {cutlass.Float32, cutlass.Float16}:
            is_valid = False
        # tested d_dtype
        if out_dtype not in {
            cutlass.Float32,
            cutlass.Float16,
            cutlass.BFloat16,
        }:
            is_valid = False
        # make sure a_dtype == b_dtype for Float16
        if a_dtype.width == 16 and a_dtype != b_dtype:
            is_valid = False
        # make sure a_dtype.width == b_dtype.width (i.e, Float8E4M3FN or Float8E5M2)
        if a_dtype.width != b_dtype.width:
            is_valid = False
        # for Float8 types, this implementation only supports k-major layout
        if (a_dtype.width == 8 and a_major != "k") or (b_dtype.width == 8 and b_major != "k"):
            is_valid = False
        return is_valid
