biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e0f60b9f052d2aa1fdd5ff2f7df922a00f26e76b
nvfp4-megamoe-kernel/dsv4/kernels/router/nvfp4_fused_router_kernel.py
biondizzle e0f60b9f05 Fix fused router: plain ints for mma_tiler + @cute.jit pattern 
Root cause of previous crash: cutlass.Int32(128) wrapping of mma_inst_shape_mn
caused _unpack_x_tuple to fail in cute.size(tiled_mma.shape_mnk, mode=[2]).

The fused_swiglu kernel uses plain Python ints for mma_tiler_mnk and
mma_inst_shape_mn — NOT cutlass.Int32. Inside @cute.jit, CuTeDSL
auto-converts plain ints to MLIR values. The Int32 wrapping was unnecessary
and actually harmful.

Pattern: same as fused_swiglu.py __call__:
- @cute.jit compiled_fn takes CuTe tensors
- _setup_attributes called inside JIT (needs MLIR context)
- cute.compile at the end
2026-06-01 10:37:15 +00:00

866 lines
41 KiB
Python
Raw Blame History

"""DSV4 NVFP4 Fused Router Kernel — Block-scaled GEMM + Activation Epilogue.
Two-phase production path:
Phase 1 (this kernel): NVFP4 block-scaled GEMM + fused sqrt(softplus) + e_bias
activation epilogue. Writes FP32 activated scores to GMEM. No intermediate
BF16 logits buffer. Pure NVFP4 + Blackwell tensor cores the entire way.
Phase 2 (activation_topk CUDA kernel): top-k + renorm on the activated scores.
The GEMM mainloop and epilogue structure follow FusedSwiGLUScaledGroupedGemmKernel
(dsv4/kernels/gemm/fused_swiglu.py) exactly, with a different activation function
(sqrt(softplus) + e_bias instead of SwiGLU) and no SwiGLU clamp.
Warp specialization (6 warps, no scheduler for dense GEMM):
Warps 0-3: Epilogue (TMEM -> register -> activation -> SMEM -> TMA store -> GMEM)
Warp 4: MMA (tcgen05.mma.block_scale with SFA/SFB in TMEM)
Warp 5: TMA load (A, B, SFA, SFB from GMEM -> SMEM)
Pipeline structure (2 pipelines):
AB pipeline: TMA (producer) -> MMA (consumer) [PipelineTmaUmma]
Acc pipeline: MMA (producer) -> Epilogue (consumer) [PipelineUmmaAsync]
The epilogue uses the proven one-way TMEM→registers→SMEM→GMEM path from the MoE
kernel. This is the same pattern that compiles and runs correctly in
FusedSwigGLUScaledGroupedGemmKernel. No SMEM top-k merge (which crashed MLIR).
"""
from __future__ import annotations
from typing import Tuple, Optional, Type, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.typing import Pointer
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils as utils
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.utils.gemm.sm100 import (
epilogue_tmem_copy_and_partition,
epilogue_smem_copy_and_partition,
transform_partitioned_tensor_layout,
)
class Nvfp4FusedRouterKernel:
"""
NVFP4 blockscaled GEMM + fused activation epilogue.
Dense (non-grouped) GEMM: [M, K] @ [K, E] -> [M, E] with NVFP4 weights.
Custom epilogue: TMEM -> registers -> sqrt(softplus(logit)) + e_bias -> SMEM -> GMEM.
Follows FusedSwiGLUScaledGroupedGemmKernel pattern exactly.
"""
def __init__(
self,
sf_vec_size: int = 16,
mma_tiler_mnk: Tuple[int, int, int] = (128, 128, 64),
cluster_shape_mnk: Tuple[int, int, int] = (1, 1, 1),
):
self.sf_vec_size = sf_vec_size
self.mma_tiler_mnk = mma_tiler_mnk
self.cluster_shape_mn = (cluster_shape_mnk[0], cluster_shape_mnk[1])
self.use_2cta_instrs = mma_tiler_mnk[0] == 256
self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
self.arch = "sm_100"
self.mma_inst_shape_mn = (mma_tiler_mnk[0], mma_tiler_mnk[1])
self.mma_inst_shape_mn_sfb = (
mma_tiler_mnk[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(mma_tiler_mnk[1], 128),
)
# 6-warp specialization (no scheduler warp for dense GEMM)
self.epilogue_warp_id = (0, 1, 2, 3)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.threads_per_warp = 32
self.threads_per_cta = self.threads_per_warp * 6
# Barrier IDs
self.cta_sync_bar_id = 1
self.epilogue_sync_bar_id = 2
self.tmem_alloc_sync_bar_id = 3
self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch)
self.occupancy = 1
self.buffer_align_bytes = 1024
def _create_tiled_mma(self, a_dtype, a_major_mode, b_major_mode, sf_dtype):
return sm100_utils.make_blockscaled_trivial_tiled_mma(
a_dtype, a_major_mode, b_major_mode, sf_dtype,
self.sf_vec_size, self.cta_group,
self.mma_inst_shape_mn,
)
def _create_tiled_mma_sfb(self, a_dtype, a_major_mode, b_major_mode, sf_dtype):
return sm100_utils.make_blockscaled_trivial_tiled_mma(
a_dtype, a_major_mode, b_major_mode, sf_dtype,
self.sf_vec_size, tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
def _setup_attributes(self, tiled_mma, tiled_mma_sfb, a_dtype, b_dtype, sf_dtype, c_dtype, c_layout):
"""Set up kernel attributes. Mirrors fused_swiglu._setup_attributes."""
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = self.mma_tiler_mnk[2] // mma_inst_shape_k
self.mma_tiler = (self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.mma_tiler_mnk[2])
self.mma_tiler_sfb = (self.mma_inst_shape_mn_sfb[0], self.mma_inst_shape_mn_sfb[1], self.mma_tiler_mnk[2])
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cta_tile_shape_mnk_sfb = (
self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler_sfb[1],
self.mma_tiler_sfb[2],
)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((self.cluster_shape_mn[0], self.cluster_shape_mn[1], 1)),
(tiled_mma.thr_id.shape,))
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout((self.cluster_shape_mn[0], self.cluster_shape_mn[1], 1)),
(tiled_mma_sfb.thr_id.shape,))
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
self.is_a_mcast = self.num_mcast_ctas_a > 1
self.is_b_mcast = self.num_mcast_ctas_b > 1
self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
# Epilogue tile (same as MoE: compute_epilogue_tile_shape for NVFP4→FP32)
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
self.use_2cta_instrs,
c_layout,
c_dtype,
)
self.epi_tile_n = cute.size(self.epi_tile[1])
# Stage counts (same as MoE)
self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
tiled_mma, self.mma_tiler, a_dtype, b_dtype,
self.epi_tile, c_dtype, c_layout, sf_dtype, self.sf_vec_size,
self.smem_capacity, self.occupancy)
# SMEM layouts
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma, self.mma_tiler, a_dtype, self.num_ab_stage)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma, self.mma_tiler, b_dtype, self.num_ab_stage)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage)
self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
c_dtype, c_layout, self.epi_tile, self.num_c_stage)
# Overlapping accumulator
self.overlapping_accum = self.cta_tile_shape_mnk[1] == 256
if self.overlapping_accum:
self.num_acc_pipeline_stages = 1
else:
self.num_acc_pipeline_stages = self.num_acc_stage
# TMEM column counts
sf_atom_mn = 32
self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage - (
self.num_sf_tmem_cols if self.overlapping_accum else 0
)
self.iter_acc_early_release_in_epilogue = (
self.num_sf_tmem_cols // self.epi_tile_n
)
# TMA load bytes
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
a_smem_0 = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
b_smem_0 = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
sfa_smem_0 = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
sfb_smem_0 = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
self.num_tma_load_bytes = (
cute.size_in_bytes(a_dtype, a_smem_0) +
cute.size_in_bytes(b_dtype, b_smem_0) +
cute.size_in_bytes(sf_dtype, sfa_smem_0) +
cute.size_in_bytes(sf_dtype, sfb_smem_0)
) * atom_thr_size
# TMEM allocation size
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)
@staticmethod
def _compute_stages(
tiled_mma, mma_tiler_mnk, a_dtype, b_dtype,
epi_tile, c_dtype, c_layout, sf_dtype, sf_vec_size,
smem_capacity, occupancy,
):
num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
num_c_stage = 2
a_smem_layout_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1)
b_smem_layout_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1)
sfa_smem_layout_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
sfb_smem_layout_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
c_smem_layout_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1)
ab_bytes_per_stage = (
cute.size_in_bytes(a_dtype, a_smem_layout_one) +
cute.size_in_bytes(b_dtype, b_smem_layout_one) +
cute.size_in_bytes(sf_dtype, sfa_smem_layout_one) +
cute.size_in_bytes(sf_dtype, sfb_smem_layout_one)
)
mbar_helpers_bytes = 1024
c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_one)
c_bytes = c_bytes_per_stage * num_c_stage
num_ab_stage = (
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
num_c_stage += (
smem_capacity
- occupancy * ab_bytes_per_stage * num_ab_stage
- occupancy * (mbar_helpers_bytes + c_bytes)
) // (occupancy * c_bytes_per_stage)
return num_acc_stage, num_ab_stage, num_c_stage
def mainloop_s2t_copy_and_partition(self, sSF, tSF, cta_group):
tCsSF_compact = cute.filter_zeros(sSF)
tCtSF_compact = cute.filter_zeros(tSF)
copy_atom_s2t = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), self.sf_dtype)
tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
thr_copy_s2t = tiled_copy_s2t.get_slice(0)
tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
# -----------------------------------------------------------------
# run() — Python entry point
# -----------------------------------------------------------------
def run(self, mat_a, mat_b, scale_a, scale_b, mat_c,
M, N, K, gsa, gsb, stream=None):
if stream is None:
stream = cuda.CUstream(0)
a_dtype = mat_a.element_type
b_dtype = mat_b.element_type
sf_dtype = scale_a.element_type
c_dtype = mat_c.element_type
a_major_mode = utils.LayoutEnum.from_tensor(mat_a).mma_major_mode()
b_major_mode = utils.LayoutEnum.from_tensor(mat_b).mma_major_mode()
c_layout = utils.LayoutEnum.from_tensor(mat_c)
self.a_dtype = a_dtype
self.b_dtype = b_dtype
self.sf_dtype = sf_dtype
self.c_dtype = c_dtype
self.a_major_mode = a_major_mode
self.b_major_mode = b_major_mode
cta_m = self.mma_tiler_mnk[0]
cta_n = self.mma_tiler_mnk[1]
num_M_tiles = (M + cta_m - 1) // cta_m
num_N_tiles = (N + cta_n - 1) // cta_n
grid = (num_M_tiles * num_N_tiles, 1, 1)
@cute.jit
def _compiled_fn(mat_a, mat_b, scale_a, scale_b, mat_c):
# Create tiled MMA and setup inside JIT context
# (same pattern as fused_swiglu.py @cute.jit __call__)
# Plain int mma_tiler values work with cute.size() inside JIT
tiled_mma = self._create_tiled_mma(a_dtype, a_major_mode, b_major_mode, sf_dtype)
tiled_mma_sfb = self._create_tiled_mma_sfb(a_dtype, a_major_mode, b_major_mode, sf_dtype)
self._setup_attributes(tiled_mma, tiled_mma_sfb, a_dtype, b_dtype, sf_dtype, c_dtype, c_layout)
# TMA atoms (inside JIT, same as fused_swiglu)
a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op, mat_a, a_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape)
b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op, mat_b, b_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape)
sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
sfa_op, scale_a, sfa_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape,
internal_type=cutlass.Uint64)
sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)
sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op, scale_b, sfb_smem_layout, self.mma_tiler_sfb, tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Uint64)
epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(), mat_c, epi_smem_layout, self.epi_tile)
tile_sched_params = utils.PersistentTileSchedulerParams(
(num_M_tiles, num_N_tiles, 1), (1, 1, 1))
self._kernel(
tiled_mma, tiled_mma_sfb,
tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b,
tma_atom_sfa, tma_tensor_sfa, tma_atom_sfb, tma_tensor_sfb,
tma_atom_c, tma_tensor_c,
self.cluster_layout_vmnk, self.cluster_layout_sfb_vmnk,
self.a_smem_layout_staged, self.b_smem_layout_staged,
self.sfa_smem_layout_staged, self.sfb_smem_layout_staged,
self.c_smem_layout_staged,
self.epi_tile,
tile_sched_params,
M, N, K, gsa, gsb,
).launch(
grid=grid, block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
stream=stream, min_blocks_per_mp=1,
)
cute.compile(_compiled_fn, mat_a, mat_b, scale_a, scale_b, mat_c)
@cute.kernel
def _kernel(self, tiled_mma, tiled_mma_sfb,
tma_atom_a, mA_mkl, tma_atom_b, mB_nkl,
tma_atom_sfa, mSFA_mkl, tma_atom_sfb, mSFB_nkl,
tma_atom_c, mC_mnl,
cluster_layout_vmnk, cluster_layout_sfb_vmnk,
a_smem_layout_staged, b_smem_layout_staged,
sfa_smem_layout_staged, sfb_smem_layout_staged,
c_smem_layout_staged,
epi_tile,
tile_sched_params,
M, N, K, gsa, gsb):
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
use_2cta = cute.size(tiled_mma.thr_id.shape) == 2
is_leader_cta = (bidx % cute.size(tiled_mma.thr_id.shape)) == 0
mma_tile_v = bidx % cute.size(tiled_mma.thr_id.shape)
cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
block_coord = cluster_layout_vmnk.get_flat_coord(cta_rank)
acc_dtype = cutlass.Float32
c_dtype = self.c_dtype
# ============================================================
# Shared storage
# ============================================================
@cute.struct
class SharedStorage:
ab_full_mbar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
acc_full_mbar: cute.struct.MemRange[cutlass.Int64, self.num_acc_pipeline_stages * 2]
tmem_dealloc_mbar: cutlass.Int64
tmem_holding: cutlass.Int32
# C staging SMEM for TMA store (same as MoE epilogue)
sC: cute.struct.Align[
cute.struct.MemRange[c_dtype, cute.cosize(c_smem_layout_staged.outer)],
self.buffer_align_bytes,
]
smem = utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
# ============================================================
# Pipelines
# ============================================================
ab_pipeline = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_full_mbar.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread,
self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1),
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
num_acc_cons = self.threads_per_warp * len(self.epilogue_warp_id) * (2 if use_2cta else 1)
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar.data_ptr(),
num_stages=self.num_acc_pipeline_stages,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_cons),
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# C pipeline for TMA store (same as MoE)
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=c_producer_group,
)
tmem = utils.TmemAllocator(
storage.tmem_holding.ptr,
barrier_for_retrieve=pipeline.NamedBarrier(
barrier_id=self.tmem_alloc_sync_bar_id,
num_threads=self.threads_per_warp * len((self.mma_warp_id, *self.epilogue_warp_id))),
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr)
cta_bar = pipeline.NamedBarrier(self.cta_sync_bar_id, self.threads_per_cta)
epi_sync_bar = pipeline.NamedBarrier(
self.epilogue_sync_bar_id,
self.threads_per_warp * len(self.epilogue_warp_id))
# SMEM tensors
sA = smem.allocate_tensor(
element_type=self.a_dtype, layout=a_smem_layout_staged.outer,
byte_alignment=128, swizzle=a_smem_layout_staged.inner)
sB = smem.allocate_tensor(
element_type=self.b_dtype, layout=b_smem_layout_staged.outer,
byte_alignment=128, swizzle=b_smem_layout_staged.inner)
sSFA = smem.allocate_tensor(
element_type=self.sf_dtype, layout=sfa_smem_layout_staged, byte_alignment=128)
sSFB = smem.allocate_tensor(
element_type=self.sf_dtype, layout=sfb_smem_layout_staged, byte_alignment=128)
sC = smem.allocate_tensor(
element_type=c_dtype, layout=c_smem_layout_staged.outer,
byte_alignment=128, swizzle=c_smem_layout_staged.inner)
# Multicast masks
a_mcast = None; b_mcast = None; sfa_mcast = None; sfb_mcast = None
if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta):
a_mcast = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_coord, mcast_mode=2)
b_mcast = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_coord, mcast_mode=1)
sfa_mcast = a_mcast
sfb_mcast = cpasync.create_tma_multicast_mask(cluster_layout_sfb_vmnk, block_coord, mcast_mode=1)
# Partition global tensors
gA = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
gB = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
gSFA = cute.local_tile(mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
gSFB = cute.local_tile(mSFB_nkl, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None))
k_tiles = cute.size(gA, mode=[3])
thr_mma = tiled_mma.get_slice(mma_tile_v)
tCgA = thr_mma.partition_A(gA)
tCgB = thr_mma.partition_B(gB)
tCgSFA = thr_mma.partition_A(gSFA)
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_v)
tCgSFB = thr_mma_sfb.partition_B(gSFB)
# TMA partitions for A/B
a_cta_l = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
tAsA, tAgA = cpasync.tma_partition(tma_atom_a, block_coord[2], a_cta_l,
cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3))
b_cta_l = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
tBsB, tBgB = cpasync.tma_partition(tma_atom_b, block_coord[1], b_cta_l,
cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3))
# TMA partitions for SFA/SFB
tAsSFA, tAgSFA = cpasync.tma_partition(tma_atom_sfa, block_coord[2], a_cta_l,
cute.group_modes(sSFA, 0, 3), cute.group_modes(tCgSFA, 0, 3))
tAsSFA = cute.filter_zeros(tAsSFA); tAgSFA = cute.filter_zeros(tAgSFA)
block_coord_sfb = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank)
sfb_cta_l = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape)
tBsSFB, tBgSFB = cpasync.tma_partition(tma_atom_sfb, block_coord_sfb[1], sfb_cta_l,
cute.group_modes(sSFB, 0, 3), cute.group_modes(tCgSFB, 0, 3))
tBsSFB = cute.filter_zeros(tBsSFB); tBgSFB = cute.filter_zeros(tBgSFB)
# TMEM accumulator
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
# Cluster arrive
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
else:
cta_bar.arrive_and_wait()
# ============================================================
# TMA WARP
# ============================================================
if warp_idx == self.tma_warp_id:
cpasync.prefetch_descriptor(tma_atom_a)
cpasync.prefetch_descriptor(tma_atom_b)
cpasync.prefetch_descriptor(tma_atom_sfa)
cpasync.prefetch_descriptor(tma_atom_sfb)
tsched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, bidx, cute.arch.grid_dim())
wt = tsched.initial_work_tile_info()
ab_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage)
while wt.is_valid_tile:
tc = wt.tile_idx
mc = (tc[0] // cute.size(tiled_mma.thr_id.shape), tc[1], tc[2])
tAgA_s = tAgA[(None, mc[0], None, mc[2])]
tBgB_s = tBgB[(None, mc[1], None, mc[2])]
tAgSFA_s = tAgSFA[(None, mc[0], None, mc[2])]
slice_n = mc[1]
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
slice_n = mc[1] // 2
tBgSFB_s = tBgSFB[(None, slice_n, None, mc[2])]
ab_ps.reset_count()
peek_ab = cutlass.Boolean(1)
if ab_ps.count < k_tiles:
peek_ab = ab_pipeline.producer_try_acquire(ab_ps)
for kt in cutlass.range(0, k_tiles, 1, unroll=1):
ab_pipeline.producer_acquire(ab_ps, peek_ab)
cute.copy(tma_atom_a, tAgA_s[(None, ab_ps.count)], tAsA[(None, ab_ps.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps), mcast_mask=a_mcast)
cute.copy(tma_atom_b, tBgB_s[(None, ab_ps.count)], tBsB[(None, ab_ps.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps), mcast_mask=b_mcast)
cute.copy(tma_atom_sfa, tAgSFA_s[(None, ab_ps.count)], tAsSFA[(None, ab_ps.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps), mcast_mask=sfa_mcast)
cute.copy(tma_atom_sfb, tBgSFB_s[(None, ab_ps.count)], tBsSFB[(None, ab_ps.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps), mcast_mask=sfb_mcast)
ab_ps.advance()
peek_ab = cutlass.Boolean(1)
if ab_ps.count < k_tiles:
peek_ab = ab_pipeline.producer_try_acquire(ab_ps)
ab_pipeline.producer_tail(ab_ps)
tsched.advance_to_next_work()
wt = tsched.get_current_work()
# ============================================================
# MMA WARP
# ============================================================
if warp_idx == self.mma_warp_id:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cta_bar.arrive_and_wait()
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
tCrA = tiled_mma.make_fragment_A(sA)
tCrB = tiled_mma.make_fragment_B(sB)
# S2T for SFA
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma, self.mma_tiler, self.sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)))
tCtSFA = cute.make_tensor(acc_tmem_ptr, tCtSFA_layout)
# S2T for SFB
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma_sfb, self.mma_tiler, self.sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)))
tCtSFB = cute.make_tensor(acc_tmem_ptr, tCtSFB_layout)
tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = \
self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA, self.cta_group)
tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = \
self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB, tcgen05.CtaGroup.ONE)
tsched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, bidx, cute.arch.grid_dim())
wt = tsched.initial_work_tile_info()
ab_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage)
acc_ps = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages)
while wt.is_valid_tile:
if is_leader_cta:
acc_pipeline.producer_acquire(acc_ps)
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_ps.phase ^ 1
else:
acc_stage_index = acc_ps.index
tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
ab_cs.reset_count()
peek_ab_full = cutlass.Boolean(1)
if ab_cs.count < k_tiles and is_leader_cta:
peek_ab_full = ab_pipeline.consumer_try_wait(ab_cs)
for kt in cutlass.range(0, k_tiles, 1, unroll=1):
if is_leader_cta:
ab_pipeline.consumer_wait(ab_cs, peek_ab_full)
s2t_stage_coord = (None, None, None, None, ab_cs.index)
cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t[s2t_stage_coord], tCtSFA_compact_s2t)
cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t[s2t_stage_coord], tCtSFB_compact_s2t)
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll=1):
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)
kb_coord = (None, None, kblock_idx, ab_cs.index)
cute.gemm(tiled_mma, tCrA[kb_coord], tCrB[kb_coord], tCtAcc, tCtAcc)
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
ab_pipeline.consumer_release(ab_cs)
ab_cs.advance()
peek_ab_full = cutlass.Boolean(1)
if ab_cs.count < k_tiles:
if is_leader_cta:
peek_ab_full = ab_pipeline.consumer_try_wait(ab_cs)
if is_leader_cta:
acc_pipeline.producer_commit(acc_ps)
acc_ps.advance()
tsched.advance_to_next_work()
wt = tsched.get_current_work()
if is_leader_cta:
acc_pipeline.producer_tail(acc_ps)
tmem.relinquish_alloc_permit()
# ============================================================
# EPILOGUE WARPS — TMEM→regs→activation→SMEM→GMEM
# Same pattern as FusedSwiGLUScaledGroupedGemmKernel.
# Activation: sqrt(softplus(logit)) + e_bias (replaces SwiGLU)
# ============================================================
if warp_idx in self.epilogue_warp_id:
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cta_bar.arrive_and_wait()
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
# TMEM → register copy (paired atoms, same as MoE)
tiled_copy_t2r, tTR_tAcc_base = epilogue_tmem_copy_and_partition(
tCtAcc_base, epi_tile, self.epilogue_warp_id, acc_dtype, use_2cta)
tTR_rAcc = tiled_copy_t2r.fragments_slice(tiled_copy_t2r, tTR_tAcc_base)
# Register tensor for activation output (same pattern as MoE)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, c_dtype)
# Register → SMEM copy (paired atoms, same as MoE)
tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition(
self, tiled_copy_t2r, tTR_rC, tidx, sC)
# TMA partition for C store
tCgC_epi = cute.flat_divide(mC_mnl, epi_tile)
bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
tma_atom_c, 0, cute.make_layout(1),
cute.group_modes(sC, 0, 2),
cute.group_modes(tCgC_epi, 0, 2))
# Tile scheduler + pipeline states
tsched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, bidx, cute.arch.grid_dim())
wt = tsched.initial_work_tile_info()
acc_cs = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_pipeline_stages)
while wt.is_valid_tile:
acc_pipeline.consumer_wait(acc_cs)
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_cs.phase
reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
else:
acc_stage_index = acc_cs.index
reverse_subtile = cutlass.Boolean(False)
tc = wt.tile_idx
mma_tile_coord_mnl = (
tc[0] // cute.size(tiled_mma.thr_id.shape), tc[1], tc[2])
bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]
tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]
tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
# Process subtiles
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
num_prev_subtiles = tsched.num_tiles_executed * subtile_cnt
for subtile_idx in cutlass.range(subtile_cnt):
real_subtile_idx = subtile_idx
if cutlass.const_expr(self.overlapping_accum):
if reverse_subtile:
real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx
# Load accumulator from TMEM to registers
tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
cute.arch.fence_view_async_tmem_load()
# Early release accumulator for overlapping case
if cutlass.const_expr(self.overlapping_accum):
if subtile_idx == self.iter_acc_early_release_in_epilogue:
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_cs)
acc_cs.advance()
# Activation: sqrt(softplus(logit * gsa * gsb))
# Global scales are applied before the activation, same as
# how MoE epilogue applies them before SwiGLU.
# The MMA output is (A * SFA) @ (B * SFB), missing gsa*gsb.
scale = cutlass.Float32(gsa * gsb)
acc_vec = tTR_rAcc.load()
for e in cutlass.range(cute.size(acc_vec), unroll=4):
logit = acc_vec[e] * scale
# softplus(x) = max(x, 0) + log(1 + exp(-|x|))
abs_x = cute.math.absf(logit)
pos = cute.math.fmax(logit, cutlass.Float32(0.0))
exp_neg = cute.math.exp(-abs_x)
sp = pos + cute.math.log(cutlass.Float32(1.0) + exp_neg)
acc_vec[e] = cute.math.sqrt(sp)
tRS_rC.store(acc_vec.to(c_dtype))
# RMEM → SMEM
c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
cute.copy(
tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]
)
cute.arch.fence_proxy(
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta)
epi_sync_bar.arrive_and_wait()
# SMEM → GMEM (TMA store)
if warp_idx == self.epilogue_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buffer)],
bSG_gC[(None, real_subtile_idx)],
)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
epi_sync_bar.arrive_and_wait()
# Release accumulator (non-overlapping case)
if cutlass.const_expr(not self.overlapping_accum):
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_cs)
acc_cs.advance()
tsched.advance_to_next_work()
wt = tsched.get_current_work()
# Cleanup
tmem.relinquish_alloc_permit()
epi_sync_bar.arrive_and_wait()
tmem.free(acc_tmem_ptr)
c_pipeline.producer_tail()
# =====================================================================
# Python entry point
# =====================================================================
def run_nvfp4_fused_router(
hidden_states: torch.Tensor, # [N, hidden_size] BF16
mat_b: torch.Tensor, # [K_packed, E_packed] uint8 NVFP4 weight
scale_b: torch.Tensor, # [K_sf, E_sf] FP8 E4M3 weight scale
gsa: float, # activation global scale
gsb_val: float, # weight global scale (weight_scale_2)
e_bias: torch.Tensor, # [num_experts] FP32
routed_scaling_factor: float,
top_k: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run the NVFP4 fused router: GEMM + activation → top-k.
Phase 1: CuTeDSL NVFP4 blockscaled GEMM + sqrt(softplus) epilogue
writes FP32 activated scores to GMEM.
Phase 2: activation_topk CUDA kernel for top-k + renorm.
Parameters
----------
hidden_states : [N, hidden_size] BF16 activation tensor
mat_b : [K_packed, E_packed] uint8 NVFP4 weight (gate projection)
scale_b : [K_sf, E_sf] FP8 E4M3 weight block scales
gsa : float, activation global scale (from checkpoint input_scale)
gsb_val : float, weight global scale (from checkpoint weight_scale_2)
e_bias : [num_experts] FP32, per-expert selection bias
routed_scaling_factor : float, post-renorm scaling
top_k : int, number of experts to select
Returns
-------
topk_weights : [N, top_k] float32
topk_ids : [N, top_k] int32
"""
N = hidden_states.shape[0] # number of tokens
hidden_size = hidden_states.shape[1]
E = mat_b.shape[0] # num_experts (N dimension of GEMM)
K = mat_b.shape[1] * 2 # K dimension (packed * 2 for FP4)
device = hidden_states.device
# Quantize activation to NVFP4
from dsv4.ops.quantize import quantize_activation_nvfp4
mat_a_bf16_packed, scale_a_fp8 = quantize_activation_nvfp4(hidden_states, gsa)
# Output tensor: FP32 activated scores [N, E]
activated_scores = torch.empty(N, E, dtype=torch.float32, device=device)
# Convert PyTorch tensors to CuTe tensors (same as gemm_runner.py pattern)
import cutlass.torch as cutlass_torch
def _to_cute(t, leading_dim=None):
ct = cutlass_torch.from_dlpack(t)
if leading_dim is not None:
return ct.mark_layout_dynamic(leading_dim=leading_dim)
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
# Determine leading dimensions from tensor shapes
# mat_a_bf16_packed: [N, K_packed] — K-major (row-major for GEMM A)
# mat_b: [E, K_packed] — K-major (col-major for GEMM B, i.e. N-major)
# Actually, for NVFP4 GEMM: A is M-major, B is N-major
# Check the existing Nvfp4Linear to see how it handles this
cute_a = _to_cute(mat_a_bf16_packed)
cute_b = _to_cute(mat_b)
cute_sfa = _to_cute(scale_a_fp8)
cute_sfb = _to_cute(scale_b)
cute_c = _to_cute(activated_scores)
# Run the CuTeDSL kernel: NVFP4 GEMM + sqrt(softplus) epilogue
kernel = Nvfp4FusedRouterKernel(
sf_vec_size=16,
mma_tiler_mnk=(128, 128, 64),
cluster_shape_mnk=(1, 1, 1),
)
kernel.run(
mat_a=cute_a,
mat_b=cute_b,
scale_a=cute_sfa,
scale_b=cute_sfb,
mat_c=cute_c,
M=N, N=E, K=K,
gsa=gsa,
gsb=gsb_val,
)
# Add e_bias (selection bias) and run top-k
# The kernel writes sqrt(softplus(logits)) in FP32
# activation_topk expects raw logits, so we pass the activated scores
# and tell it to skip the activation step
from dsv4.kernels.router._activation_topk import run_fused_activation_topk_pre_activated
out_weights = torch.empty(N, top_k, dtype=torch.float32, device=device)
out_ids = torch.empty(N, top_k, dtype=torch.int32, device=device)
run_fused_activation_topk_pre_activated(
activated_scores, e_bias, routed_scaling_factor, top_k,
out_weights, out_ids,
)
return out_weights, out_ids
Reference in New Issue View Git Blame Copy Permalink