Implementation:Mlc ai Mlc llm MoE Matmul

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

Overview

Provides optimized TIR-based matrix multiplication operators for Mixture-of-Experts (MoE) models in MLC-LLM, supporting GEMV, Group GEMM, and various dequantization strategies including integer quantization, FP8 quantization, and block-scaled FP8.

Description

This module implements a comprehensive set of matrix multiplication operators specifically designed for MoE architectures. These operators handle the unique computational pattern of MoE models where different tokens are routed to different expert weight matrices, requiring indexed access to expert-specific weights via an indptr (index pointer) tensor.

The module provides the following operators:

• gemv -- Basic GEMV (General Matrix-Vector multiply) for single-token MoE inference. Computes the product of input x (shape: 1 or experts_per_tok x in_features) with expert weights w (shape: local_experts x out_features x in_features), routed by indptr. Uses TIR thread binding for GPU parallelism across experts.

• dequantize_gemv -- GEMV with on-the-fly integer dequantization. The weight tensor is stored in a packed integer format (e.g., int3, int4 packed into uint32). Dequantization extracts individual quantized values using bit shifting and masking, then applies per-group scale factors. Supports arbitrary quantization bit widths and group sizes.

• dequantize_float8_gemv -- GEMV with FP8 dequantization. Supports both float8_e5m2 and float8_e4m3fn formats with optional per-tensor scaling. Handles both native FP8 storage (num_elem_per_storage == 1) and packed storage, using tir.reinterpret for type casting. Provides two code paths: with scale and without scale.

