Implementation:Sgl project Sglang GEMM Interface

Knowledge Sources	Sgl_project_Sglang
Domains	GPU Kernels, GEMM, Quantization, LLM Inference
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for GEMM (General Matrix Multiply) and quantization kernels supporting FP8, INT8, FP4, AWQ, GPTQ, Marlin, QServe, and DeepSeek V3 precision formats.

Description

gemm.py provides a comprehensive set of matrix multiplication and quantization wrappers covering all quantization formats used in production LLM serving. The module enables memory-efficient inference across diverse model architectures by supporting formats from FP8 down to INT4.

Scaled Matrix Multiplication:

• fp8_scaled_mm -- FP8 per-tensor scaled matrix multiply with optional bias, returning results in the specified output dtype.
• int8_scaled_mm -- INT8 scaled matrix multiply with per-tensor scales and optional bias.
• fp8_blockwise_scaled_mm -- FP8 blockwise (per-group) scaled matrix multiply for finer-grained quantization.
• bmm_fp8 -- Batched FP8 matrix multiply using cuBLAS, with an internal workspace buffer cached via _get_cache_buf (32MB default).
• cutlass_scaled_fp4_mm -- CUTLASS-based FP4 scaled matrix multiply with block scales and global alpha.

Quantization Routines:

• sgl_per_token_group_quant_8bit -- Per-token group quantization to 8-bit (FP8 or INT8) with configurable group size, epsilon, clipping range, and UE8M0 scale format. Supports a V2 variant enabled via SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2 env var that adds fused SiLU-and-mul and masked M support.
• sgl_per_tensor_quant_fp8 -- Per-tensor FP8 quantization with static or dynamic scaling.
• sgl_per_token_quant_fp8 -- Per-token FP8 quantization.
• scaled_fp4_quant -- Quantizes input to FP4 with swizzled block scales in FP8-E4M3 format (packed as uint8 pairs).

AWQ/GPTQ Kernels:

• awq_dequantize -- Dequantizes AWQ quantized weights given scales and zero points.
• gptq_marlin_gemm -- High-performance GPTQ Marlin GEMM supporting multiple quantization types (via ScalarType), with optional global scale, zeros, group indices, and FP32 reduce.
• gptq_gemm -- Standard GPTQ GEMM with configurable bit width and shuffling.
• gptq_shuffle -- In-place shuffling of GPTQ quantized weights for optimized memory access.

DeepSeek V3 Kernels:

• dsv3_router_gemm -- Optimized GEMM for DeepSeek V3 router computation.
• dsv3_fused_a_gemm -- Fused activation + GEMM for DeepSeek V3 architecture.

QServe Kernels:

• qserve_w4a8_per_chn_gemm -- W4A8 per-channel GEMM from the QServe framework.
• qserve_w4a8_per_group_gemm -- W4A8 per-group GEMM from the QServe framework.

FP4 Expert Quantization:

• scaled_fp4_experts_quant -- Quantizes packed MoE inputs to FP4 with expert-specific offsets and block scale offsets. Supports expert mapping via row shuffling.
• scaled_fp4_grouped_quant -- Quantizes grouped tensors to FP4 for flashinfer grouped_gemm_nt_masked, with swizzled scale layout.
• silu_and_mul_scaled_fp4_grouped_quant -- Fused SiLU-and-mul with FP4 quantization for grouped MoE inputs.

Utility:

• shuffle_rows -- Permutes rows of a tensor according to a destination-to-source mapping.
• Legacy aliases sgl_per_token_group_quant_fp8 and sgl_per_token_group_quant_int8 point to sgl_per_token_group_quant_8bit.

Usage

Use these functions for quantized model inference where weight matrices and/or activations are stored in reduced precision formats. The choice of function depends on the quantization scheme used by the model (GPTQ, AWQ, FP8, FP4, etc.).

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/gemm.py
• Lines: 1-564

Signature

def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor) -> torch.ByteTensor:

def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):

def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):

def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype):

def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None) -> torch.Tensor:

def dsv3_fused_a_gemm(mat_a, mat_b, output=None) -> torch.Tensor:

