Implementation:Sgl project Sglang CUTLASS MoE

Knowledge Sources	Sgl_project_Sglang
Domains	Kernel, Mixture of Experts, Quantization
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for CUTLASS-based W4A8 (4-bit weight, 8-bit activation) Mixture-of-Experts grouped matrix multiplication.

Description

The cutlass_moe.py module provides two functions for CUTLASS-based quantized MoE inference. get_cutlass_w4a8_moe_mm_data prepares all routing data needed for MoE grouped GEMM from topk_ids (token-expert mapping), computing expert_offsets (indices marking where each expert begins), problem_sizes (M x N x K for each expert), input_permutation (to shuffle inputs before GEMM), and output_permutation (to restore output ordering). cutlass_w4a8_moe_mm performs the actual grouped matrix multiplication between int4 packed weights and FP8 (float_e4m3) activations using CUTLASS Hopper kernels. It supports per-group scaling via a_scales and b_scales tensors, with configurable chunk_size for block-wise scaling. Both functions delegate to torch.ops.sgl_kernel.* C++ ops.

Usage

Use these functions for quantized MoE inference on NVIDIA Hopper (H100) GPUs, where expert weights are stored in W4A8 format (4-bit weights with 8-bit activations) to minimize memory bandwidth and maximize throughput.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/cutlass_moe.py
• Lines: 1-112

Signature

def get_cutlass_w4a8_moe_mm_data(
    topk_ids: torch.Tensor,
    expert_offsets: torch.Tensor,
    problem_sizes1: torch.Tensor,
    problem_sizes2: torch.Tensor,
    input_permutation: torch.Tensor,
    output_permutation: torch.Tensor,
    num_experts: int,
    n: int,
    k: int,
): ...

def cutlass_w4a8_moe_mm(
    d: torch.Tensor,
    a: torch.Tensor,
    b: torch.Tensor,
    a_scales: torch.Tensor,
    b_scales: torch.Tensor,
    experts_offsets: torch.Tensor,
    problem_sizes: torch.Tensor,
    a_strides: torch.Tensor,
    b_strides: torch.Tensor,
    d_strides: torch.Tensor,
    s_strides: torch.Tensor,
    chunk_size: int = 128,
    topk: int = 8,
): ...

Import

from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data

I/O Contract

Inputs

Name	Type	Required	Description
topk_ids	torch.Tensor	Yes	Token-to-expert mapping from MoE router
expert_offsets	torch.Tensor	Yes	Pre-allocated tensor for expert offset indices
problem_sizes1	torch.Tensor	Yes	Pre-allocated tensor for first MM problem sizes
problem_sizes2	torch.Tensor	Yes	Pre-allocated tensor for second MM problem sizes
input_permutation	torch.Tensor	Yes	Pre-allocated tensor for input shuffle permutation
output_permutation	torch.Tensor	Yes	Pre-allocated tensor for output shuffle permutation
d	torch.Tensor	Yes	Output matrices of shape [total_m, total_n]
a	torch.Tensor	Yes	Activation matrices in FP8 format, shape [total_m, K]
b	torch.Tensor	Yes	Weight matrices in packed int4, shape [E, N, K//2]
a_scales	torch.Tensor	Yes	Scale factors for activations
b_scales	torch.Tensor	Yes	Scale factors for weights, shape [E, K//512, N*8]
chunk_size	int	No	Elements per scale value, default 128
topk	int	No	Number of experts per token, default 8

Outputs

Name	Type	Description
(in-place)	-	get_cutlass_w4a8_moe_mm_data writes to expert_offsets, problem_sizes, and permutation tensors
(in-place)	-	cutlass_w4a8_moe_mm writes result to output tensor d

Usage Examples

from sgl_kernel.cutlass_moe import get_cutlass_w4a8_moe_mm_data, cutlass_w4a8_moe_mm

# Step 1: Prepare MoE routing data
get_cutlass_w4a8_moe_mm_data(
    topk_ids, expert_offsets, problem_sizes1, problem_sizes2,
    input_permutation, output_permutation,
    num_experts=8, n=hidden_dim, k=input_dim
)

# Step 2: Execute the grouped GEMM
cutlass_w4a8_moe_mm(
    d, a, b, a_scales, b_scales,
    expert_offsets, problem_sizes,
    a_strides, b_strides, d_strides, s_strides,
    chunk_size=128, topk=8
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime