# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""
MoE Block-Scaled Grouped GEMM Kernel with GLU (SwiGLU/GeGLU) + Hadamard Transform Fusion.

Supports:
    - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler)
    - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout
    - BF16/F16 D output with GLU activation
    - Optional C output (pre-activation GLU output)
    - AMAX reduction for calibration
    - GLU activation fusion (SwiGLU / GeGLU)

Warp assignment (8 epilogue warps, pingpong):
    warps 0-3  : ACT warps  — TMEM→reg, alpha scale, GLU activation, C/D store
    warps 4-7  : RHT store warps — post-RHT amax from D SMEM
    warp  8    : MMA warp
    warp  9    : TMA load warp
    warp  10   : Scheduler warp (MoEPersistentTileScheduler)
    warp  11   : Bias load warp (optional)

sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt)
    Validity: tile_info[0] >= 0  (expert_idx == -1 signals end)
"""

from typing import Type, Tuple, Union, Optional

import cuda.bindings.driver as cuda

import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass.cute.nvgpu import OperandMajorMode
import cutlass.utils as utils
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.typing import Float32, Int32, AddressSpace
from ..moe_persistent_scheduler import (
    MoEPersistentTileScheduler,
    MoESchedulerParams,
    MoEWorkTileInfo,
)
from ..moe_utils import (
    compute_expert_token_range,
    MoEWeightMode,
    TensormapWorkspace,
    store_tma_desc,
)
from .hadamard_utils import (
    hadamard_setup,
    hadamard_compute,
    hadamard_smem_transpose_fwht_amax,
    hadamard_transpose_in,
    hadamard_transpose_out_amax,
    HADAMARD_SIZE,
    M_PER_CLUSTER,
)
from ..moe_sched_extension import (
    DiscreteWeightScaledGemmSchedExtension,
    ContiguousAndConsistentGroupedGemmSchedExtension,
)
from ..moe_kernel_helpers import (
    fmin,
    fmax,
    atomic_max_float32,
    silu_f32,
    silu_f32_geglu_scaled,
    compute_grid,
    get_dtype_rcp_limits,
    can_implement,
    amax_reduction_per_thread,
)


class BlockScaledMoEGroupedGemmGluHadamardKernel:
    """Block-scaled MoE grouped GEMM with GLU activation and Hadamard transform fusion.

    Always uses pingpong epilogue (8 epilogue warps: 4 ACT + 4 RHT-store).
    D output is BF16 or F16 only (no FP8/FP4, no SFD).

    :param sf_vec_size: Scalefactor vector size.
    :param mma_tiler_mn: Shape of MMA tile (M, N).
    :param cluster_shape_mn: Cluster dimensions (M, N).
    :param expert_cnt: Number of experts (compile-time constant).
    :param weight_mode: Dense or Discrete weight layout.
    :param use_dynamic_sched: Use dynamic tile scheduling.
    :param act_func: Activation function ('swiglu', 'geglu', or 'srelu').
    :param enable_bias: Enable bias addition.
    """

    FIX_PAD_SIZE = 256

    @staticmethod
    def can_implement(
        ab_dtype: Type[cutlass.Numeric],
        sf_dtype: Type[cutlass.Numeric],
        sf_vec_size: int,
        acc_dtype: Type[cutlass.Numeric],
        d_dtype: Type[cutlass.Numeric],
        use_2cta_instrs: bool,
        mma_tiler_mn: Tuple[int, int],
        cluster_shape_mn: Tuple[int, int],
        m: int,
        n: int,
        k: int,
        l: int,
        a_major: str,
        b_major: str,
        cd_major: str,
        m_aligned: int,
    ) -> bool:
        # speical requirements for hadamard fusion
        if not use_2cta_instrs or mma_tiler_mn[0] != 256 or mma_tiler_mn[1] != 256:
            return False
        return can_implement(
            ab_dtype,
            sf_dtype,
            sf_vec_size,
            acc_dtype,
            d_dtype,
            use_2cta_instrs,
            mma_tiler_mn,
            cluster_shape_mn,
            m,
            n,
            k,
            l,
            a_major,
            b_major,
            cd_major,
            m_aligned,
            fix_pad_size=BlockScaledMoEGroupedGemmGluHadamardKernel.FIX_PAD_SIZE,
        )

    def __init__(
        self,
        sf_vec_size: int,
        acc_dtype: Type[cutlass.Numeric],
        use_2cta_instrs: bool,
        mma_tiler_mn: Tuple[int, int],
        cluster_shape_mn: Tuple[int, int],
        vectorized_f32: bool,
        expert_cnt: int,
        weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE,
        use_dynamic_sched: bool = False,
        act_func: str = "swiglu",
        situ_beta1: float = 4.0,
        enable_bias: bool = False,
        use_tmem_post_rht_amax: bool = False,
    ):
        mma_tile_m = mma_tiler_mn[0]
        if self.FIX_PAD_SIZE % mma_tile_m != 0:
            raise ValueError(f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}).")
        if expert_cnt > 1024:
            raise ValueError("Expert count > 1024 is not supported.")
        if not isinstance(weight_mode, MoEWeightMode):
            raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}")

        self.sf_vec_size = sf_vec_size
        self.expert_cnt = expert_cnt
        self.acc_dtype: Type[cutlass.Numeric] = acc_dtype
        self.use_2cta_instrs = use_2cta_instrs
        self.cluster_shape_mn = cluster_shape_mn
        self.mma_tiler = (*mma_tiler_mn, 1)
        self.weight_mode = weight_mode
        self.use_dynamic_sched = use_dynamic_sched
        self.enable_bias = enable_bias
        self.use_tmem_post_rht_amax = use_tmem_post_rht_amax

        # Always use pingpong epilogue for Hadamard
        self.epilogue_pingpong = True
        # Always delay TMA store acquire sync for Hadamard
        self.delay_tma_store_acquire_sync = True

        self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE

        self.occupancy = 1
        self.threads_per_warp = 32

        # Warp assignments: 8 epilogue warps (4 ACT + 4 RHT-store)
        self.epilog_warp_id = (0, 1, 2, 3, 4, 5, 6, 7)
        self.epilog_act_warp_id = (0, 1, 2, 3)
        self.epilog_rht_store_warp_id = (4, 5, 6, 7)
        self.mma_warp_id = 8
        self.tma_warp_id = 9
        self.sched_warp_id = 10
        self.bias_load_warp_id = 11 if enable_bias else None

        self.epilogue_warp_group_size = len(self.epilog_act_warp_id)  # = 4

        all_warps = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id, self.sched_warp_id]
        warps_wo_sched = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id]
        if enable_bias:
            all_warps.append(self.bias_load_warp_id)
            warps_wo_sched.append(self.bias_load_warp_id)
        self.threads_per_cta = self.threads_per_warp * len(all_warps)
        self.threads_wo_sched = self.threads_per_warp * len(warps_wo_sched)

        # Named barriers
        self.cta_sync_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=self.threads_per_cta,
        )
        self.epilog_sync_barrier = pipeline.NamedBarrier(
            barrier_id=2,
            num_threads=32 * len(self.epilog_warp_id),
        )
        tmem_alloc_warp_ids = self.epilog_warp_id if self.use_tmem_post_rht_amax else self.epilog_act_warp_id
        self.tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=3,
            num_threads=32 * len((self.mma_warp_id, *tmem_alloc_warp_ids)),
        )
        self.sched_sync_barrier = pipeline.NamedBarrier(
            barrier_id=4,
            num_threads=self.threads_per_warp,
        )
        # Pingpong barriers (group 0 = ACT warps, group 1 = RHT store warps)
        self.epilog_sync_barrier_group0 = pipeline.NamedBarrier(
            barrier_id=5,
            num_threads=32 * self.epilogue_warp_group_size,
        )
        self.epilog_sync_barrier_group1 = pipeline.NamedBarrier(
            barrier_id=6,
            num_threads=32 * self.epilogue_warp_group_size,
        )

        self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
        SM100_TMEM_CAPACITY_COLUMNS = 512
        self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS

        self.vectorized_f32 = vectorized_f32

        # Amax: ACT warps reduce the post-GLU output.
        self.num_epilog_warps = len(self.epilog_act_warp_id)  # = 4

        self.act_func = act_func
        self.situ_beta1 = float(situ_beta1)
        if act_func not in ["swiglu", "geglu", "situglu", "srelu"]:
            raise ValueError(f"Invalid activation function: {act_func}")

    def _setup_attributes(self):
        """Set up configurations dependent on GEMM inputs (called inside __call__)."""

        self.mma_inst_shape_mn = (
            self.mma_tiler[0],
            self.mma_tiler[1],
        )
        self.mma_inst_shape_mn_sfb = (
            self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
            cute.round_up(self.mma_inst_shape_mn[1], 128),
        )

        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.b_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            self.cta_group,
            self.mma_inst_shape_mn,
        )

        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.b_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            cute.nvgpu.tcgen05.CtaGroup.ONE,
            self.mma_inst_shape_mn_sfb,
        )

        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
        mma_inst_tile_k = 4
        self.mma_tiler = (
            self.mma_tiler[0],
            self.mma_tiler[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )

        self.mma_tiler_sfb = (
            self.mma_inst_shape_mn_sfb[0],
            self.mma_inst_shape_mn_sfb[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )

        self.cta_tile_shape_mnk = (
            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler[1],
            self.mma_tiler[2],
        )
        self.cta_tile_shape_mnk_sfb = (
            self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler_sfb[1],
            self.mma_tiler_sfb[2],
        )

        d_tile_n = self.mma_inst_shape_mn[1] if self.act_func == "srelu" else self.mma_inst_shape_mn[1] // 2
        self.mma_tiler_d = (
            self.mma_inst_shape_mn[0],
            d_tile_n,
            mma_inst_shape_k * mma_inst_tile_k,
        )
        self.cta_tile_shape_mnk_d = (
            self.mma_tiler_d[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler_d[1],
            self.mma_tiler_d[2],
        )

        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((*self.cluster_shape_mn, 1)),
            (tiled_mma.thr_id.shape,),
        )
        self.cluster_layout_sfb_vmnk = cute.tiled_divide(
            cute.make_layout((*self.cluster_shape_mn, 1)),
            (tiled_mma_sfb.thr_id.shape,),
        )

        self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
        self.is_a_mcast = self.num_mcast_ctas_a > 1
        self.is_b_mcast = self.num_mcast_ctas_b > 1

        self.epi_tile = (128, 32)
        self.epi_tile_cnt = (
            self.cta_tile_shape_mnk_d[0] // self.epi_tile[0],
            self.cta_tile_shape_mnk_d[1] // self.epi_tile[1],
        )
        self.epi_tile_c = self.epi_tile if self.act_func == "srelu" else (128, 64)

        (
            self.num_acc_stage,
            self.num_ab_stage,
            self.num_c_stage,
            self.num_d_stage,
            self.num_tile_stage,
            self.num_bias_stage,
            self.num_pingpong_stage,
        ) = self._compute_stages(
            tiled_mma,
            self.mma_tiler,
            self.a_dtype,
            self.b_dtype,
            self.epi_tile,
            self.epi_tile_c,
            self.c_dtype,
            self.c_layout,
            self.d_dtype,
            self.d_layout,
            self.sf_dtype,
            self.sf_vec_size,
            self.num_smem_capacity,
            self.occupancy,
            self.bias_dtype if self.enable_bias else None,
            self.use_tmem_post_rht_amax,
        )

        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
            tiled_mma,
            self.mma_tiler,
            self.a_dtype,
            self.num_ab_stage,
        )
        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
            tiled_mma,
            self.mma_tiler,
            self.b_dtype,
            self.num_ab_stage,
        )
        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma,
            self.mma_tiler,
            self.sf_vec_size,
            self.num_ab_stage,
        )
        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma,
            self.mma_tiler,
            self.sf_vec_size,
            self.num_ab_stage,
        )
        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            self.c_dtype,
            self.c_layout,
            self.epi_tile_c,
            self.num_c_stage,
        )
        self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            self.d_dtype,
            self.d_layout,
            self.epi_tile,
            self.num_d_stage,
        )

        if self.enable_bias:
            self.bias_smem_layout_staged = cute.make_layout(
                (self.mma_tiler[1], self.num_bias_stage),
                stride=(1, self.mma_tiler[1]),
            )
        else:
            self.bias_smem_layout_staged = cute.make_layout((1, 1))

        self.overlapping_accum = self.num_acc_stage == 1 and self.mma_tiler[1] == 256

        sf_atom_mn = 32
        self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
        self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
        self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
        self.num_rht_tmem_cols = HADAMARD_SIZE if self.use_tmem_post_rht_amax else 0
        self.num_reserved_tmem_cols = self.num_sf_tmem_cols + self.num_rht_tmem_cols
        self.num_accumulator_tmem_cols = (
            self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_reserved_tmem_cols
        )

        self.epi_tile_n_required = cute.size(self.epi_tile[1]) if self.act_func == "srelu" else 2 * cute.size(self.epi_tile[1])
        self.iter_acc_early_release_in_epilogue = (self.num_reserved_tmem_cols + self.epi_tile_n_required - 1) // self.epi_tile_n_required - 1
        if self.act_func != "srelu":
            self.iter_acc_early_release_in_epilogue = self.iter_acc_early_release_in_epilogue * 2

    def get_desc_workspace_bytes(self) -> int:
        """Return descriptor workspace size in bytes."""
        if self.weight_mode == MoEWeightMode.DISCRETE:
            from ..moe_utils import DiscreteWeightTensormapConstructor

            return DiscreteWeightTensormapConstructor.get_workspace_size(self.expert_cnt)
        return 0

    def get_workspace_bytes(self) -> int:
        """Return total workspace size in bytes."""
        desc_workspace_bytes = self.get_desc_workspace_bytes()
        dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0
        return desc_workspace_bytes + dynamic_sched_bytes

    @cute.jit
    def _get_sched_counter_ptr(self, workspace_ptr):
        counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes()
        return cute.make_ptr(
            cutlass.Int32,
            counter_addr,
            AddressSpace.gmem,
            assumed_align=4,
        )

    @cute.kernel
    def helper_kernel(
        self,
        ptrs_b: cute.Pointer,
        ptrs_sfb: cute.Pointer,
        n: Int32,
        k: Int32,
        b_stride_size: cutlass.Int64,
        b_major_mode: cutlass.Constexpr,
        workspace_ptr,
        tiled_mma_arg: cute.TiledMma,
        tiled_mma_sfb_arg: cute.TiledMma,
        b_smem_layout_arg,
        sfb_smem_layout_arg,
        cluster_layout_vmnk_shape_arg: cutlass.Constexpr,
        cluster_layout_sfb_vmnk_shape_arg: cutlass.Constexpr,
    ):
        """Pre-main-kernel: build per-expert TMA descriptors (discrete mode) and/or reset sched counter."""
        expert_idx = cute.arch.block_idx()[0]

        if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE):
            b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id)
            sfb_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma_arg.thr_id)

            # Read per-expert base addresses from the pointer arrays
            b_ptr_tensor = cute.make_tensor(
                cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8),
                cute.make_layout((self.expert_cnt,)),
            )
            sfb_ptr_tensor = cute.make_tensor(
                cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8),
                cute.make_layout((self.expert_cnt,)),
            )

            c1 = cutlass.Int32(1)
            c0 = cutlass.Int64(0)
            c1_64 = 1
            if cutlass.const_expr(b_major_mode == OperandMajorMode.K):
                stride_n = b_stride_size
                stride_k = c1_64
            else:
                stride_n = c1_64
                stride_k = b_stride_size

            b_ptr_val = b_ptr_tensor[expert_idx]
            b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem)
            b_expert = cute.make_tensor(
                b_ptr,
                cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)),
            )
            tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B(
                b_tma_op_arg,
                b_expert,
                b_smem_layout_arg,
                self.mma_tiler,
                tiled_mma_arg,
                cluster_layout_vmnk_shape_arg,
            )

            workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"])
            store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx))

            sfb_ptr_val = sfb_ptr_tensor[expert_idx]
            sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem)
            sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size)
            sfb_expert = cute.make_tensor(sfb_ptr, sfb_layout)
            tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B(
                sfb_tma_op_arg,
                sfb_expert,
                sfb_smem_layout_arg,
                self.mma_tiler_sfb,
                tiled_mma_sfb_arg,
                cluster_layout_sfb_vmnk_shape_arg,
                internal_type=cutlass.Uint64,
            )
            store_tma_desc(tma_atom_sfb, workspace.get_ptr("sfb", expert_idx))

        if cutlass.const_expr(self.use_dynamic_sched):
            if expert_idx == cutlass.Int32(0):
                sched_counter = cute.make_tensor(
                    self._get_sched_counter_ptr(workspace_ptr),
                    cute.make_layout(1),
                )
                sched_counter[0] = cutlass.Int32(0)

    @cute.jit
    def __call__(
        self,
        a: cute.Tensor,
        b,  # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[]
        sfa: cute.Tensor,
        sfb,  # Dense: cute.Tensor         | Discrete: cute.Pointer to int64[]
        n: Int32,  # Ignored for dense mode
        k: Int32,  # Ignored for dense mode
        b_stride_size: cutlass.Int64,  # Ignored for dense mode
        b_major_mode: cutlass.Constexpr,  # Ignored for dense mode
        workspace_ptr,
        c: cute.Tensor,
        d: cute.Tensor,
        amax_tensor: Optional[cute.Tensor],
        post_rht_amax_tensor: Optional[cute.Tensor],
        padded_offsets: cute.Tensor,
        alpha: cute.Tensor,
        prob: cute.Tensor,
        hadamard_tensor: Optional[cute.Tensor],
        bias: Optional[cute.Tensor],
        max_active_clusters: cutlass.Constexpr,
        stream: cuda.CUstream,
        epilogue_op: cutlass.Constexpr = lambda x: x,
        linear_offset: cutlass.Float32 = 0.0,
        situ_beta1: cutlass.Float32 = 4.0,
        situ_beta2: cutlass.Float32 = 25.0,
    ):
        """Execute the MoE GEMM + GLU + Hadamard kernel.

        Dense mode: ``b`` and ``sfb`` are 3-D cute.Tensor (N, K, L).
        Discrete mode: ``b`` and ``sfb`` are cute.Pointer to device int64[]
        arrays of per-expert base addresses.

        ``situ_beta1`` and ``situ_beta2`` configure SiTU-GLU:

            out = prob * beta1 * tanh(gate / beta1) * sigmoid(gate)
                  * beta2 * tanh(up / beta2)

        They are ignored unless ``act_func == "situglu"``.
        """
        self.a_dtype: Type[cutlass.Numeric] = a.element_type
        self.b_dtype: Type[cutlass.Numeric] = a.element_type
        self.c_dtype: Type[cutlass.Numeric] = c.element_type
        self.d_dtype: Type[cutlass.Numeric] = d.element_type
        self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type
        self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16
        self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode()
        self.c_layout = utils.LayoutEnum.from_tensor(c)
        self.d_layout = utils.LayoutEnum.from_tensor(d)

        if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE):
            self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode()
        else:
            self.b_major_mode = b_major_mode

        if cutlass.const_expr(self.a_dtype != self.b_dtype):
            raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")

        self._setup_attributes()

        # ---- B / SFB setup (mode-dependent) ----
        b_from_call_arg = b
        sfb_from_call_arg = sfb
        if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE):
            sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, self.sf_vec_size)
            sfb = cute.make_tensor(sfb.iterator, sfb_layout)
        else:
            c1 = cutlass.Int32(1)
            c0 = cutlass.Int64(0)
            c1_64 = 1
            if cutlass.const_expr(b_major_mode == OperandMajorMode.K):
                b_template_stride = (b_stride_size, c1_64, c0)
            else:
                b_template_stride = (c1_64, b_stride_size, c0)
            b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride)
            b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16)
            b = cute.make_tensor(b_ptr_typed, b_template_layout)

            sfb_ptr_typed = cute.make_ptr(self.sf_dtype, sfb.toint(), AddressSpace.gmem, assumed_align=16)
            sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size)
            sfb = cute.make_tensor(sfb_ptr_typed, sfb_layout)

        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size)
        sfa = cute.make_tensor(sfa.iterator, sfa_layout)

        self.generate_amax = amax_tensor is not None
        self.generate_post_rht_amax = post_rht_amax_tensor is not None

        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.b_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            self.cta_group,
            self.mma_inst_shape_mn,
        )
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.b_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            cute.nvgpu.tcgen05.CtaGroup.ONE,
            self.mma_inst_shape_mn_sfb,
        )
        atom_thr_size = cute.size(tiled_mma.thr_id.shape)

        # TMA load A
        a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
            a_op,
            a,
            a_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )

        # TMA load B
        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
            b_op,
            b,
            b_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )

        # TMA load SFA
        sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
        sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
        tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
            sfa_op,
            sfa,
            sfa_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=cutlass.Int16,
        )

        # TMA load SFB
        sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)
        sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
        tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
            sfb_op,
            sfb,
            sfb_smem_layout,
            self.mma_tiler_sfb,
            tiled_mma_sfb,
            self.cluster_layout_sfb_vmnk.shape,
            internal_type=cutlass.Uint64,
        )

        if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
            x = tma_tensor_sfb.stride[0][1]
            y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4)
            new_shape = (
                (tma_tensor_sfb.shape[0][0], ((2, 2), y)),
                tma_tensor_sfb.shape[1],
                tma_tensor_sfb.shape[2],
            )
            x_times_3 = 3 * x
            new_stride = (
                (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)),
                tma_tensor_sfb.stride[1],
                tma_tensor_sfb.stride[2],
            )
            tma_tensor_sfb = cute.make_tensor(
                tma_tensor_sfb.iterator,
                cute.make_layout(new_shape, stride=new_stride),
            )

        a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
        b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
        sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
        sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
        self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size

        # TMA store C
        c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            c,
            c_smem_layout,
            self.epi_tile_c,
        )

        # TMA store D
        d_smem_layout = cute.slice_(self.d_smem_layout_staged, (None, None, 0))
        tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            d,
            d_smem_layout,
            self.epi_tile,
        )

        # ---- Helper kernel (discrete TMA desc init + dynamic sched counter reset) ----
        _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched)
        if cutlass.const_expr(_need_helper):
            _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1
            _helper_args = (
                b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem),
                sfb_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem),
                n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0),
                k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0),
                b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0),
                b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode,
                workspace_ptr,
                tiled_mma,
                tiled_mma_sfb,
                b_smem_layout,
                sfb_smem_layout,
                self.cluster_layout_vmnk.shape,
                self.cluster_layout_sfb_vmnk.shape,
            )
            self.helper_kernel(*_helper_args).launch(
                grid=(_helper_grid_x, 1, 1),
                block=(1, 1, 1),
                stream=stream,
                min_blocks_per_mp=1,
            )

        # ---- Grid computation via MoE scheduler ----
        if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE):
            b_n, b_k, b_l = cute.shape(b)
            sched_expert_shape = (self.expert_cnt, b_n, b_k)
        else:
            sched_expert_shape = (self.expert_cnt, n, k)

        sched_params = MoESchedulerParams(
            scenario="2Dx3D",
            expert_shape=sched_expert_shape,
            cta_tile_shape_mnk=self.cta_tile_shape_mnk,
            cluster_shape_mn=self.cluster_shape_mn,
            use_dynamic_sched=self.use_dynamic_sched,
        )
        self.sched_params, grid = compute_grid(
            sched_params,
            max_active_clusters,
            self.use_2cta_instrs,
        )

        self.buffer_align_bytes = 1024

        # ---- Shared storage ----
        SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched)

        @cute.struct
        class SharedStorage:
            ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
            acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
            pingpong_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_pingpong_stage * 2]
            if cutlass.const_expr(self.enable_bias):
                bias_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_bias_stage * 2]
            if cutlass.const_expr(self.use_tmem_post_rht_amax):
                hadamard_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
                hadamard_prerequisite_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
            scheduler: SchedulerStorage
            tmem_dealloc_mbar_ptr: cutlass.Int64
            tmem_holding_buf: cutlass.Int32
            sC: cute.struct.Align[
                cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)],
                self.buffer_align_bytes,
            ]
            sD: cute.struct.Align[
                cute.struct.MemRange[self.d_dtype, cute.cosize(self.d_smem_layout_staged.outer)],
                self.buffer_align_bytes,
            ]
            sA: cute.struct.Align[
                cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)],
                self.buffer_align_bytes,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)],
                self.buffer_align_bytes,
            ]
            sSFA: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)],
                self.buffer_align_bytes,
            ]
            sSFB: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)],
                self.buffer_align_bytes,
            ]
            if cutlass.const_expr(self.enable_bias):
                sBias: cute.struct.Align[
                    cute.struct.MemRange[self.bias_dtype, cute.cosize(self.bias_smem_layout_staged)],
                    16,
                ]
            if cutlass.const_expr(self.use_tmem_post_rht_amax):
                sHadamard: cute.struct.Align[cute.struct.MemRange[cutlass.BFloat16, HADAMARD_SIZE * HADAMARD_SIZE], 16]
            sAmax: cute.struct.Align[
                cute.struct.MemRange[cutlass.Float32, self.num_epilog_warps],
                1,
            ]
            sPostRhtAmax: cute.struct.Align[
                cute.struct.MemRange[cutlass.Float32, self.num_epilog_warps],
                1,
            ]

        self.shared_storage = SharedStorage

        # Launch main kernel
        self.kernel(
            tiled_mma,
            tiled_mma_sfb,
            tma_atom_a,
            tma_tensor_a,
            tma_atom_b,
            tma_tensor_b,
            tma_atom_sfa,
            tma_tensor_sfa,
            tma_atom_sfb,
            tma_tensor_sfb,
            tma_atom_c,
            tma_tensor_c,
            tma_atom_d,
            tma_tensor_d,
            amax_tensor,
            post_rht_amax_tensor,
            padded_offsets,
            alpha,
            bias,
            prob,
            hadamard_tensor,
            workspace_ptr,
            self.cluster_layout_vmnk,
            self.cluster_layout_sfb_vmnk,
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.sfa_smem_layout_staged,
            self.sfb_smem_layout_staged,
            self.c_smem_layout_staged,
            self.d_smem_layout_staged,
            self.bias_smem_layout_staged,
            self.epi_tile,
            self.sched_params,
            epilogue_op,
            linear_offset,
            situ_beta1,
            situ_beta2,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=(*self.cluster_shape_mn, 1),
            max_number_threads=[self.threads_per_cta, 1, 1],
            smem=self.shared_storage.size_in_bytes(),
            stream=stream,
            min_blocks_per_mp=1,
        )
        return

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @cute.jit
    def _make_extension(self, workspace_ptr):
        if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE):
            desc_workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"])
            return DiscreteWeightScaledGemmSchedExtension(
                tensormap_ctor=desc_workspace,
                sf_vec_size=self.sf_vec_size,
            )
        else:
            return ContiguousAndConsistentGroupedGemmSchedExtension(
                sf_vec_size=self.sf_vec_size,
            )

    def mainloop_s2t_copy_and_partition(self, sSF, tSF):
        tCsSF_compact = cute.filter_zeros(sSF)
        tCtSF_compact = cute.filter_zeros(tSF)
        copy_atom_s2t = cute.make_copy_atom(
            tcgen05.Cp4x32x128bOp(self.cta_group),
            self.sf_dtype,
        )
        tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
        thr_copy_s2t = tiled_copy_s2t.get_slice(0)
        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
        tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
        return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t

    @cute.jit
    def amax_reduction_per_thread(self, vec_fp32, amax_fp32):
        vec_fp32_ssa = vec_fp32.load()
        import cutlass._mlir.dialects.math as _math

        abs_acc_values_ir = _math.absf(vec_fp32_ssa.ir_value())
        abs_acc_values = type(vec_fp32_ssa)(abs_acc_values_ir, vec_fp32_ssa.shape, vec_fp32_ssa.dtype)
        subtile_amax = abs_acc_values.reduce(cute.ReductionOp.MAX, cutlass.Float32(0.0), 0)
        return cute.arch.fmax(amax_fp32, subtile_amax)

    @cute.jit
    def amax_reduction_per_warp_and_cta(self, amax_fp32, warp_idx, amax_smem, amax_gmem):
        warp_amax = cute.arch.warp_redux_sync(value=amax_fp32, kind="fmax", mask_and_clamp=0xFFFFFFFF, nan=True)
        if cute.arch.lane_idx() == 0:
            amax_smem[warp_idx & 0x3] = cutlass.Float32(warp_amax)
        self.epilog_sync_barrier_group0.arrive_and_wait()
        if warp_idx == self.epilog_act_warp_id[0] and cute.arch.lane_idx() == 0:
            block_amax = cutlass.Float32(0.0)
            for i in cutlass.range(self.num_epilog_warps):
                block_amax = cute.arch.fmax(block_amax, amax_smem[i])
            _ = atomic_max_float32(ptr=amax_gmem, value=block_amax)

    @cute.jit
    def post_rht_amax_reduction_per_warp_and_cta(self, amax_fp32, warp_idx, amax_smem, amax_gmem):
        warp_amax = cute.arch.warp_redux_sync(value=amax_fp32, kind="fmax", mask_and_clamp=0xFFFFFFFF, nan=True)
        if cute.arch.lane_idx() == 0:
            amax_smem[warp_idx & 0x3] = cutlass.Float32(warp_amax)
        self.epilog_sync_barrier_group1.arrive_and_wait()
        if warp_idx == self.epilog_rht_store_warp_id[0] and cute.arch.lane_idx() == 0:
            block_amax = cutlass.Float32(0.0)
            for i in cutlass.range(self.num_epilog_warps):
                block_amax = cute.arch.fmax(block_amax, amax_smem[i])
            block_amax = block_amax * cutlass.Float32(0.25)
            _ = atomic_max_float32(ptr=amax_gmem, value=block_amax)

    @cute.jit
    def store_c(
        self,
        tiled_copy_r2s,
        tma_atom_c,
        warp_idx,
        tTR_rAcc,
        tTR_rAcc_up,
        tRS_rC,
        tRS_sC,
        bSG_gC,
        bSG_sC,
        c_pipeline,
        prev_subtile_idx,
        real_subtile_idx,
    ):
        c_buffer = prev_subtile_idx % self.num_c_stage
        tRS_rC.store(tTR_rAcc.load().to(self.c_dtype))
        cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)])
        tRS_rC.store(tTR_rAcc_up.load().to(self.c_dtype))
        cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 1, c_buffer)])
        cute.arch.fence_proxy("async.shared", space="cta")
        self.epilog_sync_barrier_group0.arrive_and_wait()
        if warp_idx == self.epilog_act_warp_id[0]:
            cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)])
            c_pipeline.producer_commit()
            if not cutlass.const_expr(self.delay_tma_store_acquire_sync):
                c_pipeline.producer_acquire()
        if not cutlass.const_expr(self.delay_tma_store_acquire_sync):
            self.epilog_sync_barrier_group0.arrive_and_wait()

    @cute.jit
    def store_c_unary(
        self,
        tiled_copy_r2s,
        tma_atom_c,
        warp_idx,
        tTR_rAcc,
        tRS_rC,
        tRS_sC,
        bSG_gC,
        bSG_sC,
        c_pipeline,
        prev_subtile_idx,
        real_subtile_idx,
    ):
        c_buffer = prev_subtile_idx % self.num_c_stage
        tRS_rC.store(tTR_rAcc.load().to(self.c_dtype))
        cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)])
        cute.arch.fence_proxy("async.shared", space="cta")
        self.epilog_sync_barrier_group0.arrive_and_wait()
        if warp_idx == self.epilog_act_warp_id[0]:
            cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)])
            c_pipeline.producer_commit()
            if not cutlass.const_expr(self.delay_tma_store_acquire_sync):
                c_pipeline.producer_acquire()
        if not cutlass.const_expr(self.delay_tma_store_acquire_sync):
            self.epilog_sync_barrier_group0.arrive_and_wait()

    @cute.jit
    def geglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset=1.0):
        if cutlass.const_expr(self.vectorized_f32):
            LOG2_E = cutlass.Float32(1.4426950408889634)
            for i in cutlass.range_constexpr(0, cute.size(tCompute), 2):
                scaled_gate_0, scaled_gate_1 = cute.arch.mul_packed_f32x2(
                    (acc_vec_gate[i], acc_vec_gate[i + 1]),
                    (1.702, 1.702),
                    rnd="rn",
                    ftz=False,
                )
                tCompute_log2e = cute.arch.mul_packed_f32x2(
                    (scaled_gate_0, scaled_gate_1),
                    (-LOG2_E, -LOG2_E),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2(
                    (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)),
                    (1.0, 1.0),
                )
                tCompute[i] = cute.arch.rcp_approx(tCompute[i])
                tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1])
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (acc_vec_gate[i], acc_vec_gate[i + 1]),
                    rnd="rn",
                    ftz=False,
                )
                up0, up1 = cute.arch.add_packed_f32x2(
                    (linear_offset, linear_offset),
                    (acc_vec_up[i], acc_vec_up[i + 1]),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (up0, up1),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (mProb, mProb),
                    rnd="rn",
                    ftz=False,
                )
        else:
            for i in cutlass.range_constexpr(cute.size(tCompute)):
                tCompute[i] = (acc_vec_up[i] + linear_offset) * silu_f32_geglu_scaled(acc_vec_gate[i], fastmath=True)
                tCompute[i] = tCompute[i] * mProb

    @cute.jit
    def swiglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb):
        if cutlass.const_expr(self.vectorized_f32):
            LOG2_E = cutlass.Float32(1.4426950408889634)
            for i in cutlass.range_constexpr(0, cute.size(tCompute), 2):
                tCompute_log2e = cute.arch.mul_packed_f32x2(
                    (acc_vec_gate[i], acc_vec_gate[i + 1]),
                    (-LOG2_E, -LOG2_E),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.add_packed_f32x2(
                    (cute.math.exp2(tCompute_log2e[0], fastmath=True), cute.math.exp2(tCompute_log2e[1], fastmath=True)),
                    (1.0, 1.0),
                )
                tCompute[i] = cute.arch.rcp_approx(tCompute[i])
                tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1])
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (acc_vec_gate[i], acc_vec_gate[i + 1]),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (acc_vec_up[i], acc_vec_up[i + 1]),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (mProb, mProb),
                    rnd="rn",
                    ftz=False,
                )
        else:
            for i in cutlass.range_constexpr(cute.size(tCompute)):
                tCompute[i] = acc_vec_up[i] * silu_f32(acc_vec_gate[i], fastmath=True)
                tCompute[i] = tCompute[i] * mProb

    @cute.jit
    def situglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb, beta1, beta2):
        beta1_rcp = cutlass.Float32(1.0 / self.situ_beta1)
        beta2_rcp = cute.arch.rcp_approx(beta2)
        beta_product = beta1 * beta2
        for i in cutlass.range_constexpr(cute.size(tCompute)):
            gate = acc_vec_gate[i]
            up = acc_vec_up[i]
            gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True)
            up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True)
            if cutlass.const_expr(self.situ_beta1 == 4.0):
                # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2).
                sigmoid = cutlass.Float32(0.5) + gate_tanh * cute.arch.rcp_approx(cutlass.Float32(1.0) + gate_tanh * gate_tanh)
            else:
                sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True))
            tCompute[i] = beta_product * gate_tanh * sigmoid * up_tanh * mProb

    @cute.jit
    def srelu_act(self, tCompute, acc_vec, mProb):
        acc_relu = cute.where(acc_vec > 0, acc_vec, cute.full_like(acc_vec, 0))
        if cutlass.const_expr(self.vectorized_f32):
            for i in cutlass.range_constexpr(0, cute.size(tCompute), 2):
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (acc_relu[i], acc_relu[i + 1]),
                    (acc_relu[i], acc_relu[i + 1]),
                    rnd="rn",
                    ftz=False,
                )
                tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(
                    (tCompute[i], tCompute[i + 1]),
                    (mProb, mProb),
                    rnd="rn",
                    ftz=False,
                )
        else:
            for i in cutlass.range_constexpr(cute.size(tCompute)):
                tCompute[i] = acc_relu[i] * acc_relu[i] * mProb

    @cute.jit
    def query_hadamard_tmem_a_ptr(self, tile_idx, reverse_subtile, tmem_ptr):
        hadamard_tmem_offset = self.num_accumulator_tmem_cols + self.num_sf_tmem_cols
        return cute.recast_ptr(tmem_ptr + hadamard_tmem_offset, dtype=self.acc_dtype)

    @cute.jit
    def query_hadamard_tmem_acc_ptr(self, tile_idx, reverse_subtile, tmem_ptr):
        hadamard_tmem_offset = self.num_accumulator_tmem_cols + self.num_sf_tmem_cols
        return cute.recast_ptr(tmem_ptr + hadamard_tmem_offset, dtype=self.acc_dtype)

    def epilog_tmem_copy_and_partition(self, tidx, tAcc, gD_mnl, epi_tile, use_2cta_instrs):
        copy_atom_t2r = sm100_utils.get_tmem_load_op(
            self.cta_tile_shape_mnk,
            self.d_layout,
            self.d_dtype,
            self.acc_dtype,
            epi_tile,
            use_2cta_instrs,
        )
        tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])
        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
        tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
        gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile)
        tTR_gC = thr_copy_t2r.partition_D(gD_mnl_epi)
        tTR_rAcc_gate = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype)
        tTR_rAcc_up = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype)
        return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_gate, tTR_rAcc_up

    def epilog_smem_copy_and_partition(self, tiled_copy_t2r, tTR_rC, tidx, sD):
        copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, self.d_dtype, self.acc_dtype, tiled_copy_t2r)
        tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
        tRS_sD = thr_copy_r2s.partition_D(sD)
        tRS_rD = tiled_copy_r2s.retile(tTR_rC)
        return tiled_copy_r2s, tRS_rD, tRS_sD

    def epilog_gmem_copy_and_partition(self, tidx, atom, gD_mnl, epi_tile, sD):
        gD_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile)
        tma_atom_d = atom
        sD_for_tma_partition = cute.group_modes(sD, 0, 2)
        gD_for_tma_partition = cute.group_modes(gD_epi, 0, 2)
        bSG_sD, bSG_gD = cpasync.tma_partition(
            tma_atom_d,
            0,
            cute.make_layout(1),
            sD_for_tma_partition,
            gD_for_tma_partition,
        )
        return tma_atom_d, bSG_sD, bSG_gD

    @staticmethod
    def _compute_stages(
        tiled_mma,
        mma_tiler_mnk,
        a_dtype,
        b_dtype,
        epi_tile,
        epi_tile_c,
        c_dtype,
        c_layout,
        d_dtype,
        d_layout,
        sf_dtype,
        sf_vec_size,
        num_smem_capacity,
        occupancy,
        bias_dtype,
        use_tmem_post_rht_amax,
    ):
        num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
        num_c_stage = 1
        num_d_stage = 1
        num_tile_stage = 2
        num_pingpong_stage = mma_tiler_mnk[1] // epi_tile_c[1]

        a_smem_layout_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1)
        b_smem_layout_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1)
        sfa_smem_layout_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
        sfb_smem_layout_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
        c_smem_layout_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile_c, 1)
        d_smem_layout_one = sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1)

        ab_bytes_per_stage = (
            cute.size_in_bytes(a_dtype, a_smem_layout_one)
            + cute.size_in_bytes(b_dtype, b_smem_layout_one)
            + cute.size_in_bytes(sf_dtype, sfa_smem_layout_one)
            + cute.size_in_bytes(sf_dtype, sfb_smem_layout_one)
        )
        mbar_helpers_bytes = 1024
        # sInfo is in SchedulerStorage, not here, so use 4-int sInfo
        sinfo_bytes = 4 * 4 * num_tile_stage
        c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_one) * num_c_stage
        d_bytes = cute.size_in_bytes(d_dtype, d_smem_layout_one) * num_d_stage
        hadamard_bytes = (cutlass.BFloat16.width // 8) * HADAMARD_SIZE * HADAMARD_SIZE if use_tmem_post_rht_amax else 0
        # sAmax and sPostRhtAmax for the 4 ACT/RHT epilogue warps.
        amax_bytes = 2 * 4 * 4  # two sets of 4 Float32 values

        if bias_dtype is not None:
            num_bias_stage = 2
            bias_bytes = mma_tiler_mnk[1] * num_bias_stage * (bias_dtype.width // 8)
        else:
            num_bias_stage = 0
            bias_bytes = 0

        epi_bytes = c_bytes + d_bytes + amax_bytes + bias_bytes
        epi_bytes += hadamard_bytes

        num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage

        return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage, num_pingpong_stage

    # GPU device kernel
    @cute.kernel
    def kernel(
        self,
        tiled_mma: cute.TiledMma,
        tiled_mma_sfb: cute.TiledMma,
        tma_atom_a: cute.CopyAtom,
        mA_mkl: cute.Tensor,
        tma_atom_b: cute.CopyAtom,
        mB_nkl: cute.Tensor,
        tma_atom_sfa: cute.CopyAtom,
        mSFA_mkl: cute.Tensor,
        tma_atom_sfb: cute.CopyAtom,
        mSFB_nkl: cute.Tensor,
        tma_atom_c: cute.CopyAtom,
        mC_mnl: cute.Tensor,
        tma_atom_d: cute.CopyAtom,
        mD_mnl: cute.Tensor,
        mAmax_tensor: Optional[cute.Tensor],
        mPostRhtAmax_tensor: Optional[cute.Tensor],
        padded_offsets: cute.Tensor,
        alpha: cute.Tensor,
        mBias_nl: Optional[cute.Tensor],
        prob: cute.Tensor,
        hadamard_tensor: Optional[cute.Tensor],
        workspace_ptr,
        cluster_layout_vmnk: cute.Layout,
        cluster_layout_sfb_vmnk: cute.Layout,
        a_smem_layout_staged: cute.ComposedLayout,
        b_smem_layout_staged: cute.ComposedLayout,
        sfa_smem_layout_staged: cute.Layout,
        sfb_smem_layout_staged: cute.Layout,
        c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
        d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
        bias_smem_layout_staged: cute.Layout,
        epi_tile: cute.Tile,
        sched_params: MoESchedulerParams,
        epilogue_op: cutlass.Constexpr,
        linear_offset: cutlass.Float32 = 0.0,
        situ_beta1: cutlass.Float32 = 4.0,
        situ_beta2: cutlass.Float32 = 25.0,
    ):
        """GPU device kernel: MoE persistent GEMM + GLU + Hadamard (pingpong epilogue)."""
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        lane_idx = cute.arch.lane_idx()
        cta_rank = cute.arch.block_in_cluster_idx()[0]

        total_token = padded_offsets[self.expert_cnt - 1]

        # Prefetch TMA descriptors
        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_sfa)
            if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE):
                cpasync.prefetch_descriptor(tma_atom_b)
                cpasync.prefetch_descriptor(tma_atom_sfb)
            cpasync.prefetch_descriptor(tma_atom_c)
            cpasync.prefetch_descriptor(tma_atom_d)

        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2

        # CTA coordinates
        bidx, bidy, bidz = cute.arch.block_idx()
        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
        is_leader_cta = mma_tile_coord_v == 0
        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
        block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster)
        tidx, _, _ = cute.arch.thread_idx()

        # Shared memory allocation
        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)
        sched_storage = storage.scheduler

        # AB pipeline
        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
        ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer)
        ab_pipeline = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_mbar_ptr.data_ptr(),
            num_stages=self.num_ab_stage,
            producer_group=ab_pipeline_producer_group,
            consumer_group=ab_pipeline_consumer_group,
            tx_count=self.num_tma_load_bytes,
            cta_layout_vmnk=cluster_layout_vmnk,
        )

        # ACC pipeline
        acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_acc_consumer_threads = len(self.epilog_act_warp_id) * (2 if use_2cta_instrs else 1)
        acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads)
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=acc_pipeline_producer_group,
            consumer_group=acc_pipeline_consumer_group,
            cta_layout_vmnk=cluster_layout_vmnk,
        )

        if cutlass.const_expr(self.use_tmem_post_rht_amax):
            # Hadamard SMEM tensor
            hadamard_b_layout_staged = cute.make_layout((16, 8), stride=(1, 16))
            sHadamard = storage.sHadamard.get_tensor(hadamard_b_layout_staged)
            hadamard_h_iter = hadamard_tensor.iterator + (cutlass.Int32(128) if cta_rank == 1 else cutlass.Int32(0))
            hadamard_tensor_local = cute.make_tensor(
                cute.make_ptr(
                    cutlass.BFloat16,
                    hadamard_h_iter.toint(),
                    hadamard_h_iter.memspace,
                    assumed_align=16,
                ),
                hadamard_b_layout_staged,
            )

            # Hadamard UMMA pipeline (1 stage)
            hadamard_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
            num_hadamard_consumer_threads = 128 * (2 if use_2cta_instrs else 1)
            hadamard_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_hadamard_consumer_threads)
            hadamard_producer, hadamard_consumer = pipeline.PipelineUmmaAsync.create(
                barrier_storage=storage.hadamard_mbar_ptr.data_ptr(),
                num_stages=1,
                producer_group=hadamard_pipeline_producer_group,
                consumer_group=hadamard_pipeline_consumer_group,
                cta_layout_vmnk=cluster_layout_vmnk,
            ).make_participants()

            # Hadamard prerequisite pipeline (cross-CTA sync, 1 stage)
            num_hadamard_prerequisite_threads = 128
            hadamard_prerequisite_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_hadamard_prerequisite_threads)
            hadamard_prerequisite_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_hadamard_prerequisite_threads)
            peer_rank = cutlass.Int32(1) - cta_rank
            hadamard_prerequisite_producer, hadamard_prerequisite_consumer = pipeline.PipelineAsync.create(
                barrier_storage=storage.hadamard_prerequisite_mbar_ptr.data_ptr(),
                num_stages=1,
                producer_group=hadamard_prerequisite_pipeline_producer_group,
                consumer_group=hadamard_prerequisite_pipeline_consumer_group,
                producer_mask=peer_rank,
                defer_sync=True,
            ).make_participants()

        # Pingpong pipeline
        pingpong_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_act_warp_id) * self.threads_per_warp)
        pingpong_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_rht_store_warp_id) * self.threads_per_warp)
        pingpong_pipeline = pipeline.PipelineAsync.create(
            barrier_storage=storage.pingpong_mbar_ptr.data_ptr(),
            num_stages=self.num_pingpong_stage,
            producer_group=pingpong_producer_group,
            consumer_group=pingpong_consumer_group,
        )

        # Tile info pipeline (uses SchedulerStorage's barrier)
        tile_info_pipeline_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_per_warp * 1,
        )
        tile_info_pipeline_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_wo_sched,
        )
        tile_info_pipeline = pipeline.PipelineAsync.create(
            barrier_storage=sched_storage.tile_info_mbar.data_ptr(),
            num_stages=self.num_tile_stage,
            producer_group=tile_info_pipeline_producer_group,
            consumer_group=tile_info_pipeline_consumer_group,
        )

        # MoE persistent tile scheduler
        scheduler = MoEPersistentTileScheduler.create(
            sched_params,
            padded_offsets,
            cute.arch.block_idx(),
            cute.arch.grid_dim(),
            counter_ptr=self._get_sched_counter_ptr(workspace_ptr),
            sched_storage=sched_storage,
        )
        scheduler.internal_init()

        # Bias pipeline
        if cutlass.const_expr(self.enable_bias):
            bias_pipeline_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.threads_per_warp,
            )
            bias_pipeline_consumer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.threads_per_warp * len(self.epilog_act_warp_id),
            )
            bias_pipeline = pipeline.PipelineCpAsync.create(
                barrier_storage=storage.bias_mbar_ptr.data_ptr(),
                num_stages=self.num_bias_stage,
                producer_group=bias_pipeline_producer_group,
                consumer_group=bias_pipeline_consumer_group,
            )
            sBias = storage.sBias.get_tensor(bias_smem_layout_staged)
            gBias_nl = cute.local_tile(mBias_nl, cute.slice_(self.mma_tiler[:2], (0, None)), (None, None))

        # TMEM allocator
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=self.tmem_alloc_barrier,
            allocator_warp_id=self.epilog_act_warp_id[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr,
        )

        # Cluster arrive after barrier init
        if cute.size(self.cluster_shape_mn) > 1:
            cute.arch.cluster_arrive_relaxed()

        # SMEM tensors
        sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner)
        sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner)
        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
        sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
        sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
        amax_layout = cute.make_layout((self.num_epilog_warps,))
        sAmax = storage.sAmax.get_tensor(amax_layout)
        sPostRhtAmax = storage.sPostRhtAmax.get_tensor(amax_layout)

        # sInfo from SchedulerStorage
        info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4))
        sInfo = sched_storage.sInfo.get_tensor(info_layout)

        # Multicast masks
        a_full_mcast_mask = None
        b_full_mcast_mask = None
        sfa_full_mcast_mask = None
        sfb_full_mcast_mask = None
        if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
            a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2)
            b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1)
            sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2)
            sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1)

        # MMA fragments
        tCrA = tiled_mma.make_fragment_A(sA)
        tCrB = tiled_mma.make_fragment_B(sB)
        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
        if cutlass.const_expr(self.overlapping_accum):
            num_acc_stage_overlapped = 2
            tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage_overlapped))
            tCtAcc_fake = cute.make_tensor(
                tCtAcc_fake.iterator,
                cute.make_layout(
                    tCtAcc_fake.shape,
                    stride=(
                        tCtAcc_fake.stride[0],
                        tCtAcc_fake.stride[1],
                        tCtAcc_fake.stride[2],
                        (256 - self.num_reserved_tmem_cols) * tCtAcc_fake.stride[0][1],
                    ),
                ),
            )
        else:
            tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))

        # Cluster wait / CTA sync
        if cute.size(self.cluster_shape_mn) > 1:
            cute.arch.cluster_wait()
        else:
            self.cta_sync_barrier.arrive_and_wait()

        if total_token <= 0:
            cute.arch.nvvm.exit()

        # ---------------------------------------------------------------
        # Specialized Scheduler warp (MoEPersistentTileScheduler)
        # ---------------------------------------------------------------
        if warp_idx == self.sched_warp_id:
            work_tile_info = scheduler.initial_work_tile_info()
            tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage)

            while work_tile_info.is_valid_tile:
                tile_info_pipeline.producer_acquire(tile_info_producer_state)
                with cute.arch.elect_one():
                    sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx
                    sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx
                    sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx
                    sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt
                cute.arch.fence_proxy("async.shared", space="cta")
                self.sched_sync_barrier.arrive_and_wait()
                tile_info_pipeline.producer_commit(tile_info_producer_state)
                tile_info_producer_state.advance()
                work_tile_info = scheduler.advance_to_next_work()

            # Send invalid signal: expert_idx = -1
            tile_info_pipeline.producer_acquire(tile_info_producer_state)
            with cute.arch.elect_one():
                sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1)
                sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0)
                sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0)
                sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0)
            cute.arch.fence_proxy("async.shared", space="cta")
            self.sched_sync_barrier.arrive_and_wait()
            tile_info_pipeline.producer_commit(tile_info_producer_state)
            tile_info_producer_state.advance()
            tile_info_pipeline.producer_tail(tile_info_producer_state)

        # ---------------------------------------------------------------
        # Specialized TMA load warp
        # ---------------------------------------------------------------
        if warp_idx == self.tma_warp_id:
            ext = self._make_extension(workspace_ptr)

            ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage)
            tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage)

            tile_info = cute.make_rmem_tensor((4,), cutlass.Int32)
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            for idx in cutlass.range(4, unroll_full=True):
                tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
            is_valid_tile = tile_info[0] >= cutlass.Int32(0)
            cute.arch.fence_proxy("async.shared", space="cta")
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()

            while is_valid_tile:
                work_tile_info = MoEWorkTileInfo(
                    expert_idx=tile_info[0],
                    tile_m_idx=tile_info[1],
                    tile_n_idx=tile_info[2],
                    k_tile_cnt=tile_info[3],
                )
                k_tile_cnt = work_tile_info.k_tile_cnt
                ext.update_expert_info(padded_offsets, work_tile_info.expert_idx)

                real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info)
                real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info)
                real_sfa, _ = ext.get_gmem_tensor("sfa", mSFA_mkl, padded_offsets, work_tile_info)
                real_sfb, desc_ptr_sfb = ext.get_gmem_tensor("sfb", mSFB_nkl, padded_offsets, work_tile_info)

                gA_mkl = cute.local_tile(real_a, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
                gB_nkl = cute.local_tile(real_b, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
                gSFA_mkl = cute.local_tile(real_sfa, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
                gSFB_nkl = cute.local_tile(real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None))

                thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
                thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
                tCgA = thr_mma.partition_A(gA_mkl)
                tCgB = thr_mma.partition_B(gB_nkl)
                tCgSFA = thr_mma.partition_A(gSFA_mkl)
                tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)

                a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
                tAsA, tAgA = cpasync.tma_partition(
                    tma_atom_a,
                    block_in_cluster_coord_vmnk[2],
                    a_cta_layout,
                    cute.group_modes(sA, 0, 3),
                    cute.group_modes(tCgA, 0, 3),
                )
                b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
                tBsB, tBgB = cpasync.tma_partition(
                    tma_atom_b,
                    block_in_cluster_coord_vmnk[1],
                    b_cta_layout,
                    cute.group_modes(sB, 0, 3),
                    cute.group_modes(tCgB, 0, 3),
                )
                sfa_cta_layout = a_cta_layout
                tAsSFA, tAgSFA = cpasync.tma_partition(
                    tma_atom_sfa,
                    block_in_cluster_coord_vmnk[2],
                    sfa_cta_layout,
                    cute.group_modes(sSFA, 0, 3),
                    cute.group_modes(tCgSFA, 0, 3),
                )
                tAsSFA = cute.filter_zeros(tAsSFA)
                tAgSFA = cute.filter_zeros(tAgSFA)
                sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape)
                tBsSFB, tBgSFB = cpasync.tma_partition(
                    tma_atom_sfb,
                    block_in_cluster_coord_sfb_vmnk[1],
                    sfb_cta_layout,
                    cute.group_modes(sSFB, 0, 3),
                    cute.group_modes(tCgSFB, 0, 3),
                )
                tBsSFB = cute.filter_zeros(tBsSFB)
                tBgSFB = cute.filter_zeros(tBgSFB)

                mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)
                mma_tile_coord_n = work_tile_info.tile_n_idx
                tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)]
                tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)]
                tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)]
                slice_n = mma_tile_coord_n
                if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
                    slice_n = mma_tile_coord_n // 2
                tBgSFB_slice = tBgSFB[(None, slice_n, None, 0)]

                ab_producer_state.reset_count()
                peek_ab_empty_status = cutlass.Boolean(1)
                if ab_producer_state.count < k_tile_cnt:
                    peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)

                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    tAgA_k = tAgA_slice[(None, ab_producer_state.count)]
                    tBgB_k = tBgB_slice[(None, ab_producer_state.count)]
                    tAgSFA_k = tAgSFA_slice[(None, ab_producer_state.count)]
                    tBgSFB_k = tBgSFB_slice[(None, ab_producer_state.count)]
                    tAsA_pipe = tAsA[(None, ab_producer_state.index)]
                    tBsB_pipe = tBsB[(None, ab_producer_state.index)]
                    tAsSFA_pipe = tAsSFA[(None, ab_producer_state.index)]
                    tBsSFB_pipe = tBsSFB[(None, ab_producer_state.index)]
                    tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state)

                    ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
                    ab_producer_state_next = ab_producer_state.clone()
                    ab_producer_state_next.advance()
                    if ab_producer_state_next.count < k_tile_cnt:
                        peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state_next)
                    else:
                        peek_ab_empty_status = cutlass.Boolean(1)

                    cute.copy(tma_atom_a, tAgA_k, tAsA_pipe, tma_bar_ptr=tma_bar, mcast_mask=a_full_mcast_mask)
                    cute.copy(tma_atom_b, tBgB_k, tBsB_pipe, tma_bar_ptr=tma_bar, mcast_mask=b_full_mcast_mask, tma_desc_ptr=desc_ptr_b)
                    cute.copy(tma_atom_sfa, tAgSFA_k, tAsSFA_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfa_full_mcast_mask)
                    cute.copy(tma_atom_sfb, tBgSFB_k, tBsSFB_pipe, tma_bar_ptr=tma_bar, mcast_mask=sfb_full_mcast_mask, tma_desc_ptr=desc_ptr_sfb)

                    ab_producer_state.advance()

                # Advance to next tile
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                for idx in cutlass.range(4, unroll_full=True):
                    tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                cute.arch.fence_proxy("async.shared", space="cta")
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

            ab_pipeline.producer_tail(ab_producer_state)

        # ---------------------------------------------------------------
        # Specialized MMA warp
        # ---------------------------------------------------------------
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()
            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)

            sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + self.num_accumulator_tmem_cols, dtype=self.sf_dtype)
            tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
                tiled_mma,
                self.mma_tiler,
                self.sf_vec_size,
                cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
            )
            tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)

            sfb_tmem_ptr = cute.recast_ptr(
                acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
                dtype=self.sf_dtype,
            )
            tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
                tiled_mma,
                self.mma_tiler,
                self.sf_vec_size,
                cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
            )
            tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)

            tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
            tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)

            ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage)
            acd_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage)
            tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage)

            tile_info = cute.make_rmem_tensor((4,), cutlass.Int32)
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            for idx in cutlass.range(4, unroll_full=True):
                tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
            is_valid_tile = tile_info[0] >= cutlass.Int32(0)
            cute.arch.fence_proxy("async.shared", space="cta")
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()

            while is_valid_tile:
                k_tile_cnt = tile_info[3]
                ab_consumer_state.reset_count()
                peek_ab_full_status = cutlass.Boolean(1)
                if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
                    peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)

                acd_producer_state.reset_count()
                peek_acc_empty_status = cutlass.Boolean(1)
                if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
                    peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state)

                mma_tile_coord_mnl = (
                    tile_info[1] // cute.size(tiled_mma.thr_id.shape),
                    tile_info[2],
                    tile_info[0],
                )

                if cutlass.const_expr(self.overlapping_accum):
                    acc_stage_index = acd_producer_state.phase ^ 1
                else:
                    acc_stage_index = acd_producer_state.index

                tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
                tCtSFB_mma = tCtSFB
                if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
                    offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
                    shifted_ptr = cute.recast_ptr(
                        acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset,
                        dtype=self.sf_dtype,
                    )
                    tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
                elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
                    offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
                    shifted_ptr = cute.recast_ptr(
                        acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset,
                        dtype=self.sf_dtype,
                    )
                    tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)

                if is_leader_cta:
                    acc_pipeline.producer_acquire(acd_producer_state, peek_acc_empty_status)

                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)

                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    if is_leader_cta:
                        ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                        s2t_stage_coord = (None, None, None, None, ab_consumer_state.index)
                        tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
                        tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
                        cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_staged, tCtSFA_compact_s2t)
                        cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_staged, tCtSFB_compact_s2t)

                        num_kblocks = cute.size(tCrA, mode=[2])
                        ab_consumer_state_next = ab_consumer_state.clone()
                        ab_consumer_state_next.advance()
                        if ab_consumer_state_next.count < k_tile_cnt:
                            peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next)

                        for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
                            kblock_coord = (None, None, kblock_idx, ab_consumer_state.index)
                            sf_kblock_coord = (None, None, kblock_idx)
                            tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
                            tiled_mma.set(tcgen05.Field.SFB, tCtSFB_mma[sf_kblock_coord].iterator)
                            cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc)
                            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)

                        ab_pipeline.consumer_release(ab_consumer_state)
                        ab_consumer_state = ab_consumer_state_next

                if is_leader_cta:
                    acc_pipeline.producer_commit(acd_producer_state)

                acd_producer_state.advance()
                if acd_producer_state.count < k_tile_cnt:
                    if is_leader_cta:
                        peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state)

                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                for idx in cutlass.range(4, unroll_full=True):
                    tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                cute.arch.fence_proxy("async.shared", space="cta")
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

            acc_pipeline.producer_tail(acd_producer_state)

        # ---------------------------------------------------------------
        # Specialized bias load warp
        # ---------------------------------------------------------------
        if cutlass.const_expr(self.enable_bias):
            if warp_idx == self.bias_load_warp_id and total_token > 0:
                bias_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_bias_stage)
                tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage)

                bias_elems_per_thread = 128 // self.bias_dtype.width
                bias_g2s_atom = cute.make_copy_atom(
                    cute.nvgpu.cpasync.CopyG2SOp(),
                    self.bias_dtype,
                    num_bits_per_copy=128,
                )
                bias_g2s_tiled = cute.make_tiled_copy_tv(
                    bias_g2s_atom,
                    cute.make_layout((self.threads_per_warp,)),
                    cute.make_layout((bias_elems_per_thread,)),
                )
                thr_bias_g2s = bias_g2s_tiled.get_slice(cute.arch.lane_idx())
                tBs_sBias = thr_bias_g2s.partition_D(sBias)

                bias_n_total = mBias_nl.shape[0]
                tBpBias = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Boolean)

                tile_info = cute.make_rmem_tensor((4,), cutlass.Int32)
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                for idx in cutlass.range(4, unroll_full=True):
                    tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                cute.arch.fence_proxy("async.shared", space="cta")
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

                while is_valid_tile:
                    bias_producer_state.reset_count()
                    mma_n_coord = tile_info[2]
                    expert_idx = tile_info[0]
                    gBias_tile = gBias_nl[(None, mma_n_coord, expert_idx)]
                    tBs_gBias = thr_bias_g2s.partition_S(gBias_tile)
                    tBpBias[0] = mma_n_coord * self.mma_tiler[1] + cute.arch.lane_idx() * bias_elems_per_thread < bias_n_total
                    bias_pipeline.producer_acquire(bias_producer_state)
                    cute.copy(
                        bias_g2s_tiled,
                        tBs_gBias[(None, 0)],
                        tBs_sBias[(None, 0, bias_producer_state.index)],
                        pred=tBpBias,
                    )
                    bias_pipeline.producer_commit(bias_producer_state)
                    bias_producer_state.advance()

                    tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                    for idx in cutlass.range(4, unroll_full=True):
                        tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                    is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                    cute.arch.fence_proxy("async.shared", space="cta")
                    tile_info_pipeline.consumer_release(tile_info_consumer_state)
                    tile_info_consumer_state.advance()

                bias_pipeline.producer_tail(bias_producer_state)

        # ---------------------------------------------------------------
        # Specialized ACT epilogue warps (0-3): TMEM→regs, alpha, GLU activation,
        # C store, hadamard_in
        # ---------------------------------------------------------------
        if warp_idx < self.epilog_rht_store_warp_id[0] and total_token > 0:
            epi_tidx = tidx

            #
            # Alloc tensor memory buffer
            #
            tmem.allocate(self.num_tmem_alloc_cols)

            #
            # Bar sync for retrieve tensor memory ptr from shared memory
            #
            tmem.wait_for_alloc()

            #
            # Retrieving tensor memory ptr and make accumulator tensor
            #
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            # (MMA, MMA_M, MMA_N, STAGE)
            tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)

            #
            # Partition for epilogue (shape-only via mD_mnl for invariant setup)
            #
            thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v)
            gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None))
            tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape)

            (
                tiled_copy_t2r,
                tTR_tAcc_base,
                tTR_rAcc_gate,
                tTR_rAcc_up,
            ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs)

            tTR_rC = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.c_dtype)
            tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC)

            tTR_rD = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.d_dtype)
            tiled_copy_r2s_d, tRS_rD, tRS_sD = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rD, epi_tidx, sD)

            #
            # Create per-expert extension (for C/prob tensors inside tile loop)
            #
            epi_ext = self._make_extension(workspace_ptr)

            #
            # Persistent tile scheduling state
            #
            acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage)

            #
            # Pingpong producer state
            #
            pingpong_act_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_pingpong_stage)

            # Threads/warps participating in TMA store pipeline for C
            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.threads_per_warp * len(self.epilog_act_warp_id),
            )
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage,
                producer_group=c_producer_group,
            )

            d_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.threads_per_warp * len(self.epilog_act_warp_id),
            )
            d_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_d_stage,
                producer_group=d_producer_group,
            )

            tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage)

            if cutlass.const_expr(self.enable_bias):
                bias_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_bias_stage)
                bias_s2r_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.bias_dtype, num_bits_per_copy=128)
                tTR_rBias_gate = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype)
                tTR_rBias_up = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype)

            # Get the first tile info
            tile_info = cute.make_rmem_tensor((4,), cutlass.Int32)

            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            for idx in cutlass.range(4, unroll_full=True):
                tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
            is_valid_tile = tile_info[0] >= cutlass.Int32(0)
            cute.arch.fence_proxy("async.shared", space="cta")
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()

            num_prev_subtiles = cutlass.Int32(0)
            num_prev_d_subtiles = cutlass.Int32(0)
            while is_valid_tile:
                # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt)
                epi_work_tile_info = MoEWorkTileInfo(
                    expert_idx=tile_info[0],
                    tile_m_idx=tile_info[1],
                    tile_n_idx=tile_info[2],
                    k_tile_cnt=tile_info[3],
                )
                mma_tile_coord_mnl = (
                    epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape),
                    epi_work_tile_info.tile_n_idx,
                    cutlass.Int32(0),
                )

                expert_idx = epi_work_tile_info.expert_idx
                alpha_val = alpha[expert_idx]
                epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx)

                if cutlass.const_expr(self.generate_amax):
                    thread_tile_amax = cutlass.Float32(0.0)

                if cutlass.const_expr(self.enable_bias):
                    bias_consumer_state.reset_count()
                    bias_pipeline.consumer_wait(bias_consumer_state)
                    sBias_stage = sBias[(None, bias_consumer_state.index)]
                    if cutlass.const_expr(self.act_func == "srelu"):
                        sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(self.epi_tile[1]))
                    else:
                        sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(2 * self.epi_tile[1]))

                #
                # Get per-expert C tensor inside tile loop
                #
                real_c, _ = epi_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, epi_work_tile_info)
                gC_mnl = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None))
                thr_mma_epi_loop = tiled_mma.get_slice(mma_tile_coord_v)
                tCgC = thr_mma_epi_loop.partition_C(gC_mnl)
                _, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, self.epi_tile_c, sC)
                bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]
                bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))

                #
                # Get per-expert D tensor inside tile loop
                #
                real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info)
                gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None))
                tCgD = thr_mma_epi_loop.partition_C(gD_mnl_loop)
                _, bSG_sD, bSG_gD_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgD, epi_tile, sD)
                bSG_gD = bSG_gD_partitioned[(None, None, None, *mma_tile_coord_mnl)]
                bSG_gD = cute.group_modes(bSG_gD, 1, cute.rank(bSG_gD))

                #
                # Get per-expert prob tensor inside tile loop
                #
                real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info)
                mPosition = (
                    (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0]
                    + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape))
                    + tidx
                )
                mProb = real_prob[mPosition, 0, 0]

                #
                # Get accumulator stage index
                #
                if cutlass.const_expr(self.overlapping_accum):
                    acc_stage_index = acc_consumer_state.phase
                    reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
                else:
                    acc_stage_index = acc_consumer_state.index

                # Set tensor memory buffer for current tile
                # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE)
                tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]

                #
                # Wait for accumulator buffer full
                #
                acc_pipeline.consumer_wait(acc_consumer_state)
                tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))

                #
                # Store accumulator to global memory in subtiles
                #
                subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
                subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2
                for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1):
                    real_subtile_idx = subtile_idx if cutlass.const_expr(self.act_func == "srelu") else subtile_idx // 2
                    if cutlass.const_expr(self.overlapping_accum):
                        if reverse_subtile:
                            real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - real_subtile_idx

                    #
                    # Load accumulator from tensor memory buffer to register
                    #
                    if cutlass.const_expr(self.act_func == "srelu"):
                        tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx)]
                    else:
                        tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx * 2)]
                        tTR_tAcc_mn_up = tTR_tAcc[(None, None, None, real_subtile_idx * 2 + 1)]

                    cute.copy(tiled_copy_t2r, tTR_tAcc_mn_gate, tTR_rAcc_gate)
                    if cutlass.const_expr(self.act_func != "srelu"):
                        cute.copy(tiled_copy_t2r, tTR_tAcc_mn_up, tTR_rAcc_up)

                    #
                    # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled
                    #
                    if cutlass.const_expr(self.overlapping_accum):
                        if subtile_idx == self.iter_acc_early_release_in_epilogue:
                            cute.arch.fence_view_async_tmem_load()
                            with cute.arch.elect_one():
                                acc_pipeline.consumer_release(acc_consumer_state)
                            acc_consumer_state.advance()

                    #
                    # Apply alpha (+ bias when enabled)
                    #
                    if cutlass.const_expr(self.enable_bias):
                        if cutlass.const_expr(self.act_func == "srelu"):
                            sBias_sub = sBias_subtiles[(None, real_subtile_idx)]
                            cute.copy(bias_s2r_atom, sBias_sub, tTR_rBias_gate)
                        else:
                            sBias_pair = sBias_subtiles[(None, real_subtile_idx)]
                            sBias_sub = cute.flat_divide(sBias_pair, cute.make_layout(self.epi_tile[1]))
                            cute.copy(bias_s2r_atom, sBias_sub[(None, 0)], tTR_rBias_gate)
                        bias_vec_gate = tTR_rBias_gate.load()
                        if cutlass.const_expr(self.act_func != "srelu"):
                            cute.copy(bias_s2r_atom, sBias_sub[(None, 1)], tTR_rBias_up)
                            bias_vec_up = tTR_rBias_up.load()

                        if cutlass.const_expr(self.vectorized_f32):
                            for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2):
                                bias_gate_f32_0 = bias_vec_gate[i].to(cutlass.Float32)
                                bias_gate_f32_1 = bias_vec_gate[i + 1].to(cutlass.Float32)
                                tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.fma_packed_f32x2(
                                    (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]),
                                    (
                                        cutlass.Float32(alpha_val),
                                        cutlass.Float32(alpha_val),
                                    ),
                                    (bias_gate_f32_0, bias_gate_f32_1),
                                    rnd="rn",
                                    ftz=False,
                                )
                                if cutlass.const_expr(self.act_func != "srelu"):
                                    bias_up_f32_0 = bias_vec_up[i].to(cutlass.Float32)
                                    bias_up_f32_1 = bias_vec_up[i + 1].to(cutlass.Float32)
                                    tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.fma_packed_f32x2(
                                        (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]),
                                        (
                                            cutlass.Float32(alpha_val),
                                            cutlass.Float32(alpha_val),
                                        ),
                                        (bias_up_f32_0, bias_up_f32_1),
                                        rnd="rn",
                                        ftz=False,
                                    )
                        else:
                            for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)):
                                tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + bias_vec_gate[i].to(cutlass.Float32)
                                if cutlass.const_expr(self.act_func != "srelu"):
                                    tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + bias_vec_up[i].to(cutlass.Float32)

                        last_bias_subtile = subtile_cnt - 1 if cutlass.const_expr(self.act_func == "srelu") else subtile_cnt - 2
                        if subtile_idx == last_bias_subtile:
                            bias_pipeline.consumer_release(bias_consumer_state)
                            bias_consumer_state.advance()
                    else:
                        if cutlass.const_expr(self.vectorized_f32):
                            for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2):
                                tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.mul_packed_f32x2(
                                    (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]),
                                    (
                                        cutlass.Float32(alpha_val),
                                        cutlass.Float32(alpha_val),
                                    ),
                                    rnd="rn",
                                    ftz=False,
                                )
                                if cutlass.const_expr(self.act_func != "srelu"):
                                    tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.mul_packed_f32x2(
                                        (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]),
                                        (
                                            cutlass.Float32(alpha_val),
                                            cutlass.Float32(alpha_val),
                                        ),
                                        rnd="rn",
                                        ftz=False,
                                    )
                        else:
                            for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)):
                                tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val)
                                if cutlass.const_expr(self.act_func != "srelu"):
                                    tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val)

                    #
                    # Store pre-activation output to C tensor for residual/backward.
                    #
                    if cutlass.const_expr(self.act_func == "srelu"):
                        self.store_c_unary(
                            tiled_copy_r2s,
                            tma_atom_c,
                            warp_idx,
                            tTR_rAcc_gate,
                            tRS_rC,
                            tRS_sC,
                            bSG_gC,
                            bSG_sC,
                            c_pipeline,
                            num_prev_subtiles,
                            real_subtile_idx,
                        )
                    else:
                        self.store_c(
                            tiled_copy_r2s,
                            tma_atom_c,
                            warp_idx,
                            tTR_rAcc_gate,
                            tTR_rAcc_up,
                            tRS_rC,
                            tRS_sC,
                            bSG_gC,
                            bSG_sC,
                            c_pipeline,
                            num_prev_subtiles,
                            real_subtile_idx,
                        )
                    num_prev_subtiles = num_prev_subtiles + 1

                    #
                    # GeGLU clamp before C store
                    #
                    if cutlass.const_expr(self.act_func == "geglu"):
                        geglu_max_val = cutlass.Float32(7.0)
                        geglu_min_val = cutlass.Float32(-7.0)
                        for i in cutlass.range_constexpr(cute.size(tTR_rAcc_up)):
                            tTR_rAcc_gate[i] = fmin(tTR_rAcc_gate[i], geglu_max_val)
                            tTR_rAcc_up[i] = fmin(tTR_rAcc_up[i], geglu_max_val)
                            tTR_rAcc_up[i] = fmax(tTR_rAcc_up[i], geglu_min_val)

                    acc_vec_gate = tTR_rAcc_gate.load()

                    #
                    # Compute activation.
                    #
                    tCompute = cute.make_rmem_tensor(acc_vec_gate.shape, self.acc_dtype)
                    if cutlass.const_expr(self.act_func == "srelu"):
                        self.srelu_act(tCompute, acc_vec_gate, mProb)
                    elif cutlass.const_expr(self.act_func == "geglu"):
                        acc_vec_up = tTR_rAcc_up.load()
                        self.geglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset)
                    elif cutlass.const_expr(self.act_func == "swiglu"):
                        acc_vec_up = tTR_rAcc_up.load()
                        self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb)
                    elif cutlass.const_expr(self.act_func == "situglu"):
                        acc_vec_up = tTR_rAcc_up.load()
                        self.situglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, situ_beta1, situ_beta2)

                    if cutlass.const_expr(self.generate_amax):
                        thread_tile_amax = self.amax_reduction_per_thread(tCompute, thread_tile_amax)

                    #
                    # Store post-activation output to D.
                    #
                    acc_vec = tiled_copy_r2s_d.retile(tCompute).load()
                    tRS_rD.store(acc_vec.to(self.d_dtype))
                    d_buffer = num_prev_d_subtiles % self.num_d_stage
                    cute.copy(
                        tiled_copy_r2s_d,
                        tRS_rD,
                        tRS_sD[(None, None, None, d_buffer)],
                    )
                    cute.arch.fence_proxy("async.shared", space="cta")
                    self.epilog_sync_barrier_group0.arrive_and_wait()
                    if warp_idx == self.epilog_act_warp_id[0]:
                        cute.copy(
                            tma_atom_d,
                            bSG_sD[(None, d_buffer)],
                            bSG_gD[(None, real_subtile_idx)],
                        )
                        d_pipeline.producer_commit()

                    #
                    # Signal the RHT epilogue warps that the post-activation D tile is in SMEM.
                    #
                    pingpong_pipeline.producer_acquire(pingpong_act_producer_state)
                    pingpong_pipeline.producer_commit(pingpong_act_producer_state)
                    pingpong_act_producer_state.advance()

                    num_prev_d_subtiles = num_prev_d_subtiles + 1

                    #
                    # Delayed TMA store acquire + group sync (always enabled)
                    #
                    if cutlass.const_expr(self.delay_tma_store_acquire_sync):
                        if warp_idx == self.epilog_act_warp_id[0]:
                            d_pipeline.producer_acquire()
                            c_pipeline.producer_acquire()
                        self.epilog_sync_barrier_group0.arrive_and_wait()

                    self.epilog_sync_barrier.arrive_and_wait()

                #
                # Full epilogue barrier (ACT + RHT must both arrive)
                #
                if cutlass.const_expr(self.generate_amax):
                    gAmax = mAmax_tensor[(expert_idx, None)].iterator.llvm_ptr  # First element
                    self.amax_reduction_per_warp_and_cta(thread_tile_amax, warp_idx, sAmax, gAmax)

                self.epilog_sync_barrier.arrive_and_wait()

                #
                # Async arrive accumulator buffer empty
                #
                if cutlass.const_expr(not self.overlapping_accum):
                    with cute.arch.elect_one():
                        acc_pipeline.consumer_release(acc_consumer_state)
                    acc_consumer_state.advance()

                #
                # Advance to next tile
                #
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                for idx in cutlass.range(4, unroll_full=True):
                    tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                cute.arch.fence_proxy("async.shared", space="cta")
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

            #
            # Dealloc the tensor memory buffer
            #
            tmem.relinquish_alloc_permit()
            self.epilog_sync_barrier_group0.arrive_and_wait()
            tmem.free(tmem_ptr)
            #
            # Wait for C store / pingpong complete
            #
            c_pipeline.producer_tail()
            d_pipeline.producer_tail()
            pingpong_pipeline.producer_tail(pingpong_act_producer_state)

        # ---------------------------------------------------------------
        # Specialized RHT store warps (4-7): post-RHT amax from D SMEM
        # ---------------------------------------------------------------
        if warp_idx < self.mma_warp_id and warp_idx >= self.epilog_rht_store_warp_id[0] and total_token > 0:
            epi_tidx = tidx % 128

            if cutlass.const_expr(self.use_tmem_post_rht_amax):
                #
                # Alloc tensor memory buffer
                #
                tmem.allocate(self.num_tmem_alloc_cols)

                #
                # Bar sync for retrieve tensor memory ptr from shared memory
                #
                tmem.wait_for_alloc()

                #
                # Retrieving tensor memory ptr
                #
                tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

                #
                # Hadamard setup (loads Hadamard matrix from GMEM to SMEM)
                #
                tiled_hmma = hadamard_setup(hadamard_tensor_local, sHadamard, epi_tidx)

                #
                # Partition for epilogue (shape-only via mD_mnl)
                #
                tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
                thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v)
                gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None))
                tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape)

                (
                    tiled_copy_t2r,
                    tTR_tAcc_base,
                    tTR_rAcc_gate,
                    tTR_rAcc_up,
                ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs)

            #
            # Pingpong consumer state
            #
            pingpong_rht_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_pingpong_stage)

            #
            # Create per-expert extension (for D tensor inside tile loop)
            #
            epi_ext = self._make_extension(workspace_ptr)

            # Get the first tile info
            tile_info = cute.make_rmem_tensor((4,), cutlass.Int32)

            tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage)
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            for idx in cutlass.range(4, unroll_full=True):
                tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
            is_valid_tile = tile_info[0] >= cutlass.Int32(0)
            cute.arch.fence_proxy("async.shared", space="cta")
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()

            if cutlass.const_expr(self.use_tmem_post_rht_amax):
                #
                # Accumulator stage index (for overlapping_accum)
                #
                if cutlass.const_expr(self.overlapping_accum):
                    acc_stage_index = 0
                    reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
                else:
                    acc_stage_index = 0
                    reverse_subtile = cutlass.Boolean(False)

                num_prev_subtiles = cutlass.Int32(0)
            num_prev_d_subtiles = cutlass.Int32(0)
            while is_valid_tile:
                # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt)
                epi_work_tile_info = MoEWorkTileInfo(
                    expert_idx=tile_info[0],
                    tile_m_idx=tile_info[1],
                    tile_n_idx=tile_info[2],
                    k_tile_cnt=tile_info[3],
                )
                mma_tile_coord_mnl = (
                    epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape),
                    epi_work_tile_info.tile_n_idx,
                    cutlass.Int32(0),
                )
                expert_idx = epi_work_tile_info.expert_idx
                epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx)

                if cutlass.const_expr(self.generate_post_rht_amax):
                    thread_tile_post_rht_amax = cutlass.Float32(0.0)

                if cutlass.const_expr(self.use_tmem_post_rht_amax):
                    #
                    # Set tensor memory buffer for current tile
                    # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE)
                    #
                    tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]
                    tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))

                #
                # Store accumulator to global memory in subtiles
                #
                if cutlass.const_expr(self.use_tmem_post_rht_amax):
                    subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
                else:
                    subtile_cnt = self.cta_tile_shape_mnk[1] // cute.size(self.epi_tile[1])
                subtile_step = 1 if cutlass.const_expr(self.act_func == "srelu") else 2
                for subtile_idx in cutlass.range(0, subtile_cnt, subtile_step, unroll=1):
                    #
                    # Wait for ACT warps to finish writing the post-activation D tile to SMEM.
                    #
                    pingpong_pipeline.consumer_wait(pingpong_rht_consumer_state)

                    if cutlass.const_expr(self.generate_post_rht_amax):
                        d_buffer = num_prev_d_subtiles % self.num_d_stage
                        if cutlass.const_expr(self.use_tmem_post_rht_amax):
                            hadamard_tmem_a_ptr = self.query_hadamard_tmem_a_ptr(subtile_idx, reverse_subtile, tmem_ptr)
                            hadamard_tmem_out_ptr = self.query_hadamard_tmem_acc_ptr(subtile_idx, reverse_subtile, tmem_ptr)
                            hadamard_transpose_in(
                                sD,
                                d_buffer,
                                0,
                                hadamard_tmem_a_ptr,
                                epi_tidx,
                            )
                            hadamard_prerequisite_empty = hadamard_prerequisite_producer.acquire_and_advance()
                            hadamard_prerequisite_empty.commit()
                            hadamard_prerequisite_empty = hadamard_prerequisite_consumer.wait_and_advance()
                            hadamard_prerequisite_consumer.release(hadamard_prerequisite_empty)
                            hadamard_compute(
                                tiled_hmma,
                                hadamard_tmem_a_ptr,
                                hadamard_tmem_out_ptr,
                                storage.sHadamard,
                                (M_PER_CLUSTER, HADAMARD_SIZE),
                                epi_tidx,
                                hadamard_producer,
                            )
                            hadamard_consumer.wait_and_advance()
                            thread_tile_post_rht_amax = hadamard_transpose_out_amax(
                                hadamard_tmem_out_ptr,
                                epi_tidx,
                                thread_tile_post_rht_amax,
                            )
                            hadamard_transpose_in(
                                sD,
                                d_buffer,
                                HADAMARD_SIZE,
                                hadamard_tmem_a_ptr,
                                epi_tidx,
                            )
                            hadamard_prerequisite_empty = hadamard_prerequisite_producer.acquire_and_advance()
                            hadamard_prerequisite_empty.commit()
                            hadamard_prerequisite_empty = hadamard_prerequisite_consumer.wait_and_advance()
                            hadamard_prerequisite_consumer.release(hadamard_prerequisite_empty)
                            hadamard_compute(
                                tiled_hmma,
                                hadamard_tmem_a_ptr,
                                hadamard_tmem_out_ptr,
                                storage.sHadamard,
                                (M_PER_CLUSTER, HADAMARD_SIZE),
                                epi_tidx,
                                hadamard_producer,
                            )
                            hadamard_consumer.wait_and_advance()
                            thread_tile_post_rht_amax = hadamard_transpose_out_amax(
                                hadamard_tmem_out_ptr,
                                epi_tidx,
                                thread_tile_post_rht_amax,
                            )
                        else:
                            thread_tile_post_rht_amax = hadamard_smem_transpose_fwht_amax(
                                sD,
                                d_buffer,
                                0,
                                epi_tidx,
                                thread_tile_post_rht_amax,
                            )
                            thread_tile_post_rht_amax = hadamard_smem_transpose_fwht_amax(
                                sD,
                                d_buffer,
                                HADAMARD_SIZE,
                                epi_tidx,
                                thread_tile_post_rht_amax,
                            )

                    #
                    # Release pingpong consumer slot
                    #
                    pingpong_pipeline.consumer_release(pingpong_rht_consumer_state)
                    pingpong_rht_consumer_state.advance()
                    num_prev_d_subtiles = num_prev_d_subtiles + 1
                    self.epilog_sync_barrier.arrive_and_wait()

                #
                # Full epilogue barrier (ACT + RHT must both arrive)
                #
                if cutlass.const_expr(self.generate_post_rht_amax):
                    gPostRhtAmax = mPostRhtAmax_tensor[(expert_idx, None)].iterator.llvm_ptr
                    self.post_rht_amax_reduction_per_warp_and_cta(
                        thread_tile_post_rht_amax,
                        warp_idx,
                        sPostRhtAmax,
                        gPostRhtAmax,
                    )

                self.epilog_sync_barrier.arrive_and_wait()

                if cutlass.const_expr(self.use_tmem_post_rht_amax):
                    #
                    # Update overlapping_accum stage index for next tile
                    #
                    if cutlass.const_expr(self.overlapping_accum):
                        acc_stage_index = acc_stage_index ^ 1
                        reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
                    else:
                        acc_stage_index = 0

                #
                # Advance to next tile
                #
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                for idx in cutlass.range(4, unroll_full=True):
                    tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                is_valid_tile = tile_info[0] >= cutlass.Int32(0)
                cute.arch.fence_proxy("async.shared", space="cta")
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

            if cutlass.const_expr(self.use_tmem_post_rht_amax):
                #
                # Dealloc the tensor memory buffer
                #
                tmem.relinquish_alloc_permit()
                self.epilog_sync_barrier_group1.arrive_and_wait()
                tmem.free(tmem_ptr)
        # END OF KERNEL


class BlockScaledMoEGroupedGemmGluHadamardCompatKernel(BlockScaledMoEGroupedGemmGluHadamardKernel):
    """Compatibility adapter for the existing grouped GLU frontend wrapper shape."""

    def __init__(
        self,
        sf_vec_size: int,
        acc_dtype: Type[cutlass.Numeric],
        use_2cta_instrs: bool,
        mma_tiler_mn: Tuple[int, int],
        cluster_shape_mn: Tuple[int, int],
        vectorized_f32: bool,
        generate_sfd: bool,
        discrete_col_sfd: bool,
        expert_cnt: int,
        weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE,
        use_dynamic_sched: bool = False,
        act_func: str = "swiglu",
        enable_bias: bool = False,
    ):
        del generate_sfd, discrete_col_sfd
        super().__init__(
            sf_vec_size=sf_vec_size,
            acc_dtype=acc_dtype,
            use_2cta_instrs=use_2cta_instrs,
            mma_tiler_mn=mma_tiler_mn,
            cluster_shape_mn=cluster_shape_mn,
            vectorized_f32=vectorized_f32,
            expert_cnt=expert_cnt,
            weight_mode=weight_mode,
            use_dynamic_sched=use_dynamic_sched,
            act_func=act_func,
            enable_bias=enable_bias,
        )

    def __call__(
        self,
        a: cute.Tensor,
        b,
        sfb,
        n: Int32,
        k: Int32,
        b_stride_size: cutlass.Int64,
        b_major_mode: cutlass.Constexpr,
        workspace_ptr,
        c: cute.Tensor,
        d: cute.Tensor,
        d_col: cute.Tensor,
        sfa: cute.Tensor,
        sfd_row_tensor: Optional[cute.Tensor],
        sfd_col_tensor: Optional[cute.Tensor],
        amax_tensor: Optional[cute.Tensor],
        norm_const_tensor: Optional[cute.Tensor],
        padded_offsets: cute.Tensor,
        alpha: cute.Tensor,
        prob: cute.Tensor,
        bias: Optional[cute.Tensor],
        max_active_clusters: cutlass.Constexpr,
        stream: cuda.CUstream,
        epilogue_op: cutlass.Constexpr = lambda x: x,
        linear_offset: cutlass.Float32 = 0.0,
    ):
        del d_col, sfd_row_tensor, sfd_col_tensor, norm_const_tensor
        return super().__call__(
            a=a,
            b=b,
            sfa=sfa,
            sfb=sfb,
            n=n,
            k=k,
            b_stride_size=b_stride_size,
            b_major_mode=b_major_mode,
            workspace_ptr=workspace_ptr,
            c=c,
            d=d,
            amax_tensor=amax_tensor,
            post_rht_amax_tensor=None,
            padded_offsets=padded_offsets,
            alpha=alpha,
            prob=prob,
            hadamard_tensor=None,
            bias=bias,
            max_active_clusters=max_active_clusters,
            stream=stream,
            epilogue_op=epilogue_op,
            linear_offset=linear_offset,
        )
