Implementation:Mlc ai Mlc llm FT GEMM

Overview

The FT GEMM module provides operators that call into NVIDIA's FasterTransformer (FT) library for quantized and mixed-precision GEMM (General Matrix Multiplication) operations. It is located at python/mlc_llm/op/ft_gemm.py (136 lines).

This module exposes two external function interfaces: one for dequantize-GEMM with CutlassFpAIntB (integer-weight, float-activation matmul with optional bias and activation), and one for Mixture-of-Experts (MoE) GEMM. Both functions emit calls to external FasterTransformer kernels via TVM Relax's op.extern mechanism.

Source File

• File: python/mlc_llm/op/ft_gemm.py
• Lines: 136
• Module: mlc_llm.op.ft_gemm

Dependencies

Import	Purpose
operator	operator.mul for computing the M dimension via reduce
functools.reduce	Reduces shape dimensions to compute total row count
typing.Optional	Type annotations for optional parameters
tvm.relax.frontend.nn	Tensor types for input/output specification
tvm.relax.frontend.nn.op	op.extern for calling external library functions

Function: faster_transformer_dequantize_gemm

def faster_transformer_dequantize_gemm(
    x: nn.Tensor,
    weight: nn.Tensor,
    scale: nn.Tensor,
    bias: Optional[nn.Tensor] = None,
    activation: Optional[str] = None,
    group_size: Optional[int] = None,
):

Performs a fused dequantize-GEMM operation using FasterTransformer's CutlassFpAIntB kernel. This computes Y = activation(X @ dequantize(W) + bias) where weights are stored in quantized integer format and dequantized on-the-fly during the matrix multiply.

Parameters

Parameter	Shape	dtype	Description
x	[*m, k]	float16	Input activation tensor (arbitrary batch dimensions)
weight	[k, n // num_elem_per_storage]	varies	Quantized weight tensor (packed integer storage)
scale	[k // group_size, n]	float16	Per-group dequantization scale factors
bias	broadcastable to [*m, n]	float16	Optional bias tensor
activation	--	str	Optional fused activation: None, "identity", "relu", "gelu", or "silu"
group_size	--	int	Quantization group size; defaults to k (per-column quantization)

Assertions

The function enforces strict input requirements:

assert x.dtype == "float16" and x.ndim >= 1
assert weight.ndim == 2
assert scale.dtype == "float16" and scale.ndim == 2
assert x.shape[-1] == weight.shape[0]  # reduction dimension check
assert activation in [None, "relu", "gelu", "silu", "identity"]

Dimension Computation

m = reduce(operator.mul, x.shape[:-1], 1)  # total rows (product of batch dims)
k = x.shape[-1]                             # reduction dimension
n = scale.shape[1]                           # output dimension

The output dimension n is derived from the scale tensor's second dimension rather than from the weight tensor, since the weight tensor has a packed representation where n is divided by the number of elements per storage unit.

External Function Calls

With bias: Calls fastertransformer.gemm_fp16_int_bias with an additional bias stride parameter:

bias_stride = (
    bias.shape[-1]
    if bias and not reduce(operator.mul, bias.shape, 1) == bias.shape[-1]
    else 0
)
return op.extern(
    name="fastertransformer.gemm_fp16_int_bias",
    args=[x, weight, scale, bias, activation, m, n, k, group_size, bias_stride],
    out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"),
)

The bias_stride is set to 0 when the bias is effectively 1D (broadcast), or to bias.shape[-1] when it has a non-trivial batch dimension.

Without bias: Calls fastertransformer.gemm_fp16_int:

return op.extern(
    name="fastertransformer.gemm_fp16_int",
    args=[x, weight, scale, activation, m, n, k, group_size],
    out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"),
)

Function: faster_transformer_moe_gemm

def faster_transformer_moe_gemm(
    x: nn.Tensor,
    weight: nn.Tensor,
    total_rows_before: nn.Tensor,
):

Performs a Mixture-of-Experts (MoE) GEMM using FasterTransformer's moe_gemm_fp16_fp16 kernel. Each expert has its own weight matrix, and the kernel routes input rows to the appropriate expert based on the total_rows_before indptr-like tensor.

Parameters

Parameter	Shape	dtype	Description
x	[*m, k]	float16	Input activation tensor
weight	[num_experts, n, k]	float16	Per-expert weight matrices
total_rows_before	[num_experts]	varies	Cumulative row count per expert (indptr without the leading zero)

Assertions

assert x.dtype == "float16" and x.ndim >= 1
assert weight.dtype == "float16" and weight.ndim == 3
assert x.shape[-1] == weight.shape[-1]  # reduction dimension check

External Function Call

return op.extern(
    name="fastertransformer.moe_gemm_fp16_fp16",
    args=[x, weight, total_rows_before, m, n, k, num_experts],
    out=nn.Tensor.placeholder((*x.shape[:-1], n), dtype="float16"),
)

The dimensions are extracted as:

• m: product of all dimensions except the last (total input rows)
• num_experts: weight.shape[0]
• n: weight.shape[1] (output dimension per expert)
• k: x.shape[-1] (reduction dimension)

Design Notes

• Both functions operate exclusively in float16 precision for activations and outputs.
• The quantized GEMM function supports group-wise quantization where the group size determines how many rows share a single scale factor.
• The op.extern mechanism generates a call to a pre-registered external function that is linked at compile time, allowing these operators to leverage NVIDIA's optimized CUTLASS-based GEMM kernels.
• Output tensors are specified as placeholders with the expected shape, enabling TVM to track tensor shapes through the computation graph.

Categories

• GEMM Operators
• FasterTransformer Integration
• Quantized Inference
• Mixture of Experts
• External Kernels