• dequantize_block_scale_float8_gemv -- GEMV with block-scaled FP8 dequantization. Each block of the weight tensor has its own scale factor, enabling finer-grained quantization. The scale tensor has shape (local_experts, out_features // block_size[0], in_features // block_size[1]).

• group_gemm -- Batched Group GEMM for multi-token MoE inference. This is a highly optimized tiled matrix multiplication with a sophisticated scheduling strategy:
  • Uses a persistent-thread CTA (Cooperative Thread Array) approach with 1024 CTAs
  • Tile sizes: BLK_M=8, BLK_N=128, BLK_K=32 with TX=8, TY=32 threads
  • Implements cooperative fetching of X and W tiles to shared memory
  • Handles dynamic batch sizes and expert routing via indptr
  • Includes boundary checking for partial tiles
  • Uses TVM's schedule primitives for thread binding, vectorization, and loop unrolling (UNROLL=64)

• dequantize_group_gemm -- Combines Group GEMM with integer dequantization. Uses the same tiling and scheduling strategy as group_gemm but adds an inline dequantization step when loading weight tiles to shared memory. Supports both int32 and int64 indptr (with int64 padded for compatibility). Accumulates in float32 precision for numerical stability.

All TIR functions use op_pattern=4 (kOutEWiseFusable) to enable operator fusion in the compilation pipeline.

Usage

Use these operators in MoE model implementations within MLC-LLM. The GEMV variants (gemv, dequantize_gemv, dequantize_float8_gemv, dequantize_block_scale_float8_gemv) are used during single-token decode, while the Group GEMM variants (group_gemm, dequantize_group_gemm) are used during multi-token prefill and batch decode. The choice of operator depends on the weight storage format and quantization scheme.

Code Reference

Source Location

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

Signature

def gemv(
    x: Tensor,       # (1 or experts_per_tok, in_features)
    w: Tensor,       # (local_experts, out_features, in_features)
    indptr: Tensor,  # (1, experts_per_tok), int32
) -> Tensor: ...     # (experts_per_tok, out_features)

def dequantize_gemv(
    x: Tensor,
    w: Tensor,          # packed quantized weights
    scale: Tensor,      # per-group scales
    indptr: Tensor,
    quantize_dtype: str, # e.g., "int4"
    group_size: int,     # e.g., 32 or 128
) -> Tensor: ...

def dequantize_float8_gemv(
    x: Tensor,
    w: Tensor,
    scale: Optional[Tensor],  # per-tensor scale or None
    indptr: Tensor,
    quantize_dtype: Literal["float8_e5m2", "float8_e4m3fn"],
) -> Tensor: ...

def dequantize_block_scale_float8_gemv(
    x: Tensor,
    w: Tensor,
    w_scale: Tensor,       # per-block scales
    expert_indices: Tensor,
    block_size: Tuple[int, int],
    out_dtype: str,
) -> Tensor: ...

def group_gemm(
    x: Tensor,       # (batch_size, in_features), dynamic batch_size
    w: Tensor,       # (num_local_experts, out_features, in_features)
    indptr: Tensor,  # (num_local_experts + 1,), int32
) -> Tensor: ...     # (batch_size, out_features)

def dequantize_group_gemm(
    x: Tensor,
    w: Tensor,           # packed quantized weights
    scale: Tensor,       # per-group scales
    indptr: Tensor,
    quantize_dtype: str,
    indptr_dtype: str,   # "int32" or "int64"
    group_size: int,
) -> Tensor: ...         # (batch_size, out_features)

Import

from mlc_llm.op.moe_matmul import (
    gemv,
    dequantize_gemv,
    dequantize_float8_gemv,
    dequantize_block_scale_float8_gemv,
    group_gemm,
    dequantize_group_gemm,
)

I/O Contract

Function	Input Shape (x)	Weight Shape	Output Shape	Use Case
gemv	(1 or E_tok, K)	(E_all, N, K)	(E_tok, N)	Single-token decode, FP16/BF16 weights
dequantize_gemv	(1 or E_tok, K)	(E_all, N, K/n_pack)	(E_tok, N)	Single-token decode, int-quantized weights
dequantize_float8_gemv	(1 or E_tok, K)	(E_all, N, K/n_pack)	(E_tok, N)	Single-token decode, FP8 weights
dequantize_block_scale_float8_gemv	(1 or E_tok, K)	(E_all, N, K)	(E_tok, N)	Single-token decode, block-scaled FP8 weights
group_gemm	(B, K) dynamic	(E_all, N, K)	(B, N)	Multi-token prefill/batch, FP16/BF16 weights
dequantize_group_gemm	(B, K) dynamic	(E_all, N, K/n_pack)	(B, N)	Multi-token prefill/batch, int-quantized weights

Where: E_tok = experts_per_tok, E_all = local_experts (total), K = in_features, N = out_features, B = batch_size, n_pack = elements per storage unit.

GPU Scheduling Parameter	Value	Description
BLK_M	8	Tile size along the M (batch/row) dimension
BLK_N	128	Tile size along the N (output feature) dimension
BLK_K	32	Tile size along the K (reduction) dimension
TX	8	Thread count along the X dimension per block
TY	32	Thread count along the Y dimension per block
CTA_COUNT	1024	Number of persistent CTAs (thread blocks)
UNROLL	64	Maximum unroll step for loop optimization

Usage Examples

from mlc_llm.op.moe_matmul import gemv, group_gemm, dequantize_gemv

# Single-token MoE GEMV (decode phase)
# x: current token embedding routed to experts
# w: all expert weight matrices
# indptr: which experts are activated
output = gemv(
    x=token_embed,        # shape: (1, 4096)
    w=expert_weights,     # shape: (8, 11008, 4096)
    indptr=expert_ids,    # shape: (1, 2), e.g., [[3, 7]]
)
# output shape: (2, 11008) -- one row per activated expert

# Multi-token MoE Group GEMM (prefill phase)
# x: all tokens routed and sorted by expert
# indptr: boundaries of each expert's token group
output = group_gemm(
    x=sorted_tokens,     # shape: (total_tokens, 4096)
    w=expert_weights,    # shape: (8, 11008, 4096)
    indptr=boundaries,   # shape: (9,), e.g., [0, 5, 12, 12, 20, ...]
)
# output shape: (total_tokens, 11008)

# Single-token MoE with int4 dequantization
output = dequantize_gemv(
    x=token_embed,
    w=quantized_weights,       # int4 packed into uint32
    scale=scale_factors,
    indptr=expert_ids,
    quantize_dtype="int4",
    group_size=128,
)

Related Pages

• Implementation:Mlc_ai_Mlc_llm_CUTLASS_Ops