diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py index 31735b87..3b6a57d5 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py @@ -1,524 +1,94 @@ """ CUTLASS NVFP4 Block-Scaled GEMM — Native Blackwell SM100 kernel. -Uses CUTLASS 3.x GemmUniversalAdapter with the -MainloopSm100TmaUmmaWarpSpecializedBlockScaled dispatch policy. -This invokes the native mxf8f6f4.block_scale tensor core instruction -(tcgen05.mma) which performs E2M1 × E2M1 with UE4M3 block-16 scaling -entirely in hardware. - -The kernel is compiled as a PyTorch CUDA extension at runtime. +Uses the pre-compiled PyTorch CUDA extension (cutlass_nvfp4_gemm._C) +which invokes native mxf8f6f4.block_scale tensor core instructions. """ import os import torch -import tempfile -import shutil -from typing import Optional MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) -# --------------------------------------------------------------------------- -# CUDA kernel source — CUTLASS 3.x GemmUniversalAdapter for NVFP4 -# --------------------------------------------------------------------------- +try: + from cutlass_nvfp4_gemm import _C + _CUTLASS_AVAILABLE = True +except ImportError: + _CUTLASS_AVAILABLE = False -CUTLASS_NVFP4_GEMM_CU = r""" -#include -#include -#include -#include - -// CUTLASS includes -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" - -#include "cutlass/cutlass.h" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/thread/linear_combination.h" - -#include "cute/numeric/float8.hpp" - -using namespace cute; - -// ============================================================================ -// Type aliases for NVFP4 -// ============================================================================ -using ElementA = cutlass::float_e2m1_t; // Packed FP4 (2 per byte) -using ElementB = cutlass::float_e2m1_t; // Packed FP4 (2 per byte) -using ElementSF = cutlass::float_ue4m3_t; // Block scale factor (group_size=16) -using ElementAccum = float; // Accumulator -using ElementOutput = cutlass::bfloat16_t; // Output - -// Stride types — K-major layout for both A and B -// A: (M, K_packed) — K-major means contiguous in K -// B: (N, K_packed) — K-major means contiguous in K -using StrideA = cutlass::gemm::TagToStrideA_t; -using StrideB = cutlass::gemm::TagToStrideB_t; - -// Scale factor strides — K-major -using StrideSFA = Stride; -using StrideSFB = Stride; - -// ============================================================================ -// GEMM kernel definition using CUTLASS 3.x API -// ============================================================================ - -// Tile shape: 128x128x64 (M, N, K) — K is in E2M1 elements (32 bytes packed) -// This matches the UMMA atom shape for f8f6f4.block_scale -using TileShape = Shape<_128, _128, _64>; - -// Cluster shape -using ClusterShape = Shape<_1, _1, _1>; - -// Dispatch policy: Warp-specialized UMMA with block scaling, 2 pipeline stages -using DispatchPolicy = cutlass::gemm::collective::MainloopSm100TmaUmmaWarpSpecializedBlockScaled< - 2, // Stages - 1, // SchedulerPipelineStageCount - 1, // AccumulatorPipelineStageCount - ClusterShape, - cutlass::arch::Sm100 ->; - -// ElementPair types: (element, scale_factor) tuples -using ElementPairA = cute::tuple; -using ElementPairB = cute::tuple; - -// StridePair types: (stride_data, stride_scale) tuples -using StridePairA = cute::tuple; -using StridePairB = cute::tuple; - -// Collective mainloop -using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma< - DispatchPolicy, - TileShape, - ElementPairA, - StridePairA, - ElementPairB, - StridePairB, - /*TiledMma=*/void, // Let the builder deduce - /*GmemTiledCopyPairA=*/void, - /*SmemLayoutAtomPairA=*/void, - /*SmemCopyAtomA=*/void, - /*TransformA=*/void, - /*GmemTiledCopyPairB=*/void, - /*SmemLayoutAtomPairB=*/void, - /*SmemCopyAtomB=*/void, - /*TransformB=*/void ->; - -// Collective epilogue -using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue< - cutlass::gemm::collective::EpilogueSm100TmaUmma<2, ClusterShape, cutlass::arch::Sm100>, - TileShape, - ElementAccum, - Stride, // StrideC - ElementOutput, - Stride, // StrideD - cutlass::epilogue::thread::LinearCombination ->; - -// Gemm kernel -using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, // ProblemShape - CollectiveMainloop, - CollectiveEpilogue ->; - -// Device GEMM -using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter; - - -// ============================================================================ -// PyTorch bindings -// ============================================================================ - -torch::Tensor cutlass_nvfp4_gemm_forward( - torch::Tensor A_packed, // (M, K_packed) int8 — E2M1 packed - torch::Tensor SFA, // (M, K_sf) uint8 — UE4M3 block16 scales - torch::Tensor B_packed, // (N, K_packed) int8 — E2M1 packed - torch::Tensor SFB, // (N, K_sf) uint8 — UE4M3 block16 scales - int64_t M, int64_t N, int64_t K // Dimensions (K in E2M1 elements) -) { - auto options = torch::TensorOptions().dtype(torch::kBF16).device(A_packed.device()); - auto C = torch::zeros({M, N}, options); - - // Construct CUTLASS arguments - typename GemmDevice::Arguments args; - args.mode = cutlass::gemm::GemmUniversalMode::kGemm; - args.problem_shape = {static_cast(M), static_cast(N), static_cast(K), 1}; - args.mainloop.ptr_A = reinterpret_cast(A_packed.data_ptr()); - args.mainloop.dA = StrideA{K / 2, 1}; // K_packed stride - args.mainloop.ptr_B = reinterpret_cast(B_packed.data_ptr()); - args.mainloop.dB = StrideB{K / 2, 1}; - args.mainloop.ptr_SFA = reinterpret_cast(SFA.data_ptr()); - args.mainloop.layout_SFA = /* ... */; - args.mainloop.ptr_SFB = reinterpret_cast(SFB.data_ptr()); - args.mainloop.layout_SFB = /* ... */; - args.epilogue.ptr_C = reinterpret_cast(C.data_ptr()); - args.epilogue.dC = Stride{N, 1}; - args.epilogue.ptr_D = reinterpret_cast(C.data_ptr()); - args.epilogue.dD = Stride{N, 1}; - args.epilogue.thread = {ElementAccum(1), ElementAccum(0)}; - args.hw_info.device_id = A_packed.get_device(); - - GemmDevice gemm; - auto status = gemm.initialize(args); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("CUTLASS NVFP4 GEMM initialize failed"); - } - auto stream = c10::cuda::getCurrentCUDAStream(); - status = gemm.run(stream); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("CUTLASS NVFP4 GEMM run failed"); - } - - return C; -} - -TORCH_LIBRARY(cutlass_nvfp4, m) { - m.def("gemm_forward", &cutlass_nvfp4_gemm_forward); -} -""" - - -# --------------------------------------------------------------------------- -# CUTLASS CollectiveBuilder-based approach -# --------------------------------------------------------------------------- -# The above template approach requires explicit type deduction that's complex. -# CUTLASS 3.5+ provides CollectiveBuilder which auto-deduces all types. -# We'll use the builder pattern which is the recommended way. - -CUTLASS_NVFP4_GEMM_BUILDER_CU = r""" -#include -#include -#include -#include - -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" - -#include "cutlass/cutlass.h" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/util/packed_stride.hpp" - -#include "cute/numeric/float8.hpp" - -using namespace cute; - -// NVFP4 types -using ElementA = cutlass::float_e2m1_t; -using ElementB = cutlass::float_e2m1_t; -using ElementSF = cutlass::float_ue4m3_t; -using ElementAccum = float; -using ElementOutput = cutlass::bfloat16_t; - -// K-major layout for both A and B -using LayoutA = cutlass::layout::ColumnMajor; // K-major for A -using LayoutB = cutlass::layout::ColumnMajor; // K-major for B -using LayoutC = cutlass::layout::RowMajor; -using LayoutD = cutlass::layout::RowMajor; - -// Builder for mainloop — uses CUTLASS auto-deduction -using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm100, - cutlass::gemm::collective::OpBlockScaled, // Block-scaled MMA - ElementA, LayoutA, 2, // A type, layout, alignment (2 = 1 byte packed) - ElementB, LayoutB, 2, // B type, layout, alignment - ElementAccum, - Shape<_128, _128, _64>, // TileShape - Shape<_1, _1, _1>, // ClusterShape - cutlass::gemm::collective::StageCountAutoCarveout<0>, - cutlass::gemm::collective::KernelScheduleAuto ->::CollectiveOp; - -// Builder for epilogue -using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - cutlass::arch::Sm100, - cutlass::gemm::EpilogueDefault, - ElementAccum, - ElementOutput, LayoutC, 1, - ElementOutput, LayoutD, 1, - cutlass::gemm::collective::EpilogueScheduleAuto ->::CollectiveOp; - -// Gemm kernel -using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, - CollectiveMainloop, - CollectiveEpilogue ->; - -using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter; - - -torch::Tensor cutlass_nvfp4_gemm_forward( - torch::Tensor A_packed, // (M, K_packed) int8 — E2M1 packed - torch::Tensor SFA, // (M, K_sf) uint8 — UE4M3 block16 scales - torch::Tensor B_packed, // (N, K_packed) int8 — E2M1 packed - torch::Tensor SFB, // (N, K_sf) uint8 — UE4M3 block16 scales - int64_t M, int64_t N, int64_t K -) { - auto options = torch::TensorOptions().dtype(torch::kBF16).device(A_packed.device()); - auto C = torch::zeros({M, N}, options); - - int K_packed = K / 2; - int K_sf = K / 16; // One UE4M3 scale per 16 E2M1 elements - - // Scale factor layouts using Sm1xxBlockScaledConfig - // SFA layout: tile_to_shape(SfAtom{}, make_shape(M, K_sf), Step<_2,_1>{}) - // For the builder, we pass the scale factor pointers and strides - auto layout_SFA = cutlass::detail::Sm1xxBlockScaledConfig<16>::tile_atom_to_shape_SFA( - cute::make_shape(static_cast(M), static_cast(N), static_cast(K))); - auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig<16>::tile_atom_to_shape_SFB( - cute::make_shape(static_cast(M), static_cast(N), static_cast(K))); - - typename GemmDevice::Arguments args; - args.mode = cutlass::gemm::GemmUniversalMode::kGemm; - args.problem_shape = {static_cast(M), static_cast(N), static_cast(K), 1}; - - // Mainloop args — paired (element, scale) inputs - args.mainloop.ptr_A = reinterpret_cast(A_packed.data_ptr()); - args.mainloop.dA = cutlass::make_cute_packed_stride(Stride{}, - {static_cast(M), static_cast(K), 1}); - args.mainloop.ptr_B = reinterpret_cast(B_packed.data_ptr()); - args.mainloop.dB = cutlass::make_cute_packed_stride(Stride{}, - {static_cast(N), static_cast(K), 1}); - args.mainloop.ptr_SFA = reinterpret_cast(SFA.data_ptr()); - args.mainloop.layout_SFA = layout_SFA; - args.mainloop.ptr_SFB = reinterpret_cast(SFB.data_ptr()); - args.mainloop.layout_SFB = layout_SFB; - - // Epilogue args - args.epilogue.ptr_C = nullptr; - args.epilogue.dC = cutlass::make_cute_packed_stride(Stride{}, - {static_cast(M), static_cast(N), 1}); - args.epilogue.ptr_D = reinterpret_cast(C.data_ptr()); - args.epilogue.dD = cutlass::make_cute_packed_stride(Stride{}, - {static_cast(M), static_cast(N), 1}); - args.epilogue.thread = {ElementAccum(1), ElementAccum(0)}; - - GemmDevice gemm; - auto can_impl = GemmDevice::can_implement(args); - if (can_impl != cutlass::Status::kSuccess) { - throw std::runtime_error("CUTLASS NVFP4 GEMM: can_implement returned false"); - } - - auto status = gemm.initialize(args); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("CUTLASS NVFP4 GEMM initialize failed"); - } - auto stream = c10::cuda::getCurrentCUDAStream(); - status = gemm.run(stream); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("CUTLASS NVFP4 GEMM run failed"); - } - - return C; -} - -TORCH_LIBRARY(cutlass_nvfp4, m) { - m.def("gemm_forward", &cutlass_nvfp4_gemm_forward); -} -""" - - -# --------------------------------------------------------------------------- -# Kernel compilation and caching -# --------------------------------------------------------------------------- - -_compiled_ext = None - - -def _find_cutlass_include_dir(): - """Find CUTLASS include directory.""" - candidates = [ - "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/", - "/usr/local/lib/python3.12/dist-packages/cutlass/", - "/usr/include/cutlass/", - os.path.join(os.path.dirname(__file__), "..", "..", "3rdparty", "cutlass"), - ] - # Also check pip-installed cutlass - try: - import cutlass - candidates.insert(0, os.path.dirname(cutlass.__file__)) - except ImportError: - pass - - for c in candidates: - h = os.path.join(c, "include", "cutlass", "cutlass.h") - if os.path.exists(h): - return os.path.join(c, "include") - return None - - -def _get_compiled_extension(): - """Compile and cache the CUTLASS NVFP4 GEMM extension.""" - global _compiled_ext - if _compiled_ext is not None: - return _compiled_ext - - import torch.utils.cpp_extension as cpp_ext - - cutlass_include = _find_cutlass_include_dir() - if cutlass_include is None: - raise RuntimeError( - "Cannot find CUTLASS headers. Install cutlass or set the path. " - "Expected at /usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include/" - ) - - if MEGA_MOE_DEBUG: - print(f"[CUTLASS NVFP4] Using CUTLASS from: {cutlass_include}") - - with tempfile.TemporaryDirectory() as tmpdir: - cu_path = os.path.join(tmpdir, "cutlass_nvfp4_gemm.cu") - with open(cu_path, "w") as f: - f.write(CUTLASS_NVFP4_GEMM_BUILDER_CU) - - ext = cpp_ext.load( - name="cutlass_nvfp4", - sources=[cu_path], - extra_include_paths=[cutlass_include], - extra_cuda_cflags=[ - "-gencode=arch=compute_100a,code=sm_100a", - "--expt-relaxed-constexpr", - "-DCUTLASS_ENABLE_GEMP_OPERATION=1", - "-DCUTLASS_ARCH_SM100_ENABLED=1", - ], - extra_cflags=["-O2", "-std=c++17"], - verbose=MEGA_MOE_DEBUG, - ) - - _compiled_ext = ext - return ext - - -# --------------------------------------------------------------------------- -# Native NVFP4 GEMM API -# --------------------------------------------------------------------------- def cutlass_nvfp4_blockscaled_gemm( - A_packed: torch.Tensor, # (M, K//2) int8 — E2M1 packed - A_scales: torch.Tensor, # (M, K//16) float8_e4m3fn — UE4M3 block16 scales - B_packed: torch.Tensor, # (N, K//2) int8 — E2M1 packed - B_scales: torch.Tensor, # (N, K//16) float8_e4m3fn — UE4M3 block16 scales -) -> torch.Tensor: - """Native NVFP4 block-scaled GEMM using CUTLASS: C = A @ B^T. - - A is (M, K//2) int8 E2M1 packed with (M, K//16) UE4M3 scales. - B is (N, K//2) int8 E2M1 packed with (N, K//16) UE4M3 scales. - C is (M, N) bfloat16. - - Uses CUTLASS's MainloopSm100TmaUmmaWarpSpecializedBlockScaled - which invokes native mxf8f6f4.block_scale tensor core instructions. - """ - M = A_packed.shape[0] - K_half = A_packed.shape[1] - K = K_half * 2 - N = B_packed.shape[0] - - assert A_packed.dtype == torch.int8, f"A must be int8, got {A_packed.dtype}" - assert B_packed.dtype == torch.int8, f"B must be int8, got {B_packed.dtype}" - assert A_packed.is_cuda and B_packed.is_cuda, "Tensors must be on CUDA" - - # Convert scales to uint8 view (raw bytes of float8_e4m3fn) - if A_scales.dtype == torch.float8_e4m3fn: - A_sf_u8 = A_scales.view(torch.uint8).contiguous() - elif A_scales.dtype == torch.uint8: - A_sf_u8 = A_scales.contiguous() - else: - raise ValueError(f"A_scales must be float8_e4m3fn or uint8, got {A_scales.dtype}") - - if B_scales.dtype == torch.float8_e4m3fn: - B_sf_u8 = B_scales.view(torch.uint8).contiguous() - elif B_scales.dtype == torch.uint8: - B_sf_u8 = B_scales.contiguous() - else: - raise ValueError(f"B_scales must be float8_e4m3fn or uint8, got {B_scales.dtype}") - - try: - ext = _get_compiled_extension() - return ext.gemm_forward(A_packed, A_sf_u8, B_packed, B_sf_u8, M, N, K) - except Exception as e: - if MEGA_MOE_DEBUG: - print(f"[CUTLASS NVFP4] Kernel failed, using dequant fallback: {e}") - # Fallback: dequantize and use torch.matmul - from nvfp4_megamoe_kernel.nvfp4_dequant import unpack_e2m1_to_bf16 - A_bf16 = unpack_e2m1_to_bf16(A_packed, A_scales) - B_bf16 = unpack_e2m1_to_bf16(B_packed, B_scales) - return torch.matmul(A_bf16, B_bf16.t()) + A_packed, # (M, K_half) int8 packed E2M1 + SFA, # scale factors for A (float8_e4m3fn) + B_packed, # (N, K_half) int8 packed E2M1 + SFB, # scale factors for B (float8_e4m3fn) + M, N, K, # Problem dimensions (K in FP4 elements) +): + """Single NVFP4 block-scaled GEMM using CUTLASS.""" + if not _CUTLASS_AVAILABLE: + raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") + return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K) def cutlass_grouped_nvfp4_gemm( - x_packed: torch.Tensor, # (num_tokens, K//2) int8 — E2M1 packed - x_scales: torch.Tensor, # (num_tokens, K//16) UE4M3 scales - weights: torch.Tensor, # (E, N, K//2) int8 — per-expert E2M1 weights - weight_scales: torch.Tensor, # (E, N, K//16) UE4M3 per-expert scales - topk_ids: torch.Tensor, # (num_tokens, NUM_TOPK) int32 - topk_weights: torch.Tensor, # (num_tokens, NUM_TOPK) float32 -) -> torch.Tensor: - """Segmented grouped expert GEMM with native NVFP4 block-scaled MMA. - - For each expert, runs the CUTLASS NVFP4 GEMM on tokens routed to it. - Results are scattered back with routing weights. - - This can be optimized in the future using CUTLASS grouped GEMM - (GemmUniversalMode::kGrouped) to batch all expert GEMMs into a single - kernel launch, reducing launch overhead. - - Args: - x_packed: Packed E2M1 activations (num_tokens, K//2) - x_scales: UE4M3 block16 scales (num_tokens, K//16) - weights: Per-expert E2M1 weights (E, N, K//2) - weight_scales: Per-expert UE4M3 scales (E, N, K//16) - topk_ids: Expert assignments (num_tokens, NUM_TOPK) - topk_weights: Routing weights (num_tokens, NUM_TOPK) - - Returns: - (num_tokens, N) bfloat16 output + x_fp4, # (num_tokens, K_half) int8 packed E2M1 + x_sf, # (num_tokens, sf_k) float8_e4m3fn block scales + weights, # (E_per_rank, N, K_half) int8 packed E2M1 + weight_sf, # (E_per_rank, N, sf_k) float8_e4m3fn block scales + topk_ids, # (num_tokens, NUM_TOPK) int32 + topk_weights, # (num_tokens, NUM_TOPK) float32 +): + """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4. + + For each expert, gather the tokens routed to it, run the block-scaled GEMM, + then scatter results back with routing weights. """ - num_tokens = x_packed.shape[0] - K_half = x_packed.shape[1] - K = K_half * 2 - E = weights.shape[0] - N = weights.shape[1] - top_k = topk_ids.shape[1] - device = x_packed.device - - output = torch.zeros(num_tokens, N, dtype=torch.float32, device=device) - - # Process per expert - for e in range(E): - mask = (topk_ids == e) # (num_tokens, top_k) - if not mask.any(): + num_tokens = x_fp4.shape[0] + K_half = x_fp4.shape[1] + K = K_half * 2 # Actual K dimension (2 FP4 per byte) + N = weights.shape[1] # Output dimension + num_experts = weights.shape[0] + num_topk = topk_ids.shape[1] + + if MEGA_MOE_DEBUG: + print(f"[cutlass_grouped_gemm] tokens={num_tokens} K={K} N={N} " + f"experts={num_experts} topk={num_topk}") + + # For now, run all tokens through all experts and select using topk_ids + # This is simpler than per-expert gather and works for moderate token counts + # TODO: optimize with per-expert gather for large batch sizes + + output = torch.zeros(num_tokens, N, dtype=torch.bfloat16, device=x_fp4.device) + + for e in range(num_experts): + # Find tokens routed to this expert + expert_mask = (topk_ids == e) # (num_tokens, num_topk) + token_indices = expert_mask.any(dim=1).nonzero(as_tuple=True)[0] + + if token_indices.numel() == 0: continue - - for k_idx in range(top_k): - token_mask = mask[:, k_idx] - if not token_mask.any(): - continue - token_indices = token_mask.nonzero(as_tuple=True)[0] - - # Gather activations for this expert - x_sub_packed = x_packed[token_indices] # (n, K//2) - x_sub_scales = x_scales[token_indices] # (n, K//16) - w_packed = weights[e] # (N, K//2) - w_scales = weight_scales[e] # (N, K//16) - - # Native NVFP4 GEMM: (n, K) @ (N, K)^T → (n, N) - result = cutlass_nvfp4_blockscaled_gemm( - x_sub_packed, x_sub_scales, - w_packed, w_scales, - ) # (n, N) bfloat16 - - # Weighted scatter-add - weights_f32 = topk_weights[token_indices, k_idx].unsqueeze(-1) - output[token_indices] += result.to(torch.float32) * weights_f32 - - return output.to(torch.bfloat16) + + # Gather tokens for this expert + expert_x = x_fp4[token_indices] # (num_expert_tokens, K_half) + expert_x_sf = x_sf[token_indices] # (num_expert_tokens, sf_k) + expert_w = weights[e] # (N, K_half) + expert_w_sf = weight_sf[e] # (N, sf_k) + + M_expert = token_indices.shape[0] + + # Run CUTLASS NVFP4 block-scaled GEMM + expert_out = cutlass_nvfp4_blockscaled_gemm( + expert_x, expert_x_sf, + expert_w.unsqueeze(0).expand(1, N, K_half).reshape(N, K_half) if expert_w.dim() == 2 else expert_w, + expert_w_sf.unsqueeze(0).expand(1, N, K_half).reshape(N, K_half) if expert_w_sf.dim() == 2 else expert_w_sf, + M_expert, N, K, + ) # (M_expert, N) bfloat16 + + # Scatter back with routing weights + for t_idx, token_idx in enumerate(token_indices): + # Find which topk slot(s) route to this expert + for k_idx in range(num_topk): + if topk_ids[token_idx, k_idx] == e: + output[token_idx] += topk_weights[token_idx, k_idx] * expert_out[t_idx] + + return output