def sgl_per_token_group_quant_8bit(
    input, output_q, output_s, group_size, eps, fp8_min, fp8_max,
    scale_ue8m0=False, fuse_silu_and_mul=False, masked_m=None, enable_v2=None,
) -> None:

def sgl_per_tensor_quant_fp8(input, output_q, output_s, is_static) -> None:

def sgl_per_token_quant_fp8(input, output_q, output_s) -> None:

def cutlass_scaled_fp4_mm(a, b, block_scale_a, block_scale_b, alpha, out_dtype) -> torch.Tensor:

def scaled_fp4_quant(input, input_global_scale) -> Tuple[torch.Tensor, torch.Tensor]:

def gptq_marlin_gemm(a, c, b_q_weight, b_scales, global_scale, b_zeros, g_idx, perm,
    workspace, b_q_type, size_m, size_n, size_k, is_k_full=True,
    use_atomic_add=False, use_fp32_reduce=False, is_zp_float=False) -> torch.Tensor:

def gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit) -> torch.Tensor:

def gptq_shuffle(q_weight, q_perm, bit) -> None:

def qserve_w4a8_per_chn_gemm(in_feats, kernel, wscales, ascales, w_szs, a_ssums, out_feats=None) -> torch.Tensor:

def qserve_w4a8_per_group_gemm(in_feats, kernel, zeros, scales_i8, wscales, ascales, out_feats=None) -> torch.Tensor:

def dsv3_router_gemm(hidden_states, router_weights, out_dtype=torch.bfloat16) -> torch.Tensor:

def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):

def scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask):

def silu_and_mul_scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask):

def scaled_fp4_experts_quant(input_tensor, input_global_scale, expert_offsets,
    blockscale_offsets, topk, expert_map=None) -> tuple[torch.Tensor, torch.Tensor]:

Import

from sgl_kernel import fp8_scaled_mm, int8_scaled_mm, bmm_fp8
from sgl_kernel import gptq_marlin_gemm, sgl_per_token_group_quant_fp8

I/O Contract

Inputs

Name	Type	Required	Description
mat_a	torch.Tensor	Yes	Left matrix operand (activations)
mat_b	torch.Tensor	Yes	Right matrix operand (weights, possibly quantized)
scales_a, scales_b	torch.Tensor	Yes	Per-tensor or per-group scale factors
out_dtype	torch.dtype	Yes	Output data type (e.g., torch.bfloat16)
bias	torch.Tensor	No	Optional bias vector
b_q_type	ScalarType	Yes (GPTQ)	Quantization type descriptor for GPTQ/Marlin
group_size	int	Yes (quant)	Quantization group size for per-token-group quantization

Outputs

Name	Type	Description
result	torch.Tensor	Matrix multiplication result in specified output dtype
output_q	torch.Tensor	Quantized output tensor (for quantization routines, in-place)
output_s	torch.Tensor	Scale factors (for quantization routines, in-place)

Usage Examples

from sgl_kernel import fp8_scaled_mm, bmm_fp8, gptq_marlin_gemm
from sgl_kernel.scalar_type import scalar_types
import torch

# FP8 scaled matrix multiplication
result = fp8_scaled_mm(
    mat_a=activations_fp8,  # (M, K) in FP8
    mat_b=weights_fp8,       # (K, N) in FP8
    scales_a=scale_a,        # (1,) or per-row
    scales_b=scale_b,        # (1,) or per-col
    out_dtype=torch.bfloat16,
)

# Batched FP8 GEMM
output = bmm_fp8(
    A=batch_a,    # (B, M, K) in FP8
    B=batch_b,    # (B, K, N) in FP8
    A_scale=a_s,
    B_scale=b_s,
    dtype=torch.bfloat16,
)

# GPTQ Marlin GEMM (4-bit quantized)
result = gptq_marlin_gemm(
    a=activations,
    c=None,
    b_q_weight=packed_weights,
    b_scales=weight_scales,
    global_scale=None,
    b_zeros=weight_zeros,
    g_idx=None,
    perm=None,
    workspace=workspace_buf,
    b_q_type=scalar_types.uint4b8,
    size_m=M, size_n=N, size_k=K,
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime