"""
MoEDynamicKernel — queue-driven routed NVFP4 MoE kernel for SM120/SM121.

Ported from the b12x kernel library to FlashInfer.

This is the first dynamic fused control-plane kernel derived from the current
static implementation. It keeps the proven FC1 / activation / quant / FC2 /
scatter compute body, but replaces the resident-grid route/pack -> compute
barrier with a global ready-task queue.

Supports gated and non-gated activations selected at construction time:
  Gated (activation="silu", "gelu_tanh", "swigluoai_uninterleave"):
    FC1:     A x gate^T, A x up^T     (paired FP4 block-scaled GEMMs)
    Act:     act(gate) * up_term       (see moe_activation.gated_activation_f32)
  Non-gated (activation="relu2"):
    FC1:     A x W1^T                  (single FP4 block-scaled GEMM)
    Act:     max(0, x)^2               (squared ReLU activation)

Execution model
  Phase 0: cooperative init / clear scratch state
  Phase 1: all CTAs start as producers
           - claim routed (token, topk_slot) pairs from pair_head
           - append expert rows
           - write token_map + token_weights
           - quantize each routed token row into expert-major packed A + scales
           - publish one compute task per ready (expert, m_tile, slice_group)
             as soon as a tile is fully written
  Phase 2: CTAs that finish producing become consumers immediately
           - CTA leader pops one ready task into shared ctrl state
           - MMA warps run FC1 -> SiLU -> quant -> FC2 -> scatter for that task
           - DMA warp streams the corresponding FC1 / FC2 weights

This is intentionally conservative:
  - still one CTA per SM
  - still the static per-slice microkernel, now executed sequentially for a
    small grouped slice task
  - still one initial resident-grid barrier after init

What changes relative to the static path
  - no global route/pack -> compute barrier
  - no static scheduler in the compute steady state
  - route/pack is warp-private instead of CTA-broadcast
  - compute work is driven by a global append-only ready-task queue

This implementation remains the compatibility fallback for dynamic shapes that
the specialized gated kernel does not support.
"""

from __future__ import annotations

from typing import Tuple

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.utils as utils
import cutlass.utils.blockscaled_layout as blockscaled_utils

from cutlass.cutlass_dsl import (
    Int32,
    Int64,
    Uint8,
    Uint32,
    Uint64,
    T,
    dsl_user_op,
    extract_mlir_values,
    new_from_mlir_values,
)
from cutlass._mlir.dialects import llvm
from cutlass.cute.nvgpu import cpasync

from flashinfer.cute_dsl.utils import (
    sm120_make_smem_layout_sfa,
    sm120_make_smem_layout_sfb,
)
from flashinfer.cute_dsl.fp4_common import (
    atomic_add_global_i32,
    fabs_f32,
    fmax_f32,
    ld_shared_f32,
    ld_shared_i32_relaxed,
    rcp_approx_ftz,
    quantize_block_fp4,
    quantize_block_fp4_fast,
    get_ptr_as_int64,
    st_global_f32,
    st_global_i32,
    shared_ptr_to_u32,
    st_shared_f32,
    st_shared_u8,
    st_global_u64,
    st_global_v4_u32,
    scatter_add_v4_bf16x2,
)
from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import (
    Sm120B12xBlockScaledDenseGemmKernel as DenseGemmKernel,
)
from ..moe_activation import gated_activation_f32, is_gated_activation


_SF_VEC_SIZE = 16
_TASK_SLICE_CHUNK = 1
_PRODUCER_PAIRS_PER_WARP = 2

_WORK_ITEM_FIELDS = 5
_WORK_EXPERT = 0
_WORK_M_TILE = 1
_WORK_SLICE_BEGIN = 2
_WORK_SLICE_COUNT = 3
_WORK_VALID_ROWS = 4

_CTRL_HAS_WORK = 0
_CTRL_DONE = 1
_CTRL_WORK_BEGIN = 2


class DynamicLaunchParams:
    """Minimal runtime launch state shared between host setup and kernel code."""

    def __init__(
        self,
        row_counts: cute.Tensor,
        gate_tile_cnt: Int32,
        *,
        loc=None,
    ):
        self.row_counts = row_counts
        self.gate_tile_cnt = gate_tile_cnt
        self._loc = loc

    def __extract_mlir_values__(self):
        values, self._values_pos = [], []
        for obj in [self.row_counts, self.gate_tile_cnt]:
            obj_values = extract_mlir_values(obj)
            values += obj_values
            self._values_pos.append(len(obj_values))
        return values

    def __new_from_mlir_values__(self, values):
        obj_list = []
        for obj, n_items in zip(
            [self.row_counts, self.gate_tile_cnt],
            self._values_pos,
            strict=True,
        ):
            obj_list.append(new_from_mlir_values(obj, values[:n_items]))
            values = values[n_items:]
        return DynamicLaunchParams(*(tuple(obj_list)), loc=self._loc)


@dsl_user_op
def _st_shared_i32(addr, val, *, loc=None, ip=None):
    llvm.inline_asm(
        None,
        [Int32(addr).ir_value(loc=loc, ip=ip), Int32(val).ir_value(loc=loc, ip=ip)],
        "st.shared.s32 [$0], $1;",
        "r,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
    )


@dsl_user_op
def _ld_shared_i32(addr, *, loc=None, ip=None):
    return Int32(
        llvm.inline_asm(
            T.i32(),
            [Int32(addr).ir_value(loc=loc, ip=ip)],
            "ld.shared.s32 $0, [$1];",
            "=r,r",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
        )
    )


@dsl_user_op
def _ld_global_u64(addr, *, loc=None, ip=None):
    return Uint64(
        llvm.inline_asm(
            T.i64(),
            [Int64(addr).ir_value(loc=loc, ip=ip)],
            "ld.global.u64 $0, [$1];",
            "=l,l",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
        )
    )


@dsl_user_op
def _ld_global_acquire_i32(addr, *, loc=None, ip=None):
    return Int32(
        llvm.inline_asm(
            T.i32(),
            [Int64(addr).ir_value(loc=loc, ip=ip)],
            "ld.global.acquire.gpu.s32 $0, [$1];",
            "=r,l",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
        )
    )


@dsl_user_op
def _st_global_release_i32(addr, val, *, loc=None, ip=None):
    llvm.inline_asm(
        None,
        [Int64(addr).ir_value(loc=loc, ip=ip), Int32(val).ir_value(loc=loc, ip=ip)],
        "st.global.release.gpu.s32 [$0], $1;",
        "l,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
    )


@dsl_user_op
def _spin_wait_global_eq_i32(addr, expected, *, loc=None, ip=None):
    llvm.inline_asm(
        None,
        [
            Int64(addr).ir_value(loc=loc, ip=ip),
            Int32(expected).ir_value(loc=loc, ip=ip),
        ],
        "{\n"
        ".reg .pred %p0;\n"
        ".reg .s32 %val;\n"
        "spin_loop:\n"
        "  ld.global.acquire.gpu.s32 %val, [$0];\n"
        "  setp.eq.s32 %p0, %val, $1;\n"
        "  @%p0 bra spin_loop;\n"
        "}",
        "l,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
    )


@dsl_user_op
def _threadfence(*, loc=None, ip=None):
    llvm.inline_asm(
        None,
        [],
        "membar.gl;",
        "",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
    )


class MoEDynamicKernel:
    """Queue-driven first-pass dynamic MoE kernel."""

    def __init__(
        self,
        sf_vec_size: int,
        mma_tiler_mn: Tuple[int, int],
        *,
        input_scales_are_reciprocal: bool = False,
        fast_math: bool = False,
        activation: str = "silu",
        swiglu_alpha: float = 1.702,
        swiglu_beta: float = 1.0,
        swiglu_limit: float | None = None,
        share_input_across_experts: bool = False,
    ):
        if activation not in {"silu", "relu2", "gelu_tanh", "swigluoai_uninterleave"}:
            raise ValueError(f"unsupported activation {activation!r}")
        self._dense_cls = DenseGemmKernel
        self.acc_dtype = cutlass.Float32
        self.sf_vec_size = sf_vec_size
        self.input_scales_are_reciprocal = input_scales_are_reciprocal
        self.activation = activation
        self.is_gated = is_gated_activation(activation)
        # relu2's squared outputs need the exact quantizer and scale math.
        self.fast_math = bool(fast_math) and self.is_gated
        self.swiglu_alpha = float(swiglu_alpha)
        self.swiglu_beta = float(swiglu_beta)
        self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None
        self.share_input_across_experts = share_input_across_experts
        tile_k = sf_vec_size * 8
        self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k)
        # Scale factors come in 128-row atoms, so for sub-128 MMA tiles the
        # TMA atoms and smem are built at max(128, tile) and the kernel
        # offsets into the shared block by `*_tiles_per_block`.
        self.sa_tile_shape_mk = (max(128, mma_tiler_mn[0]), tile_k)
        self.sa_tiles_per_block = self.sa_tile_shape_mk[0] // mma_tiler_mn[0]
        self.sfa_tile_shape_mk = (max(128, mma_tiler_mn[0]), tile_k)
        self.sfa_tiles_per_block = self.sfa_tile_shape_mk[0] // mma_tiler_mn[0]
        self.sfb_tile_shape_nk = (max(128, mma_tiler_mn[1]), tile_k)
        self.sfb_tiles_per_block = self.sfb_tile_shape_nk[0] // mma_tiler_mn[1]
        self.cluster_shape_mnk = (1, 1, 1)
        self.cluster_shape_mn = (1, 1)
        self.epi_tile = (mma_tiler_mn[0], mma_tiler_mn[1])
        self.occupancy = 1
        # Smaller tiles need a matching MMA atom shape so the SF smem layout
        # and V-map stay consistent.
        if mma_tiler_mn[0] == 128:
            self.atom_shape = (2, 2, 1)
            self.num_mma_warps = 4
        elif mma_tiler_mn[0] == 64:
            self.atom_shape = (4, 2, 1)
            self.num_mma_warps = 8
        elif mma_tiler_mn[0] == 32:
            self.atom_shape = (2, 2, 1)
            self.num_mma_warps = 4
        elif mma_tiler_mn[0] == 16:
            self.atom_shape = (1, 2, 1)
            self.num_mma_warps = 2
        else:
            raise ValueError(f"unsupported dynamic MMA tile_m {mma_tiler_mn[0]}")
        self.tma_load_warp_id = self.num_mma_warps
        self.num_threads_per_warp = 32
        self.threads_per_cta = (self.num_mma_warps + 1) * self.num_threads_per_warp
        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_120")
        self.buffer_align_bytes = 1024

        self.epilog_sync_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=self.num_mma_warps * self.num_threads_per_warp,
        )
        self.pass_gate_barrier = pipeline.NamedBarrier(
            barrier_id=2,
            num_threads=self.threads_per_cta,
        )
        self.pass_final_barrier = pipeline.NamedBarrier(
            barrier_id=3,
            num_threads=self.threads_per_cta,
        )
        self.load_register_requirement = 32
        self.mma_register_requirement = 232

    def _thrfrg_SFA(self, sfa_tensor, tiled_mma):
        return self._dense_cls._thrfrg_SFA(self, sfa_tensor, tiled_mma)

    def _thrfrg_SFB(self, sfb_tensor, tiled_mma):
        return self._dense_cls._thrfrg_SFB(self, sfb_tensor, tiled_mma)

    def _get_layoutSFA_TV(self, tiled_mma):
        return self._dense_cls._get_layoutSFA_TV(self, tiled_mma)  # type: ignore[arg-type]

    def _get_layoutSFB_TV(self, tiled_mma):
        return self._dense_cls._get_layoutSFB_TV(self, tiled_mma)  # type: ignore[arg-type]

    def _setup_attributes(self, hidden_size: int):
        import cutlass.utils.blackwell_helpers as sm120_utils

        self._hidden_size = hidden_size

        mma_op = cute.nvgpu.warp.MmaMXF4NVF4Op(
            self.a_dtype,
            self.acc_dtype,
            self.sf_dtype,
        )
        atom_layout = cute.make_layout(self.atom_shape)
        permutation_mnk = sm120_utils.get_permutation_mnk(
            self.tile_shape_mnk,
            self.sf_vec_size,
            False,
        )
        self.tiled_mma = cute.make_tiled_mma(
            mma_op,
            atom_layout,
            permutation_mnk=permutation_mnk,
        )
        self.mma_atom = cute.make_mma_atom(mma_op)
        self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk)
        self.num_m_tiles = self.tile_shape_mnk[0] // (16 * self.atom_shape[0])
        self.num_n_tiles = self.tile_shape_mnk[1] // (8 * self.atom_shape[1])
        self.num_k_blocks = self.tile_shape_mnk[2] // 64

        # A/SFA smem hold the whole 128-row block for sub-128 MMA tiles (the
        # SF smem helper also rejects tile M % 64 != 0), so build all smem
        # layouts at the block shape; sub-128 tasks slice into it.
        smem_tile_shape_mnk = (
            self.sa_tile_shape_mk[0],
            self.tile_shape_mnk[1],
            self.tile_shape_mnk[2],
        )
        sfa_smem = sm120_make_smem_layout_sfa(
            self.tiled_mma,
            smem_tile_shape_mnk,
            self.sf_vec_size,
            1,
        )
        sfb_smem = sm120_make_smem_layout_sfb(
            self.tiled_mma,
            smem_tile_shape_mnk,
            self.sf_vec_size,
            1,
        )

        self.ab_stage, self.epi_stage = self._dense_cls._compute_stages(
            self.tile_shape_mnk,
            self.a_dtype,
            self.b_dtype,
            self.sf_dtype,
            sfa_smem,
            sfb_smem,
            self.epi_tile,
            cutlass.BFloat16,
            self.smem_capacity,
            self.occupancy,
        )
        # dense._compute_stages models a single B/SFB buffer, but the gated
        # path double-buffers both. Recompute the stage cap from the real
        # per-stage footprint so shapes with room keep their extra stages.
        nb = 2 if self.is_gated else 1
        n_pipe = 3 if self.is_gated else 2
        # sA smem holds the whole 128-row block, so size the stage from it.
        a_bytes = (
            self.sa_tile_shape_mk[0]
            * self.sa_tile_shape_mk[1]
            * self.a_dtype.width
            // 8
        )
        b_bytes = (
            cute.size(cute.slice_(self.tile_shape_mnk, (0, None, None)))
            * self.b_dtype.width
            // 8
        )
        sfa_bytes = (
            cute.size(cute.filter_zeros(sfa_smem).shape) * self.sf_dtype.width // 8
        )
        sfb_bytes = (
            cute.size(cute.filter_zeros(sfb_smem).shape) * self.sf_dtype.width // 8
        )
        per_stage = a_bytes + nb * b_bytes + sfa_bytes + nb * sfb_bytes + n_pipe * 2 * 8
        fixed = (
            self.tile_shape_mnk[0] * self.tile_shape_mnk[1] * 2  # sC (bf16 epi)
            + 2 * (self.num_mma_warps + 1) * 32 * 4  # route caches
            + self.tile_shape_mnk[0] * 8  # scatter caches
            + 8 * 8
            + 1024
            + 8 * 1024  # ctrl/mbar/align slack
        )
        max_fit = max(1, (self.smem_capacity - fixed) // per_stage)
        self.ab_stage = min(self.ab_stage, max_fit)
        # ab_stage must divide k_tile_cnt evenly to avoid pipeline phase mismatch.
        # _compute_stages returns the max that fits in smem, but it may not
        # divide k_tile_cnt. Round down to the nearest divisor.
        k_tile_cnt = self._hidden_size // self.tile_shape_mnk[2]
        while self.ab_stage > 1 and k_tile_cnt % self.ab_stage != 0:
            self.ab_stage -= 1
        self.epi_stage = 1
        (
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.sfa_smem_layout_staged,
            self.sfb_smem_layout_staged,
            self.epi_smem_layout_staged,
        ) = self._dense_cls._make_smem_layouts(
            smem_tile_shape_mnk,
            self.epi_tile,
            self.a_dtype,
            self.a_layout,
            self.b_dtype,
            self.b_layout,
            self.ab_stage,
            cutlass.BFloat16,
            self.c_layout,
            self.epi_stage,
            self.sf_vec_size,
            self.tiled_mma,
        )

    @cute.jit
    def _resident_grid_barrier(
        self,
        barrier_count: cute.Tensor,
        barrier_epoch: cute.Tensor,
        grid_x: Int32,
        is_cta_leader: Int32,
    ):
        cute.arch.sync_threads()
        _threadfence()
        if is_cta_leader > Int32(0):
            barrier_count_addr = get_ptr_as_int64(barrier_count, Int32(0))
            barrier_epoch_addr = get_ptr_as_int64(barrier_epoch, Int32(0))
            old_epoch = _ld_global_acquire_i32(barrier_epoch_addr)
            arrived = atomic_add_global_i32(barrier_count_addr, Int32(1))
            if arrived == grid_x - Int32(1):
                st_global_i32(barrier_count_addr, Int32(0))
                _st_global_release_i32(barrier_epoch_addr, old_epoch + Int32(1))
            else:
                _spin_wait_global_eq_i32(barrier_epoch_addr, old_epoch)
        cute.arch.sync_threads()

    @cute.jit
    def _publish_deferred_tasks(
        self,
        task_expert: cute.Tensor,
        task_valid_rows: cute.Tensor,
        gate_tile_cnt: Int32,
        slice_chunk: Int32,
        expert_idx: Int32,
        m_tile_idx: Int32,
        valid_rows: Int32,
    ):
        num_groups = (gate_tile_cnt + slice_chunk - Int32(1)) // slice_chunk
        start = m_tile_idx * num_groups

        g = Int32(0)
        while g < num_groups:
            slot = start + g
            task_expert[slot] = expert_idx
            task_valid_rows[slot] = valid_rows
            g += Int32(1)

    @cute.jit
    def _decode_materialized_work_item(
        self,
        work_item: cute.Tensor,
        task_expert: cute.Tensor,
        task_valid_rows: cute.Tensor,
        slot: Int32,
        num_groups: Int32,
        slice_chunk: Int32,
        gate_tile_cnt: Int32,
    ):
        """Decode a deferred task from its slot index.

        Slots are ordered by m-tile, then group, so only expert and
        valid-row metadata need global storage.
        """

        m_tile = slot // num_groups
        group = slot - m_tile * num_groups
        slice_begin = group * slice_chunk
        slice_count = gate_tile_cnt - slice_begin
        if slice_count > slice_chunk:
            slice_count = slice_chunk
        work_item[_WORK_EXPERT] = task_expert[slot].to(Int32)
        work_item[_WORK_M_TILE] = m_tile
        work_item[_WORK_SLICE_BEGIN] = slice_begin
        work_item[_WORK_SLICE_COUNT] = slice_count
        work_item[_WORK_VALID_ROWS] = task_valid_rows[slot].to(Int32)

    @cute.jit
    def _store_shared_work_item(
        self,
        ctrl_base_addr: Int32,
        work_item: cute.Tensor,
    ):
        for field in cutlass.range_constexpr(_WORK_ITEM_FIELDS):
            _st_shared_i32(
                ctrl_base_addr + Int32((_CTRL_WORK_BEGIN + field) * 4),
                work_item[field],
            )

    @cute.jit
    def _load_shared_work_item(
        self,
        work_item: cute.Tensor,
        ctrl_base_addr: Int32,
    ):
        for field in cutlass.range_constexpr(_WORK_ITEM_FIELDS):
            work_item[field] = _ld_shared_i32(
                ctrl_base_addr + Int32((_CTRL_WORK_BEGIN + field) * 4)
            )

    @cute.jit
    def __call__(
        self,
        a_input: cute.Tensor,  # [num_tokens, K] bf16
        topk_ids: cute.Tensor,  # [num_tokens * topk] int32
        topk_weights: cute.Tensor,  # [num_tokens * topk] float32
        packed_a: cute.Tensor,  # [rows_padded, K, 1] fp4x2 view for compute
        sfa_ptr: cute.Pointer,
        packed_a_storage: cute.Tensor,  # flat uint8 backing packed_a
        scale_storage: cute.Tensor,  # flat uint8 backing sfa_ptr
        barrier_count: cute.Tensor,  # [1] int32 (host-zeroed)
        barrier_epoch: cute.Tensor,  # [1] int32 (host-zeroed)
        pair_head: cute.Tensor,  # [1] int32
        task_head: cute.Tensor,  # [1] int32
        task_tail: cute.Tensor,  # [1] int32
        task_expert: cute.Tensor,  # [max_tasks] int32
        task_valid_rows: cute.Tensor,  # [max_tasks] int32
        b_w13: cute.Tensor,  # [2*I_tp, K, E] (gated) or [I_tp, K, E] (relu2)
        sfb_w13_ptr: cute.Pointer,  # scale factors for w13
        b_down: cute.Tensor,  # [K, I_tp, E]
        sfb_down_ptr: cute.Pointer,
        row_counts: cute.Tensor,  # expert row histogram [E]
        expert_write_rows: cute.Tensor,  # route/pack write cursors [E]
        expert_tile_base: cute.Tensor,  # compact physical-tile prefix [E + 1]
        input_global_scale: cute.Tensor,  # [E] per-expert FC1 input scale
        alpha: cute.Tensor,
        down_alpha: cute.Tensor,
        global_scale: cute.Tensor,
        scatter_output: cute.Tensor,  # [num_tokens, K]
        token_map: cute.Tensor,
        token_weights: cute.Tensor,
        max_active_clusters: cutlass.Constexpr,
        stream: cuda.CUstream,
    ):
        self.a_dtype = packed_a.element_type
        self.b_dtype = b_w13.element_type
        self.sf_dtype = sfa_ptr.dtype
        self.a_layout = utils.LayoutEnum.from_tensor(packed_a)
        self.b_layout = utils.LayoutEnum.from_tensor(b_w13)
        # Dynamic never materializes the intermediate C tensor. Preserve the
        # original row-major epilogue layout without carrying a dead memref.
        self.c_layout = utils.LayoutEnum.ROW_MAJOR

        hidden_size = a_input.shape[1]
        self._setup_attributes(hidden_size=hidden_size)

        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
            packed_a.shape, self.sf_vec_size
        )
        sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

        # SF tensor for w13 (gated: gate+up concatenated; relu2: single W1)
        sfb_w13_layout = blockscaled_utils.tile_atom_to_shape_SF(
            b_w13.shape, self.sf_vec_size
        )
        sfb_w13_tensor = cute.make_tensor(sfb_w13_ptr, sfb_w13_layout)

        # TMA descriptors
        tma_a, gA = self._dense_cls._make_tma_atoms_and_tensors(
            packed_a,
            self.a_smem_layout_staged,
            self.sa_tile_shape_mk,
            1,
        )
        tma_sfa, gSFA = self._dense_cls._make_tma_atoms_and_tensors(
            sfa_tensor,
            self.sfa_smem_layout_staged,
            self.sfa_tile_shape_mk,
            1,
            internal_type=cutlass.Int16,
        )
        # Single TMA descriptor over w13. Gated: [2*I_tp, K, E]; relu2: [I_tp, K, E].
        # Up tiles at N=0..I_tp/tile_N-1, gate tiles at N=I_tp/tile_N..2*I_tp/tile_N-1.
        tma_b_w13, gB_w13 = self._dense_cls._make_tma_atoms_and_tensors(
            b_w13,
            self.b_smem_layout_staged,
            (self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
            1,
        )
        tma_sfb_w13, gSFB_w13 = self._dense_cls._make_tma_atoms_and_tensors(
            sfb_w13_tensor,
            self.sfb_smem_layout_staged,
            self.sfb_tile_shape_nk,
            1,
            internal_type=cutlass.Int16,
        )
        # B_down TMA
        sfb_down_layout = blockscaled_utils.tile_atom_to_shape_SF(
            b_down.shape, self.sf_vec_size
        )
        sfb_down_tensor = cute.make_tensor(sfb_down_ptr, sfb_down_layout)
        tma_b_down, gB_down = self._dense_cls._make_tma_atoms_and_tensors(
            b_down,
            self.b_smem_layout_staged,
            (self.tile_shape_mnk[1], self.tile_shape_mnk[2]),
            1,
        )
        tma_sfb_down, gSFB_down = self._dense_cls._make_tma_atoms_and_tensors(
            sfb_down_tensor,
            self.sfb_smem_layout_staged,
            self.sfb_tile_shape_nk,
            1,
            internal_type=cutlass.Int16,
        )

        gate_tile_cnt = Int32(b_w13.shape[0]) // Int32(self.tile_shape_mnk[1])
        if self.is_gated:
            gate_tile_cnt = gate_tile_cnt // Int32(2)
        launch_params = DynamicLaunchParams(row_counts, gate_tile_cnt)
        grid = (*self.cluster_shape_mn, max_active_clusters)
        self.kernel(
            a_input,
            topk_ids,
            topk_weights,
            packed_a_storage,
            scale_storage,
            barrier_count,
            barrier_epoch,
            pair_head,
            task_head,
            task_tail,
            task_expert,
            task_valid_rows,
            tma_a,
            gA,
            tma_sfa,
            gSFA,
            tma_b_w13,
            gB_w13,
            tma_sfb_w13,
            gSFB_w13,
            tma_b_down,
            gB_down,
            tma_sfb_down,
            gSFB_down,
            self.tiled_mma,
            self.mma_atom,
            self.cta_layout_mnk,
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.sfa_smem_layout_staged,
            self.sfb_smem_layout_staged,
            self.epi_smem_layout_staged,
            launch_params,
            expert_write_rows,
            expert_tile_base,
            input_global_scale,
            alpha,
            down_alpha,
            global_scale,
            scatter_output,
            token_map,
            token_weights,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=[1, 1, 1],
            # A regular launch beside other stream work can admit only part
            # of the grid, deadlocking the software grid barriers below.
            cooperative=True,
            stream=stream,
        )

    @cute.kernel
    def kernel(
        self,
        a_input: cute.Tensor,
        topk_ids: cute.Tensor,
        topk_weights: cute.Tensor,
        packed_a_storage: cute.Tensor,
        scale_storage: cute.Tensor,
        barrier_count: cute.Tensor,
        barrier_epoch: cute.Tensor,
        pair_head: cute.Tensor,
        task_head: cute.Tensor,
        task_tail: cute.Tensor,
        task_expert: cute.Tensor,
        task_valid_rows: cute.Tensor,
        tma_a: cute.CopyAtom,
        mA: cute.Tensor,
        tma_sfa: cute.CopyAtom,
        mSFA: cute.Tensor,
        tma_b_w13: cute.CopyAtom,
        mB_w13: cute.Tensor,
        tma_sfb_w13: cute.CopyAtom,
        mSFB_w13: cute.Tensor,
        tma_b_down: cute.CopyAtom,
        mB_down: cute.Tensor,
        tma_sfb_down: cute.CopyAtom,
        mSFB_down: cute.Tensor,
        tiled_mma: cute.TiledMma,
        mma_atom: cute.MmaAtom,
        cta_layout_mnk: cute.Layout,
        a_smem_staged: cute.ComposedLayout,
        b_smem_staged: cute.ComposedLayout,
        sfa_smem_staged: cute.Layout,
        sfb_smem_staged: cute.Layout,
        epi_smem_staged: cute.ComposedLayout,
        launch_params: DynamicLaunchParams,
        expert_write_rows: cute.Tensor,
        expert_tile_base: cute.Tensor,
        input_global_scale: cute.Tensor,
        alpha: cute.Tensor,
        down_alpha: cute.Tensor,
        global_scale: cute.Tensor,
        scatter_output: cute.Tensor,
        token_map: cute.Tensor,
        token_weights: cute.Tensor,
    ):
        """Kernel entry point."""
        from cutlass.cute.nvgpu.warp.mma import Field as WarpField

        tidx, _, _ = cute.arch.thread_idx()
        _bidx, _, bidz = cute.arch.block_idx()
        _, _, gdim_z = cute.arch.grid_dim()
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        is_cta_leader = Int32(Int32(tidx) == Int32(0))

        if warp_idx == 0:
            cpasync.prefetch_descriptor(tma_a)
            cpasync.prefetch_descriptor(tma_sfa)
            cpasync.prefetch_descriptor(tma_b_w13)
            cpasync.prefetch_descriptor(tma_sfb_w13)
            cpasync.prefetch_descriptor(tma_b_down)
            cpasync.prefetch_descriptor(tma_sfb_down)

        cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
        cluster_coord = cta_layout_mnk.get_flat_coord(cta_rank)

        a_smem_one = cute.slice_(a_smem_staged, (None, None, 0))
        b_smem_one = cute.slice_(b_smem_staged, (None, None, 0))
        sfa_smem_one = cute.slice_(sfa_smem_staged, (None, None, 0))
        sfb_smem_one = cute.slice_(sfb_smem_staged, (None, None, 0))
        tma_copy_bytes = (
            cute.size_in_bytes(self.a_dtype, a_smem_one)
            + cute.size_in_bytes(self.b_dtype, b_smem_one)
            + cute.size_in_bytes(self.sf_dtype, sfa_smem_one)
            + cute.size_in_bytes(self.sf_dtype, sfb_smem_one)
        )
        phase2_tma_copy_bytes = cute.size_in_bytes(
            self.b_dtype, b_smem_one
        ) + cute.size_in_bytes(self.sf_dtype, sfb_smem_one)

        smem = cutlass.utils.SmemAllocator()

        @cute.struct
        class StorageGated:
            # ctrl layout (8 x Int32, accessed via raw shared memory PTX):
            #   [0] has_task     [4] done          [8]  expert_idx
            #   [12] m_tile_idx  [16] slice_begin   [20] slice_count
            #   [24] valid_rows  [28] batch_base
            ctrl: cute.struct.MemRange[cutlass.Int32, 8]
            route_phys_rows: cute.struct.MemRange[
                cutlass.Int32, (self.num_mma_warps + 1) * 32
            ]
            route_expert_ids: cute.struct.MemRange[
                cutlass.Int32, (self.num_mma_warps + 1) * 32
            ]
            pipeline_array: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2]
            up_pipeline_array: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2]
            phase2_pipeline_array: cute.struct.MemRange[
                cutlass.Int64, self.ab_stage * 2
            ]
            scatter_tok_cache: cute.struct.MemRange[
                cutlass.Int32, self.tile_shape_mnk[0]
            ]
            scatter_weight_cache: cute.struct.MemRange[
                cutlass.Float32, self.tile_shape_mnk[0]
            ]
            sA: cute.struct.Align[
                cute.struct.MemRange[self.a_dtype, cute.cosize(a_smem_staged)],
                self.buffer_align_bytes,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(b_smem_staged)],
                self.buffer_align_bytes,
            ]
            sB_up: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(b_smem_staged)],
                self.buffer_align_bytes,
            ]
            sSFA: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(sfa_smem_staged)],
                self.buffer_align_bytes,
            ]
            sSFB: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(sfb_smem_staged)],
                self.buffer_align_bytes,
            ]
            sSFB_up: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(sfb_smem_staged)],
                self.buffer_align_bytes,
            ]
            sC: cute.struct.Align[
                cute.struct.MemRange[cutlass.BFloat16, cute.cosize(epi_smem_staged)],
                self.buffer_align_bytes,
            ]

        @cute.struct
        class StorageRelu2:
            # ctrl layout (8 x Int32, accessed via raw shared memory PTX):
            #   [0] has_task     [4] done          [8]  expert_idx
            #   [12] m_tile_idx  [16] slice_begin   [20] slice_count
            #   [24] valid_rows  [28] batch_base
            ctrl: cute.struct.MemRange[cutlass.Int32, 8]
            route_phys_rows: cute.struct.MemRange[
                cutlass.Int32, (self.num_mma_warps + 1) * 32
            ]
            route_expert_ids: cute.struct.MemRange[
                cutlass.Int32, (self.num_mma_warps + 1) * 32
            ]
            pipeline_array: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2]
            phase2_pipeline_array: cute.struct.MemRange[
                cutlass.Int64, self.ab_stage * 2
            ]
            scatter_tok_cache: cute.struct.MemRange[
                cutlass.Int32, self.tile_shape_mnk[0]
            ]
            scatter_weight_cache: cute.struct.MemRange[
                cutlass.Float32, self.tile_shape_mnk[0]
            ]
            sA: cute.struct.Align[
                cute.struct.MemRange[self.a_dtype, cute.cosize(a_smem_staged)],
                self.buffer_align_bytes,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(b_smem_staged)],
                self.buffer_align_bytes,
            ]
            sSFA: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(sfa_smem_staged)],
                self.buffer_align_bytes,
            ]
            sSFB: cute.struct.Align[
                cute.struct.MemRange[self.sf_dtype, cute.cosize(sfb_smem_staged)],
                self.buffer_align_bytes,
            ]
            sC: cute.struct.Align[
                cute.struct.MemRange[cutlass.BFloat16, cute.cosize(epi_smem_staged)],
                self.buffer_align_bytes,
            ]

        storage = smem.allocate(StorageGated if self.is_gated else StorageRelu2)

        prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        cons_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, self.num_mma_warps
        )
        cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape))
        ml_pipeline = pipeline.PipelineTmaAsync.create(
            num_stages=self.ab_stage,
            producer_group=prod_group,
            consumer_group=cons_group,
            tx_count=tma_copy_bytes,
            barrier_storage=storage.pipeline_array.data_ptr(),
            cta_layout_vmnk=cta_layout_vmnk,
        )
        up_pipeline = (
            pipeline.PipelineTmaAsync.create(
                num_stages=self.ab_stage,
                producer_group=prod_group,
                consumer_group=cons_group,
                tx_count=tma_copy_bytes,
                barrier_storage=storage.up_pipeline_array.data_ptr(),
                cta_layout_vmnk=cta_layout_vmnk,
            )
            if self.is_gated
            else ml_pipeline
        )
        phase2_pipeline = pipeline.PipelineTmaAsync.create(
            num_stages=self.ab_stage,
            producer_group=prod_group,
            consumer_group=cons_group,
            tx_count=phase2_tma_copy_bytes,
            barrier_storage=storage.phase2_pipeline_array.data_ptr(),
            cta_layout_vmnk=cta_layout_vmnk,
        )

        cute.arch.sync_threads()

        sA = storage.sA.get_tensor(a_smem_staged.outer, swizzle=a_smem_staged.inner)
        sB = storage.sB.get_tensor(b_smem_staged.outer, swizzle=b_smem_staged.inner)
        sB_up = (
            storage.sB_up.get_tensor(b_smem_staged.outer, swizzle=b_smem_staged.inner)
            if self.is_gated
            else sB
        )
        cute.recast_tensor(sA, cutlass.Uint8)
        cute.recast_tensor(sB, cutlass.Uint8)
        cute.recast_tensor(sB_up, cutlass.Uint8)
        sSFA = storage.sSFA.get_tensor(sfa_smem_staged)
        sSFB = storage.sSFB.get_tensor(sfb_smem_staged)
        sSFB_up = storage.sSFB_up.get_tensor(sfb_smem_staged) if self.is_gated else sSFB
        cute.recast_tensor(sSFA, cutlass.Uint8)
        cute.recast_tensor(sSFB, cutlass.Uint8)
        cute.recast_tensor(sSFB_up, cutlass.Uint8)
        sC = storage.sC.get_tensor(
            epi_smem_staged.outer,
            swizzle=epi_smem_staged.inner,
        )
        sfa_base_addr = shared_ptr_to_u32(storage.sSFA.data_ptr())
        ctrl_base_addr = shared_ptr_to_u32(storage.ctrl.data_ptr())
        route_phys_rows_addr = shared_ptr_to_u32(storage.route_phys_rows.data_ptr())
        route_expert_ids_addr = shared_ptr_to_u32(storage.route_expert_ids.data_ptr())
        scatter_tok_base_addr = shared_ptr_to_u32(storage.scatter_tok_cache.data_ptr())
        scatter_weight_base_addr = shared_ptr_to_u32(
            storage.scatter_weight_cache.data_ptr()
        )

        num_tokens = Int32(a_input.shape[0])
        cols = Int32(a_input.shape[1])
        scatter_base = scatter_output.iterator.toint()
        row_counts = launch_params.row_counts
        num_experts = Int32(row_counts.shape[0])
        sf_blocks_per_row = cols // Int32(16)
        output_bytes_per_row = cols // Int32(2)
        cols_u32 = cols // Int32(2)
        scatter_output_u32 = cute.recast_tensor(scatter_output, cutlass.Uint32)
        total_pairs = Int32(topk_ids.shape[0])
        num_topk = total_pairs // num_tokens
        flat_tid = Int32(bidz) * Int32(self.threads_per_cta) + Int32(tidx)
        flat_stride = Int32(gdim_z) * Int32(self.threads_per_cta)
        num_k_tiles = (cols + Int32(63)) // Int32(64)
        route_gate_tile_cnt = launch_params.gate_tile_cnt
        task_slice_chunk = Int32(_TASK_SLICE_CHUNK)
        materialized_num_groups = (
            route_gate_tile_cnt + task_slice_chunk - Int32(1)
        ) // task_slice_chunk

        # Phase 0: cooperative init. Zero routing state, queue state, and
        # output. Task metadata slots are overwritten before the second grid
        # barrier and need no clear.
        i = flat_tid
        while i < num_experts:
            row_counts[i] = Int32(0)
            expert_write_rows[i] = Int32(0)
            i += flat_stride
        if flat_tid < num_experts + Int32(1):
            expert_tile_base[flat_tid] = Int32(0)

        scatter_total_u32 = Int32(scatter_output.shape[0]) * cols_u32
        scatter_vecs = scatter_total_u32 // Int32(4)
        zero_u32 = Uint32(0)
        zv = flat_tid
        while zv < scatter_vecs:
            st_global_v4_u32(
                scatter_base + Int64(zv) * Int64(16),
                zero_u32,
                zero_u32,
                zero_u32,
                zero_u32,
            )
            zv += flat_stride

        j = scatter_vecs * Int32(4) + flat_tid
        while j < scatter_total_u32:
            scatter_output_u32[j // cols_u32, j % cols_u32] = Uint32(0)
            j += flat_stride

        if flat_tid == Int32(0):
            pair_head[Int32(0)] = Int32(0)
            task_head[Int32(0)] = Int32(0)
            task_tail[Int32(0)] = Int32(0)

        cute.arch.sync_threads()
        self._resident_grid_barrier(
            barrier_count,
            barrier_epoch,
            Int32(gdim_z),
            is_cta_leader,
        )

        # Phase 1: histogram routed rows per expert.
        hist_idx = flat_tid
        while hist_idx < total_pairs:
            expert_id = topk_ids[hist_idx].to(Int32)
            atomic_add_global_i32(get_ptr_as_int64(row_counts, expert_id), Int32(1))
            hist_idx += flat_stride

        self._resident_grid_barrier(
            barrier_count,
            barrier_epoch,
            Int32(gdim_z),
            is_cta_leader,
        )

        if flat_tid == Int32(0):
            tile_acc = Int32(0)
            expert_idx = Int32(0)
            while expert_idx < num_experts:
                expert_tile_base[expert_idx] = tile_acc
                rows = row_counts[expert_idx]
                tile_acc += (rows + Int32(self.tile_shape_mnk[0]) - Int32(1)) // Int32(
                    self.tile_shape_mnk[0]
                )
                expert_idx += Int32(1)
            expert_tile_base[num_experts] = tile_acc

        self._resident_grid_barrier(
            barrier_count,
            barrier_epoch,
            Int32(gdim_z),
            is_cta_leader,
        )

        # Phase 2: warp-private route/pack producers into compact physical tiles.
        lane_id = Int32(tidx) & Int32(31)
        num_cta_warps = Int32(self.num_mma_warps + 1)
        producer_batch_pairs = num_cta_warps * Int32(_PRODUCER_PAIRS_PER_WARP)
        shared_input_gs_value = cutlass.Float32(0.0)
        if cutlass.const_expr(self.share_input_across_experts):
            shared_input_gs_value = input_global_scale[Int32(0)].to(cutlass.Float32)
            if (
                self.input_scales_are_reciprocal
                and shared_input_gs_value != cutlass.Float32(0.0)
            ):
                if self.fast_math:
                    shared_input_gs_value = rcp_approx_ftz(shared_input_gs_value)
                else:
                    shared_input_gs_value = cutlass.Float32(1.0) / shared_input_gs_value
        pair_idx = Int32(0)
        expert_id = Int32(0)
        token_idx = Int32(0)
        weight = cutlass.Float32(0.0)
        row = Int32(0)
        phys_tile = Int32(0)
        phys_row = Int32(0)
        produce_active = Int32(1)
        while produce_active > Int32(0):
            batch_base = Int32(0)
            if is_cta_leader > Int32(0):
                claim_count = producer_batch_pairs
                if cutlass.const_expr(self.share_input_across_experts):
                    claim_count = num_cta_warps
                batch_base = atomic_add_global_i32(
                    get_ptr_as_int64(pair_head, Int32(0)),
                    claim_count,
                )
                _st_shared_i32(ctrl_base_addr + Int32(28), batch_base)
            cute.arch.sync_threads()
            batch_base = _ld_shared_i32(ctrl_base_addr + Int32(28))
            producer_limit = total_pairs
            if cutlass.const_expr(self.share_input_across_experts):
                producer_limit = num_tokens
            if batch_base >= producer_limit:
                produce_active = Int32(0)
            else:
                if cutlass.const_expr(self.share_input_across_experts):
                    token_idx = batch_base + warp_idx
                    if token_idx < num_tokens:
                        route_slot_base = warp_idx * Int32(32)
                        if lane_id == Int32(0):
                            topk_slot = Int32(0)
                            while topk_slot < num_topk:
                                pair_idx = token_idx * num_topk + topk_slot
                                expert_id = topk_ids[pair_idx].to(Int32)
                                weight = topk_weights[pair_idx].to(cutlass.Float32)
                                row = atomic_add_global_i32(
                                    get_ptr_as_int64(expert_write_rows, expert_id),
                                    Int32(1),
                                )
                                phys_tile = expert_tile_base[expert_id] + row // Int32(
                                    self.tile_shape_mnk[0]
                                )
                                phys_row = phys_tile * Int32(
                                    self.tile_shape_mnk[0]
                                ) + row % Int32(self.tile_shape_mnk[0])
                                st_global_i32(
                                    get_ptr_as_int64(token_map, phys_row), token_idx
                                )
                                st_global_f32(
                                    get_ptr_as_int64(token_weights, phys_row), weight
                                )
                                slot = route_slot_base + topk_slot
                                _st_shared_i32(
                                    route_phys_rows_addr + slot * Int32(4), phys_row
                                )
                                _st_shared_i32(
                                    route_expert_ids_addr + slot * Int32(4), expert_id
                                )
                                topk_slot += Int32(1)
                        cute.arch.sync_warp()

                        gs_value = shared_input_gs_value
                        if num_topk == Int32(8):
                            route_output_base = cute.make_rmem_tensor((8,), Int32)
                            route_scale_base = cute.make_rmem_tensor((8,), Int32)
                            for cache_slot in cutlass.range_constexpr(8):
                                slot = route_slot_base + Int32(cache_slot)
                                phys_row = _ld_shared_i32(
                                    route_phys_rows_addr + slot * Int32(4)
                                )
                                route_output_base[cache_slot] = (
                                    phys_row * output_bytes_per_row
                                )
                                # Scale storage is tiled in 128-row SF atoms,
                                # independently of the MMA tile.
                                sf_atom = phys_row >> Int32(7)
                                sf_row = phys_row & Int32(127)
                                route_scale_base[cache_slot] = (
                                    sf_atom * num_k_tiles * Int32(32 * 4 * 4)
                                    + (sf_row % Int32(32)) * Int32(4 * 4)
                                    + (sf_row // Int32(32)) * Int32(4)
                                )

                            sf_idx = lane_id
                            while sf_idx < sf_blocks_per_row:
                                block_start = sf_idx * Int32(16)
                                values = cute.make_rmem_tensor((16,), cutlass.Float32)
                                block_max = cutlass.Float32(0.0)
                                for elem_idx in cutlass.range_constexpr(16):
                                    value = cutlass.Float32(
                                        a_input[
                                            token_idx, block_start + Int32(elem_idx)
                                        ]
                                    )
                                    values[elem_idx] = value
                                    block_max = fmax_f32(block_max, fabs_f32(value))
                                packed64 = Uint64(0)
                                scale_byte = Uint8(0)
                                if self.fast_math:
                                    packed64, scale_byte = quantize_block_fp4_fast(
                                        values, block_max, gs_value
                                    )
                                else:
                                    packed64, scale_byte = quantize_block_fp4(
                                        values, block_max, gs_value
                                    )

                                k_tile_idx = sf_idx // Int32(4)
                                scale_k_base = k_tile_idx * Int32(32 * 4 * 4) + (
                                    sf_idx % Int32(4)
                                )
                                for cache_slot in cutlass.range_constexpr(8):
                                    output_offset = route_output_base[
                                        cache_slot
                                    ] + sf_idx * Int32(8)
                                    st_global_u64(
                                        get_ptr_as_int64(
                                            packed_a_storage, output_offset
                                        ),
                                        packed64,
                                    )
                                    scale_storage[
                                        route_scale_base[cache_slot] + scale_k_base
                                    ] = scale_byte
                                sf_idx += Int32(32)
                        else:
                            sf_idx = lane_id
                            while sf_idx < sf_blocks_per_row:
                                block_start = sf_idx * Int32(16)
                                values = cute.make_rmem_tensor((16,), cutlass.Float32)
                                block_max = cutlass.Float32(0.0)
                                for elem_idx in cutlass.range_constexpr(16):
                                    value = cutlass.Float32(
                                        a_input[
                                            token_idx, block_start + Int32(elem_idx)
                                        ]
                                    )
                                    values[elem_idx] = value
                                    block_max = fmax_f32(block_max, fabs_f32(value))
                                packed64 = Uint64(0)
                                scale_byte = Uint8(0)
                                if self.fast_math:
                                    packed64, scale_byte = quantize_block_fp4_fast(
                                        values, block_max, gs_value
                                    )
                                else:
                                    packed64, scale_byte = quantize_block_fp4(
                                        values, block_max, gs_value
                                    )

                                topk_slot = Int32(0)
                                while topk_slot < num_topk:
                                    slot = route_slot_base + topk_slot
                                    phys_row = _ld_shared_i32(
                                        route_phys_rows_addr + slot * Int32(4)
                                    )
                                    output_offset = (
                                        phys_row * output_bytes_per_row
                                        + sf_idx * Int32(8)
                                    )
                                    st_global_u64(
                                        get_ptr_as_int64(
                                            packed_a_storage, output_offset
                                        ),
                                        packed64,
                                    )
                                    # Scale storage is tiled in 128-row SF
                                    # atoms, not MMA tiles.
                                    k_tile_idx = sf_idx // Int32(4)
                                    sf_atom = phys_row >> Int32(7)
                                    sf_row = phys_row & Int32(127)
                                    outer_m_idx = sf_row % Int32(32)
                                    inner_m_idx = sf_row // Int32(32)
                                    inner_k_idx = sf_idx % Int32(4)
                                    scale_offset = (
                                        sf_atom * num_k_tiles * Int32(32 * 4 * 4)
                                        + k_tile_idx * Int32(32 * 4 * 4)
                                        + outer_m_idx * Int32(4 * 4)
                                        + inner_m_idx * Int32(4)
                                        + inner_k_idx
                                    )
                                    scale_storage[scale_offset] = scale_byte
                                    topk_slot += Int32(1)
                                sf_idx += Int32(32)

                else:
                    warp_item = Int32(0)
                    while warp_item < Int32(_PRODUCER_PAIRS_PER_WARP):
                        pair_idx = batch_base + warp_idx + warp_item * num_cta_warps
                        expert_id = Int32(0)
                        token_idx = Int32(0)
                        weight = cutlass.Float32(0.0)
                        row = Int32(0)
                        phys_tile = Int32(0)
                        if pair_idx < total_pairs:
                            expert_id = topk_ids[pair_idx].to(Int32)
                            token_idx = pair_idx // num_topk
                            weight = topk_weights[pair_idx].to(cutlass.Float32)

                            if lane_id == Int32(0):
                                row = atomic_add_global_i32(
                                    get_ptr_as_int64(expert_write_rows, expert_id),
                                    Int32(1),
                                )
                                phys_tile = expert_tile_base[expert_id] + row // Int32(
                                    self.tile_shape_mnk[0]
                                )
                                phys_row = phys_tile * Int32(
                                    self.tile_shape_mnk[0]
                                ) + row % Int32(self.tile_shape_mnk[0])
                                st_global_i32(
                                    get_ptr_as_int64(token_map, phys_row), token_idx
                                )
                                st_global_f32(
                                    get_ptr_as_int64(token_weights, phys_row), weight
                                )

                            row = cute.arch.shuffle_sync(row, Int32(0))
                            phys_tile = cute.arch.shuffle_sync(phys_tile, Int32(0))
                            expert_id = cute.arch.shuffle_sync(expert_id, Int32(0))
                            token_idx = cute.arch.shuffle_sync(token_idx, Int32(0))

                            gs_value = input_global_scale[expert_id].to(cutlass.Float32)
                            if (
                                self.input_scales_are_reciprocal
                                and gs_value != cutlass.Float32(0.0)
                            ):
                                if self.fast_math:
                                    gs_value = rcp_approx_ftz(gs_value)
                                else:
                                    gs_value = cutlass.Float32(1.0) / gs_value
                            sf_idx = lane_id
                            while sf_idx < sf_blocks_per_row:
                                block_start = sf_idx * Int32(16)
                                values = cute.make_rmem_tensor((16,), cutlass.Float32)
                                block_max = cutlass.Float32(0.0)
                                for elem_idx in cutlass.range_constexpr(16):
                                    value = cutlass.Float32(
                                        a_input[
                                            token_idx, block_start + Int32(elem_idx)
                                        ]
                                    )
                                    values[elem_idx] = value
                                    block_max = fmax_f32(block_max, fabs_f32(value))
                                packed64 = Uint64(0)
                                scale_byte = Uint8(0)
                                if self.fast_math:
                                    packed64, scale_byte = quantize_block_fp4_fast(
                                        values, block_max, gs_value
                                    )
                                else:
                                    packed64, scale_byte = quantize_block_fp4(
                                        values, block_max, gs_value
                                    )

                                output_offset = (
                                    phys_tile * Int32(self.tile_shape_mnk[0])
                                    + row % Int32(self.tile_shape_mnk[0])
                                ) * output_bytes_per_row + sf_idx * Int32(8)
                                st_global_u64(
                                    get_ptr_as_int64(packed_a_storage, output_offset),
                                    packed64,
                                )

                                # Scale storage is tiled in 128-row SF atoms,
                                # not MMA tiles.
                                k_tile_idx = sf_idx // Int32(4)
                                inner_k_idx = sf_idx % Int32(4)
                                phys_row = phys_tile * Int32(
                                    self.tile_shape_mnk[0]
                                ) + row % Int32(self.tile_shape_mnk[0])
                                sf_atom = phys_row >> Int32(7)
                                sf_row = phys_row & Int32(127)
                                scale_offset = (
                                    sf_atom * num_k_tiles * Int32(32 * 4 * 4)
                                    + k_tile_idx * Int32(32 * 4 * 4)
                                    + (sf_row % Int32(32)) * Int32(4 * 4)
                                    + (sf_row // Int32(32)) * Int32(4)
                                    + inner_k_idx
                                )
                                scale_storage[scale_offset] = scale_byte
                                sf_idx += Int32(32)

                        warp_item += Int32(1)

        cute.arch.sync_threads()
        # Conservative publish fence before the last-producer CTA flushes any
        # partial tiles. All producer threads in the CTA must have ordered
        # their global writes before lane 0 can publish work.
        _threadfence()
        cute.arch.sync_threads()

        # Rendezvous once, publish every physical tile, then consume the
        # finished queue.
        self._resident_grid_barrier(
            barrier_count,
            barrier_epoch,
            Int32(gdim_z),
            is_cta_leader,
        )

        if is_cta_leader > Int32(0):
            expert_flush = Int32(bidz)
            while expert_flush < num_experts:
                rows_remaining = row_counts[expert_flush]
                m_tile_offset = Int32(0)
                while rows_remaining > Int32(0):
                    valid_rows = rows_remaining
                    if valid_rows > Int32(self.tile_shape_mnk[0]):
                        valid_rows = Int32(self.tile_shape_mnk[0])
                    self._publish_deferred_tasks(
                        task_expert,
                        task_valid_rows,
                        route_gate_tile_cnt,
                        task_slice_chunk,
                        expert_flush,
                        expert_tile_base[expert_flush] + m_tile_offset,
                        valid_rows,
                    )
                    rows_remaining -= Int32(self.tile_shape_mnk[0])
                    m_tile_offset += Int32(1)
                expert_flush += Int32(gdim_z)

        if flat_tid == Int32(0):
            st_global_i32(
                get_ptr_as_int64(task_tail, Int32(0)),
                expert_tile_base[num_experts] * materialized_num_groups,
            )

        self._resident_grid_barrier(
            barrier_count,
            barrier_epoch,
            Int32(gdim_z),
            is_cta_leader,
        )

        gA = cute.local_tile(mA, self.sa_tile_shape_mk, (None, None, None))
        # Tiled view over w13.
        # Gated: [2*I_tp, K, E] packed as [up, gate] across N.
        #   Up tiles: N-indices 0..gate_tile_cnt-1
        #   Gate tiles: N-indices gate_tile_cnt..2*gate_tile_cnt-1
        # ReLU2: [I_tp, K, E] single FC1 pass, gate_tile_cnt == intermediate_tile_cnt.
        gB_w13_tiled = cute.local_tile(
            mB_w13,
            cute.slice_(self.tile_shape_mnk, (0, None, None)),
            (None, None, None),
        )
        # A/SF tiles use the 128-row block shape; for sub-128 MMA tiles one
        # block backs `sa_tiles_per_block` MMA tiles (offset applied below).
        gSFA = cute.local_tile(mSFA, self.sfa_tile_shape_mk, (None, None, None))
        gSFB_w13_tiled = cute.local_tile(
            mSFB_w13, self.sfb_tile_shape_nk, (None, None, None)
        )
        thr_mma = tiled_mma.get_slice(tidx)

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

        tAsA, tAgA = cpasync.tma_partition(
            tma_a,
            a_cta_crd,
            a_cta_layout,
            cute.group_modes(sA, 0, 2),
            cute.group_modes(gA, 0, 2),
        )
        tAsSFA, tAgSFA = cpasync.tma_partition(
            tma_sfa,
            a_cta_crd,
            a_cta_layout,
            cute.group_modes(sSFA, 0, 2),
            cute.group_modes(gSFA, 0, 2),
        )
        tAsSFA = cute.filter_zeros(tAsSFA)
        tAgSFA = cute.filter_zeros(tAgSFA)

        # w13 TMA partition (gated: gate+up concatenated; relu2: single W1)
        tBsB_w13, tBgB_w13 = cpasync.tma_partition(
            tma_b_w13,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sB, 0, 2),
            cute.group_modes(gB_w13_tiled, 0, 2),
        )
        tBsB_w13_up, _ = cpasync.tma_partition(
            tma_b_w13,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sB_up, 0, 2),
            cute.group_modes(gB_w13_tiled, 0, 2),
        )
        tBsSFB_w13, tBgSFB_w13 = cpasync.tma_partition(
            tma_sfb_w13,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sSFB, 0, 2),
            cute.group_modes(gSFB_w13_tiled, 0, 2),
        )
        tBsSFB_w13_up, _ = cpasync.tma_partition(
            tma_sfb_w13,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sSFB_up, 0, 2),
            cute.group_modes(gSFB_w13_tiled, 0, 2),
        )
        tBsB_w13_up = cute.filter_zeros(tBsB_w13_up)
        tBsSFB_w13 = cute.filter_zeros(tBsSFB_w13)
        tBgSFB_w13 = cute.filter_zeros(tBgSFB_w13)
        tBsSFB_w13_up = cute.filter_zeros(tBsSFB_w13_up)

        # B_down TMA partitions
        gB_down = cute.local_tile(
            mB_down,
            cute.slice_(self.tile_shape_mnk, (0, None, None)),
            (None, None, None),
        )
        gSFB_down = cute.local_tile(
            mSFB_down,
            cute.slice_(self.tile_shape_mnk, (0, None, None)),
            (None, None, None),
        )
        tBsB_down, tBgB_down = cpasync.tma_partition(
            tma_b_down,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sB, 0, 2),
            cute.group_modes(gB_down, 0, 2),
        )
        tBsSFB_down, tBgSFB_down = cpasync.tma_partition(
            tma_sfb_down,
            b_cta_crd,
            b_cta_layout,
            cute.group_modes(sSFB, 0, 2),
            cute.group_modes(gSFB_down, 0, 2),
        )
        tBsSFB_down = cute.filter_zeros(tBsSFB_down)
        tBgSFB_down = cute.filter_zeros(tBgSFB_down)

        # MMA fragment partitions
        # sA/sSFA hold the whole 128-row block; slice to the tile_m sub-tile
        # the V-map expects (the per-task offset is applied at consumption).
        if cutlass.const_expr(self.sa_tiles_per_block > 1):
            sA_part = cute.local_tile(
                sA,
                cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                (Int32(0), 0, None),
            )
        else:
            sA_part = sA
        tCsA = thr_mma.partition_A(sA_part)
        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
        if cutlass.const_expr(self.sfa_tiles_per_block > 1):
            sSFA_part = cute.local_tile(
                sSFA,
                cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                (Int32(0), 0, None),
            )
        else:
            sSFA_part = sSFA
        tCrSFA = self._dense_cls._partition_fragment_SFA(
            self,  # type: ignore[arg-type]
            sSFA_part[None, None, 0],
            thr_mma,
            tidx,
        )
        tCsB = thr_mma.partition_B(sB)
        tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
        tCrSFB = self._dense_cls._partition_fragment_SFB(
            self,  # type: ignore[arg-type]
            sSFB[None, None, 0],
            thr_mma,
            tidx,
        )

        tCsC_for_shape = thr_mma.partition_C(sC[None, None, 0])
        epi_m_scale = self.tile_shape_mnk[0] // self.epi_tile[0]
        sub_shape = tCsC_for_shape.shape[:3]
        acc_shape = (sub_shape[0], sub_shape[1] * epi_m_scale, sub_shape[2])
        gate_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
        up_acc = (
            cute.make_rmem_tensor(acc_shape, self.acc_dtype)
            if self.is_gated
            else gate_acc
        )

        k_tile_cnt = cute.size(gA, mode=[3])
        fc1_k_tile_cnt = k_tile_cnt
        # Gated: w13 has 2*I_tp/tile_N N-tiles. Gate = second half, up = first half.
        # ReLU2: w13 has I_tp/tile_N N-tiles. Single FC1 pass, no split.
        intermediate_tile_cnt = cute.size(gB_w13_tiled, mode=[2])
        gate_tile_cnt = (
            intermediate_tile_cnt // Int32(2)
            if self.is_gated
            else intermediate_tile_cnt
        )
        output_tile_cnt = cute.size(gB_down, mode=[2])

        prod_state = pipeline.make_pipeline_state(
            pipeline.PipelineUserType.Producer, self.ab_stage
        )
        cons_state = pipeline.make_pipeline_state(
            pipeline.PipelineUserType.Consumer, self.ab_stage
        )
        up_prod_state = (
            pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.ab_stage
            )
            if self.is_gated
            else prod_state
        )
        up_cons_state = (
            pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.ab_stage
            )
            if self.is_gated
            else cons_state
        )
        phase2_prod_state = pipeline.make_pipeline_state(
            pipeline.PipelineUserType.Producer, self.ab_stage
        )
        phase2_cons_state = pipeline.make_pipeline_state(
            pipeline.PipelineUserType.Consumer, self.ab_stage
        )

        num_k_blocks = cute.size(tCrA, mode=[2])

        atom_ld_A = cute.make_copy_atom(
            cute.nvgpu.warp.LdMatrix8x8x16bOp(self.a_layout.is_m_major_a(), 4),
            self.a_dtype,
        )
        atom_ld_B = cute.make_copy_atom(
            cute.nvgpu.warp.LdMatrix8x8x16bOp(self.b_layout.is_n_major_b(), 4),
            self.b_dtype,
        )
        smem_copy_A = cute.make_tiled_copy_A(atom_ld_A, tiled_mma)
        smem_copy_B = cute.make_tiled_copy_B(atom_ld_B, tiled_mma)
        atom_ld_SF = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.sf_dtype)
        smem_copy_SFA = cute.make_tiled_copy(
            atom_ld_SF,
            self._dense_cls._get_layoutSFA_TV(self, tiled_mma),  # type: ignore[arg-type]
            (
                cute.size(tiled_mma.permutation_mnk[0]),
                cute.size(tiled_mma.permutation_mnk[2]),
            ),
        )
        smem_copy_SFB = cute.make_tiled_copy(
            atom_ld_SF,
            self._dense_cls._get_layoutSFB_TV(self, tiled_mma),  # type: ignore[arg-type]
            (
                cute.size(tiled_mma.permutation_mnk[1]),
                cute.size(tiled_mma.permutation_mnk[2]),
            ),
        )

        thr_ld_A = smem_copy_A.get_slice(tidx)
        thr_ld_B = smem_copy_B.get_slice(tidx)
        csA = thr_ld_A.partition_S(sA_part)
        crA = thr_ld_A.retile(tCrA)
        csB = thr_ld_B.partition_S(sB)
        csB_up = thr_ld_B.partition_S(sB_up)
        crB = thr_ld_B.retile(tCrB)

        thr_ld_SFA = smem_copy_SFA.get_slice(tidx)
        thr_ld_SFB = smem_copy_SFB.get_slice(tidx)
        csSFA = thr_ld_SFA.partition_S(sSFA_part)
        crSFA = thr_ld_SFA.retile(tCrSFA)
        csSFB = thr_ld_SFB.partition_S(sSFB)
        csSFB_up = thr_ld_SFB.partition_S(sSFB_up)
        crSFB = thr_ld_SFB.retile(tCrSFB)

        # ===================================================================
        # Per-warp setup for the consumer steady state
        # ===================================================================
        if warp_idx < self.num_mma_warps:
            cute.arch.setmaxregister_increase(self.mma_register_requirement)
        elif warp_idx == self.tma_load_warp_id:
            cute.arch.setmaxregister_decrease(self.load_register_requirement)

        # ===================================================================
        # Consumer steady state: pop one ready task per CTA, then let
        # the MMA warps and DMA warp cooperate on that task.
        # ===================================================================
        work_item = cute.make_rmem_tensor((_WORK_ITEM_FIELDS,), Int32)
        consumer_live = Int32(1)
        while consumer_live > Int32(0):
            if is_cta_leader > Int32(0):
                _st_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4), Int32(0))
                _st_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4), Int32(0))
                slot = atomic_add_global_i32(
                    get_ptr_as_int64(task_head, Int32(0)), Int32(1)
                )
                # task_tail is final after the pre-consume grid barrier, so an
                # out-of-range slot is a definitive done signal.
                tail = _ld_global_acquire_i32(get_ptr_as_int64(task_tail, Int32(0)))
                if slot < tail:
                    self._decode_materialized_work_item(
                        work_item,
                        task_expert,
                        task_valid_rows,
                        slot,
                        materialized_num_groups,
                        task_slice_chunk,
                        route_gate_tile_cnt,
                    )
                    self._store_shared_work_item(ctrl_base_addr, work_item)
                    _st_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4), Int32(1))
                else:
                    _st_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4), Int32(1))
            cute.arch.sync_threads()

            has_task = _ld_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4))
            is_done = _ld_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4))
            if has_task > Int32(0):
                self._load_shared_work_item(work_item, ctrl_base_addr)
                task_m_tile_idx_cache = work_item[_WORK_M_TILE]
                task_valid_rows_cache = work_item[_WORK_VALID_ROWS]
                tile_m_base_cache = task_m_tile_idx_cache * Int32(
                    self.tile_shape_mnk[0]
                )
                cache_row = Int32(tidx)
                while cache_row < Int32(self.tile_shape_mnk[0]):
                    tok = Int32(0)
                    wv = cutlass.Float32(0.0)
                    if cache_row < task_valid_rows_cache:
                        global_row_cache = tile_m_base_cache + cache_row
                        tok = token_map[global_row_cache].to(Int32)
                        wv = token_weights[global_row_cache].to(cutlass.Float32)
                    _st_shared_i32(scatter_tok_base_addr + cache_row * Int32(4), tok)
                    st_shared_f32(scatter_weight_base_addr + cache_row * Int32(4), wv)
                    cache_row += Int32(self.threads_per_cta)
                cute.arch.sync_threads()
            if has_task == Int32(0):
                if is_done > Int32(0):
                    consumer_live = Int32(0)
            elif warp_idx < self.num_mma_warps:
                task_expert_idx = work_item[_WORK_EXPERT]
                task_m_tile_idx = work_item[_WORK_M_TILE]
                task_slice_begin_idx = work_item[_WORK_SLICE_BEGIN]
                task_slice_count_val = work_item[_WORK_SLICE_COUNT]
                task_valid_rows_val = work_item[_WORK_VALID_ROWS]

                alpha_value = alpha[task_expert_idx].to(cutlass.Float32)
                valid_rows = task_valid_rows_val

                # FC1's activation rows sit at offset (task_m_tile_idx %
                # sfa_tiles_per_block) within the shared 128-row block; FC2
                # re-slices at offset 0 since its intermediate is written there.
                if cutlass.const_expr(self.sfa_tiles_per_block > 1):
                    _fc1_off = task_m_tile_idx % Int32(self.sfa_tiles_per_block)
                    _sA_il = cute.local_tile(
                        sA,
                        cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                        (_fc1_off, 0, None),
                    )
                    tCrA = tiled_mma.make_fragment_A(
                        thr_mma.partition_A(_sA_il)[None, None, None, 0]
                    )
                    csA = thr_ld_A.partition_S(_sA_il)
                    crA = thr_ld_A.retile(tCrA)
                    _sSFA_il = cute.local_tile(
                        sSFA,
                        cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                        (_fc1_off, 0, None),
                    )
                    tCrSFA = self._dense_cls._partition_fragment_SFA(
                        self,  # type: ignore[arg-type]
                        _sSFA_il[None, None, 0],
                        thr_mma,
                        tidx,
                    )
                    csSFA = thr_ld_SFA.partition_S(_sSFA_il)
                    crSFA = thr_ld_SFA.retile(tCrSFA)

                _is_m_major = self.c_layout.is_m_major_c()
                copy_atom_r2s = cute.make_copy_atom(
                    cute.nvgpu.CopyUniversalOp(),
                    cutlass.BFloat16,
                )
                copy_atom_C = cute.make_copy_atom(
                    cute.nvgpu.warp.StMatrix8x8x16bOp(_is_m_major, 2),
                    cutlass.BFloat16,
                )
                tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
                tiled_copy_r2s = cute.make_tiled_copy_S(
                    copy_atom_r2s, tiled_copy_C_Atom
                )

                thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
                tRS_sD = thr_copy_r2s.partition_D(sC)
                tRS_rGate = tiled_copy_r2s.retile(gate_acc)
                tRS_rUp = tiled_copy_r2s.retile(up_acc)

                rD_shape = cute.shape(thr_copy_r2s.partition_S(sC))
                tRS_rD_layout = cute.make_layout(rD_shape[:3])
                tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype)
                tRS_rD_out = cute.make_rmem_tensor(
                    tRS_rD_layout.shape, cutlass.BFloat16
                )

                mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rGate, mode=[1])
                mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rGate, mode=[2])
                epi_buffer = Int32(0)

                down_alpha_value = down_alpha[task_expert_idx].to(cutlass.Float32)
                down_acc = cute.make_rmem_tensor(acc_shape, self.acc_dtype)

                epi_rest_m = self.tile_shape_mnk[0] // self.epi_tile[0]
                MmaMPerEpiM = self.epi_tile[0] // mma_tile_m
                MmaNPerEpiN = self.epi_tile[1] // mma_tile_n
                fc2_m_tiles = cute.size(tCrA, mode=[1])
                fc2_n_tiles = cute.size(tCrB, mode=[1])

                fc1_m_tiles = cute.size(tCrA, mode=[1])
                fc1_n_tiles = cute.size(tCrB, mode=[1])
                slice_idx = Int32(0)
                while slice_idx < task_slice_count_val:
                    # ============================================================
                    # PHASE A: FC1 for this slice (gate + up)
                    # ============================================================

                    # Gate GEMM (inlined to avoid @cute.jit pass-by-value for acc)
                    fz_crSFA = cute.filter_zeros(crSFA)
                    fz_crSFB = cute.filter_zeros(crSFB)
                    gate_acc.fill(0.0)
                    cons_state.reset_count()
                    peek = ml_pipeline.consumer_try_wait(cons_state)
                    ml_pipeline.consumer_wait(cons_state, peek)
                    csA_p = csA[None, None, None, cons_state.index]
                    csB_p = csB[None, None, None, cons_state.index]
                    csSFA_p = csSFA[None, None, None, cons_state.index]
                    csSFB_p = csSFB[None, None, None, cons_state.index]
                    cute.copy(smem_copy_A, csA_p[None, None, 0], crA[None, None, 0])
                    cute.copy(smem_copy_B, csB_p[None, None, 0], crB[None, None, 0])
                    fz_csSFA_p = cute.filter_zeros(csSFA_p)
                    fz_csSFB_p = cute.filter_zeros(csSFB_p)
                    cute.copy(
                        smem_copy_SFA,
                        fz_csSFA_p[None, None, 0],
                        fz_crSFA[None, None, 0],
                    )
                    cute.copy(
                        smem_copy_SFB,
                        fz_csSFB_p[None, None, 0],
                        fz_crSFB[None, None, 0],
                    )
                    for _k_tile in range(0, fc1_k_tile_cnt - 1, 1, unroll=4):  # type: ignore[call-overload]
                        for k_block_idx in cutlass.range_constexpr(num_k_blocks):
                            k_next = (
                                0
                                if k_block_idx + 1 == num_k_blocks
                                else k_block_idx + 1
                            )
                            if k_block_idx == num_k_blocks - 1:
                                ml_pipeline.consumer_release(cons_state)
                                cons_state.advance()
                                peek = ml_pipeline.consumer_try_wait(cons_state)
                                csA_p = csA[None, None, None, cons_state.index]
                                csB_p = csB[None, None, None, cons_state.index]
                                csSFA_p = csSFA[None, None, None, cons_state.index]
                                csSFB_p = csSFB[None, None, None, cons_state.index]
                                fz_csSFA_p = cute.filter_zeros(csSFA_p)
                                fz_csSFB_p = cute.filter_zeros(csSFB_p)
                                ml_pipeline.consumer_wait(cons_state, peek)
                            for _mt in cutlass.range_constexpr(fc1_m_tiles):
                                for _nt in cutlass.range_constexpr(fc1_n_tiles):
                                    mma_atom.set(
                                        WarpField.SFA,
                                        tCrSFA[None, _mt, k_block_idx].iterator,
                                    )
                                    mma_atom.set(
                                        WarpField.SFB,
                                        tCrSFB[None, _nt, k_block_idx].iterator,
                                    )
                                    cute.gemm(
                                        mma_atom,
                                        gate_acc[None, _mt, _nt],
                                        tCrA[None, _mt, k_block_idx],
                                        tCrB[None, _nt, k_block_idx],
                                        gate_acc[None, _mt, _nt],
                                    )
                            cute.copy(
                                smem_copy_A,
                                csA_p[None, None, k_next],
                                crA[None, None, k_next],
                            )
                            cute.copy(
                                smem_copy_B,
                                csB_p[None, None, k_next],
                                crB[None, None, k_next],
                            )
                            fz_csSFA_cur = cute.filter_zeros(
                                csSFA[None, None, None, cons_state.index]
                            )
                            fz_csSFB_cur = cute.filter_zeros(
                                csSFB[None, None, None, cons_state.index]
                            )
                            cute.copy(
                                smem_copy_SFA,
                                fz_csSFA_cur[None, None, k_next],
                                fz_crSFA[None, None, k_next],
                            )
                            cute.copy(
                                smem_copy_SFB,
                                fz_csSFB_cur[None, None, k_next],
                                fz_crSFB[None, None, k_next],
                            )
                    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
                        k_next = (
                            0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1
                        )
                        if k_block_idx == num_k_blocks - 1:
                            ml_pipeline.consumer_release(cons_state)
                            cons_state.advance()
                        if k_next > 0 and fc1_k_tile_cnt > Int32(0):
                            cute.copy(
                                smem_copy_A,
                                csA_p[None, None, k_next],
                                crA[None, None, k_next],
                            )
                            cute.copy(
                                smem_copy_B,
                                csB_p[None, None, k_next],
                                crB[None, None, k_next],
                            )
                            cute.copy(
                                smem_copy_SFA,
                                fz_csSFA_p[None, None, k_next],
                                fz_crSFA[None, None, k_next],
                            )
                            cute.copy(
                                smem_copy_SFB,
                                fz_csSFB_p[None, None, k_next],
                                fz_crSFB[None, None, k_next],
                            )
                        for _mt in cutlass.range_constexpr(fc1_m_tiles):
                            for _nt in cutlass.range_constexpr(fc1_n_tiles):
                                mma_atom.set(
                                    WarpField.SFA,
                                    tCrSFA[None, _mt, k_block_idx].iterator,
                                )
                                mma_atom.set(
                                    WarpField.SFB,
                                    tCrSFB[None, _nt, k_block_idx].iterator,
                                )
                                cute.gemm(
                                    mma_atom,
                                    gate_acc[None, _mt, _nt],
                                    tCrA[None, _mt, k_block_idx],
                                    tCrB[None, _nt, k_block_idx],
                                    gate_acc[None, _mt, _nt],
                                )
                    # Signal FC1 gate/only completion before the DMA warp
                    # reuses the shared A/gate buffers for the next pass.
                    self.pass_gate_barrier.arrive_unaligned()

                    if cutlass.const_expr(self.is_gated):
                        # Up GEMM (inlined, same pattern)
                        up_acc.fill(0.0)
                        up_cons_state.reset_count()
                        peek = up_pipeline.consumer_try_wait(up_cons_state)
                        up_pipeline.consumer_wait(up_cons_state, peek)
                        csA_p = csA[None, None, None, up_cons_state.index]
                        csB_p = csB_up[None, None, None, up_cons_state.index]
                        csSFA_p = csSFA[None, None, None, up_cons_state.index]
                        csSFB_p = csSFB_up[None, None, None, up_cons_state.index]
                        cute.copy(smem_copy_A, csA_p[None, None, 0], crA[None, None, 0])
                        cute.copy(smem_copy_B, csB_p[None, None, 0], crB[None, None, 0])
                        fz_csSFA_p = cute.filter_zeros(csSFA_p)
                        fz_csSFB_p = cute.filter_zeros(csSFB_p)
                        cute.copy(
                            smem_copy_SFA,
                            fz_csSFA_p[None, None, 0],
                            fz_crSFA[None, None, 0],
                        )
                        cute.copy(
                            smem_copy_SFB,
                            fz_csSFB_p[None, None, 0],
                            fz_crSFB[None, None, 0],
                        )
                        for _k_tile in range(0, fc1_k_tile_cnt - 1, 1, unroll=4):  # type: ignore[call-overload]
                            for k_block_idx in cutlass.range_constexpr(num_k_blocks):
                                k_next = (
                                    0
                                    if k_block_idx + 1 == num_k_blocks
                                    else k_block_idx + 1
                                )
                                if k_block_idx == num_k_blocks - 1:
                                    up_pipeline.consumer_release(up_cons_state)
                                    up_cons_state.advance()
                                    peek = up_pipeline.consumer_try_wait(up_cons_state)
                                    csA_p = csA[None, None, None, up_cons_state.index]
                                    csB_p = csB_up[
                                        None, None, None, up_cons_state.index
                                    ]
                                    csSFA_p = csSFA[
                                        None, None, None, up_cons_state.index
                                    ]
                                    csSFB_p = csSFB_up[
                                        None, None, None, up_cons_state.index
                                    ]
                                    fz_csSFA_p = cute.filter_zeros(csSFA_p)
                                    fz_csSFB_p = cute.filter_zeros(csSFB_p)
                                    up_pipeline.consumer_wait(up_cons_state, peek)
                                for _mt in cutlass.range_constexpr(fc1_m_tiles):
                                    for _nt in cutlass.range_constexpr(fc1_n_tiles):
                                        mma_atom.set(
                                            WarpField.SFA,
                                            tCrSFA[None, _mt, k_block_idx].iterator,
                                        )
                                        mma_atom.set(
                                            WarpField.SFB,
                                            tCrSFB[None, _nt, k_block_idx].iterator,
                                        )
                                        cute.gemm(
                                            mma_atom,
                                            up_acc[None, _mt, _nt],
                                            tCrA[None, _mt, k_block_idx],
                                            tCrB[None, _nt, k_block_idx],
                                            up_acc[None, _mt, _nt],
                                        )
                                cute.copy(
                                    smem_copy_A,
                                    csA_p[None, None, k_next],
                                    crA[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_B,
                                    csB_p[None, None, k_next],
                                    crB[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_SFA,
                                    fz_csSFA_p[None, None, k_next],
                                    fz_crSFA[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_SFB,
                                    fz_csSFB_p[None, None, k_next],
                                    fz_crSFB[None, None, k_next],
                                )
                        for k_block_idx in cutlass.range_constexpr(num_k_blocks):
                            k_next = (
                                0
                                if k_block_idx + 1 == num_k_blocks
                                else k_block_idx + 1
                            )
                            if k_block_idx == num_k_blocks - 1:
                                up_pipeline.consumer_release(up_cons_state)
                                up_cons_state.advance()
                            if k_next > 0 and fc1_k_tile_cnt > Int32(0):
                                cute.copy(
                                    smem_copy_A,
                                    csA_p[None, None, k_next],
                                    crA[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_B,
                                    csB_p[None, None, k_next],
                                    crB[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_SFA,
                                    fz_csSFA_p[None, None, k_next],
                                    fz_crSFA[None, None, k_next],
                                )
                                cute.copy(
                                    smem_copy_SFB,
                                    fz_csSFB_p[None, None, k_next],
                                    fz_crSFB[None, None, k_next],
                                )
                            for _mt in cutlass.range_constexpr(fc1_m_tiles):
                                for _nt in cutlass.range_constexpr(fc1_n_tiles):
                                    mma_atom.set(
                                        WarpField.SFA,
                                        tCrSFA[None, _mt, k_block_idx].iterator,
                                    )
                                    mma_atom.set(
                                        WarpField.SFB,
                                        tCrSFB[None, _nt, k_block_idx].iterator,
                                    )
                                    cute.gemm(
                                        mma_atom,
                                        up_acc[None, _mt, _nt],
                                        tCrA[None, _mt, k_block_idx],
                                        tCrB[None, _nt, k_block_idx],
                                        up_acc[None, _mt, _nt],
                                    )
                    # Activation + quant into sA
                    sA_u8 = cute.recast_tensor(sA[None, None, 0], cutlass.Uint8)
                    packed_cols = Int32(self.tile_shape_mnk[2] // 2)
                    sf_blocks_per_row = Int32(self.tile_shape_mnk[2] // 16)
                    gs_value = global_scale[task_expert_idx].to(cutlass.Float32)
                    if self.input_scales_are_reciprocal and gs_value != cutlass.Float32(
                        0.0
                    ):
                        if self.fast_math:
                            gs_value = rcp_approx_ftz(gs_value)
                        else:
                            gs_value = cutlass.Float32(1.0) / gs_value

                    for epi_m in cutlass.range_constexpr(epi_rest_m):
                        epi_m_valid = valid_rows - Int32(epi_m) * Int32(
                            self.epi_tile[0]
                        )
                        silu_epi_buffer = Int32(epi_m) % cute.size(tRS_sD, mode=[3])
                        if epi_m_valid > Int32(0):
                            for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN):
                                for mma_m_in_epi in cutlass.range_constexpr(
                                    MmaMPerEpiM
                                ):
                                    mma_m = epi_m * MmaMPerEpiM + mma_m_in_epi
                                    mma_n = mma_n_in_epi
                                    tRS_rD_slice = tRS_rD[
                                        (None, mma_m_in_epi, mma_n_in_epi)
                                    ]
                                    gate_slice = tRS_rGate[(None, mma_m, mma_n)]
                                    if cutlass.const_expr(self.is_gated):
                                        up_slice = tRS_rUp[(None, mma_m, mma_n)]
                                        for elem_idx in cutlass.range_constexpr(
                                            cute.size(tRS_rD_slice)
                                        ):
                                            g = alpha_value * gate_slice[elem_idx]
                                            u = alpha_value * up_slice[elem_idx]
                                            tRS_rD_slice[elem_idx] = (
                                                gated_activation_f32(
                                                    g,
                                                    u,
                                                    activation=self.activation,
                                                    limit=self.swiglu_limit,
                                                    alpha=self.swiglu_alpha,
                                                    beta=self.swiglu_beta,
                                                    fast_math=self.fast_math,
                                                )
                                            )
                                    else:
                                        for elem_idx in cutlass.range_constexpr(
                                            cute.size(tRS_rD_slice)
                                        ):
                                            g = alpha_value * gate_slice[elem_idx]
                                            relu_g = fmax_f32(g, cutlass.Float32(0.0))
                                            tRS_rD_slice[elem_idx] = relu_g * relu_g

                            acc_vec = tRS_rD.load()
                            acc_vec = acc_vec.to(cutlass.BFloat16)
                            tRS_rD_out.store(acc_vec)
                            cute.copy(
                                tiled_copy_r2s,
                                tRS_rD_out,
                                tRS_sD[(None, None, None, silu_epi_buffer)],
                            )
                            cute.arch.fence_proxy("async.shared", space="cta")
                        self.epilog_sync_barrier.arrive_and_wait()

                        rows_offset = Int32(epi_m) * Int32(self.epi_tile[0])
                        epi_rows = epi_m_valid
                        if epi_rows > Int32(self.epi_tile[0]):
                            epi_rows = Int32(self.epi_tile[0])
                        if epi_rows < Int32(0):
                            epi_rows = Int32(0)
                        quant_idx = Int32(tidx)
                        while quant_idx < epi_rows * sf_blocks_per_row:
                            local_row = quant_idx // sf_blocks_per_row
                            row = rows_offset + local_row
                            sf_block = quant_idx - local_row * sf_blocks_per_row
                            block_start = sf_block * Int32(16)

                            values = cute.make_rmem_tensor((16,), cutlass.Float32)
                            block_max = cutlass.Float32(0.0)
                            for elem_idx in cutlass.range_constexpr(16):
                                value = cutlass.Float32(
                                    sC[
                                        local_row,
                                        block_start + elem_idx,
                                        silu_epi_buffer,
                                    ]
                                )
                                values[elem_idx] = value
                                block_max = fmax_f32(block_max, fabs_f32(value))

                            packed64 = Uint64(0)
                            scale_byte = Uint8(0)
                            if self.fast_math:
                                packed64, scale_byte = quantize_block_fp4_fast(
                                    values, block_max, gs_value
                                )
                            else:
                                packed64, scale_byte = quantize_block_fp4(
                                    values, block_max, gs_value
                                )
                            packed_base = sf_block << Int32(3)
                            dst_pcol = row & Int32(63)
                            xor_bits = ((dst_pcol >> Int32(1)) & Int32(0x3)) << Int32(4)
                            row_high = row >> Int32(6)
                            for byte_idx in cutlass.range_constexpr(8):
                                src_pcol = packed_base + Int32(byte_idx)
                                dst_row = ((src_pcol ^ xor_bits) << Int32(1)) + row_high
                                dst_flat = dst_row * packed_cols + dst_pcol
                                byte_val = Uint8(
                                    (packed64 >> Uint64(byte_idx * 8)) & Uint64(0xFF)
                                )
                                sA_u8[dst_flat] = byte_val

                            outer_m_idx = row % Int32(32)
                            inner_m_idx = row // Int32(32)
                            inner_k_idx = sf_block % Int32(4)
                            k_tile_idx = sf_block // Int32(4)
                            sf_raw_idx = (
                                k_tile_idx * Int32(32 * 4 * 4)
                                + outer_m_idx * Int32(4 * 4)
                                + inner_m_idx * Int32(4)
                                + inner_k_idx
                            )
                            st_shared_u8(sfa_base_addr + sf_raw_idx, scale_byte)
                            quant_idx += Int32(
                                self.num_mma_warps * self.num_threads_per_warp
                            )

                    cute.arch.fence_proxy("async.shared", space="cta")
                    # epilog_sync: MMA-only barrier. DMA warp doesn't need to wait
                    # for quant — it only loads B_down into sB (separate buffer).
                    # This allows DMA to prefetch B_down tiles earlier.
                    self.epilog_sync_barrier.arrive_and_wait()

                    # ============================================================
                    # PHASE B: Sweep ALL FC2 output tiles using cached sA
                    # No CTA-wide barrier needed here: gate is done with sB/sSFB
                    # (barrier at line 925 ensured that), up uses sB_up/sSFB_up,
                    # and DMA's B_down loads into sB/sSFB don't conflict with
                    # MMA's SiLU+quant on sC/sA/sSFA. The phase2_pipeline
                    # handles B_down availability for FC2 GEMM.
                    # ============================================================
                    scatter_N = Int32(scatter_output.shape[1])
                    lane_id = Int32(tidx) & Int32(31)
                    warp_in_tile = Int32(tidx) >> Int32(5)
                    warp_m_base = (warp_in_tile >> Int32(1)) * Int32(64)
                    warp_n_base = (warp_in_tile & Int32(1)) * Int32(64)

                    # FC2's intermediate was quant-written to the head of the
                    # shared 128-row block, so phase B re-slices at offset 0
                    # rather than FC1's per-task offset.
                    if cutlass.const_expr(self.sfa_tiles_per_block > 1):
                        _sA_p2 = cute.local_tile(
                            sA,
                            cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                            (Int32(0), 0, None),
                        )
                        tCrA = tiled_mma.make_fragment_A(
                            thr_mma.partition_A(_sA_p2)[None, None, None, 0]
                        )
                        csA = thr_ld_A.partition_S(_sA_p2)
                        crA = thr_ld_A.retile(tCrA)
                        _sSFA_p2 = cute.local_tile(
                            sSFA,
                            cute.slice_(self.tile_shape_mnk, (None, 0, None)),
                            (Int32(0), 0, None),
                        )
                        tCrSFA = self._dense_cls._partition_fragment_SFA(
                            self,  # type: ignore[arg-type]
                            _sSFA_p2[None, None, 0],
                            thr_mma,
                            tidx,
                        )
                        csSFA = thr_ld_SFA.partition_S(_sSFA_p2)
                        crSFA = thr_ld_SFA.retile(tCrSFA)

                    csA_phase2 = csA[None, None, None, 0]
                    csSFA_phase2 = csSFA[None, None, None, 0]

                    expert_idx = task_expert_idx

                    # Consume all output tiles continuously from phase2_pipeline.

                    # Hoist A-side register loads: sA is constant across all
                    # FC2 output tiles (quantized intermediate). Load crA and
                    # crSFA for all k-blocks once, reuse for all 32 tiles.
                    fz_crSFA_p2 = cute.filter_zeros(crSFA)
                    cute.copy(
                        smem_copy_A, csA_phase2[None, None, 0], crA[None, None, 0]
                    )
                    fz_csSFA_p2 = cute.filter_zeros(csSFA_phase2)
                    cute.copy(
                        smem_copy_SFA,
                        fz_csSFA_p2[None, None, 0],
                        fz_crSFA_p2[None, None, 0],
                    )
                    for _kb_pre in cutlass.range_constexpr(num_k_blocks - 1):
                        k_pre = _kb_pre + 1
                        cute.copy(
                            smem_copy_A,
                            csA_phase2[None, None, k_pre],
                            crA[None, None, k_pre],
                        )
                        cute.copy(
                            smem_copy_SFA,
                            fz_csSFA_p2[None, None, k_pre],
                            fz_crSFA_p2[None, None, k_pre],
                        )

                    phase2_cons_state.reset_count()
                    for output_tile_idx in range(0, output_tile_cnt, 1, unroll=4):  # type: ignore[call-overload]
                        phase2_peek = phase2_pipeline.consumer_try_wait(
                            phase2_cons_state
                        )
                        phase2_pipeline.consumer_wait(phase2_cons_state, phase2_peek)
                        csB_phase2 = csB[None, None, None, phase2_cons_state.index]
                        csSFB_phase2 = csSFB[None, None, None, phase2_cons_state.index]

                        # Only load B-side (B_down changes per output tile; A is hoisted)
                        cute.copy(
                            smem_copy_B, csB_phase2[None, None, 0], crB[None, None, 0]
                        )
                        f2 = cute.filter_zeros(csSFB_phase2)
                        f4 = cute.filter_zeros(crSFB)
                        cute.copy(smem_copy_SFB, f2[None, None, 0], f4[None, None, 0])

                        down_acc.fill(0.0)
                        for k_block_idx in cutlass.range_constexpr(num_k_blocks):
                            k_next = (
                                0
                                if k_block_idx + 1 == num_k_blocks
                                else k_block_idx + 1
                            )
                            if k_block_idx == num_k_blocks - 1:
                                phase2_pipeline.consumer_release(phase2_cons_state)
                                phase2_cons_state.advance()
                            if k_next > 0:
                                # Only B-side for next k-block (A already in registers)
                                cute.copy(
                                    smem_copy_B,
                                    csB_phase2[None, None, k_next],
                                    crB[None, None, k_next],
                                )
                                f2 = cute.filter_zeros(csSFB_phase2)
                                f4 = cute.filter_zeros(crSFB)
                                cute.copy(
                                    smem_copy_SFB,
                                    f2[None, None, k_next],
                                    f4[None, None, k_next],
                                )
                            for _mt in cutlass.range_constexpr(fc2_m_tiles):
                                for _nt in cutlass.range_constexpr(fc2_n_tiles):
                                    mma_atom.set(
                                        WarpField.SFA,
                                        tCrSFA[None, _mt, k_block_idx].iterator,
                                    )
                                    mma_atom.set(
                                        WarpField.SFB,
                                        tCrSFB[None, _nt, k_block_idx].iterator,
                                    )
                                    cute.gemm(
                                        mma_atom,
                                        down_acc[None, _mt, _nt],
                                        tCrA[None, _mt, k_block_idx],
                                        tCrB[None, _nt, k_block_idx],
                                        down_acc[None, _mt, _nt],
                                    )

                        # Scatter using precomputed metadata (no redundant gmem loads)
                        tile_n_base_cur = output_tile_idx * Int32(
                            self.tile_shape_mnk[1]
                        )
                        for epi_m in cutlass.range_constexpr(epi_rest_m):
                            for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN):
                                for mma_m_in_epi in cutlass.range_constexpr(
                                    MmaMPerEpiM
                                ):
                                    mma_n = mma_n_in_epi
                                    mma_m = epi_m * MmaMPerEpiM + mma_m_in_epi
                                    tRS_rD_slice = tRS_rD[
                                        (None, mma_m_in_epi, mma_n_in_epi)
                                    ]
                                    down_epi_acc_slice = down_acc[(None, mma_m, mma_n)]
                                    for elem_idx in cutlass.range_constexpr(
                                        cute.size(tRS_rD_slice)
                                    ):
                                        tRS_rD_slice[elem_idx] = (
                                            down_alpha_value
                                            * down_epi_acc_slice[elem_idx]
                                        )

                            acc_vec = tRS_rD.load()
                            acc_vec = acc_vec.to(cutlass.BFloat16)
                            tRS_rD_out.store(acc_vec)
                            epi_buffer = Int32(epi_m) % cute.size(tRS_sD, mode=[3])
                            cute.copy(
                                tiled_copy_r2s,
                                tRS_rD_out,
                                tRS_sD[(None, None, None, epi_buffer)],
                            )
                            cute.arch.fence_proxy("async.shared", space="cta")
                            # Vector scatter reads wider spans from sC than the
                            # scalar path, so wait for all MMA-warps' stores.
                            self.epilog_sync_barrier.arrive_and_wait()
                            rows_offset = Int32(epi_m) * Int32(self.epi_tile[0])

                            # Per-warp scatter: each warp scatters its own quadrant
                            # of sC (64 M-rows x 64 N-cols).
                            warp_epi_rows = valid_rows - rows_offset - warp_m_base
                            if warp_epi_rows > Int32(64):
                                warp_epi_rows = Int32(64)
                            if warp_epi_rows < Int32(0):
                                warp_epi_rows = Int32(0)

                            tile_vec_cols = Int32(64) // Int32(8)
                            vec_idx = lane_id
                            while vec_idx < warp_epi_rows * tile_vec_cols:
                                local_row = vec_idx // tile_vec_cols
                                local_vec_col = vec_idx - local_row * tile_vec_cols
                                local_col = warp_n_base + local_vec_col * Int32(8)
                                global_col = tile_n_base_cur + local_col
                                cached_row = rows_offset + warp_m_base + local_row
                                tok = ld_shared_i32_relaxed(
                                    scatter_tok_base_addr + cached_row * Int32(4)
                                )
                                wv = ld_shared_f32(
                                    scatter_weight_base_addr + cached_row * Int32(4)
                                )
                                sc_v0 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col,
                                        epi_buffer,
                                    ]
                                )
                                sc_v1 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(1),
                                        epi_buffer,
                                    ]
                                )
                                sc_v2 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(2),
                                        epi_buffer,
                                    ]
                                )
                                sc_v3 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(3),
                                        epi_buffer,
                                    ]
                                )
                                sc_v4 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(4),
                                        epi_buffer,
                                    ]
                                )
                                sc_v5 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(5),
                                        epi_buffer,
                                    ]
                                )
                                sc_v6 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(6),
                                        epi_buffer,
                                    ]
                                )
                                sc_v7 = cutlass.Float32(
                                    sC[
                                        warp_m_base + local_row,
                                        local_col + Int32(7),
                                        epi_buffer,
                                    ]
                                )
                                scatter_add_v4_bf16x2(
                                    get_ptr_as_int64(
                                        scatter_output, tok * scatter_N + global_col
                                    ),
                                    wv * sc_v0,
                                    wv * sc_v1,
                                    wv * sc_v2,
                                    wv * sc_v3,
                                    wv * sc_v4,
                                    wv * sc_v5,
                                    wv * sc_v6,
                                    wv * sc_v7,
                                )
                                vec_idx += Int32(self.num_threads_per_warp)

                            # Post-scatter barrier: needed to ensure all warps
                            # finish scatter before next output tile's pipeline ops
                            # (pipeline consumer is collective across all MMA warps).
                            self.epilog_sync_barrier.arrive_and_wait()

                    # Signal that FC2/scatter no longer needs sA, so the DMA
                    # warp may start the next slice/task's FC1 loads.
                    self.pass_final_barrier.arrive_unaligned()
                    slice_idx += Int32(1)

            elif warp_idx == self.tma_load_warp_id:
                task_expert_idx = work_item[_WORK_EXPERT]
                task_m_tile_idx = work_item[_WORK_M_TILE]
                task_slice_begin_idx = work_item[_WORK_SLICE_BEGIN]
                task_slice_count_val = work_item[_WORK_SLICE_COUNT]

                # gA/gSFA are tiled in 128-row blocks; the fragment partition
                # selects the within-block sub-tile.
                sa_block_idx = task_m_tile_idx // Int32(self.sa_tiles_per_block)
                tAgA_mk = tAgA[(None, sa_block_idx, None, Int32(0))]
                sfa_block_idx = task_m_tile_idx // Int32(self.sfa_tiles_per_block)
                tAgSFA_mk = tAgSFA[(None, sfa_block_idx, None, Int32(0))]
                slice_idx = Int32(0)
                while slice_idx < task_slice_count_val:
                    intermediate_slice = task_slice_begin_idx + slice_idx

                    # FC1 producer slice. Gated activation packs [up, gate]
                    # across N; relu2 uses a single FC1 pass.
                    tBgB_w13_up_nk = tBgB_w13[
                        (None, intermediate_slice, None, task_expert_idx)
                    ]
                    tBgSFB_w13_up_nk = tBgSFB_w13[
                        (None, intermediate_slice, None, task_expert_idx)
                    ]
                    gate_slice_idx = (
                        intermediate_slice + gate_tile_cnt
                        if self.is_gated
                        else intermediate_slice
                    )
                    tBgB_w13_gate_nk = tBgB_w13[
                        (None, gate_slice_idx, None, task_expert_idx)
                    ]
                    tBgSFB_w13_gate_nk = tBgSFB_w13[
                        (None, gate_slice_idx, None, task_expert_idx)
                    ]

                    # ---- FC1 gate/only pass ----
                    prod_state.reset_count()
                    for k_tile in range(0, fc1_k_tile_cnt, 1, unroll=4):  # type: ignore[call-overload]
                        ml_pipeline.producer_acquire(prod_state)
                        cute.copy(
                            tma_a,
                            tAgA_mk[(None, k_tile)],
                            tAsA[(None, prod_state.index)],
                            tma_bar_ptr=ml_pipeline.producer_get_barrier(prod_state),
                        )
                        cute.copy(
                            tma_b_w13,
                            tBgB_w13_gate_nk[(None, k_tile)],
                            tBsB_w13[(None, prod_state.index)],
                            tma_bar_ptr=ml_pipeline.producer_get_barrier(prod_state),
                        )
                        cute.copy(
                            tma_sfa,
                            tAgSFA_mk[(None, k_tile)],
                            tAsSFA[(None, prod_state.index)],
                            tma_bar_ptr=ml_pipeline.producer_get_barrier(prod_state),
                        )
                        cute.copy(
                            tma_sfb_w13,
                            tBgSFB_w13_gate_nk[(None, k_tile)],
                            tBsSFB_w13[(None, prod_state.index)],
                            tma_bar_ptr=ml_pipeline.producer_get_barrier(prod_state),
                        )
                        ml_pipeline.producer_commit(prod_state)
                        prod_state.advance()

                    # Wait for the MMA warps to finish the FC1 gate/only pass
                    # before reusing sA/sB/sSFA/sSFB for up/FC2 staging.
                    self.pass_gate_barrier.wait_unaligned()

                    if cutlass.const_expr(self.is_gated):
                        # ---- FC1 up pass ----
                        up_prod_state.reset_count()
                        for k_tile in range(0, fc1_k_tile_cnt, 1, unroll=4):  # type: ignore[call-overload]
                            up_pipeline.producer_acquire(up_prod_state)
                            cute.copy(
                                tma_a,
                                tAgA_mk[(None, k_tile)],
                                tAsA[(None, up_prod_state.index)],
                                tma_bar_ptr=up_pipeline.producer_get_barrier(
                                    up_prod_state
                                ),
                            )
                            cute.copy(
                                tma_b_w13,
                                tBgB_w13_up_nk[(None, k_tile)],
                                tBsB_w13_up[(None, up_prod_state.index)],
                                tma_bar_ptr=up_pipeline.producer_get_barrier(
                                    up_prod_state
                                ),
                            )
                            cute.copy(
                                tma_sfa,
                                tAgSFA_mk[(None, k_tile)],
                                tAsSFA[(None, up_prod_state.index)],
                                tma_bar_ptr=up_pipeline.producer_get_barrier(
                                    up_prod_state
                                ),
                            )
                            cute.copy(
                                tma_sfb_w13,
                                tBgSFB_w13_up_nk[(None, k_tile)],
                                tBsSFB_w13_up[(None, up_prod_state.index)],
                                tma_bar_ptr=up_pipeline.producer_get_barrier(
                                    up_prod_state
                                ),
                            )
                            up_pipeline.producer_commit(up_prod_state)
                            up_prod_state.advance()

                    # ---- FC2 B_down loads: continuous pipeline ----
                    # No barrier needed: sB/sSFB are free (gate done, up uses
                    # sB_up/sSFB_up). phase2_pipeline handles data availability.
                    # intermediate_slice selects the K-tile of GEMM2 (FC1 output N-tile
                    # = GEMM2 K-tile since intermediate dim is the reduction dim).
                    # Load ALL FC2 tiles continuously once stage1 no longer needs
                    # the gate staging buffers.
                    phase2_prod_state.reset_count()
                    for output_tile_idx in range(0, output_tile_cnt, 1, unroll=4):  # type: ignore[call-overload]
                        phase2_pipeline.producer_acquire(phase2_prod_state)
                        cute.copy(
                            tma_b_down,
                            tBgB_down[
                                (
                                    None,
                                    output_tile_idx,
                                    intermediate_slice,
                                    task_expert_idx,
                                )
                            ],
                            tBsB_down[(None, phase2_prod_state.index)],
                            tma_bar_ptr=phase2_pipeline.producer_get_barrier(
                                phase2_prod_state
                            ),
                        )
                        cute.copy(
                            tma_sfb_down,
                            tBgSFB_down[
                                (
                                    None,
                                    output_tile_idx,
                                    intermediate_slice,
                                    task_expert_idx,
                                )
                            ],
                            tBsSFB_down[(None, phase2_prod_state.index)],
                            tma_bar_ptr=phase2_pipeline.producer_get_barrier(
                                phase2_prod_state
                            ),
                        )
                        phase2_pipeline.producer_commit(phase2_prod_state)
                        phase2_prod_state.advance()

                    # Ensure MMA warps finish FC2/scatter before DMA starts the
                    # next slice/task's FC1 loads into shared A buffers.
                    self.pass_final_barrier.wait_unaligned()
                    slice_idx += Int32(1)

        if warp_idx == self.tma_load_warp_id:
            ml_pipeline.producer_tail(prod_state)
            if cutlass.const_expr(self.is_gated):
                up_pipeline.producer_tail(up_prod_state)
            phase2_pipeline.producer_tail(phase2_prod_state)
        return


__all__ = ["MoEDynamicKernel"]
