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

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

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

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

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

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

"""SM100 grouped GEMM for NVFP4 expert weights and BF16 activations.

Each routed-row tile selects one expert through FlashInfer's MoE sort metadata.
The kernel decodes each 16-value E2M1 weight block and its E4M3 scale to BF16,
then executes BF16 tensor-core MMA with FP32 accumulation.
"""

from math import ceil, log2
from typing import Optional, Union

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils as utils
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
import cutlass.utils.mixed_input_helpers as mixed_input_utils
from cutlass.utils.mixed_input_helpers import TransformMode
from cutlass.cute.nvgpu import cpasync, tcgen05

from flashinfer.tllm_enums import (
    ActivationType,
    DEFAULT_SWIGLU_ALPHA,
    DEFAULT_SWIGLU_BETA,
    DEFAULT_SWIGLU_LIMIT,
)

from .moe_w4a16_utils import decode_nvfp4_fragment_to_bf16
from .utils import (
    blk_reduce_bf16,
    fmin,
    griddepcontrol_launch_dependents,
    griddepcontrol_wait,
)


_NVFP4_SCALE_GRANULARITY_M = 1
_NVFP4_SCALE_GRANULARITY_K = 16
_TMA_CACHE_HINT_EVICT_FIRST = 0x12F0000000000000
_TMA_CACHE_HINT_EVICT_LAST = 0x14F0000000000000


