biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replace kernel.py with thin wrapper around pre-compiled _C extension

Browse Source
This commit is contained in:
biondizzle
2026-05-14 09:25:56 +00:00
parent 33e5d67326
commit 25cbc85afe
1 changed files with 76 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files

582
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
View File

@@ -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 <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
// 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<cutlass::layout::ColumnMajor>;
using StrideB = cutlass::gemm::TagToStrideB_t<cutlass::layout::ColumnMajor>;
// Scale factor strides — K-major
using StrideSFA = Stride<int64_t, _1, int64_t>;
using StrideSFB = Stride<int64_t, _1, int64_t>;
// ============================================================================
// 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<ElementA, ElementSF>;
using ElementPairB = cute::tuple<ElementB, ElementSF>;
// StridePair types: (stride_data, stride_scale) tuples
using StridePairA = cute::tuple<StrideA, StrideSFA>;
using StridePairB = cute::tuple<StrideB, StrideSFB>;
// 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<int64_t, _1, int64_t>, // StrideC
ElementOutput,
Stride<int64_t, _1, int64_t>, // StrideD
cutlass::epilogue::thread::LinearCombination<ElementOutput, 1, ElementAccum, ElementAccum>
>;
// Gemm kernel
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int,int,int,int>, // ProblemShape
CollectiveMainloop,
CollectiveEpilogue
>;
// Device GEMM
using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// ============================================================================
// 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<int>(M), static_cast<int>(N), static_cast<int>(K), 1};
args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(A_packed.data_ptr<int8_t>());
args.mainloop.dA = StrideA{K / 2, 1}; // K_packed stride
args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(B_packed.data_ptr<int8_t>());
args.mainloop.dB = StrideB{K / 2, 1};
args.mainloop.ptr_SFA = reinterpret_cast<ElementSF const*>(SFA.data_ptr<uint8_t>());
args.mainloop.layout_SFA = /* ... */;
args.mainloop.ptr_SFB = reinterpret_cast<ElementSF const*>(SFB.data_ptr<uint8_t>());
args.mainloop.layout_SFB = /* ... */;
args.epilogue.ptr_C = reinterpret_cast<ElementOutput const*>(C.data_ptr<cutlass::bfloat16_t>());
args.epilogue.dC = Stride<int64_t, _1, int64_t>{N, 1};
args.epilogue.ptr_D = reinterpret_cast<ElementOutput*>(C.data_ptr<cutlass::bfloat16_t>());
args.epilogue.dD = Stride<int64_t, _1, int64_t>{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 <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#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<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue
>;
using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
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<int>(M), static_cast<int>(N), static_cast<int>(K)));
auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig<16>::tile_atom_to_shape_SFB(
cute::make_shape(static_cast<int>(M), static_cast<int>(N), static_cast<int>(K)));
typename GemmDevice::Arguments args;
args.mode = cutlass::gemm::GemmUniversalMode::kGemm;
args.problem_shape = {static_cast<int>(M), static_cast<int>(N), static_cast<int>(K), 1};
// Mainloop args — paired (element, scale) inputs
args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(A_packed.data_ptr<int8_t>());
args.mainloop.dA = cutlass::make_cute_packed_stride(Stride<int64_t, _1, int64_t>{},
{static_cast<int>(M), static_cast<int>(K), 1});
args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(B_packed.data_ptr<int8_t>());
args.mainloop.dB = cutlass::make_cute_packed_stride(Stride<int64_t, _1, int64_t>{},
{static_cast<int>(N), static_cast<int>(K), 1});
args.mainloop.ptr_SFA = reinterpret_cast<ElementSF const*>(SFA.data_ptr<uint8_t>());
args.mainloop.layout_SFA = layout_SFA;
args.mainloop.ptr_SFB = reinterpret_cast<ElementSF const*>(SFB.data_ptr<uint8_t>());
args.mainloop.layout_SFB = layout_SFB;
// Epilogue args
args.epilogue.ptr_C = nullptr;
args.epilogue.dC = cutlass::make_cute_packed_stride(Stride<int64_t, _1, int64_t>{},
{static_cast<int>(M), static_cast<int>(N), 1});
args.epilogue.ptr_D = reinterpret_cast<ElementOutput*>(C.data_ptr<cutlass::bfloat16_t>());
args.epilogue.dD = cutlass::make_cute_packed_stride(Stride<int64_t, _1, int64_t>{},
{static_cast<int>(M), static_cast<int>(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