Implementation:Mlc ai Mlc llm CUTLASS Ops

Knowledge Sources	Mlc_ai_Mlc_llm
Domains	Machine Learning, GPU Computing, Quantization, Linear Algebra
Last Updated	2026-02-09 19:00 GMT

Overview

Provides Python-level wrappers for CUTLASS-based GEMM operations in MLC-LLM, supporting FP8 quantized, block-scaled, batched, and group GEMM kernels via TVM's extern operator mechanism.

Description

This module defines a collection of functions that invoke NVIDIA CUTLASS (CUDA Templates for Linear Algebra Subroutines) library kernels through TVM Relax's extern operator interface. These operators are critical for high-performance inference, particularly when using FP8 quantization for reduced memory footprint and increased throughput.

The module provides the following operators:

• group_gemm -- CUTLASS group GEMM for Mixture-of-Experts (MoE) models. Computes grouped matrix multiplications where each group corresponds to an expert. Supports float16, bfloat16, and multiple FP8 format combinations (e5m2/e4m3). Requires a 4MB workspace buffer. The weight tensor has shape (num_groups, n, k) and the indptr tensor maps input rows to expert groups.

• fp8_gemm -- Standard FP8 GEMM with per-tensor scaling. Supports input tensors of arbitrary batch dimensions (ndim >= 2) with a 2D weight matrix and a scalar scale factor. Dispatches to format-specific CUTLASS kernels based on the combination of input/weight/output dtypes.

• fp8_groupwise_scaled_gemm -- Block-scaled FP8 GEMM with per-block (128x128) quantization scaling factors. Both input and weight have separate scale tensors. Requires float8_e4m3fn for both x and weight, float32 scales, and float16 or bfloat16 output. Calls the "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn" extern.

• fp8_groupwise_scaled_bmm -- Batched version of block-scaled FP8 GEMM. All input tensors have a leading batch dimension (ndim == 3). Same quantization constraints as fp8_groupwise_scaled_gemm but operates across a batch of independent matrix multiplications.

• fp8_groupwise_scaled_group_gemm -- Combines block-scaled FP8 quantization with group GEMM for MoE architectures. Supports dynamic input shapes, multi-expert weight tensors of shape (num_experts, n, k), per-block scales for both input and weight, and int64 indptr for expert routing. Handles reshaping for inputs with more than 2 dimensions.

All functions allocate a 4MB (4096 * 1024 bytes) uint8 workspace buffer for CUTLASS internal use and produce output via TVM's op.extern call with a placeholder output tensor.

Usage

Use these operators when deploying models with FP8 quantization on NVIDIA GPUs that support the FP8 data type (Hopper architecture and newer). They are called internally by MLC-LLM model implementations -- particularly MoE models and quantized dense models -- during the TVM Relax compilation pipeline.

Code Reference

Source Location

• Repository: Mlc_ai_Mlc_llm
• File: python/mlc_llm/op/cutlass.py

Signature

def group_gemm(
    x: nn.Tensor,          # shape: (m, k)
    weight: nn.Tensor,      # shape: (num_groups, n, k)
    indptr: nn.Tensor,      # shape: (num_groups,), dtype=int64
    scale: Optional[nn.Tensor] = None,
    weight_dtype: Optional[str] = None,
    out_dtype: Optional[str] = None,
) -> nn.Tensor: ...         # shape: (m, n)

def fp8_gemm(
    x: nn.Tensor,           # shape: (m, k) or higher dim
    weight: nn.Tensor,      # shape: (n, k)
    scale: nn.Tensor,       # shape: (1,)
    weight_dtype: Optional[str] = None,
    out_dtype: Optional[str] = None,
) -> nn.Tensor: ...         # shape: (*batch, n)