class Sm100W4A16GroupedGemmKernel:
    """Warp-specialized grouped GEMM for the W4A16 MoE pipeline."""

    def __init__(
        self,
        acc_dtype: type[cutlass.Numeric],
        use_2cta_instrs: bool,
        mma_tiler_mnk: tuple[int, int, int],
        cluster_shape_mn: tuple[int, int],
        group_count: int,
        activation_type: Optional[int],
        swiglu_alpha: float,
        swiglu_beta: float,
        swiglu_limit: float,
        use_fused_finalize: bool,
        enable_pdl: bool,
        use_clc_scheduler: bool,
        raster_along_m: bool,
        transform_fragment_size: int,
    ):
        """Initialize the W4A16 grouped GEMM configuration."""
        self.group_count = group_count
        self.acc_dtype = acc_dtype
        self.use_2cta_instrs = use_2cta_instrs
        self.cluster_shape_mn = cluster_shape_mn
        self.mma_tiler = mma_tiler_mnk
        if activation_type not in (
            None,
            ActivationType.Swiglu.value,
            ActivationType.Relu2.value,
        ):
            raise ValueError(
                f"unsupported W4A16 epilogue activation: {activation_type}"
            )
        self.fuse_activation = activation_type is not None
        self.use_clc_scheduler = use_clc_scheduler
        self.raster_along_m = raster_along_m
        self.transform_fragment_size = transform_fragment_size
        self.gated = activation_type == ActivationType.Swiglu.value
        self.swiglu_alpha = swiglu_alpha
        self.swiglu_beta = swiglu_beta
        self.swiglu_limit = swiglu_limit
        self.parameterized_swiglu = (
            swiglu_alpha != DEFAULT_SWIGLU_ALPHA
            or swiglu_beta != DEFAULT_SWIGLU_BETA
            or swiglu_limit != DEFAULT_SWIGLU_LIMIT
        )
        self.output_m_factor = 2 if self.gated else 1
        if self.gated and use_fused_finalize:
            raise ValueError(
                "W4A16 fused finalize is only implemented for the non-gated epilogue"
            )
        self.use_fused_finalize = use_fused_finalize
        self.enable_pdl = enable_pdl
        self.cta_group = (
            tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
        )
        self.epilog_warp_id = (0, 1, 2, 3)
        self.mma_warp_id = 4
        self.schedule_warp_id = 5
        self.weight_tma_warp_id = 6
        self.activation_tma_warp_id = 7
        self.num_transform_warpgroups = 2
        self.transform_warp_id = tuple(range(8, 8 + 4 * self.num_transform_warpgroups))
        self.num_transform_warps = len(self.transform_warp_id)
        # Keep the role-control warpgroup on one absolute register target.
        # Transform warps retain the 128-register launch default.
        self.num_regs_epilogue_warps = 176
        self.num_regs_generic_warps = 80
        self.threads_per_cta = 32 * (
            max(
                (
                    self.mma_warp_id,
                    self.schedule_warp_id,
                    self.weight_tma_warp_id,
                    self.activation_tma_warp_id,
                    *self.epilog_warp_id,
                    *self.transform_warp_id,
                )
            )
            + 1
        )

        self.epilog_sync_barrier = pipeline.NamedBarrier(
            1, 32 * len(self.epilog_warp_id)
        )
        self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta)
        self.sched_sync_barrier = pipeline.NamedBarrier(5, 32)

        self.smem_buffer_align_bytes = 1024

    def _setup_attributes(self):
        """Derive static MMA, pipeline, and storage layouts from operand types."""
        self.transform_a_source = mixed_input_utils.get_transform_a_source(
            self.a_major_mode
        )
        tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.mma_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.acc_dtype,
            self.cta_group,
            self.mma_tiler[:2],
            self.transform_a_source,
        )
        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.mma_tiler_c = (
            self.mma_tiler[0] // self.output_m_factor,
            self.mma_tiler[1],
            self.mma_tiler[2],
        )
        self.cta_tile_shape_mnk_c = (
            self.cta_tile_shape_mnk[0] // self.output_m_factor,
            self.cta_tile_shape_mnk[1],
            self.cta_tile_shape_mnk[2],
        )
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((*self.cluster_shape_mn, 1)),
            (tiled_mma.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

        default_epi_tile = sm100_utils.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk,
            self.use_2cta_instrs,
            self.c_layout,
            self.c_dtype,
        )
        self.epi_tile = (
            (self.cta_tile_shape_mnk_c[0], default_epi_tile[1])
            if self.gated
            else default_epi_tile
        )
        self.epi_tile_n = cute.size(self.epi_tile[1])

        (
            self.num_load2trans_stage,
            self.num_trans2mma_stage,
            self.num_acc_stage,
            self.num_c_stage,
            self.num_tile_info_stage,
            self.num_acc_tmem_cols,
            self.num_a_tmem_cols,
        ) = self._compute_stages_and_tmem_cols(
            tiled_mma,
            self.mma_tiler,
            self.cta_tile_shape_mnk,
            self.epi_tile,
            self.a_dtype,
            self.b_dtype,
            self.c_dtype,
            self.c_layout,
            self.transform_a_source,
            self.smem_buffer_align_bytes,
            self.use_fused_finalize,
            self.use_clc_scheduler,
        )

        # TMEM allocation requires power-of-two column alignment.
        self.num_tmem_alloc_cols = cute.round_up(
            self.num_acc_tmem_cols + self.num_a_tmem_cols,
            cute.arch.get_min_tmem_alloc_cols("sm_100"),
        )
        self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols)))
        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            self.c_dtype,
            self.c_layout,
            self.epi_tile,
            self.num_c_stage,
        )
        (
            self.smem_layout_a,
            self.smem_layout_a_transform,
            self.smem_layout_b,
        ) = mixed_input_utils.compute_smem_layout(
            tiled_mma,
            self.mma_tiler,
            self.a_dtype,
            self.b_dtype,
            self.num_load2trans_stage,
            self.num_trans2mma_stage,
        )
        scale_tile_shape = (
            self.cta_tile_shape_mnk[0],
            self.mma_tiler[2],
        )
        # The transform and TMA partitions use different logical views of the
        # same block-scaled shared-memory storage.
        scale_smem_layout = blockscaled_utils.make_smem_layout_sf(
            scale_tile_shape,
            _NVFP4_SCALE_GRANULARITY_K,
            self.num_load2trans_stage,
        )
        scale_smem_layout_per_stage = cute.slice_(scale_smem_layout, (None, None, 0))
        trivial_swizzle = cute.make_swizzle(0, 4, 3)
        self.smem_layout_scale_per_stage = cute.make_composed_layout(
            trivial_swizzle, 0, scale_smem_layout_per_stage
        )
        self.smem_layout_scale = cute.make_composed_layout(
            trivial_swizzle, 0, scale_smem_layout
        )
        self.smem_layout_scale_tma = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma,
            self.mma_tiler,
            _NVFP4_SCALE_GRANULARITY_K,
            self.num_load2trans_stage,
        )

    @cute.jit
    def _setup_transform_partitions(
        self,
        tiled_mma: cute.TiledMma,
        copy_atom_a_input: cute.CopyAtom,
        copy_atom_a_transform: cute.CopyAtom,
        sA_transform_input: cute.Tensor,
        a_smem_layout_transform: cute.ComposedLayout,
        sA_transform: cute.Tensor,
        tCsS: cute.Tensor,
        accumulators: cute.Tensor,
        transform_tidx: cutlass.Int32,
        transform_group_idx: cutlass.Int32,
    ):
        if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM):
            tmem_ptr_transform = cute.recast_ptr(
                accumulators.iterator + self.num_acc_tmem_cols,
                dtype=self.mma_dtype,
            )
            transform_destination = cute.make_tensor(
                tmem_ptr_transform,
                tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout,
            )
        else:
            transform_destination = sA_transform

        src_copy_a, dst_copy_a, tAsA_input, tAsA_transform = (
            mixed_input_utils.transform_partition(
                self.transform_a_source,
                TransformMode.ConvertScale,
                copy_atom_a_input,
                copy_atom_a_transform,
                sA_transform_input,
                transform_destination,
                transform_tidx,
            )
        )
        num_transform_blocks = cute.size(tAsA_input, mode=[3])
        assert num_transform_blocks % self.num_transform_warpgroups == 0, (
            "A blocks must divide evenly across transform warpgroups"
        )
        blocks_per_transform_group = (
            num_transform_blocks // self.num_transform_warpgroups
        )
        transform_group_tiler = (
            None,
            None,
            None,
            blocks_per_transform_group,
            None,
        )
        transform_group_coord = (
            None,
            None,
            None,
            (None, transform_group_idx),
            None,
        )
        tAsA_input = cute.logical_divide(tAsA_input, transform_group_tiler)[
            transform_group_coord
        ]
        tAsA_transform = cute.logical_divide(tAsA_transform, transform_group_tiler)[
            transform_group_coord
        ]

        a_stage_coord = (None,) * (cute.rank(tAsA_input) - 1) + (0,)
        tArA = cute.make_rmem_tensor(
            tAsA_input[a_stage_coord].shape,
            tAsA_input.element_type,
        )
        tArA_transform = cute.make_rmem_tensor(
            tAsA_input[a_stage_coord].shape,
            self.mma_dtype,
        )

        smem_thr_copy_S = src_copy_a.get_slice(transform_tidx)
        tSsS_trans = smem_thr_copy_S.partition_S(tCsS)
        tSsS_trans = cute.logical_divide(tSsS_trans, transform_group_tiler)[
            transform_group_coord
        ]
        scale_stage_coord = (None,) * (cute.rank(tSsS_trans) - 1) + (0,)
        tSsS_layout_per_stage = tSsS_trans[scale_stage_coord].layout
        tSrS_copy = cute.make_rmem_tensor(
            cute.filter_zeros(tSsS_layout_per_stage).shape,
            self.a_scale_dtype,
        )
        tSrS = cute.make_tensor(
            tSrS_copy.iterator,
            cute.make_layout(
                tSsS_layout_per_stage.shape,
                stride=tSrS_copy.layout.stride,
            ),
        )
        assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), (
            "tSrS and tArA have different leading dimension"
        )
        assert cute.size(tSrS) == cute.size(tArA), "tSrS and tArA have different shape"

        transform_tiler_size = min(
            cute.size(cute.coalesce(tAsA_input.layout), mode=[0]),
            self.transform_fragment_size,
        )
        assert transform_tiler_size % _NVFP4_SCALE_GRANULARITY_K == 0, (
            "transform fragment must cover whole NVFP4 scale blocks"
        )
        transform_tiler = cute.make_layout(transform_tiler_size)
        tArA_load = cute.flat_divide(tArA, transform_tiler)
        tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load))
        scale_tiler = cute.make_layout(
            transform_tiler_size // _NVFP4_SCALE_GRANULARITY_K
        )
        tSrS_load = cute.flat_divide(tSrS_copy, scale_tiler)
        tSrS_load = cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load))
        tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler)
        tArA_transform_store = cute.group_modes(
            tArA_transform_store,
            1,
            cute.rank(tArA_transform_store),
        )
        return (
            dst_copy_a,
            tAsA_input,
            tAsA_transform,
            tSsS_trans,
            tSrS_copy,
            tSrS_load,
            tArA_load,
            tArA_transform,
            tArA_transform_store,
            transform_tiler,
        )

    @cute.jit
    def _transform_tile(
        self,
        a_load2trans_pipeline: pipeline.PipelineTmaAsync,
        trans2mma_pipeline: pipeline.PipelineAsyncUmma,
        a_load2trans_consumer_state: pipeline.PipelineState,
        trans2mma_producer_state: pipeline.PipelineState,
        dst_copy_a: cute.CopyAtom,
        tAsA_input: cute.Tensor,
        tAsA_transform: cute.Tensor,
        tSsS_trans: cute.Tensor,
        tSrS_copy: cute.Tensor,
        tSrS_load: cute.Tensor,
        tArA_load: cute.Tensor,
        tArA_transform: cute.Tensor,
        tArA_transform_store: cute.Tensor,
        transform_tiler: cute.Layout,
        k_tile_cnt: cutlass.Int32,
    ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]:
        a_load2trans_consumer_state.reset_count()
        peek_load2trans_full_status = cutlass.Boolean(1)
        if a_load2trans_consumer_state.count < k_tile_cnt:
            peek_load2trans_full_status = a_load2trans_pipeline.consumer_try_wait(
                a_load2trans_consumer_state
            )
        trans2mma_producer_state.reset_count()
        peek_trans2mma_empty_status = cutlass.Boolean(1)
        if trans2mma_producer_state.count < k_tile_cnt:
            peek_trans2mma_empty_status = trans2mma_pipeline.producer_try_acquire(
                trans2mma_producer_state
            )

        for _k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
            a_load2trans_pipeline.consumer_wait(
                a_load2trans_consumer_state,
                peek_load2trans_full_status,
            )
            a_stage_coord = (None,) * (cute.rank(tAsA_input) - 1) + (
                a_load2trans_consumer_state.index,
            )
            tAsA_input_slice = tAsA_input[a_stage_coord]
            tAsA_input_slice = cute.flat_divide(tAsA_input_slice, transform_tiler)
            tAsA_input_slice = cute.group_modes(
                tAsA_input_slice,
                1,
                cute.rank(tAsA_input_slice),
            )
            trans2mma_pipeline.producer_acquire(
                trans2mma_producer_state,
                peek_trans2mma_empty_status,
            )
            scale_stage_coord = (None,) * (cute.rank(tSsS_trans) - 1) + (
                a_load2trans_consumer_state.index,
            )
            tSsS_slice = tSsS_trans[scale_stage_coord]
            tSsS_slice_filtered = cute.make_tensor(
                tSsS_slice.iterator,
                cute.filter_zeros(tSsS_slice.layout),
            )
            cute.autovec_copy(tSsS_slice_filtered, tSrS_copy)
            cur_a_load2trans_consumer_state = a_load2trans_consumer_state.clone()
            for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])):
                cute.autovec_copy(
                    tAsA_input_slice[(None, idx)],
                    tArA_load[(None, idx)],
                )
                if cutlass.const_expr(idx == cute.size(tArA_load, mode=[1]) - 1):
                    a_load2trans_consumer_state.advance()
                    if a_load2trans_consumer_state.count < k_tile_cnt:
                        peek_load2trans_full_status = (
                            a_load2trans_pipeline.consumer_try_wait(
                                a_load2trans_consumer_state
                            )
                        )
                packed_fragment = cute.recast_tensor(
                    tArA_load[(None, idx)], cutlass.Uint8
                ).load()
                scale_slice = tSrS_load[(None, idx)]
                scale_slice_compact = cute.make_tensor(
                    scale_slice.iterator,
                    cute.filter_zeros(scale_slice.layout),
                )
                scale_fragment = cute.recast_tensor(
                    scale_slice_compact, cutlass.Uint8
                ).load()
                tensor_transformed = decode_nvfp4_fragment_to_bf16(
                    packed_fragment,
                    scale_fragment,
                )
                tArA_transform_store[(None, idx)].store(tensor_transformed)

            a_transform_stage_coord = (None,) * (cute.rank(tAsA_transform) - 1) + (
                trans2mma_producer_state.index,
            )
            mixed_input_utils.store_transformed_a(
                tArA_transform,
                tAsA_transform[a_transform_stage_coord],
                dst_copy_a,
            )
            if cutlass.const_expr(
                self.transform_a_source == tcgen05.OperandSource.TMEM
            ):
                cute.arch.fence_view_async_tmem_store()
            else:
                cute.arch.fence_proxy("async.shared", space="cta")
            a_load2trans_pipeline.consumer_release(cur_a_load2trans_consumer_state)
            trans2mma_pipeline.producer_commit(trans2mma_producer_state)
            trans2mma_producer_state.advance()
            if trans2mma_producer_state.count < k_tile_cnt:
                peek_trans2mma_empty_status = trans2mma_pipeline.producer_try_acquire(
                    trans2mma_producer_state
                )

        return a_load2trans_consumer_state, trans2mma_producer_state

    @cute.jit
    def wrapper(
        self,
        weight_ptr: cute.Pointer,
        weight_sf_ptr: cute.Pointer,
        activation_ptr: cute.Pointer,
        tile_idx_to_expert_idx_ptr: cute.Pointer,
        tile_idx_to_mn_limit_ptr: cute.Pointer,
        num_non_exiting_tiles_ptr: cute.Pointer,
        alpha_ptr: cute.Pointer,
        output_ptr: cute.Pointer,
        permuted_idx_to_expanded_idx_ptr: Optional[cute.Pointer],
        token_final_scales_ptr: Optional[cute.Pointer],
        m: cutlass.Int64,
        n: cutlass.Int64,
        k: cutlass.Int64,
        num_tokens: cutlass.Int64,
        top_k: cutlass.Int64,
        max_active_clusters: cutlass.Constexpr,
        stream: cuda.CUstream,
    ):
        """Build logical tensors over FlashInfer's packed NVFP4 storage."""
        weights = cute.make_tensor(
            weight_ptr,
            cute.make_ordered_layout((m, k, self.group_count), order=(1, 0, 2)),
        )
        weight_sf = cute.make_tensor(
            weight_sf_ptr,
            blockscaled_utils.tile_atom_to_shape_SF(
                weights.shape, _NVFP4_SCALE_GRANULARITY_K
            ),
        )
        activations = cute.make_tensor(
            activation_ptr,
            cute.make_ordered_layout((n, k, 1), order=(1, 0, 2)),
        )
        output_m = m // self.output_m_factor
        output_layout = cute.make_ordered_layout((output_m, n, 1), order=(0, 1, 2))
        output = cute.make_tensor(output_ptr, output_layout)
        final_output = (
            cute.make_tensor(
                output_ptr,
                cute.make_ordered_layout((output_m, num_tokens, 1), order=(0, 1, 2)),
            )
            if cutlass.const_expr(self.use_fused_finalize)
            else output
        )
        num_route_tiles = n // self.mma_tiler[1]
        tile_idx_to_expert_idx = cute.make_tensor(
            tile_idx_to_expert_idx_ptr, cute.make_layout((num_route_tiles,))
        )
        tile_idx_to_mn_limit = cute.make_tensor(
            tile_idx_to_mn_limit_ptr, cute.make_layout((num_route_tiles,))
        )
        num_non_exiting_tiles = cute.make_tensor(
            num_non_exiting_tiles_ptr, cute.make_layout((1,))
        )
        alpha = cute.make_tensor(alpha_ptr, cute.make_layout((self.group_count,)))
        permuted_idx_to_expanded_idx = (
            cute.make_tensor(permuted_idx_to_expanded_idx_ptr, cute.make_layout((n,)))
            if cutlass.const_expr(self.use_fused_finalize)
            else None
        )
        token_final_scales = (
            cute.make_tensor(
                token_final_scales_ptr,
                cute.make_ordered_layout((num_tokens, top_k), order=(1, 0)),
            )
            if cutlass.const_expr(self.use_fused_finalize)
            else None
        )
        return self(
            weights,
            weight_sf,
            activations,
            tile_idx_to_expert_idx,
            tile_idx_to_mn_limit,
            num_non_exiting_tiles,
            alpha,
            output,
            final_output,
            permuted_idx_to_expanded_idx,
            token_final_scales,
            max_active_clusters,
            stream,
        )

    @cute.jit
    def __call__(
        self,
        a: cute.Tensor,
        a_scale: cute.Tensor,
        b: cute.Tensor,
        tile_idx_to_expert_idx: cute.Tensor,
        tile_idx_to_mn_limit: cute.Tensor,
        num_non_exiting_tiles: cute.Tensor,
        alpha: cute.Tensor,
        c: cute.Tensor,
        final_output: cute.Tensor,
        permuted_idx_to_expanded_idx: Optional[cute.Tensor],
        token_final_scales: Optional[cute.Tensor],
        max_active_clusters: cutlass.Constexpr,
        stream: cuda.CUstream,
    ):
        """Configure and launch the grouped GEMM."""
        self.a_dtype: type[cutlass.Numeric] = a.element_type
        self.a_scale_dtype: type[cutlass.Numeric] = a_scale.element_type
        self.b_dtype: type[cutlass.Numeric] = b.element_type
        self.c_dtype: type[cutlass.Numeric] = c.element_type
        self.mma_dtype = self.b_dtype

        self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode()
        self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode()
        self.c_layout = utils.LayoutEnum.from_tensor(c)

        if cutlass.const_expr(
            utils.LayoutEnum.from_tensor(a_scale).mma_major_mode()
            != cute.nvgpu.OperandMajorMode.K
        ):
            raise ValueError("scale_major_mode must be K-major")

        self._setup_attributes()

        tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.mma_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.acc_dtype,
            self.cta_group,
            self.mma_tiler[:2],
            self.transform_a_source,
        )
        a_op = mixed_input_utils.get_tma_atom_kind(
            self.is_a_mcast, self.use_2cta_instrs, False
        )
        b_op = mixed_input_utils.get_tma_atom_kind(
            self.is_b_mcast, self.use_2cta_instrs, True
        )
        smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0))
        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
            a_op,
            a,
            smem_layout_a_per_stage,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=(
                cutlass.TFloat32 if a.element_type is cutlass.Float32 else None
            ),
        )

        scale_smem_layout_for_tma = cute.slice_(
            self.smem_layout_scale_tma, (None, None, None, 0)
        )
        tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A(
            a_op,
            a_scale,
            scale_smem_layout_for_tma,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=cutlass.Int16,
        )

        smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0))
        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
            b_op,
            b,
            smem_layout_b_per_stage,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=(
                cutlass.TFloat32 if b.element_type is cutlass.Float32 else None
            ),
        )

        a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage)
        b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage)
        a_scale_copy_size = cute.size_in_bytes(
            self.a_scale_dtype, self.smem_layout_scale_per_stage
        )

        self.num_tma_load_bytes_a_scale = a_copy_size + a_scale_copy_size
        self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape)
        self.tile_sched_params, grid = self._compute_grid(
            c,
            self.cta_tile_shape_mnk_c,
            self.cluster_shape_mn,
            max_active_clusters,
            self.use_clc_scheduler,
            self.raster_along_m,
        )

        epi_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,
            epi_smem_layout,
            self.epi_tile,
        )

        @cute.struct
        class SharedStorage:
            # Routed-tile scheduling metadata.
            tile_info: cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_info_stage]
            a_load2trans_full_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_load2trans_stage
            ]
            a_load2trans_empty_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_load2trans_stage
            ]
            a_trans2mma_full_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_trans2mma_stage
            ]
            a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_trans2mma_stage
            ]
            b_load2mma_full_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_load2trans_stage
            ]
            b_load2mma_empty_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_load2trans_stage
            ]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            tile_info_full_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_tile_info_stage
            ]
            tile_info_empty_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.num_tile_info_stage
            ]
            clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
            clc_response: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 4], 16]
            tmem_dealloc_mbar: cutlass.Int64
            tmem_holding_buf: cutlass.Int32

        self.shared_storage = SharedStorage

        self.kernel(
            tiled_mma,
            tma_atom_a,
            tma_tensor_a,
            tma_atom_scale,
            tma_tensor_scale,
            tma_atom_b,
            tma_tensor_b,
            tma_atom_c,
            tma_tensor_c,
            c,
            final_output,
            tile_idx_to_expert_idx,
            tile_idx_to_mn_limit,
            num_non_exiting_tiles,
            alpha,
            permuted_idx_to_expanded_idx,
            token_final_scales,
            self.group_count,
            self.cluster_layout_vmnk,
            self.smem_layout_a,
            self.smem_layout_scale,
            self.smem_layout_scale_tma,
            self.smem_layout_a_transform,
            self.smem_layout_b,
            self.c_smem_layout_staged,
            self.epi_tile,
            self.tile_sched_params,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=(*self.cluster_shape_mn, 1),
            min_blocks_per_mp=1,
            stream=stream,
            use_pdl=self.enable_pdl,
        )

    @cute.kernel
    def kernel(
        self,
        tiled_mma: cute.TiledMma,
        tma_atom_a: cute.CopyAtom,
        mA_mkl: cute.Tensor,
        tma_atom_s: cute.CopyAtom,
        mS_mkl: cute.Tensor,
        tma_atom_b: cute.CopyAtom,
        mB_nkl: cute.Tensor,
        tma_atom_c: cute.CopyAtom,
        mC_mnl: cute.Tensor,
        tensor_c: cute.Tensor,
        final_output: cute.Tensor,
        tile_idx_to_expert_idx: cute.Tensor,
        tile_idx_to_mn_limit: cute.Tensor,
        num_non_exiting_tiles: cute.Tensor,
        alpha: cute.Tensor,
        permuted_idx_to_expanded_idx: Optional[cute.Tensor],
        token_final_scales: Optional[cute.Tensor],
        group_count: cutlass.Constexpr[int],
        cluster_layout_vmnk: cute.Layout,
        a_smem_layout: cute.ComposedLayout,
        scale_smem_layout: cute.ComposedLayout,
        scale_tma_smem_layout: cute.Layout,
        a_smem_layout_transform: cute.ComposedLayout,
        b_smem_layout: cute.ComposedLayout,
        c_smem_layout_staged: cute.ComposedLayout,
        epi_tile: cute.Tile,
        tile_sched_params: Union[
            utils.ClcDynamicPersistentTileSchedulerParams,
            utils.PersistentTileSchedulerParams,
        ],
    ):
        """Run the persistent W4A16 grouped GEMM."""
        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
        tidx, _, _ = cute.arch.thread_idx()
        bidx, bidy, bidz = cute.arch.block_idx()
        # Prefetch TMA descriptors
        if warp_idx == self.epilog_warp_id[0]:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b)
            cpasync.prefetch_descriptor(tma_atom_s)
            cpasync.prefetch_descriptor(tma_atom_c)

        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
        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()
        )
        is_first_cta_in_cluster = cta_rank_in_cluster == 0
        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
            cta_rank_in_cluster
        )
        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)
        clc_response_ptr = storage.clc_response.data_ptr()

        # Initialize load2transform pipeline for the weight operand.
        transform_thread_idx = (
            tidx - 32 * self.transform_warp_id[0]
            if tidx >= 32 * self.transform_warp_id[0]
            else tidx
        )
        a_load2trans_pipeline = pipeline.PipelineTmaAsync.create(
            barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(),
            num_stages=self.num_load2trans_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.num_mcast_ctas_a * self.num_transform_warps,
            ),
            tx_count=self.num_tma_load_bytes_a_scale,
            cta_layout_vmnk=cluster_layout_vmnk,
            tidx=transform_thread_idx,
            mcast_mode_mn=(1, 0),  # multicast for A will only happen on the M-mode
            defer_sync=True,
        )
        # Initialize transform2mma pipeline, which tracks the dependencies between the transformation
        # of A and MMA's consumption of transformed A
        cta_v_size = cute.size(cluster_layout_vmnk, mode=[0])
        trans2mma_pipeline = pipeline.PipelineAsyncUmma.create(
            barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(),
            num_stages=self.num_trans2mma_stage,
            producer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                32 * self.num_transform_warps * cta_v_size,
            ),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )
        # Initialize pipeline for tensor B load to MMA
        # MMA warp informs TMA warp to proceed to load next tile of B tensor
        b_load2mma_pipeline = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(),
            num_stages=self.num_load2trans_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread, self.num_mcast_ctas_b
            ),
            tx_count=self.num_tma_load_bytes_b,
            cta_layout_vmnk=cluster_layout_vmnk,
            mcast_mode_mn=(0, 1),  # multicast for B will only happen on the N-mode
            defer_sync=True,
        )
        # Initialize accumulator pipeline, which tracks the dependencies between
        # MMA's computation of accumulators and epilogue warps' consumption of accumulators
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id)
            ),
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )
        # Initialize tile info pipeline, which tracks the dependencies between
        # tile scheduling warp and other warps
        # Every warp except the scheduler consumes routed-tile metadata.
        num_tile_info_pipeline_consumer_threads = self.threads_per_cta - 32
        tile_info_pipeline = pipeline.PipelineAsync.create(
            barrier_storage=storage.tile_info_full_mbar_ptr.data_ptr(),
            num_stages=self.num_tile_info_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * 1),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                num_tile_info_pipeline_consumer_threads,
            ),
            defer_sync=True,
        )
        clc_pipeline = None
        if cutlass.const_expr(self.use_clc_scheduler):
            clc_pipeline = pipeline.PipelineClcFetchAsync.create(
                barrier_storage=storage.clc_mbar_ptr.data_ptr(),
                num_stages=1,
                producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
                consumer_group=pipeline.CooperativeGroup(
                    pipeline.Agent.Thread,
                    32 * cute.size(self.cluster_shape_mn),
                ),
                tx_count=16,
                cta_layout_vmnk=cluster_layout_vmnk,
                defer_sync=True,
            )

        # Tensor memory dealloc barrier init
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=self.tmem_ptr_sync_barrier,
            allocator_warp_id=self.epilog_warp_id[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
        )

        # Cluster arrive after barrier init
        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)

        # Setup smem tensor A/scale/B/C
        sC = smem.allocate_tensor(
            element_type=self.c_dtype,
            layout=c_smem_layout_staged.outer,
            byte_alignment=self.smem_buffer_align_bytes,
            swizzle=c_smem_layout_staged.inner,
        )
        # Fused finalize reuses the first C stage as a linear
        # [route, hidden] tile for descriptor-free bulk reduction. The bulk
        # reduction consumes raw contiguous bytes, so drop sC's pointer swizzle.
        sFinalize = (
            cute.make_tensor(
                cute.recast_ptr(sC.iterator),
                cute.make_layout(
                    (self.epi_tile_n, self.cta_tile_shape_mnk[0]),
                    stride=(self.cta_tile_shape_mnk[0], 1),
                ),
            )
            if cutlass.const_expr(self.use_fused_finalize)
            else None
        )
        sA_input = smem.allocate_tensor(
            element_type=self.a_dtype,
            layout=a_smem_layout.outer,
            byte_alignment=self.smem_buffer_align_bytes,
            swizzle=a_smem_layout.inner,
        )
        sS_input = smem.allocate_tensor(
            element_type=self.a_scale_dtype,
            layout=scale_smem_layout.outer,
            byte_alignment=self.smem_buffer_align_bytes,
            swizzle=scale_smem_layout.inner,
        )
        sS_tma = cute.make_tensor(sS_input.iterator, scale_tma_smem_layout)
        sA_transform_input = sA_input
        sS_transform_input = sS_input
        if cutlass.const_expr(self.gated):
            # Pair 16-row up/gate slices during decode so each epilogue warp
            # owns both operands in its 32 TMEM lanes.
            gated_row_layout = cute.make_layout(((16, 2), 4), stride=((1, 64), 16))
            sA_transform_input = cute.composition(
                sA_input,
                ((gated_row_layout, None), None, None, None),
            )
            sS_transform_input = cute.composition(
                sS_input,
                ((gated_row_layout, None), None, None),
            )
        sB_input = smem.allocate_tensor(
            element_type=self.b_dtype,
            layout=b_smem_layout.outer,
            byte_alignment=self.smem_buffer_align_bytes,
            swizzle=b_smem_layout.inner,
        )
        sA_transform = None
        # Get smem tensor for transformed A when transform_a_source is SMEM
        if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM):
            sA_transform = smem.allocate_tensor(
                element_type=self.mma_dtype,
                layout=a_smem_layout_transform.outer,
                byte_alignment=self.smem_buffer_align_bytes,
                swizzle=a_smem_layout_transform.inner,
            )
        sFinalizeScale = (
            smem.allocate_tensor(
                element_type=cutlass.Float32,
                layout=cute.make_layout((self.epi_tile_n,)),
                byte_alignment=16,
            )
            if cutlass.const_expr(self.use_fused_finalize)
            else None
        )
        sTile_info = storage.tile_info.get_tensor(
            cute.make_layout((4, self.num_tile_info_stage), stride=(1, 4))
        )

        # Compute multicast mask for A/B buffer full
        a_full_mcast_mask = None
        b_full_mcast_mask = None
        s_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
            )
            # Scale tensor shares the same multicast mask as the A tensor
            s_full_mcast_mask = a_full_mcast_mask
            b_full_mcast_mask = cpasync.create_tma_multicast_mask(
                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
            )

        # local_tile partition global tensors
        # (bM, bK, loopM, loopK, loopL)
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
        )
        # (bM, bK, loopM, loopK, loopL)
        gS_mkl = cute.local_tile(
            mS_mkl,
            cute.slice_(self.mma_tiler, (None, 0, None)),
            (None, None, None),
        )
        # (bN, bK, loopN, loopK, loopL)
        gB_nkl = cute.local_tile(
            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
        )
        # (bM, bN, loopM, loopN, loopL)
        gC_mnl = cute.local_tile(
            mC_mnl,
            cute.slice_(self.mma_tiler_c, (None, None, 0)),
            (None, None, None),
        )
        gC_mnl_simt = cute.local_tile(
            tensor_c,
            cute.slice_(self.mma_tiler_c, (None, None, 0)),
            (None, None, None),
        )
        k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))

        # Partition global tensor for TiledMMA_A/B/C
        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
        # (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
        tCgA = thr_mma.partition_A(gA_mkl)
        # (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
        tCgS = thr_mma.partition_A(gS_mkl)
        # (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
        tCgB = thr_mma.partition_B(gB_nkl)
        # (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
        tCgC = thr_mma.partition_C(gC_mnl)
        tCgC_simt = thr_mma.partition_C(gC_mnl_simt)

        # Setup copy atom to load A from shared memory for further transformation
        copy_atom_a_input = cute.make_copy_atom(
            cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32
        )
        # Split complete NVFP4 blocks between two transform warpgroups.
        copy_atom_a_transform = cute.make_copy_atom(
            tcgen05.St32x32bOp(
                tcgen05.Repetition(8),
                tcgen05.Unpack.NONE,
            ),
            self.mma_dtype,
        )

        # Partition global/shared tensor for TMA load A/B
        # TMA load A partition_S/D
        a_cta_layout = cute.make_layout(
            cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
        )
        # ((atom_v, rest_v), STAGE)
        # ((atom_v, rest_v), loopM, loopK, loopL)
        tAsA, tAgA = cpasync.tma_partition(
            tma_atom_a,
            block_in_cluster_coord_vmnk[2],
            a_cta_layout,
            cute.group_modes(sA_input, 0, 3),
            cute.group_modes(tCgA, 0, 3),
        )

        thr_mma_leader_cta = tiled_mma.get_slice(0)
        # (MMA, MMA_M, MMA_K, STAGE)
        tCsS = thr_mma_leader_cta.partition_A(sS_transform_input)
        tSsS, tSgS = cpasync.tma_partition(
            tma_atom_s,
            block_in_cluster_coord_vmnk[2],
            a_cta_layout,
            cute.group_modes(sS_tma, 0, 3),
            cute.group_modes(tCgS, 0, 3),
        )
        tSsS = cute.filter_zeros(tSsS)
        tSgS = cute.filter_zeros(tSgS)

        # TMA load B partition_S/D
        b_cta_layout = cute.make_layout(
            cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
        )
        # ((atom_v, rest_v), STAGE)
        # ((atom_v, rest_v), loopM, loopK, loopL)
        tBsB, tBgB = cpasync.tma_partition(
            tma_atom_b,
            block_in_cluster_coord_vmnk[1],
            b_cta_layout,
            cute.group_modes(sB_input, 0, 3),
            cute.group_modes(tCgB, 0, 3),
        )

        # (MMA, MMA_N, MMA_K, STAGE)
        tCrB = tiled_mma.make_fragment_B(sB_input)
        # (MMA, MMA_M, MMA_N)
        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
        tCtAcc_fake = tiled_mma.make_fragment_C(
            cute.append(acc_shape, self.num_acc_stage)
        )

        # Cluster wait before TMEM alloc and ensure pipelines are ready.
        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)

        # TMEM allocation
        tmem.allocate(self.num_tmem_alloc_cols)
        tmem.wait_for_alloc()

        # Tile scheduling is independent of the preceding grid.
        if warp_idx == self.schedule_warp_id:
            cute.arch.setmaxregister_decrease(self.num_regs_generic_warps)
            tile_info_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_tile_info_stage
            )
            num_non_exiting_tiles_value = num_non_exiting_tiles[0]
            if cutlass.const_expr(self.use_clc_scheduler):
                tile_sched = utils.ClcDynamicPersistentTileScheduler.create(
                    tile_sched_params,
                    (bidx, bidy, bidz),
                    cute.arch.grid_dim(),
                    clc_response_ptr,
                )
                work_tile = tile_sched.initial_work_tile_info()
                clc_producer_state = pipeline.make_pipeline_state(
                    pipeline.PipelineUserType.ProducerConsumer, 1
                )
                clc_consumer_state = pipeline.make_pipeline_state(
                    pipeline.PipelineUserType.Consumer, 1
                )
                not_last_tile = cutlass.Boolean(1)
                while not_last_tile:
                    tile_coord_mnl = work_tile.tile_idx
                    route_tile_idx = tile_coord_mnl[1]
                    is_routed_tile = work_tile.is_valid_tile and (
                        route_tile_idx < num_non_exiting_tiles_value
                    )

                    if is_routed_tile or not work_tile.is_valid_tile:
                        tile_info_pipeline.producer_acquire(tile_info_producer_state)

                    if work_tile.is_valid_tile and is_first_cta_in_cluster:
                        clc_pipeline.producer_acquire(clc_producer_state)
                        tile_sched.advance_to_next_work(
                            clc_pipeline.producer_get_barrier(clc_producer_state)
                        )
                        clc_producer_state.advance()

                    if is_routed_tile or not work_tile.is_valid_tile:
                        cur_sTile_info = sTile_info[
                            (None, tile_info_producer_state.index)
                        ]
                        route_start = cutlass.Int32(-1)
                        group_idx = cutlass.Int32(group_count)
                        distance_to_boundary = cutlass.Int32(-1)
                        if is_routed_tile:
                            route_start = route_tile_idx * self.cta_tile_shape_mnk[1]
                            group_idx = tile_idx_to_expert_idx[route_tile_idx]
                            distance_to_boundary = (
                                tile_idx_to_mn_limit[route_tile_idx] - route_start
                            )
                        with cute.arch.elect_one():
                            cur_sTile_info[0] = tile_coord_mnl[0]
                            cur_sTile_info[1] = route_start
                            cur_sTile_info[2] = group_idx
                            cur_sTile_info[3] = distance_to_boundary
                        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()

                    if work_tile.is_valid_tile:
                        clc_pipeline.consumer_wait(clc_consumer_state)
                        work_tile = tile_sched.get_current_work()
                        clc_pipeline.consumer_release(clc_consumer_state)
                        clc_consumer_state.advance()
                    else:
                        not_last_tile = cutlass.Boolean(0)

                tile_info_pipeline.producer_tail(tile_info_producer_state)
                if is_first_cta_in_cluster:
                    clc_pipeline.producer_tail(clc_producer_state)
            else:
                tile_sched = utils.StaticPersistentRuntimeTileScheduler.create(
                    tile_sched_params,
                    (bidx, bidy, bidz),
                    cute.arch.grid_dim(),
                    inner_mode=0 if self.raster_along_m else 1,
                )
                work_tile = tile_sched.initial_work_tile_info()
                not_last_tile = cutlass.Boolean(1)
                if cutlass.const_expr(self.raster_along_m):
                    while not_last_tile:
                        tile_info_pipeline.producer_acquire(tile_info_producer_state)
                        cluster_tile_coord_mnl = work_tile.tile_idx
                        cta_tile_coord_m = (
                            cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0]
                            + block_in_cluster_coord_vmnk[1]
                            * cute.size(tiled_mma.thr_id.shape)
                            + block_in_cluster_coord_vmnk[0]
                        )
                        cta_tile_offset_n = block_in_cluster_coord_vmnk[2]
                        route_tile_idx = (
                            cluster_tile_coord_mnl[1] * self.cluster_shape_mn[1]
                            + cta_tile_offset_n
                        )
                        not_last_tile = work_tile.is_valid_tile and (
                            route_tile_idx < num_non_exiting_tiles_value
                        )
                        route_start = cutlass.Int32(-1)
                        group_idx = cutlass.Int32(group_count)
                        distance_to_boundary = cutlass.Int32(-1)
                        if not_last_tile:
                            route_start = route_tile_idx * self.cta_tile_shape_mnk[1]
                            group_idx = tile_idx_to_expert_idx[route_tile_idx]
                            distance_to_boundary = (
                                tile_idx_to_mn_limit[route_tile_idx] - route_start
                            )
                        cur_sTile_info = sTile_info[
                            (None, tile_info_producer_state.index)
                        ]
                        with cute.arch.elect_one():
                            cur_sTile_info[0] = cta_tile_coord_m
                            cur_sTile_info[1] = route_start
                            cur_sTile_info[2] = group_idx
                            cur_sTile_info[3] = distance_to_boundary
                        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_sched.advance_to_next_work()
                        work_tile = tile_sched.get_current_work()
                else:
                    while not_last_tile:
                        cluster_tile_coord_mnl = work_tile.tile_idx
                        cta_tile_coord_m = (
                            cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0]
                            + block_in_cluster_coord_vmnk[1]
                            * cute.size(tiled_mma.thr_id.shape)
                            + block_in_cluster_coord_vmnk[0]
                        )
                        cta_tile_offset_n = block_in_cluster_coord_vmnk[2]
                        route_tile_idx = (
                            cluster_tile_coord_mnl[1] * self.cluster_shape_mn[1]
                            + cta_tile_offset_n
                        )
                        # The runtime scheduler leaves its outer mode unbounded.
                        # With N as the inner mode, terminate on the finite M
                        # extent and skip padded routed tiles within each M tile.
                        has_work = work_tile.is_valid_tile and (
                            cluster_tile_coord_mnl[0]
                            < tile_sched_params.problem_layout_ncluster_mnl.shape[0]
                        )
                        is_routed_tile = has_work and (
                            route_tile_idx < num_non_exiting_tiles_value
                        )

                        if is_routed_tile or not has_work:
                            tile_info_pipeline.producer_acquire(
                                tile_info_producer_state
                            )
                            route_start = cutlass.Int32(-1)
                            group_idx = cutlass.Int32(group_count)
                            distance_to_boundary = cutlass.Int32(-1)
                            if is_routed_tile:
                                route_start = (
                                    route_tile_idx * self.cta_tile_shape_mnk[1]
                                )
                                group_idx = tile_idx_to_expert_idx[route_tile_idx]
                                distance_to_boundary = (
                                    tile_idx_to_mn_limit[route_tile_idx] - route_start
                                )
                            cur_sTile_info = sTile_info[
                                (None, tile_info_producer_state.index)
                            ]
                            with cute.arch.elect_one():
                                cur_sTile_info[0] = cta_tile_coord_m
                                cur_sTile_info[1] = route_start
                                cur_sTile_info[2] = group_idx
                                cur_sTile_info[3] = distance_to_boundary
                            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()

                        if has_work:
                            tile_sched.advance_to_next_work()
                            work_tile = tile_sched.get_current_work()
                        else:
                            not_last_tile = cutlass.Boolean(0)

                tile_info_pipeline.producer_tail(tile_info_producer_state)

        # The activation TMA warp is the only role that waits on the preceding grid.
        if warp_idx == self.activation_tma_warp_id:
            cute.arch.setmaxregister_decrease(self.num_regs_generic_warps)
            tile_info_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_tile_info_stage
            )
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                group_count, sTile_info[(None, tile_info_consumer_state.index)]
            )
            cute.arch.fence_proxy(
                "async.shared",
                space="cta",
            )
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()
            b_load2mma_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_load2trans_stage
            )
            if cutlass.const_expr(self.enable_pdl):
                griddepcontrol_wait()

            while work_tile.is_valid_tile:
                coord_n_offset = (
                    (work_tile.coord_n, 0)
                    if cutlass.const_expr(
                        self.b_major_mode == cute.nvgpu.OperandMajorMode.MN
                    )
                    else (0, work_tile.coord_n)
                )
                tBgB_slice = cute.make_tensor(
                    (
                        tBgB.iterator[0] + coord_n_offset[0],
                        coord_n_offset[1] + tBgB.iterator[1],
                    ),
                    cute.slice_(tBgB.layout, (None, 0, None, 0)),
                )

                b_load2mma_producer_state.reset_count()
                for _k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state)
                    cute.copy(
                        tma_atom_b,
                        tBgB_slice[(None, b_load2mma_producer_state.count)],
                        tBsB[(None, b_load2mma_producer_state.index)],
                        tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier(
                            b_load2mma_producer_state
                        ),
                        mcast_mask=b_full_mcast_mask,
                        cache_policy=cutlass.Int64(_TMA_CACHE_HINT_EVICT_LAST),
                    )
                    b_load2mma_pipeline.producer_commit(b_load2mma_producer_state)
                    b_load2mma_producer_state.advance()

                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                    group_count, sTile_info[(None, tile_info_consumer_state.index)]
                )
                cute.arch.fence_proxy(
                    "async.shared",
                    space="cta",
                )
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()
            b_load2mma_pipeline.producer_tail(b_load2mma_producer_state)

        # Static weight and scale loads can run before the activation is ready.
        if warp_idx == self.weight_tma_warp_id:
            cute.arch.setmaxregister_decrease(self.num_regs_generic_warps)
            tile_info_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_tile_info_stage
            )
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                group_count, sTile_info[(None, tile_info_consumer_state.index)]
            )
            cute.arch.fence_proxy(
                "async.shared",
                space="cta",
            )
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()
            a_load2trans_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_load2trans_stage
            )
            if cutlass.const_expr(self.enable_pdl and self.fuse_activation):
                with cute.arch.elect_one():
                    griddepcontrol_launch_dependents()
            while work_tile.is_valid_tile:
                tAgA_slice = tAgA[
                    (
                        None,
                        work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape),
                        None,
                        work_tile.group_idx,
                    )
                ]
                tSgS_slice = tSgS[
                    (
                        None,
                        work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape),
                        None,
                        work_tile.group_idx,
                    )
                ]
                # Filter zeros in rest mode
                rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout)
                tSgS_slice_filtered = cute.make_tensor(
                    tSgS_slice.iterator,
                    cute.make_layout(
                        (tSgS_slice.layout[0].shape, rest_filtered.shape),
                        stride=(tSgS_slice.layout[0].stride, rest_filtered.stride),
                    ),
                )

                a_load2trans_producer_state.reset_count()
                peek_load2trans_empty_status = cutlass.Boolean(1)
                if a_load2trans_producer_state.count < k_tile_cnt:
                    peek_load2trans_empty_status = (
                        a_load2trans_pipeline.producer_try_acquire(
                            a_load2trans_producer_state
                        )
                    )
                for _k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    a_load2trans_pipeline.producer_acquire(
                        a_load2trans_producer_state,
                        peek_load2trans_empty_status,
                    )
                    cute.copy(
                        tma_atom_a,
                        tAgA_slice[(None, a_load2trans_producer_state.count)],
                        tAsA[(None, a_load2trans_producer_state.index)],
                        tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier(
                            a_load2trans_producer_state
                        ),
                        mcast_mask=a_full_mcast_mask,
                        cache_policy=cutlass.Int64(_TMA_CACHE_HINT_EVICT_FIRST),
                    )
                    cute.copy(
                        tma_atom_s,
                        tSgS_slice_filtered[(None, a_load2trans_producer_state.count)],
                        tSsS[(None, a_load2trans_producer_state.index)],
                        tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier(
                            a_load2trans_producer_state
                        ),
                        mcast_mask=s_full_mcast_mask,
                    )
                    a_load2trans_pipeline.producer_commit(a_load2trans_producer_state)
                    a_load2trans_producer_state.advance()
                    if a_load2trans_producer_state.count < k_tile_cnt:
                        peek_load2trans_empty_status = (
                            a_load2trans_pipeline.producer_try_acquire(
                                a_load2trans_producer_state
                            )
                        )

                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                    group_count, sTile_info[(None, tile_info_consumer_state.index)]
                )
                cute.arch.fence_proxy(
                    "async.shared",
                    space="cta",
                )
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()

            a_load2trans_pipeline.producer_tail(a_load2trans_producer_state)

        # Specialized transform warps
        if warp_idx >= self.transform_warp_id[0]:
            transform_group_idx = (warp_idx - self.transform_warp_id[0]) // 4
            transform_local_tidx = tidx - 32 * (
                self.transform_warp_id[0] + transform_group_idx * 4
            )
            # Get the pointer to the TMEM buffer
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)

            (
                dst_copy_a,
                tAsA_input,
                tAsA_transform,
                tSsS_trans,
                tSrS_copy,
                tSrS_load,
                tArA_load,
                tArA_transform,
                tArA_transform_store,
                transform_tiler,
            ) = self._setup_transform_partitions(
                tiled_mma,
                copy_atom_a_input,
                copy_atom_a_transform,
                sA_transform_input,
                a_smem_layout_transform,
                sA_transform,
                tCsS,
                accumulators,
                transform_local_tidx,
                transform_group_idx,
            )

            tile_info_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_tile_info_stage
            )
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                group_count, sTile_info[(None, tile_info_consumer_state.index)]
            )
            cute.arch.fence_proxy(
                "async.shared",
                space="cta",
            )
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()
            a_load2trans_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer,
                self.num_load2trans_stage,
            )
            trans2mma_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer,
                self.num_trans2mma_stage,
            )
            while work_tile.is_valid_tile:
                (
                    a_load2trans_consumer_state,
                    trans2mma_producer_state,
                ) = self._transform_tile(
                    a_load2trans_pipeline,
                    trans2mma_pipeline,
                    a_load2trans_consumer_state,
                    trans2mma_producer_state,
                    dst_copy_a,
                    tAsA_input,
                    tAsA_transform,
                    tSsS_trans,
                    tSrS_copy,
                    tSrS_load,
                    tArA_load,
                    tArA_transform,
                    tArA_transform_store,
                    transform_tiler,
                    k_tile_cnt,
                )
                # Advance to next tile
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                    group_count, sTile_info[(None, tile_info_consumer_state.index)]
                )
                cute.arch.fence_proxy(
                    "async.shared",
                    space="cta",
                )
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()
            # Wait a_transform buffer empty
            trans2mma_pipeline.producer_tail(trans2mma_producer_state)

        # Specialized MMA warp
        if warp_idx == self.mma_warp_id:
            cute.arch.setmaxregister_decrease(self.num_regs_generic_warps)
            # Get the pointer to the TMEM buffer
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
            if cutlass.const_expr(
                self.transform_a_source == tcgen05.OperandSource.TMEM
            ):
                tmem_ptr_transform = cute.recast_ptr(
                    accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype
                )
                tCrA = cute.make_tensor(
                    tmem_ptr_transform,
                    tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout,
                )
            else:
                tCrA = tiled_mma.make_fragment_A(sA_transform)
            tCtAcc_base = accumulators
            # Persistent tile scheduling loop
            tile_info_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_tile_info_stage
            )
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                group_count, sTile_info[(None, tile_info_consumer_state.index)]
            )
            cute.arch.fence_proxy(
                "async.shared",
                space="cta",
            )
            tile_info_pipeline.consumer_release(tile_info_consumer_state)
            tile_info_consumer_state.advance()
            trans2mma_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage
            )
            b_load2mma_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_load2trans_stage
            )
            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage
            )
            while work_tile.is_valid_tile:
                # (MMA, MMA_M, MMA_N)
                tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
                b_load2mma_consumer_state.reset_count()
                trans2mma_consumer_state.reset_count()
                peek_trans2mma_full_status = cutlass.Boolean(1)
                if is_leader_cta:
                    if trans2mma_consumer_state.count < k_tile_cnt:
                        peek_trans2mma_full_status = (
                            trans2mma_pipeline.consumer_try_wait(
                                trans2mma_consumer_state
                            )
                        )
                    acc_pipeline.producer_acquire(acc_producer_state)

                    tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
                    # Mma mainloop
                    for _k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                        trans2mma_pipeline.consumer_wait(
                            trans2mma_consumer_state, peek_trans2mma_full_status
                        )
                        b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state)
                        num_kblocks = cute.size(tCrA, mode=[2])
                        for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
                            kblock_coord_a = (
                                None,
                                None,
                                kblock_idx,
                                trans2mma_consumer_state.index,
                            )
                            kblock_coord_b = (
                                None,
                                None,
                                kblock_idx,
                                b_load2mma_consumer_state.index,
                            )

                            cute.gemm(
                                tiled_mma,
                                tCtAcc,
                                tCrA[kblock_coord_a],
                                tCrB[kblock_coord_b],
                                tCtAcc,
                            )
                            # Enable accumulate on tCtAcc after first kblock
                            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                        trans2mma_pipeline.consumer_release(trans2mma_consumer_state)
                        b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state)
                        trans2mma_consumer_state.advance()
                        b_load2mma_consumer_state.advance()
                        peek_trans2mma_full_status = cutlass.Boolean(1)
                        if trans2mma_consumer_state.count < k_tile_cnt:
                            peek_trans2mma_full_status = (
                                trans2mma_pipeline.consumer_try_wait(
                                    trans2mma_consumer_state
                                )
                            )
                    # Async arrive accumulator buffer full
                    acc_pipeline.producer_commit(acc_producer_state)
                acc_producer_state.advance()

                # Advance to next tile
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                    group_count, sTile_info[(None, tile_info_consumer_state.index)]
                )
                cute.arch.fence_proxy(
                    "async.shared",
                    space="cta",
                )
                tile_info_pipeline.consumer_release(tile_info_consumer_state)
                tile_info_consumer_state.advance()
            # Wait for accumulator buffer empty
            acc_pipeline.producer_tail(acc_producer_state)

        # Specialized epilogue warps
        if warp_idx < self.mma_warp_id:
            cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps)
            epi_tidx = tidx
            # Get the pointer to the TMEM buffer
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
            tCtAcc_base = accumulators
            # Partition for epilogue
            tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = (
                mixed_input_utils.epilog_tmem_copy_and_partition(
                    self.cta_tile_shape_mnk,
                    self.c_layout,
                    self.c_dtype,
                    self.acc_dtype,
                    epi_tidx,
                    tCtAcc_base,
                    tCgC,
                    epi_tile,
                    self.use_2cta_instrs,
                )
            )
            gated_tTR_tAcc_base = None
            gated_tiled_copy_t2r = None
            gated_tTR_identity_base = None
            gated_tTR_rAcc = None
            if cutlass.const_expr(self.gated):
                gated_epi_tile = (
                    self.cta_tile_shape_mnk[0],
                    self.epi_tile_n,
                )
                gated_copy_atom_t2r = sm100_utils.get_tmem_load_op(
                    self.cta_tile_shape_mnk,
                    self.c_layout,
                    self.c_dtype,
                    self.acc_dtype,
                    gated_epi_tile,
                    self.use_2cta_instrs,
                )
                gated_tAcc_epi = cute.flat_divide(
                    tCtAcc_base[((None, None), 0, 0, None)],
                    gated_epi_tile,
                )
                gated_tiled_copy_t2r = tcgen05.make_tmem_copy(
                    gated_copy_atom_t2r,
                    gated_tAcc_epi[(None, None, 0, 0, 0)],
                )
                gated_thr_copy_t2r = gated_tiled_copy_t2r.get_slice(epi_tidx)
                gated_tTR_tAcc_base = gated_thr_copy_t2r.partition_S(gated_tAcc_epi)
                gated_identity = cute.make_identity_tensor(self.cta_tile_shape_mnk[:2])
                gated_identity_epi = cute.flat_divide(gated_identity, gated_epi_tile)
                gated_tTR_identity_base = gated_thr_copy_t2r.partition_D(
                    gated_identity_epi
                )
                gated_tTR_rAcc = cute.make_rmem_tensor(
                    gated_tTR_identity_base[(None, None, None, 0, 0)].shape,
                    self.acc_dtype,
                )
            tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
            tiled_copy_r2s, tRS_rC, tRS_sC = (
                mixed_input_utils.epilog_smem_copy_and_partition(
                    self.c_layout,
                    self.c_dtype,
                    self.acc_dtype,
                    tiled_copy_t2r,
                    tTR_rC,
                    epi_tidx,
                    sC,
                )
            )
            (
                tma_atom_c,
                bSG_sC,
                bSG_gC_partitioned,
                simt_atom,
                tTR_gC_partitioned,
            ) = mixed_input_utils.epilog_gmem_copy_and_partition(
                self.c_dtype,
                epi_tidx,
                tma_atom_c,
                tiled_copy_t2r,
                tCgC,
                tCgC_simt,
                epi_tile,
                sC,
            )
            if cutlass.const_expr(self.gated):
                gC_epi = cute.local_tile(mC_mnl, epi_tile, (None, None, None))
                bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
                    tma_atom_c,
                    0,
                    cute.make_layout(1),
                    cute.group_modes(sC, 0, 2),
                    cute.group_modes(gC_epi, 0, 2),
                )

            # Predicates
            thr_mapping = cute.make_identity_tensor(
                (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1])
            )
            thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile)
            thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx)
            m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn)
            m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset))

            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage
            )

            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                32 * len(self.epilog_warp_id),
            )
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage,
                producer_group=c_producer_group,
            )

            # Persistent tile scheduling loop
            tile_info_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_tile_info_stage
            )
            tile_info_pipeline.consumer_wait(tile_info_consumer_state)
            work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                group_count, sTile_info[(None, tile_info_consumer_state.index)]
            )
            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)
            while work_tile.is_valid_tile:
                alpha_val = alpha[work_tile.group_idx]
                if cutlass.const_expr(self.gated):
                    bSG_gC = bSG_gC_partitioned[(None, work_tile.cta_coord_m, None, 0)]
                else:
                    bSG_gC = bSG_gC_partitioned[
                        (
                            None,
                            None,
                            None,
                            work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape),
                            0,
                            0,
                        )
                    ]
                    tma_store_offset_coord = (
                        (work_tile.coord_n, 0)
                        if cutlass.const_expr(self.c_layout.is_n_major_c())
                        else (0, work_tile.coord_n)
                    )
                    bSG_gC = cute.make_tensor(
                        (
                            tma_store_offset_coord[0] + bSG_gC.iterator[0],
                            tma_store_offset_coord[1] + bSG_gC.iterator[1],
                        ),
                        bSG_gC.layout,
                    )
                tTR_gC = tTR_gC_partitioned[
                    (
                        None,
                        None,
                        None,
                        None,
                        None,
                        work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape),
                        0,
                        0,
                    )
                ]
                tTR_gC = cute.make_tensor(
                    tTR_gC.iterator + (work_tile.coord_n * tensor_c.layout.stride[1]),
                    tTR_gC.layout,
                )

                tTR_tAcc = tTR_tAcc_base[
                    (None, None, None, None, None, acc_consumer_state.index)
                ]
                # Wait for accumulator buffer full
                acc_pipeline.consumer_wait(acc_consumer_state)
                if cutlass.const_expr(not self.gated):
                    bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))

                if cutlass.const_expr(self.gated):
                    gated_tTR_tAcc = gated_tTR_tAcc_base[
                        (None, None, None, None, None, acc_consumer_state.index)
                    ]
                    gated_subtile_cnt = cute.size(gated_tTR_tAcc.shape, mode=[4])
                    gated_alpha_f32 = cutlass.Float32(alpha_val)
                    swiglu_alpha = cutlass.Float32(self.swiglu_alpha)
                    swiglu_beta = cutlass.Float32(self.swiglu_beta)
                    swiglu_limit = cutlass.Float32(self.swiglu_limit)
                    neg_swiglu_alpha_log2_e = -(
                        swiglu_alpha * cutlass.Float32(1.4426950408889634)
                    )
                    for gated_subtile_idx in cutlass.range(gated_subtile_cnt):
                        gated_tTR_tAcc_mn = gated_tTR_tAcc[
                            (None, None, None, 0, gated_subtile_idx)
                        ]
                        cute.copy(
                            gated_tiled_copy_t2r,
                            gated_tTR_tAcc_mn,
                            gated_tTR_rAcc,
                        )
                        if gated_subtile_idx == gated_subtile_cnt - 1:
                            cute.arch.fence_view_async_tmem_load()
                            with cute.arch.elect_one():
                                acc_pipeline.consumer_release(acc_consumer_state)
                            acc_consumer_state.advance()
                        num_prev_subtiles += 1
                        c_buffer = num_prev_subtiles % self.num_c_stage
                        gated_identity_mn = gated_tTR_identity_base[
                            (None, None, None, 0, gated_subtile_idx)
                        ]
                        up = cute.coalesce(
                            cute.flatten(gated_tTR_rAcc[(None, 0, None)])
                        ).load()
                        gate = cute.coalesce(
                            cute.flatten(gated_tTR_rAcc[(None, 1, None)])
                        ).load()
                        up_coords = cute.coalesce(
                            cute.flatten(gated_identity_mn[(None, 0, None)])
                        )
                        up_size = cute.size(gated_tTR_rAcc[(None, 0, None)])
                        for i in cutlass.range_constexpr(0, up_size, 2):
                            up_pair = cute.arch.mul_packed_f32x2(
                                (up[i], up[i + 1]),
                                (gated_alpha_f32, gated_alpha_f32),
                            )
                            gate_pair = cute.arch.mul_packed_f32x2(
                                (gate[i], gate[i + 1]),
                                (gated_alpha_f32, gated_alpha_f32),
                            )
                            if cutlass.const_expr(self.parameterized_swiglu):
                                gate_pair = (
                                    fmin(gate_pair[0], swiglu_limit, nan=True),
                                    fmin(gate_pair[1], swiglu_limit, nan=True),
                                )
                                up_pair = (
                                    -fmin(
                                        -fmin(up_pair[0], swiglu_limit, nan=True),
                                        swiglu_limit,
                                        nan=True,
                                    ),
                                    -fmin(
                                        -fmin(up_pair[1], swiglu_limit, nan=True),
                                        swiglu_limit,
                                        nan=True,
                                    ),
                                )
                            gate_log2e = cute.arch.mul_packed_f32x2(
                                gate_pair,
                                (
                                    neg_swiglu_alpha_log2_e,
                                    neg_swiglu_alpha_log2_e,
                                ),
                            )
                            sigmoid = cute.arch.add_packed_f32x2(
                                (
                                    cute.math.exp2(gate_log2e[0], fastmath=True),
                                    cute.math.exp2(gate_log2e[1], fastmath=True),
                                ),
                                (1.0, 1.0),
                            )
                            sigmoid = (
                                cute.arch.rcp_approx(sigmoid[0]),
                                cute.arch.rcp_approx(sigmoid[1]),
                            )
                            silu = cute.arch.mul_packed_f32x2(gate_pair, sigmoid)
                            if cutlass.const_expr(self.parameterized_swiglu):
                                up_pair = cute.arch.add_packed_f32x2(
                                    up_pair, (swiglu_beta, swiglu_beta)
                                )
                            result = cute.arch.mul_packed_f32x2(up_pair, silu)

                            coord_0 = up_coords[i]
                            coord_1 = up_coords[i + 1]
                            # Fold paired TMEM rows back into the reduced-M
                            # output tile and keep the staged N coordinate local.
                            output_m_0 = coord_0[0] // 32 * 16 + coord_0[0] % 16
                            output_m_1 = coord_1[0] // 32 * 16 + coord_1[0] % 16
                            output_n_0 = coord_0[1] % self.epi_tile_n
                            output_n_1 = coord_1[1] % self.epi_tile_n
                            sC[(output_m_0, output_n_0, c_buffer)] = self.c_dtype(
                                result[0]
                            )
                            sC[(output_m_1, output_n_1, c_buffer)] = self.c_dtype(
                                result[1]
                            )

                        cute.arch.fence_proxy("async.shared", space="cta")
                        self.epilog_sync_barrier.arrive_and_wait()
                        if warp_idx == self.epilog_warp_id[0]:
                            cute.copy(
                                tma_atom_c,
                                bSG_sC[(None, c_buffer)],
                                bSG_gC[
                                    (
                                        None,
                                        work_tile.coord_n // self.epi_tile_n
                                        + gated_subtile_idx,
                                    )
                                ],
                            )
                            c_pipeline.producer_commit()
                            c_pipeline.producer_acquire()
                        self.epilog_sync_barrier.arrive_and_wait()

                tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
                tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC))
                tma_distance_to_boundary = work_tile.distance_to_boundary

                # Store accumulator to global memory in subtiles
                subtile_cnt = (
                    0
                    if cutlass.const_expr(self.gated)
                    else cute.size(tTR_tAcc.shape, mode=[3])
                )
                for subtile_idx in cutlass.range(subtile_cnt):
                    # Load accumulator from tensor memory buffer to register
                    tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
                    cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
                    if subtile_idx == subtile_cnt - 1:
                        cute.arch.fence_view_async_tmem_load()
                        with cute.arch.elect_one():
                            acc_pipeline.consumer_release(acc_consumer_state)
                        acc_consumer_state.advance()

                    if cutlass.const_expr(self.fuse_activation):
                        alpha_f32 = cutlass.Float32(alpha_val)
                        acc_up = tTR_rAcc.load()
                        for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2):
                            value = cute.arch.mul_packed_f32x2(
                                (acc_up[i], acc_up[i + 1]),
                                (alpha_f32, alpha_f32),
                            )
                            relu = (
                                cute.arch.fmax(value[0], cutlass.Float32(0.0)),
                                cute.arch.fmax(value[1], cutlass.Float32(0.0)),
                            )
                            tTR_rAcc[i], tTR_rAcc[i + 1] = cute.arch.mul_packed_f32x2(
                                relu, relu
                            )
                    if cutlass.const_expr(self.use_fused_finalize):
                        finalize_thr_slice = m_thr_offset[
                            (None, None, None, subtile_idx)
                        ]
                        top_k = token_final_scales.shape[1]
                        finalize_scale_route = epi_tidx
                        finalize_scale_route_in_tile = (
                            subtile_idx * self.epi_tile_n + finalize_scale_route
                        )
                        if finalize_scale_route < self.epi_tile_n:
                            finalize_scale_permuted_row = (
                                work_tile.coord_n + finalize_scale_route_in_tile
                            )
                            finalize_scale_expanded_idx = permuted_idx_to_expanded_idx[
                                finalize_scale_permuted_row
                            ]
                            finalize_scale_safe_idx = cutlass.max(
                                finalize_scale_expanded_idx, cutlass.Int32(0)
                            )
                            finalize_scale_token_idx = finalize_scale_safe_idx // top_k
                            finalize_scale_topk_idx = finalize_scale_safe_idx % top_k
                            finalize_scale_is_valid = cutlass.Int32(
                                finalize_scale_route_in_tile
                                < work_tile.distance_to_boundary
                            )
                            finalize_token_scale = token_final_scales[
                                (
                                    finalize_scale_token_idx * finalize_scale_is_valid,
                                    finalize_scale_topk_idx,
                                )
                            ]
                            sFinalizeScale[finalize_scale_route] = (
                                cutlass.Float32(alpha_val)
                                * cutlass.Float32(finalize_token_scale)
                                * cutlass.Float32(finalize_scale_is_valid)
                            )

                        cute.arch.fence_proxy("async.shared", space="cta")
                        self.epilog_sync_barrier.arrive_and_wait()
                        for i in cutlass.range(cute.size(tTR_rC), unroll_full=True):
                            finalize_route = finalize_thr_slice[(i)][1]
                            finalize_value = sFinalizeScale[
                                finalize_route % self.epi_tile_n
                            ] * cutlass.Float32(tTR_rAcc[i])
                            sFinalize[
                                (
                                    finalize_route % self.epi_tile_n,
                                    finalize_thr_slice[(i)][0],
                                )
                            ] = self.c_dtype(finalize_value)

                        cute.arch.fence_proxy("async.shared", space="cta")
                        self.epilog_sync_barrier.arrive_and_wait()

                        reduce_route = epi_tidx
                        reduce_route_in_tile = (
                            subtile_idx * self.epi_tile_n + reduce_route
                        )
                        if (
                            reduce_route < self.epi_tile_n
                            and reduce_route_in_tile < work_tile.distance_to_boundary
                        ):
                            reduce_permuted_row = (
                                work_tile.coord_n + reduce_route_in_tile
                            )
                            reduce_expanded_idx = permuted_idx_to_expanded_idx[
                                reduce_permuted_row
                            ]
                            reduce_token_idx = reduce_expanded_idx // top_k
                            hidden_base = (
                                work_tile.cta_coord_m * self.cta_tile_shape_mnk[0]
                            )
                            valid_elements = (
                                cutlass.Int64(final_output.shape[0]) - hidden_base
                            )
                            if valid_elements > 0:
                                scatter_out = cute.domain_offset(
                                    (hidden_base, reduce_token_idx, 0), final_output
                                )
                                copy_elements = cutlass.Int32(
                                    cutlass.min(
                                        cutlass.Int64(self.cta_tile_shape_mnk[0]),
                                        valid_elements,
                                    )
                                )
                                blk_reduce_bf16(
                                    scatter_out,
                                    sFinalize[(reduce_route, None)],
                                    copy_elements * (self.c_dtype.width // 8),
                                )

                        cute.arch.cp_async_bulk_commit_group()
                        cute.arch.cp_async_bulk_wait_group(0, read=True)
                        self.epilog_sync_barrier.arrive_and_wait()
                    elif tma_distance_to_boundary >= self.cta_tile_shape_mnk[1]:
                        # Convert to C type
                        acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
                        if cutlass.const_expr(not self.fuse_activation):
                            acc_vec = cutlass.Float32(alpha_val) * acc_vec
                        acc_vec = acc_vec.to(self.c_dtype)
                        tRS_rC.store(acc_vec)
                        num_prev_subtiles += 1
                        c_buffer = num_prev_subtiles % self.num_c_stage
                        # Store C to shared memory
                        cute.copy(
                            tiled_copy_r2s,
                            tRS_rC,
                            tRS_sC[(None, None, None, c_buffer)],
                        )
                        # Fence and barrier to make sure shared memory store is visible to TMA store
                        cute.arch.fence_proxy(
                            "async.shared",
                            space="cta",
                        )
                        self.epilog_sync_barrier.arrive_and_wait()
                        # TMA store C to global memory
                        if warp_idx == self.epilog_warp_id[0]:
                            cute.copy(
                                tma_atom_c,
                                bSG_sC[(None, c_buffer)],
                                bSG_gC[(None, subtile_idx)],
                            )
                            c_pipeline.producer_commit()
                            c_pipeline.producer_acquire()
                        self.epilog_sync_barrier.arrive_and_wait()
                    else:
                        # Convert to C type
                        acc_vec = tTR_rAcc.load()
                        if cutlass.const_expr(not self.fuse_activation):
                            acc_vec = cutlass.Float32(alpha_val) * acc_vec
                        acc_vec = acc_vec.to(self.c_dtype)
                        tTR_rC.store(acc_vec)
                        # Compute predicate for SIMT store
                        tCpC = cute.make_rmem_tensor(
                            cute.make_layout(tTR_rC.shape),
                            cutlass.Boolean,
                        )
                        m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)]
                        for i in cutlass.range(cute.size(tCpC), unroll_full=True):
                            tCpC[i] = (
                                m_thr_slice[(i)][0]
                                + work_tile.cta_coord_m * self.cta_tile_shape_mnk_c[0]
                                < cute.size(tensor_c.shape[0])
                            ) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary)
                        # Store C to global memory
                        cute.copy(
                            simt_atom,
                            cute.flatten(tTR_rC),
                            cute.flatten(tTR_gC[(None, None, None, subtile_idx)]),
                            pred=cute.flatten(tCpC),
                        )
                # Advance to next tile
                tile_info_pipeline.consumer_wait(tile_info_consumer_state)
                work_tile = mixed_input_utils.make_contiguous_group_work_tile_info(
                    group_count, sTile_info[(None, tile_info_consumer_state.index)]
                )
                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.arrive_and_wait()
            tmem.free(tmem_ptr)
            if cutlass.const_expr(not self.use_fused_finalize):
                c_pipeline.producer_tail()

        if cutlass.const_expr(self.enable_pdl and not self.fuse_activation):
            if warp_idx == self.mma_warp_id:
                with cute.arch.elect_one():
                    griddepcontrol_launch_dependents()

    @staticmethod
    def _compute_stages_and_tmem_cols(
        tiled_mma: cute.TiledMma,
        mma_tiler_mnk: tuple[int, int, int],
        cta_tile_shape_mnk: tuple[int, int, int],
        epi_tile: cute.Tile,
        a_dtype: type[cutlass.Numeric],
        b_dtype: type[cutlass.Numeric],
        c_dtype: type[cutlass.Numeric],
        c_layout: utils.LayoutEnum,
        transform_a_source: tcgen05.OperandSource,
        smem_buffer_align_bytes: int,
        use_fused_finalize: bool,
        use_clc_scheduler: bool,
    ) -> tuple[int, int, int, int, int, int, int]:
        """Fit the load, transform, accumulator, and epilogue pipelines."""
        # Compute tmem columns required for accumulator
        acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
        tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1))
        num_tmem_acc_col_per_stage = utils.get_num_tmem_alloc_cols(tCtAcc_stage1, True)
        # Heuristic to decide the number of stages for accumulator
        sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100")
        accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage
        if transform_a_source == tcgen05.OperandSource.TMEM:
            if num_tmem_acc_col_per_stage < 256:
                accumulator_stage_count = 2
            else:
                accumulator_stage_count = 1
        # transformed A in 16bit, thus 1 tmem column could hold 2 elements
        num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width
        num_tmem_cols_a_per_stage = cute.round_up(
            (
                cta_tile_shape_mnk[2] // num_elts_per_tmem_col
                if transform_a_source == tcgen05.OperandSource.TMEM
                else 0
            ),
            4,
        )

        bytes_per_pipeline_stage = 16
        # Limit how far CLC can claim work ahead of the consumer pipelines.
        num_tile_info_stage = 1 if use_clc_scheduler else 3
        tile_info_bytes = (
            cute.size_in_bytes(cute.Int32, cute.make_layout((4, num_tile_info_stage)))
            + bytes_per_pipeline_stage * num_tile_info_stage
        )

        c_stage_count = 1 if use_fused_finalize else 2
        c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
            c_dtype,
            c_layout,
            epi_tile,
            1,
        )
        c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
        c_bytes = c_bytes_per_stage * c_stage_count
        finalize_metadata_bytes = (
            cute.size_in_bytes(
                cutlass.Float32,
                cute.make_layout((cute.size(epi_tile[1]),)),
            )
            if use_fused_finalize
            else 0
        )

        smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
        a_scale_k_mode = max(cta_tile_shape_mnk[2] // _NVFP4_SCALE_GRANULARITY_K, 1)
        a_scale_m_mode = max(cta_tile_shape_mnk[0] // _NVFP4_SCALE_GRANULARITY_M, 1)
        a_scale_bytes_per_stage = cute.round_up(
            cute.size_in_bytes(
                tiled_mma.op.a_dtype,
                cute.make_layout((a_scale_m_mode, a_scale_k_mode)),
            ),
            smem_buffer_align_bytes,
        )
        carveout_smem_bytes = (
            bytes_per_pipeline_stage * accumulator_stage_count
            + c_bytes
            + finalize_metadata_bytes
            + tile_info_bytes
        )

        # Compute transform stages if A is in TMEM
        num_tmem_acc_cols = cute.round_up(
            accumulator_stage_count * num_tmem_acc_col_per_stage, 4
        )

        transform2mma_stage_count_a_source_tmem_potential = (
            (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage
            if transform_a_source == tcgen05.OperandSource.TMEM
            else -1
        )
        if (
            transform_a_source == tcgen05.OperandSource.TMEM
            and transform2mma_stage_count_a_source_tmem_potential <= 0
        ):
            raise ValueError("Not enough TMEM capacity for selected tile size")
        a_load_bytes_per_stage = cute.round_up(
            cute.size_in_bytes(
                a_dtype,
                cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])),
            ),
            smem_buffer_align_bytes,
        )
        b_load_bytes_per_stage = cute.round_up(
            cute.size_in_bytes(
                b_dtype,
                cute.make_layout(
                    (
                        cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id),
                        cta_tile_shape_mnk[2],
                    )
                ),
            ),
            smem_buffer_align_bytes,
        )
        ab_scale_load_bytes_per_stage = (
            a_load_bytes_per_stage
            + b_load_bytes_per_stage
            + a_scale_bytes_per_stage
            + 2 * bytes_per_pipeline_stage
        )
        a_transform_bytes_per_stage = (
            cute.round_up(
                cute.size_in_bytes(
                    tiled_mma.op.a_dtype,
                    cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])),
                ),
                smem_buffer_align_bytes,
            )
            if transform_a_source == tcgen05.OperandSource.SMEM
            else 0
        )

        a_transform_bytes_per_stage = (
            a_transform_bytes_per_stage + bytes_per_pipeline_stage
        )
        transform2mma_stage_count_a_source_smem_potential = (
            smem_capacity - carveout_smem_bytes
        ) // (ab_scale_load_bytes_per_stage + a_transform_bytes_per_stage)
        transform2mma_stage_count = (
            min(
                transform2mma_stage_count_a_source_tmem_potential,
                transform2mma_stage_count_a_source_smem_potential,
            )
            if transform_a_source == tcgen05.OperandSource.TMEM
            else transform2mma_stage_count_a_source_smem_potential
        )
        load2transform_stage_count = (
            smem_capacity
            - carveout_smem_bytes
            - (transform2mma_stage_count * a_transform_bytes_per_stage)
        ) // ab_scale_load_bytes_per_stage
        if (
            load2transform_stage_count < 2
            or transform2mma_stage_count < 2
            or accumulator_stage_count < 1
        ):
            raise ValueError("Not enough SMEM or TMEM capacity for selected tile size")
        num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage
        # Fused finalize reuses its single C stage as reduction scratch.
        if not use_fused_finalize:
            c_stage_count += (
                smem_capacity
                - carveout_smem_bytes
                - load2transform_stage_count * ab_scale_load_bytes_per_stage
                - transform2mma_stage_count * a_transform_bytes_per_stage
            ) // c_bytes_per_stage

        return (
            load2transform_stage_count,
            transform2mma_stage_count,
            accumulator_stage_count,
            c_stage_count,
            num_tile_info_stage,
            num_tmem_acc_cols,
            num_tmem_a_cols,
        )

    @staticmethod
    def _compute_grid(
        c: cute.Tensor,
        cta_tile_shape_mnk: tuple[int, int, int],
        cluster_shape_mn: tuple[int, int],
        max_active_clusters: cutlass.Constexpr,
        use_clc_scheduler: bool,
        raster_along_m: bool,
    ) -> tuple[
        Union[
            utils.ClcDynamicPersistentTileSchedulerParams,
            utils.PersistentTileSchedulerParams,
        ],
        tuple[int, int, int],
    ]:
        """Build the persistent scheduler and launch grid."""
        c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
        gc = cute.zipped_divide(c, tiler=c_shape)
        num_ctas_mnl = gc[(0, (None, None, None))].shape
        cluster_shape_mnl = (*cluster_shape_mn, 1)

        if use_clc_scheduler:
            tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams(
                num_ctas_mnl, cluster_shape_mnl
            )
            grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(
                tile_sched_params
            )
        else:
            tile_sched_params = utils.PersistentTileSchedulerParams(
                num_ctas_mnl,
                cluster_shape_mnl,
                raster_along_m=raster_along_m,
            )
            grid = (cluster_shape_mn[0], cluster_shape_mn[1], max_active_clusters)

        return tile_sched_params, grid

    @staticmethod
    def can_implement(
        mnkl: tuple[int, int, int, int],
        a_dtype: type[cutlass.Numeric],
        b_dtype: type[cutlass.Numeric],
        c_dtype: type[cutlass.Numeric],
        a_major: str,
        b_major: str,
        c_major: str,
        mma_tiler: tuple[int, int, int],
        cluster_shape_mn: tuple[int, int],
        use_2cta_instrs: bool,
    ) -> bool:
        """Return whether the requested layout and tile shape are supported."""
        m, n, k, _ = mnkl

        if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape(
            mma_tiler, cluster_shape_mn, use_2cta_instrs
        ):
            return False
        # Unlike the generic mixed-input kernel, W4A16 consumes multiple
        # 16-value NVFP4 scale blocks within each MMA K tile.
        if (
            k % _NVFP4_SCALE_GRANULARITY_K != 0
            or mma_tiler[2] % _NVFP4_SCALE_GRANULARITY_K != 0
        ):
            return False
        if not mixed_input_utils.is_valid_tensor_alignment(
            m,
            n,
            k,
            a_dtype,
            b_dtype,
            c_dtype,
            b_dtype,
            a_major,
            b_major,
            c_major,
            mma_tiler,
            use_2cta_instrs,
            cluster_shape_mn,
            _NVFP4_SCALE_GRANULARITY_M,
            _NVFP4_SCALE_GRANULARITY_K,
        ):
            return False
        return True
