From 90b2581dfeb9e56eaea3aa308b993315c4c9a798 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 06:40:21 +0000 Subject: [PATCH] feat: NVFP4 fused router CuTeDSL kernel (WIP) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-kernel NVFP4 block-scaled GEMM + fused sqrt(softplus) + top-k epilogue. Avoids materializing intermediate FP32 logits to GMEM. Architecture: 6-warp specialization - Warp 5 (TMA): Load A, B, SFA, SFB from GMEM → SMEM - Warp 4 (MMA): NVFP4 block-scaled GEMM → FP32 accumulator in TMEM - Warps 0-3 (EPI): TMEM → registers → sqrt(softplus) + bias + top-k → GMEM Epilogue maintains per-thread min-heap across N subtiles, then merges all 128 threads' heaps in SMEM for final top-k selection. Mirrors Sm100BlockScaledPersistentDenseGemmKernel structure for TMA/MMA/SFA/SFB handling, with custom top-k epilogue replacing the standard SwiGLU + TMA store path. NOTE: This is WIP — needs compilation testing on B200. Several API details (tiled_mma_sfb, cluster_layout_sfb_vmnk) need to be passed through the kernel parameters properly. --- .../router/nvfp4_fused_router_kernel.py | 734 ++++++++++++++++++ 1 file changed, 734 insertions(+) create mode 100644 dsv4/kernels/router/nvfp4_fused_router_kernel.py diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py new file mode 100644 index 00000000..eaa06fe6 --- /dev/null +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -0,0 +1,734 @@ +"""DSV4 NVFP4 Fused Router Kernel — Blackwell SM100. + +Fuses the NVFP4 block-scaled GEMM with the sqrt(softplus) + e_bias + top-k +epilogue into a single kernel launch. Avoids materializing the intermediate +(N, E) FP32 logits tensor to global memory. + +Architecture (6-warp specialization): + Warp 5 (TMA): Load A [M,K] and B [K,N] tiles GMEM → SMEM, plus scale factors + Warp 4 (MMA): NVFP4 block-scaled GEMM, FP32 accumulator → TMEM + Warps 0-3 (EPI): TMEM → registers → sqrt(softplus) + bias + top-k heap → GMEM + +The epilogue accumulates a per-thread min-heap across all subtiles. +After all subtiles for a row are processed, thread 0 of warp 0 merges +all heaps in SMEM, sorts, renormalizes, and writes the final (k=6) +weights and expert IDs to global memory. + +Math (DSV4 §2.1): + logit = X @ W_gate (NVFP4 block-scaled GEMM, FP32 accumulator) + act = sqrt(softplus(logit)) softplus(x) = max(x,0) + log(1+exp(-|x|)) + score = act + e_bias[e] + ids = argtopk(score, k=6) min-heap, lower index wins ties + w = (act[ids] / sum(act[ids])) * scaling + +NVFP4 GEMM details: + - A operand: FP4 (quantized from BF16 activation), SFA in TMEM + - B operand: FP4 (from checkpoint), SFB in TMEM + - Accumulator: FP32 in TMEM + - Global scales: gsa (activation) and gsb (weight) applied in epilogue +""" + +from __future__ import annotations + +from typing import Tuple, Optional + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +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.torch as cutlass_torch + + +LOG2_E = 1.44269504088896340736 + + +class Nvfp4FusedRouterKernel: + """NVFP4 block-scaled GEMM + fused sqrt(softplus)/top-k router epilogue. + + Single-kernel replacement for the two-kernel path: + Nvfp4Linear (NVFP4 GEMM) → activation_topk CUDA kernel + + The fusion eliminates the intermediate FP32 logits write to GMEM + and the subsequent read-back. For decode (1 token, 384 experts), + the savings are small (1.5KB), but for large-batch prefill the + bandwidth savings and reduced kernel launch overhead are significant. + """ + + def __init__(self, mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1), top_k=6): + # Data types + self.a_dtype = cutlass.Float4E2M1FN # FP4 activation (quantized from BF16) + self.b_dtype = cutlass.Float4E2M1FN # FP4 weight + self.sf_dtype = cutlass.Float8E4M3FN # Scale factors (E4M3) + self.acc_dtype = cutlass.Float32 + self.c_dtype = cutlass.Float32 # Accumulator for topk + + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + self.top_k = top_k + + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + self.mma_kind = tcgen05.mma.Kind.BLOCK_SCALE + + # Warp layout (6 warps: 4 epi + 1 MMA + 1 TMA) + self.epilog_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 # 192 + + # Barrier IDs + self.cta_sync_bar_id = 1 + self.epilog_sync_bar_id = 2 + self.tmem_alloc_sync_bar_id = 3 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.occupancy = 1 + self.buffer_align_bytes = 1024 + + # ---------------------------------------------------------------- + # MMA setup — mirrors Sm100BlockScaledPersistentDenseGemmKernel + # ---------------------------------------------------------------- + def _create_tiled_mma(self): + """Create the tiled MMA for NVFP4 block-scaled GEMM.""" + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, self.a_major_mode, self.b_major_mode, + self.acc_dtype, self.cta_group, self.mma_tiler_mn, + self.mma_kind, self.sf_dtype, + ) + + def _setup_attributes(self): + self._tiled_mma = self._create_tiled_mma() + mma_inst_shape_k = cute.size(self._tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = (*self.mma_tiler_mn, mma_inst_shape_k * mma_inst_tile_k) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(self._tiled_mma.thr_id.shape), + self.mma_tiler[1], self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self._tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self._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 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, self.use_2cta_instrs, + layout_d=utils.LayoutEnum.ROW_MAJOR, elem_ty_d=self.c_dtype, + layout_c=None, elem_ty_c=None, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + self.num_ab_stage = 2 + self.num_acc_stage = 1 + self.overlapping_accum = False + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + self._tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + self._tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage) + + # Scale factor SMEM layouts + self.sfa_smem_layout = sm100_utils.make_smem_layout_sfa( + self._tiled_mma, self.mma_tiler, self.sf_dtype) + self.sfb_smem_layout = sm100_utils.make_smem_layout_sfb( + self._tiled_mma, self.mma_tiler, self.sf_dtype) + + acc_shape = self._tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = self._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) + + def mainloop_s2t_copy_and_partition( + self, sSF: cute.Tensor, tSF: cute.Tensor, + ) -> tuple: + """SMEM → TMEM copy partition for scale factors (mirrors dense.py).""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), self.sf_dtype) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def mainloop_s2t_copy_and_partition_sfb( + self, sSF: cute.Tensor, tSF: cute.Tensor, + ) -> tuple: + """SMEM → TMEM copy partition for SFB (uses tiled_mma_sfb).""" + return self.mainloop_s2t_copy_and_partition(sSF, tSF) + + # ---------------------------------------------------------------- + def epilog_tmem_copy_and_partition(self, epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta): + """TMEM → register copy partition. Same as dense GEMM.""" + epi_thr_idx = epi_tidx % (self.threads_per_warp * len(self.epilog_warp_id)) + tiled_copy_t2r, tTR_tAcc, tTR_rAcc = sm100_utils.epilogue_tmem_copy_and_partition( + epi_thr_idx, tCtAcc_base, epi_tile, self._tiled_mma, self.c_dtype, use_2cta, + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + # ---------------------------------------------------------------- + # Public API + # ---------------------------------------------------------------- + def run( + self, + mat_a, # (M, K//2) FP4 activation (quantized) + mat_b, # (K//2, N) FP4 weight + scale_a, # (M, K//16) E4M3 activation scale factors + scale_b, # (K//16, N) E4M3 weight scale factors + expert_offsets, # (1,) int32 — [M] for single-group + global_scale_a, # (1,) FP32 — gsa + global_scale_b, # (1,) FP32 — gsb + e_bias, # (N,) FP32 — per-expert bias + out_weights, # (M, top_k) FP32 — output weights + out_ids, # (M, top_k) int32 — output expert IDs + M, N, K, # Problem dimensions + scaling, # routed_scaling_factor + top_k, # k=6 + stream=None, + ): + if stream is None: + stream = cuda.CUstream(0) + + @cute.jit + def _compiled_fn(mat_a, mat_b, scale_a, scale_b, expert_offsets, + global_scale_a, global_scale_b, e_bias, out_weights, out_ids): + # Infer major modes + self.a_major_mode = utils.LayoutEnum.from_tensor(mat_a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(mat_b).mma_major_mode() + self._setup_attributes() + tiled_mma = self._tiled_mma + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Compute TMA load bytes for pipeline setup + a_smem_0 = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + a_copy = cute.size_in_bytes(self.a_dtype, a_smem_0) + b_smem_0 = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + b_copy = cute.size_in_bytes(self.b_dtype, b_smem_0) + # Scale factor sizes + sfa_smem_0 = cute.slice_(self.sfa_smem_layout, (None, None, 0)) + sfa_copy = cute.size_in_bytes(self.sf_dtype, sfa_smem_0) + sfb_smem_0 = cute.slice_(self.sfb_smem_layout, (None, None, 0)) + sfb_copy = cute.size_in_bytes(self.sf_dtype, sfb_smem_0) + self.num_tma_load_bytes = (a_copy + b_copy + sfa_copy + sfb_copy) * atom_thr_size + + # Make TMA atoms for A, B, SFA, SFB + a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, mat_a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) + + b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, mat_b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape) + + # Scale factor TMA atoms (same pattern as dense GEMM) + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem = 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, self.mma_tiler, tiled_mma, + self.cluster_layout_vmnk.shape, internal_type=cutlass.Int16) + + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem = 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, self.mma_tiler, self._tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Int16) + + num_M_tiles = cute.ceil_div(M, self.cta_tile_shape_mnk[0]) + num_N_tiles = cute.ceil_div(N, self.cta_tile_shape_mnk[1]) + L = 1 + grid = (num_M_tiles * num_N_tiles, 1, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + (cutlass.Int32(num_M_tiles), cutlass.Int32(num_N_tiles), cutlass.Int32(L)), + (*self.cluster_shape_mn, 1)) + + self._kernel( + tiled_mma, + tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b, + tma_atom_sfa, tma_tensor_sfa, tma_atom_sfb, tma_tensor_sfb, + 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.epi_tile, + e_bias, out_weights, out_ids, + expert_offsets, global_scale_a, global_scale_b, + tile_sched_params, + M, N, K, top_k, scaling, + ).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, expert_offsets, + global_scale_a, global_scale_b, e_bias, out_weights, out_ids) + + # ================================================================ + # KERNEL + # ================================================================ + @cute.kernel + def _kernel( + self, tiled_mma, + tma_atom_a, mA_mkl, tma_atom_b, mB_nkl, + tma_atom_sfa, mSFA_mkl, tma_atom_sfb, mSFB_nkl, + cluster_layout_vmnk, cluster_layout_sfb_vmnk, + a_smem_layout_staged, b_smem_layout_staged, + sfa_smem_layout_staged, sfb_smem_layout_staged, + epi_tile, + e_bias_tensor, out_w_tensor, out_id_tensor, + expert_offsets, gsa_tensor, gsb_tensor, + tile_sched_params, M, N, K, top_k, routed_scaling_factor, + ): + 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 + 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) + + # ============================================================== + # Shared storage + # ============================================================== + @cute.struct + class SharedStorage: + ab_full_mbar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding: cutlass.Int32 + # Top-k heap SMEM: 128 threads * 6 entries * (score + index + act) + heap_scores: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, 4*32*6], 128] + heap_indices: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 4*32*6], 128] + heap_acts: cute.struct.Align[cute.struct.MemRange[cutlass.Float32, 4*32*6], 128] + sA: cute.struct.Align[cute.struct.MemRange[self.a_dtype, cute.cosize(a_smem_layout_staged.outer)], self.buffer_align_bytes] + sB: cute.struct.Align[cute.struct.MemRange[self.b_dtype, cute.cosize(b_smem_layout_staged.outer)], self.buffer_align_bytes] + sSFA: cute.struct.Align[cute.struct.MemRange[self.sf_dtype, cute.cosize(sfa_smem_layout.outer)], self.buffer_align_bytes] + sSFB: cute.struct.Align[cute.struct.MemRange[self.sf_dtype, cute.cosize(sfb_smem_layout.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) + + num_acc_cons = self.threads_per_warp * len(self.epilog_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_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_cons), + cta_layout_vmnk=cluster_layout_vmnk) + + 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.epilog_warp_id))), + allocator_warp_id=self.epilog_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_bar = pipeline.NamedBarrier(self.epilog_sync_bar_id, + self.threads_per_warp * len(self.epilog_warp_id)) + + # ============================================================== + # SMEM tensors + # ============================================================== + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout.outer, swizzle=sfa_smem_layout.inner) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout.outer, swizzle=sfb_smem_layout.inner) + + # Multicast masks + a_mcast = None; b_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) + if cutlass.const_expr(self.is_sfb_mcast or use_2cta): + sfb_mcast = cpasync.create_tma_multicast_mask(cluster_layout_sfb_vmnk, block_coord, mcast_mode=1) + + # Partition globals + 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, (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_SFA(gSFA); tCgSFB = thr_mma.partition_SFB(gSFB) + + 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)) + + # SFA/SFB TMA partition (same pattern as dense GEMM) + 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)) + 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[1], sfb_cta_l, + cute.group_modes(sSFB,0,3), cute.group_modes(tCgSFB,0,3)) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + tCrSFA = tiled_mma.make_fragment_SFA(sSFA) + tCrSFB = tiled_mma.make_fragment_SFB(sSFB) + + 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)) + + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # ============================================================== + # TMA WARP (5) — load A, B, SFA, SFB tiles from GMEM → SMEM + # ============================================================== + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + 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]) + tA_s = tAgA[(None, mc[0], None, mc[2])] + tB_s = tBgB[(None, mc[1], None, mc[2])] + tSFA_s = tAgSFA[(None, mc[0], None, mc[2])] + tSFB_s = tBgSFB[(None, mc[1], None, mc[2])] + for kt in cutlass.range(0, k_tiles, 1, unroll=1): + ab_pipeline.producer_acquire(ab_ps) + cute.copy(tma_atom_a, tA_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, tB_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, tSFA_s[(None, ab_ps.count)], tAsSFA[(None, ab_ps.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps)) + cute.copy(tma_atom_sfb, tSFB_s[(None, ab_ps.count)], tBsSFB[(None, ab_ps.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_ps)) + ab_ps.advance() + ab_pipeline.producer_tail(ab_ps) + tsched.advance_to_next_work(); wt = tsched.get_current_work() + + # ============================================================== + # MMA WARP (4) — NVFP4 block-scaled GEMM (mirrors dense.py mainloop) + # ============================================================== + 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() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc = tCtAcc_base[(None, None, None, 0)] + # SFA/SFB SMEM → TMEM copy atoms + tiled_copy_s2t_sfa, tCsSFA, tCtSFA = self.mainloop_s2t_copy_and_partition(sSFA, self._tiled_mma) + tiled_copy_s2t_sfb, tCsSFB, tCtSFB = self.mainloop_s2t_copy_and_partition_sfb(sSFB, self._tiled_mma_sfb) + + 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_stage) + while wt.is_valid_tile: + cute.clear(tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for kt in cutlass.range(0, k_tiles, 1, unroll=1): + ab_pipeline.consumer_wait(ab_cs) + # Copy A, B from SMEM → register fragments + cute.copy(tiled_mma.partition_A(sA[(None,None,None,ab_cs.index)]), tCrA) + cute.copy(tiled_mma.partition_B(sB[(None,None,None,ab_cs.index)]), tCrB) + # Copy SFA, SFB from SMEM → TMEM + s2t_stage = (None, None, None, None, ab_cs.index) + cute.copy(tiled_copy_s2t_sfa, tCsSFA[s2t_stage], tCtSFA) + cute.copy(tiled_copy_s2t_sfb, tCsSFB[s2t_stage], tCtSFB) + # GEMM with block-scaled MMA + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = (None, None, kblock_idx, ab_cs.index) + sf_kblock = (None, None, kblock_idx) + tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock].iterator) + tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock].iterator) + cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + ab_pipeline.consumer_release(ab_cs); ab_cs.advance() + acc_pipeline.producer_commit(acc_ps); acc_ps.advance() + tsched.advance_to_next_work(); wt = tsched.get_current_work() + acc_pipeline.producer_tail(acc_ps) + tmem.relinquish_alloc_permit() + + # ============================================================== + # EPILOGUE WARPS (0-3) — TMEM → sqrt(softplus) + top-k → GMEM + # ============================================================== + if warp_idx in self.epilog_warp_id: + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + cta_bar.arrive_and_wait() + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc0 = tCtAcc_base[(None, None, None, 0)] + + # TMEM → register copy + epi_n = self.epi_tile_n + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(epi_n)), self.acc_dtype) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tCtAcc0) + sfw_idx = tidx % (self.threads_per_warp * len(self.epilog_warp_id)) + thr_ld = tiled_tmem_load.get_slice(sfw_idx) + tS = thr_ld.partition_S(tCtAcc0) + tD = thr_ld.partition_D(tS) + + # Identity tensor for expert index mapping + cAcc = cute.make_identity_tensor((self.cta_tile_shape_mnk[0], N)) + tCcAcc = tiled_mma.get_slice(mma_tile_v).partition_C(cAcc) + cFlat = cute.flatten(tCcAcc) + rFlat = cute.flatten(tD) + + # Per-thread register heap (top_k=6 entries) + hs = [cutlass.Float32(-1e30)] * 6 + hi = [cutlass.Int32(-1)] * 6 + ha = [cutlass.Float32(0.0)] * 6 + + # Read global scales (same for all tiles) + gsa_val = gsa_tensor[0] + gsb_val = gsb_tensor[0] + + 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_stage) + + while wt.is_valid_tile: + acc_pipeline.consumer_wait(acc_cs) + # Reset heap + for i in cutlass.range(6, unroll=1): + hs[i] = cutlass.Float32(-1e30) + hi[i] = cutlass.Int32(-1) + ha[i] = cutlass.Float32(0.0) + + # Process subtiles + for subtile in cutlass.range(1, unroll=1): + cute.copy(tiled_tmem_load, tS, tD) + + elem_cnt = cute.size(rFlat) + for e in cutlass.range(elem_cnt, unroll=4): + logit_raw = rFlat[e] + # Apply global scales: the GEMM output is + # sum(A_sf * B_sf) which needs * gsa * gsb + logit = logit_raw * gsa_val * gsb_val + coord = cFlat[e] + e_idx = coord[1] + + # sqrt(softplus(logit)) + abs_x = cute.math.absf(logit) + pos = cute.where(logit > cutlass.Float32(0.0), logit, cutlass.Float32(0.0)) + exp_neg = cute.math.exp(-abs_x) + sp = pos + cute.math.log(cutlass.Float32(1.0) + exp_neg) + act = cute.math.sqrt(sp) + + # score = act + bias + score = act + e_bias_tensor[e_idx] + + # Min-heap push: root = hs[0] (smallest) + do_push = (score > hs[0]) or (score == hs[0] and e_idx < hi[0]) + if do_push: + hs[0] = score; hi[0] = e_idx; ha[0] = act + # Sift down (k=6, fully unrolled) + root = 0 + _done = cutlass.Bool(False) + while root < 3 and not _done: + left = 2*root+1; right = 2*root+2 + smallest = root + if left < 6: + if hs[left] < hs[smallest] or (hs[left] == hs[smallest] and hi[left] > hi[smallest]): + smallest = left + if right < 6: + if hs[right] < hs[smallest] or (hs[right] == hs[smallest] and hi[right] > hi[smallest]): + smallest = right + if smallest == root: + _done = cutlass.Bool(True) + if not _done: + ts_ = hs[root]; ti_ = hi[root]; ta_ = ha[root] + hs[root] = hs[smallest]; hi[root] = hi[smallest]; ha[root] = ha[smallest] + hs[smallest] = ts_; hi[smallest] = ti_; ha[smallest] = ta_ + root = smallest + + # Write heap to shared memory for merge + tid = (warp_idx * 32 + tidx) + base = tid * 6 + for i in cutlass.range(6, unroll=1): + storage.heap_scores.data_ptr()[base + i] = hs[i] + storage.heap_indices.data_ptr()[base + i] = hi[i] + storage.heap_acts.data_ptr()[base + i] = ha[i] + + epi_bar.arrive_and_wait() + + # Thread 0 of warp 0 does the final merge + store + if warp_idx == 0 and tidx == 0: + # Initialize final heap from thread 0 + fs = list(hs); fi = list(hi); fa = list(ha) + # Merge all 128 threads + for t in cutlass.range(1, 128, unroll=1): + for i in cutlass.range(6, unroll=1): + cs = storage.heap_scores.data_ptr()[t*6+i] + ci = storage.heap_indices.data_ptr()[t*6+i] + ca = storage.heap_acts.data_ptr()[t*6+i] + if ci >= 0: + if cs > fs[0] or (cs == fs[0] and ci < fi[0]): + fs[0] = cs; fi[0] = ci; fa[0] = ca + # Sift down + r = 0 + _done2 = cutlass.Bool(False) + while r < 3 and not _done2: + l = 2*r+1; ri = 2*r+2; sm = r + if l < 6: + if fs[l] < fs[sm] or (fs[l] == fs[sm] and fi[l] > fi[sm]): + sm = l + if ri < 6: + if fs[ri] < fs[sm] or (fs[ri] == fs[sm] and fi[ri] > fi[sm]): + sm = ri + if sm == r: + _done2 = cutlass.Bool(True) + else: + ts_=fs[r]; ti_=fi[r]; ta_=fa[r] + fs[r]=fs[sm]; fi[r]=fi[sm]; fa[r]=fa[sm] + fs[sm]=ts_; fi[sm]=ti_; fa[sm]=ta_ + r = sm + + # Sort descending (selection sort, k=6) + sorted_s = [cutlass.Float32(-1e30)]*6 + sorted_i = [cutlass.Int32(-1)]*6 + sorted_a = [cutlass.Float32(0.0)]*6 + for i in cutlass.range(6, unroll=1): + best = 0 + for j in cutlass.range(1, 6, unroll=1): + if fs[j] > fs[best] or (fs[j] == fs[best] and fi[j] < fi[best]): + best = j + sorted_s[i] = fs[best]; sorted_i[i] = fi[best]; sorted_a[i] = fa[best] + fs[best] = cutlass.Float32(-1e30) + + # Renormalize + act_sum = sorted_a[0] + sorted_a[1] + sorted_a[2] + sorted_a[3] + sorted_a[4] + sorted_a[5] + inv_sum = cutlass.Float32(1.0) / act_sum + sc = cutlass.Float32(routed_scaling_factor) + + # Get tile coordinates for output indexing + tc = wt.tile_idx + row_base = tc[0] // cute.size(tiled_mma.thr_id.shape) * self.cta_tile_shape_mnk[0] + + # Store to GMEM + for i in cutlass.range(6, unroll=1): + out_w_tensor[row_base + 0, i] = sorted_a[i] * inv_sum * sc + out_id_tensor[row_base + 0, i] = sorted_i[i] + + epi_bar.arrive_and_wait() + + 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_bar.arrive_and_wait() + tmem.free(tmem_ptr) + + +# ================================================================ +# Python wrapper — mirrors Nvfp4Linear but uses fused kernel +# ================================================================ +def run_nvfp4_fused_router( + hidden_states: torch.Tensor, # [M, K] BF16 + mat_b: torch.Tensor, # [K//2, N_packed] FP4 weight + scale_a: torch.Tensor, # [M, K//16] E4M3 activation SF + scale_b: torch.Tensor, # [K//16, N] E4M3 weight SF + gsa: float, # activation global scale + gsb: float, # weight global scale + e_bias: torch.Tensor, # [E] FP32 + routed_scaling_factor: float, + top_k: int = 6, + out_weights: Optional[torch.Tensor] = None, # [M, top_k] FP32 + out_ids: Optional[torch.Tensor] = None, # [M, top_k] int32 +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the NVFP4 fused router kernel. + + Combines NVFP4 block-scaled GEMM + sqrt(softplus) + top-k + into a single kernel launch. + """ + M = hidden_states.shape[0] + K = hidden_states.shape[1] + N = scale_b.shape[1] # num_experts from weight scale shape + device = hidden_states.device + + if out_weights is None: + out_weights = torch.empty(M, top_k, dtype=torch.float32, device=device) + if out_ids is None: + out_ids = torch.empty(M, top_k, dtype=torch.int32, device=device) + + # Expert offsets: single group of M tokens + expert_offsets = torch.tensor([M], dtype=torch.int32, device=device) + # Global scales as 1-element tensors + gsa_t = torch.tensor([gsa], dtype=torch.float32, device=device) + gsb_t = torch.tensor([gsb], dtype=torch.float32, device=device) + + # Quantize activation to FP4 (same as Nvfp4Linear) + from dsv4.ops.quantize import quantize_activation_nvfp4 + x_fp4, x_sf = quantize_activation_nvfp4(hidden_states, gsa) + + # Create CuTe tensors + a_tensor = cutlass_torch.make_tensor(x_fp4, shape=x_fp4.shape) + b_tensor = cutlass_torch.make_tensor(mat_b, shape=mat_b.shape) + sfa_tensor = cutlass_torch.make_tensor(x_sf, shape=x_sf.shape) + sfb_tensor = cutlass_torch.make_tensor(scale_b, shape=scale_b.shape) + e_bias_ct = cutlass_torch.make_tensor(e_bias, shape=e_bias.shape) + out_w_ct = cutlass_torch.make_tensor(out_weights, shape=out_weights.shape) + out_id_ct = cutlass_torch.make_tensor(out_ids, shape=out_ids.shape) + eo_ct = cutlass_torch.make_tensor(expert_offsets, shape=expert_offsets.shape) + gsa_ct = cutlass_torch.make_tensor(gsa_t, shape=gsa_t.shape) + gsb_ct = cutlass_torch.make_tensor(gsb_t, shape=gsb_t.shape) + + kernel = Nvfp4FusedRouterKernel(top_k=top_k) + kernel.run( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, + eo_ct, gsa_ct, gsb_ct, + e_bias_ct, out_w_ct, out_id_ct, + M, N, K, routed_scaling_factor, top_k, + ) + return out_weights, out_ids \ No newline at end of file