def fp8_groupwise_scaled_gemm(
    x: nn.Tensor,           # shape: (m, k), float8_e4m3fn
    x_scale: nn.Tensor,     # shape: (k // 128, m), float32
    weight: nn.Tensor,      # shape: (n, k), float8_e4m3fn
    weight_scale: nn.Tensor,# shape: (n // 128, k // 128), float32
    block_size: Tuple[int, int],  # must be (128, 128)
    out_dtype: str,          # "float16" or "bfloat16"
) -> nn.Tensor: ...         # shape: (m, n)

def fp8_groupwise_scaled_bmm(
    x: nn.Tensor,           # shape: (b, m, k)
    x_scale: nn.Tensor,     # shape: (b, k // 128, m)
    weight: nn.Tensor,      # shape: (b, n, k)
    weight_scale: nn.Tensor,# shape: (b, n // 128, k // 128)
    block_size: Tuple[int, int],
    out_dtype: str,
) -> nn.Tensor: ...         # shape: (b, m, n)

def fp8_groupwise_scaled_group_gemm(
    x: nn.Tensor,           # shape: (m, k)
    x_scale: nn.Tensor,     # shape: (m, k // 128)
    weight: nn.Tensor,      # shape: (num_experts, n, k)
    weight_scale: nn.Tensor,# shape: (num_experts, n // 128, k // 128)
    indptr: nn.Tensor,      # shape: (num_experts,), int64
    block_size: Tuple[int, int],
    out_dtype: str,
) -> nn.Tensor: ...         # shape: (m, n)

Import

from mlc_llm.op.cutlass import (
    group_gemm,
    fp8_gemm,
    fp8_groupwise_scaled_gemm,
    fp8_groupwise_scaled_bmm,
    fp8_groupwise_scaled_group_gemm,
)

I/O Contract

Function	Input Dtypes	Output Dtype	CUTLASS Extern Name
group_gemm	fp16/fp16 or bf16/bf16	same as input	cutlass.group_gemm
group_gemm	e5m2/e5m2	fp16	cutlass.group_gemm_e5m2_e5m2_fp16
group_gemm	e4m3/e5m2	fp16	cutlass.group_gemm_e4m3_e5m2_fp16
group_gemm	e4m3/e4m3	fp16	cutlass.group_gemm_e4m3_e4m3_fp16
fp8_gemm	e5m2/e5m2	fp16	cutlass.gemm_e5m2_e5m2_fp16
fp8_gemm	e4m3/e5m2	fp16	cutlass.gemm_e5m2_e4m3_fp16
fp8_gemm	e4m3/e4m3	fp16	cutlass.gemm_e4m3_e4m3_fp16
fp8_groupwise_scaled_gemm	e4m3/e4m3	fp16 or bf16	cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn
fp8_groupwise_scaled_bmm	e4m3/e4m3	fp16 or bf16	cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn
fp8_groupwise_scaled_group_gemm	e4m3/e4m3	fp16 or bf16	cutlass.groupwise_scaled_group_gemm_e4m3fn_e4m3fn

Constraint	Description
Block size	fp8_groupwise_scaled_* functions require block_size == (128, 128)
Workspace	All functions allocate 4MB (4096 * 1024 bytes) uint8 workspace
Scale dtype	Block-scaled functions require float32 for both x_scale and weight_scale
indptr dtype	group_gemm and group_gemm variants require int64 indptr

Usage Examples

from mlc_llm.op.cutlass import fp8_gemm, group_gemm

# FP8 GEMM with per-tensor scaling
output = fp8_gemm(
    x=input_tensor,         # shape (batch, seq_len, hidden_dim), float8_e4m3fn
    weight=weight_matrix,   # shape (out_features, hidden_dim), float8_e4m3fn
    scale=scale_tensor,     # shape (1,), float32
    out_dtype="float16",
)

# Group GEMM for MoE expert computation
output = group_gemm(
    x=routed_input,         # shape (total_tokens, hidden_dim), float16
    weight=expert_weights,  # shape (num_experts, out_features, hidden_dim), float16
    indptr=expert_indptr,   # shape (num_experts,), int64
)

Related Pages

• Implementation:Mlc_ai_Mlc_llm_MoE_Matmul