chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add mxfp8 support for online quantization, Triton dense linear, and CUTLASS MoE (#17449)

Browse Source
This commit is contained in:
Ziang Li
2026-01-29 05:33:57 -08:00
committed by GitHub
parent cfa09d311c
commit 3c9cc44ff5
9 changed files with 725 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/advanced_features/server_arguments.md
View File

@@ -109,7 +109,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| Argument | Description | Defaults | Options |
| --- | --- | --- | --- |
| `--dtype` | Data type for model weights and activations. * "auto" will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. * "half" for FP16. Recommended for AWQ quantization. * "float16" is the same as "half". * "bfloat16" for a balance between precision and range. * "float" is shorthand for FP32 precision. * "float32" for FP32 precision. | `auto` | `auto`, `half`, `float16`, `bfloat16`, `float`, `float32` |
| `--quantization` | The quantization method. | `None` | `awq`, `fp8`, `gptq`, `marlin`, `gptq_marlin`, `awq_marlin`, `bitsandbytes`, `gguf`, `modelopt`, `modelopt_fp8`, `modelopt_fp4`, `petit_nvfp4`, `w8a8_int8`, `w8a8_fp8`, `moe_wna16`, `qoq`, `w4afp8`, `mxfp4`, `auto-round`, `compressed-tensors`, `modelslim`, `quark_int4fp8_moe` |
| `--quantization` | The quantization method. | `None` | `awq`, `fp8`, `gptq`, `marlin`, `gptq_marlin`, `awq_marlin`, `bitsandbytes`, `gguf`, `modelopt`, `modelopt_fp8`, `modelopt_fp4`, `petit_nvfp4`, `w8a8_int8`, `w8a8_fp8`, `moe_wna16`, `qoq`, `w4afp8`, `mxfp4`, `mxfp8`, `auto-round`, `compressed-tensors`, `modelslim`, `quark_int4fp8_moe` |
| `--quantization-param-path` | Path to the JSON file containing the KV cache scaling factors. This should generally be supplied, when KV cache dtype is FP8. Otherwise, KV cache scaling factors default to 1.0, which may cause accuracy issues. | `None` | Type: Optional[str] |
| `--kv-cache-dtype` | Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+ | `auto` | `auto`, `fp8_e5m2`, `fp8_e4m3`, `bf16`, `bfloat16`, `fp4_e2m1` |
| `--enable-fp32-lm-head` | If set, the LM head outputs (logits) are in FP32. | `False` | bool flag (set to enable) |

105
python/sglang/srt/layers/moe/cutlass_moe.py
View File

@@ -5,7 +5,7 @@ from typing import Optional, Tuple
import torch
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams
from sglang.srt.utils import is_cuda, is_sm90_supported
from sglang.srt.utils import is_cuda, is_sm90_supported, is_sm100_supported
_is_cuda = is_cuda()
if _is_cuda:
@@ -13,6 +13,8 @@ if _is_cuda:
apply_shuffle_mul_sum,
cutlass_fp4_group_mm,
es_fp8_blockwise_scaled_grouped_mm,
es_sm100_mxfp8_blockscaled_grouped_mm,
es_sm100_mxfp8_blockscaled_grouped_quant,
fp8_blockwise_scaled_grouped_mm,
prepare_moe_input,
scaled_fp4_experts_quant,
@@ -43,6 +45,7 @@ def cutlass_fused_experts_fp8(
problem_sizes1: torch.Tensor,
problem_sizes2: torch.Tensor,
use_fp8_blockscale: bool = True,
use_mxfp8: bool = False,
output: Optional[torch.Tensor] = None,
enable_es: Tuple[bool, bool] = (False, False),
) -> torch.Tensor:
@@ -99,6 +102,8 @@ def cutlass_fused_experts_fp8(
b_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert.
use_fp8_blockscale (bool, optional): Flag indicating usage of FP8 with
block scaling. Currently, only `True` is supported. Defaults to `True`.
use_mxfp8 (bool, optional): Flag indicating usage of MXFP8 (UE8M0 scales)
with SM100 expert-specialization kernels. Defaults to `False`.
output (torch.Tensor, optional): Output tensor. If not provided, a new tensor will be created.
enable_es (tuple(bool, bool)): Flag indicating usage of expert specialization kernel for (up-projection, down-projection)
Returns:
@@ -137,6 +142,44 @@ def cutlass_fused_experts_fp8(
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
if use_mxfp8:
assert es_up and es_down, "MXFP8 requires expert-specialization for both GEMMs"
assert is_sm100_supported(), "MXFP8 requires SM100"
assert k % 32 == 0, "MXFP8 requires hidden size to be divisible by 32"
assert n % 32 == 0, "MXFP8 requires intermediate size to be divisible by 32"
assert w1_scale.dtype == torch.uint8, "MXFP8 w1_scale must be uint8"
assert w2_scale.dtype == torch.uint8, "MXFP8 w2_scale must be uint8"
expected_w1_scale_shape = (
num_experts,
w1_q.shape[1] // 32,
w1_q.shape[2],
)
expected_w2_scale_shape = (
num_experts,
w2_q.shape[1] // 32,
w2_q.shape[2],
)
assert (
w1_scale.shape == expected_w1_scale_shape
), f"MXFP8 w1_scale must be {expected_w1_scale_shape}, got {w1_scale.shape}"
assert (
w2_scale.shape == expected_w2_scale_shape
), f"MXFP8 w2_scale must be {expected_w2_scale_shape}, got {w2_scale.shape}"
mxfp8_blockscale_align = 128
total_tokens = m * topk
nonzero_experts = min(num_experts, total_tokens)
max_total = total_tokens + (mxfp8_blockscale_align - 1) * nonzero_experts
max_blockscale = (
(max_total + mxfp8_blockscale_align - 1) // mxfp8_blockscale_align
) * mxfp8_blockscale_align
blockscale_offsets = None
if use_mxfp8 and (es_up or es_down):
blockscale_offsets = torch.empty(
(num_experts + 1,), dtype=torch.int32, device=device
)
prepare_moe_input(
topk_ids,
expert_offsets,
@@ -147,11 +190,27 @@ def cutlass_fused_experts_fp8(
num_experts,
n,
k,
blockscale_offsets,
)
a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
if use_mxfp8 and es_up:
rep_a = shuffle_rows(a, a_map, (m * topk, k))
rep_a_q = torch.empty_like(rep_a, dtype=torch.float8_e4m3fn)
rep_a1_scales = torch.empty(
(max_blockscale, k // 32), dtype=torch.uint8, device=device
)
es_sm100_mxfp8_blockscaled_grouped_quant(
rep_a,
problem_sizes1,
expert_offsets[:-1],
blockscale_offsets[:-1],
rep_a_q,
rep_a1_scales,
)
else:
a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
@@ -173,6 +232,17 @@ def cutlass_fused_experts_fp8(
expert_offsets[:-1],
workspace,
)
elif use_mxfp8 and es_up:
es_sm100_mxfp8_blockscaled_grouped_mm(
c1,
rep_a_q,
w1_q,
rep_a1_scales,
w1_scale,
problem_sizes1,
expert_offsets[:-1],
blockscale_offsets[:-1],
)
else:
fp8_blockwise_scaled_grouped_mm(
c1,
@@ -198,7 +268,21 @@ def cutlass_fused_experts_fp8(
intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
silu_and_mul(c1, intermediate)
intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
if use_mxfp8 and es_down:
intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn)
a2_scale = torch.empty(
(max_blockscale, n // 32), dtype=torch.uint8, device=device
)
es_sm100_mxfp8_blockscaled_grouped_quant(
intermediate,
problem_sizes2,
expert_offsets[:-1],
blockscale_offsets[:-1],
intemediate_q,
a2_scale,
)
else:
intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
if is_sm90_supported() and es_down:
es_fp8_blockwise_scaled_grouped_mm(
@@ -214,6 +298,17 @@ def cutlass_fused_experts_fp8(
expert_offsets[:-1],
workspace,
)
elif use_mxfp8 and es_down:
es_sm100_mxfp8_blockscaled_grouped_mm(
c2,
intemediate_q,
w2_q,
a2_scale,
w2_scale,
problem_sizes2,
expert_offsets[:-1],
blockscale_offsets[:-1],
)
else:
fp8_blockwise_scaled_grouped_mm(
c2,

1
python/sglang/srt/layers/quantization/__init__.py
View File

@@ -52,6 +52,7 @@ if TYPE_CHECKING:
# Base quantization methods
BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
"fp8": Fp8Config,
"mxfp8": Fp8Config,
"blockwise_int8": BlockInt8Config,
"modelopt": ModelOptFp8Config, # Auto-detect, defaults to FP8
"modelopt_fp8": ModelOptFp8Config,

254
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -47,8 +47,10 @@ from sglang.srt.layers.quantization.fp8_utils import (
cutlass_fp8_supported,
dispatch_w8a8_block_fp8_linear,
input_to_float8,
mxfp8_group_quantize,
normalize_e4m3fn_to_e4m3fnuz,
requant_weight_ue8m0_inplace,
triton_mxfp8_blockscaled_linear,
)
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
from sglang.srt.layers.quantization.marlin_utils_fp8 import (
@@ -112,6 +114,7 @@ class Fp8Config(QuantizationConfig):
activation_scheme: str = "dynamic",
ignored_layers: Optional[List[str]] = None,
weight_block_size: List[int] = None,
use_mxfp8: bool = False,
) -> None:
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
if is_checkpoint_fp8_serialized:
@@ -120,6 +123,7 @@ class Fp8Config(QuantizationConfig):
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
self.activation_scheme = activation_scheme
self.ignored_layers = ignored_layers or []
self.use_mxfp8 = use_mxfp8
if weight_block_size is not None:
if not is_checkpoint_fp8_serialized:
raise ValueError(
@@ -133,19 +137,22 @@ class Fp8Config(QuantizationConfig):
raise ValueError(
f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme."
)
if self.use_mxfp8:
if weight_block_size is None:
weight_block_size = [1, 32]
elif weight_block_size != [1, 32]:
raise ValueError("MXFP8 requires weight_block_size=[1, 32].")
self.weight_block_size = weight_block_size
@classmethod
def get_name(cls) -> str:
return "fp8"
def get_name(self) -> str:
return "mxfp8" if self.use_mxfp8 else "fp8"
@classmethod
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
return [torch.bfloat16, torch.half]
@classmethod
def get_min_capability(cls) -> int:
return 80
def get_min_capability(self) -> int:
return 100 if self.use_mxfp8 else 80
@classmethod
def get_config_filenames(cls) -> List[str]:
@@ -154,7 +161,8 @@ class Fp8Config(QuantizationConfig):
@classmethod
def from_config(cls, config: Dict[str, Any]) -> Fp8Config:
quant_method = cls.get_from_keys(config, ["quant_method"])
is_checkpoint_fp8_serialized = "fp8" in quant_method
use_mxfp8 = "mxfp8" in quant_method
is_checkpoint_fp8_serialized = ("fp8" in quant_method) or use_mxfp8
activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
ignored_layers = cls.get_from_keys_or(
config, ["ignored_layers", "modules_to_not_convert"], None
@@ -163,11 +171,17 @@ class Fp8Config(QuantizationConfig):
# hack for ministral
ignored_layers = [layer.replace("model.", "") for layer in ignored_layers]
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
if use_mxfp8 and weight_block_size is not None:
logger.warning(
"MXFP8 ignoring incoming weight_block_size in config.json; it is fixed to [1, 32]."
)
weight_block_size = [1, 32]
return cls(
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
activation_scheme=activation_scheme,
ignored_layers=ignored_layers,
weight_block_size=weight_block_size,
use_mxfp8=use_mxfp8,
)
def get_quant_method(
@@ -223,7 +237,10 @@ class Fp8LinearMethod(LinearMethodBase):
auto_enable = can_auto_enable_marlin_fp8()
self.use_marlin = force_marlin or auto_enable
self.block_quant = self.quant_config.weight_block_size is not None
self.use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
self.block_quant = (
self.use_mxfp8 or self.quant_config.weight_block_size is not None
)
self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
self.is_checkpoint_fp8_serialized = (
self.quant_config.is_checkpoint_fp8_serialized
@@ -324,18 +341,25 @@ class Fp8LinearMethod(LinearMethodBase):
assert self.quant_config.activation_scheme == "dynamic"
elif hasattr(self.quant_config, "linear_activation_scheme"):
assert self.quant_config.linear_activation_scheme == "dynamic"
if self.use_mxfp8 and not self.is_checkpoint_fp8_serialized:
raise ValueError(
"MXFP8 requires fp8-serialized checkpoint for linear layers."
)
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.empty
scale = BlockQuantScaleParameter(
data=torch.empty(
data=scale_init(
(output_size_per_partition + block_n - 1) // block_n,
(input_size_per_partition + block_k - 1) // block_k,
dtype=torch.float32,
dtype=scale_dtype,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
scale.format_ue8m0 = False
scale[:] = torch.finfo(torch.float32).min
scale.format_ue8m0 = self.use_mxfp8
if scale_dtype != torch.uint8:
scale[:] = torch.finfo(torch.float32).min
layer.register_parameter("weight_scale_inv", scale)
else:
scale = PerTensorScaleParameter(
@@ -382,6 +406,15 @@ class Fp8LinearMethod(LinearMethodBase):
layer.weight_scale_inv.data, requires_grad=False
)
return
elif self.use_mxfp8:
if not self.is_checkpoint_fp8_serialized:
self._quantize_mxfp8_weights(layer)
return
# MXFP8 scales are stored as UE8M0 uint8; no requantization here.
# Keep parameter object to preserve weight_loader attrs for hot reload.
layer.weight_scale_inv.requires_grad_(False)
layer.weight_scale_inv.format_ue8m0 = True
return
else:
# For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
from sglang.srt.layers.quantization.fp8_utils import (
@@ -414,6 +447,23 @@ class Fp8LinearMethod(LinearMethodBase):
layer.weight.data = weight.data
layer.weight_scale_inv.data = weight_scale.data
def _quantize_mxfp8_weights(self, layer: Module) -> None:
weight = layer.weight.data
qweight, weight_scale = mxfp8_group_quantize(weight)
# Keep parameter objects to preserve weight_loader attrs for hot reload.
layer.weight.data = qweight
layer.weight.requires_grad_(False)
if hasattr(layer, "weight_scale_inv") and layer.weight_scale_inv is not None:
layer.weight_scale_inv.data = weight_scale
layer.weight_scale_inv.requires_grad_(False)
else:
# First-time online MXFP8 quantization (no serialized scales).
layer.register_parameter(
"weight_scale_inv", Parameter(weight_scale, requires_grad=False)
)
layer.weight_scale_inv.format_ue8m0 = True
layer.input_scale = None
def process_weights_after_loading(self, layer: Module) -> None:
if self.block_quant:
self.process_weights_after_loading_block_quant(layer)
@@ -543,6 +593,23 @@ class Fp8LinearMethod(LinearMethodBase):
bias=bias,
)
if self.use_mxfp8:
if isinstance(x, tuple):
return triton_mxfp8_blockscaled_linear(
input=x[0],
weight=layer.weight,
weight_scale=layer.weight_scale_inv,
input_scale=x[1],
bias=bias,
)
return triton_mxfp8_blockscaled_linear(
input=x,
weight=layer.weight,
weight_scale=layer.weight_scale_inv,
input_scale=None,
bias=bias,
)
if self.block_quant:
if use_intel_amx_backend(layer):
return torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
@@ -600,7 +667,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
def __init__(self, quant_config: Fp8Config):
self.quant_config = quant_config
self.block_quant = self.quant_config.weight_block_size is not None
self.use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
self.block_quant = (
self.use_mxfp8 or self.quant_config.weight_block_size is not None
)
if get_moe_runner_backend().is_cutlass():
assert (
cutlass_fp8_supported()
@@ -708,27 +778,29 @@ class Fp8MoEMethod(FusedMoEMethodBase):
# WEIGHT_SCALES
if self.block_quant:
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.ones
w13_weight_scale = torch.nn.Parameter(
torch.ones(
scale_init(
num_experts,
2 * ((intermediate_size_per_partition + block_n - 1) // block_n),
(hidden_size + block_k - 1) // block_k,
dtype=torch.float32,
dtype=scale_dtype,
),
requires_grad=False,
)
w2_weight_scale = torch.nn.Parameter(
torch.ones(
scale_init(
num_experts,
(hidden_size + block_n - 1) // block_n,
(intermediate_size_per_partition + block_k - 1) // block_k,
dtype=torch.float32,
dtype=scale_dtype,
),
requires_grad=False,
)
# w13_weight and w2_weight are always requanted together
w13_weight_scale.format_ue8m0 = False
w2_weight_scale.format_ue8m0 = False
w13_weight_scale.format_ue8m0 = self.use_mxfp8
w2_weight_scale.format_ue8m0 = self.use_mxfp8
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
assert self.quant_config.activation_scheme == "dynamic"
@@ -856,6 +928,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
_is_cpu_amx_available
), "Fp8MoEMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
elif self.use_mxfp8:
self._process_mxfp8_moe_weights(
layer, quantize=not self.quant_config.is_checkpoint_fp8_serialized
)
else:
# For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
@@ -888,6 +964,143 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
def _process_mxfp8_moe_weights(self, layer: Module, quantize: bool = True) -> None:
if not (_is_cuda and is_sm100_supported()):
raise RuntimeError("MXFP8 MoE quantization requires SM100.")
def _quantize_and_swizzle_with_cutlass_es_kernel(weight: torch.Tensor):
from sgl_kernel import es_sm100_mxfp8_blockscaled_grouped_quant
weight = weight.contiguous()
num_experts, m, k = weight.shape
assert k % 32 == 0, f"{k=} must be divisible by 32 for MXFP8"
weight_flat = weight.view(-1, k).contiguous()
problem_sizes = torch.empty(
(num_experts, 3), dtype=torch.int32, device=weight.device
)
problem_sizes[:, 0] = m
problem_sizes[:, 1] = 0
problem_sizes[:, 2] = k
expert_offsets = torch.arange(
0, num_experts * m, m, dtype=torch.int32, device=weight.device
)
aligned_m = ((m + 127) // 128) * 128
blockscale_offsets = torch.arange(
0,
num_experts * aligned_m,
aligned_m,
dtype=torch.int32,
device=weight.device,
)
qweight = torch.empty_like(weight_flat, dtype=torch.float8_e4m3fn)
scale = torch.empty(
(num_experts * aligned_m, k // 32),
dtype=torch.uint8,
device=weight.device,
)
es_sm100_mxfp8_blockscaled_grouped_quant(
weight_flat,
problem_sizes,
expert_offsets,
blockscale_offsets,
qweight,
scale,
)
qweight = qweight.view_as(weight)
scale = scale.view(num_experts, aligned_m, k // 32)
if aligned_m != m:
scale = scale[:, :m, :]
return qweight, scale
def _swizzle_mxfp8_sf(scale, num_warps):
from triton_kernels.tensor import convert_layout, wrap_torch_tensor
from triton_kernels.tensor_details import layout
scale_layout, scale_layout_opts = (
layout.make_default_matmul_mxfp4_w_scale_layout(
mx_axis=1, num_warps=num_warps
)
)
scale = scale.transpose(-2, -1)
scale = convert_layout(
wrap_torch_tensor(scale), scale_layout, **scale_layout_opts
)
return scale
def _swizzle_with_triton_kernel(
weight_shape: tuple[int, int, int], scale: torch.Tensor
):
num_experts, m, k = weight_shape
aligned_m = ((m + 127) // 128) * 128
scale = scale.view(num_experts, aligned_m, k // 32)
num_warps = 8
scale = _swizzle_mxfp8_sf(scale, num_warps)
scale = scale.data.view(num_experts, aligned_m, k // 32)
return scale
def _quantize_and_swizzle_with_triton_kernel(weight: torch.Tensor):
weight = weight.contiguous()
_, _, k = weight.shape
assert k % 32 == 0, f"{k=} must be divisible by 32 for MXFP8"
weight_flat = weight.view(-1, k).contiguous()
qweight, scale = mxfp8_group_quantize(weight_flat)
qweight = qweight.view_as(weight)
scale = _swizzle_with_triton_kernel(weight.shape, scale)
return qweight, scale
if quantize:
if get_moe_runner_backend().is_cutlass():
w13_q, w13_s = _quantize_and_swizzle_with_cutlass_es_kernel(
layer.w13_weight.data
)
w2_q, w2_s = _quantize_and_swizzle_with_cutlass_es_kernel(
layer.w2_weight.data
)
else:
w13_q, w13_s = _quantize_and_swizzle_with_triton_kernel(
layer.w13_weight.data
)
w2_q, w2_s = _quantize_and_swizzle_with_triton_kernel(
layer.w2_weight.data
)
else:
w13_q = layer.w13_weight.data
w2_q = layer.w2_weight.data
w13_s = _swizzle_with_triton_kernel(
layer.w13_weight.data.shape, layer.w13_weight_scale_inv.data
)
w2_s = _swizzle_with_triton_kernel(
layer.w2_weight.data.shape, layer.w2_weight_scale_inv.data
)
# Keep parameter objects to preserve weight_loader attrs for hot reload.
# Prefer in-place copy; rebind only when shape/dtype changes (online quantize).
def _copy_or_rebind(param: Parameter, new_value: torch.Tensor) -> None:
if (
param.data.shape == new_value.shape
and param.data.dtype == new_value.dtype
):
param.data.copy_(new_value)
else:
param.data = new_value
_copy_or_rebind(layer.w13_weight, w13_q)
_copy_or_rebind(layer.w2_weight, w2_q)
_copy_or_rebind(layer.w13_weight_scale_inv, w13_s)
_copy_or_rebind(layer.w2_weight_scale_inv, w2_s)
layer.w13_weight.requires_grad_(False)
layer.w2_weight.requires_grad_(False)
layer.w13_weight_scale_inv.requires_grad_(False)
layer.w2_weight_scale_inv.requires_grad_(False)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
layer.w13_input_scale = None
layer.w2_input_scale = None
def process_weights_after_loading(self, layer: Module) -> None:
if _is_hip and _use_hip_int4:
self.process_weights_hip_int4(layer)
@@ -1173,6 +1386,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
symm_output = torch.empty_like(x)
topk_weights, topk_ids, _ = dispatch_output.topk_output
use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
output = cutlass_fused_experts_fp8(
x,
layer.w13_weight.transpose(1, 2),
@@ -1195,7 +1409,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
self.problem_sizes1,
self.problem_sizes2,
use_fp8_blockscale=True,
use_mxfp8=use_mxfp8,
output=symm_output,
enable_es=(use_mxfp8, use_mxfp8),
)
return StandardCombineInput(hidden_states=output)

139
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -23,6 +23,11 @@ import torch
import triton
import triton.language as tl
try:
from triton.tools.tensor_descriptor import TensorDescriptor
except:
pass
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.utils import (
ceil_align,
@@ -1175,6 +1180,140 @@ def w8a8_block_fp8_matmul(
)
# Copied and adapted from https://github.com/triton-lang/triton/blob/main/python/tutorials/10-block-scaled-matmul.py
@triton.jit
def _mxfp8_block_scaled_matmul_kernel( #
a_desc, #
a_scale_desc, #
b_desc, #
b_scale_desc, #
c_desc, #
M: tl.constexpr, #
N: tl.constexpr, #
K: tl.constexpr, #
output_type: tl.constexpr, #
BLOCK_M: tl.constexpr, #
BLOCK_N: tl.constexpr, #
BLOCK_K: tl.constexpr, #
rep_m: tl.constexpr, #
rep_n: tl.constexpr, #
rep_k: tl.constexpr, #
NUM_STAGES: tl.constexpr, #
): #
if output_type == 0:
output_dtype = tl.float32
elif output_type == 1:
output_dtype = tl.float16
elif output_type == 2:
output_dtype = tl.bfloat16
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_M)
pid_m = pid % num_pid_m
pid_n = pid // num_pid_m
offs_am = pid_m * BLOCK_M
offs_bn = pid_n * BLOCK_N
offs_k_a = 0
offs_k_b = 0
offs_scale_m = pid_m * rep_m
offs_scale_n = pid_n * rep_n
offs_scale_k = 0
VEC_SIZE: tl.constexpr = 32
accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES):
a = a_desc.load([offs_am, offs_k_a])
b = b_desc.load([offs_bn, offs_k_b])
scale_a = a_scale_desc.load([0, offs_scale_m, offs_scale_k, 0, 0])
scale_b = b_scale_desc.load([0, offs_scale_n, offs_scale_k, 0, 0])
scale_a = (
scale_a.reshape(rep_m, rep_k, 32, 4, 4)
.trans(0, 3, 2, 1, 4)
.reshape(BLOCK_M, BLOCK_K // VEC_SIZE)
)
scale_b = (
scale_b.reshape(rep_n, rep_k, 32, 4, 4)
.trans(0, 3, 2, 1, 4)
.reshape(BLOCK_N, BLOCK_K // VEC_SIZE)
)
accumulator = tl.dot_scaled(
a, scale_a, "e4m3", b.T, scale_b, "e4m3", accumulator
)
offs_k_a += BLOCK_K
offs_k_b += BLOCK_K
offs_scale_k += rep_k
c_desc.store([offs_am, offs_bn], accumulator.to(output_dtype))
# Copied and adapted from https://github.com/triton-lang/triton/blob/main/python/tutorials/10-block-scaled-matmul.py
def mxfp8_block_scaled_matmul_triton(
a: torch.Tensor,
a_scale: torch.Tensor,
b: torch.Tensor,
b_scale: torch.Tensor,
output_dtype: torch.dtype,
*,
block_m: int = 128,
block_n: int = 256,
block_k: int = 128,
num_stages: int = 4,
) -> torch.Tensor:
"""Block-scaled matmul for MXFP8 using Triton dot_scaled."""
M, K = a.shape
N, K_b = b.shape
assert K == K_b
if output_dtype == torch.float32:
output_type = 0
elif output_dtype == torch.float16:
output_type = 1
elif output_dtype == torch.bfloat16:
output_type = 2
else:
raise ValueError(f"Unsupported output dtype: {output_dtype}")
rep_m = block_m // 128
rep_n = block_n // 128
rep_k = block_k // 32 // 4
a_desc = TensorDescriptor.from_tensor(a, [block_m, block_k])
b_desc = TensorDescriptor.from_tensor(b, [block_n, block_k])
scale_block_shape = [1, rep_m, rep_k, 2, 256]
a_scale_desc = TensorDescriptor.from_tensor(a_scale, block_shape=scale_block_shape)
scale_block_shape = [1, rep_n, rep_k, 2, 256]
b_scale_desc = TensorDescriptor.from_tensor(b_scale, block_shape=scale_block_shape)
output = torch.empty((M, N), dtype=output_dtype, device=a.device)
c_desc = TensorDescriptor.from_tensor(output, [block_m, block_n])
grid = (triton.cdiv(M, block_m) * triton.cdiv(N, block_n), 1)
_mxfp8_block_scaled_matmul_kernel[grid](
a_desc,
a_scale_desc,
b_desc,
b_scale_desc,
c_desc,
M,
N,
K,
output_type,
block_m,
block_n,
block_k,
rep_m,
rep_n,
rep_k,
num_stages,
)
return output
@triton.jit
def _per_tensor_quant_mla_fp8_stage1(
x_ptr,

127
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -19,6 +19,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
fp8_dtype,
fp8_max,
is_fp8_fnuz,
mxfp8_block_scaled_matmul_triton,
per_token_group_quant_fp8,
scaled_fp8_quant,
sglang_per_token_quant_fp8,
@@ -38,6 +39,7 @@ from sglang.srt.utils import (
is_flashinfer_available,
is_hip,
is_sm90_supported,
is_sm100_supported,
offloader,
)
@@ -536,6 +538,131 @@ def triton_w8a8_block_fp8_linear(
return output.to(dtype=input_2d.dtype).view(*output_shape)
@lru_cache(maxsize=1)
def _get_triton_mxfp8_downcast():
try:
from triton_kernels.numerics_details.mxfp import downcast_to_mxfp
except Exception as err:
raise RuntimeError(
"MXFP8 quantization requires triton_kernels with MXFP8 support."
) from err
return downcast_to_mxfp
def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""Quantize a 2D contiguous tensor to MXFP8 with UE8M0 scales per group (32)."""
assert x.dim() == 2, f"Expected 2D input, got {x.dim()}D"
assert x.is_contiguous(), "MXFP8 quantization requires a contiguous 2D tensor."
_, k = x.shape
assert k % 32 == 0, f"{k=} must be divisible by 32"
downcast_to_mxfp = _get_triton_mxfp8_downcast()
q_input, scale_u8 = downcast_to_mxfp(x, torch.float8_e4m3fn, axis=1)
return q_input.contiguous(), scale_u8.contiguous()
def _pack_mxfp8_scales(scale_u8: torch.Tensor) -> torch.Tensor:
# Pack (M, K//32) UE8M0 scales into the layout expected by tl.dot_scaled.
assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D"
scale_u8 = scale_u8.contiguous()
m, k_groups = scale_u8.shape
assert (
k_groups % 4 == 0
), f"{k_groups=} must be divisible by 4 (K must be multiple of 128)"
scale_m = ceil_div(m, 128)
if m % 128 != 0:
pad_rows = scale_m * 128 - m
pad = torch.full(
(pad_rows, k_groups),
127,
dtype=scale_u8.dtype,
device=scale_u8.device,
)
scale_u8 = torch.cat([scale_u8, pad], dim=0)
scale_k = k_groups // 4
scale_u8 = scale_u8.view(scale_m, 128, scale_k, 4)
scale_u8 = scale_u8.view(scale_m, 4, 32, scale_k, 4)
packed = scale_u8.permute(0, 3, 2, 1, 4).contiguous()
return packed.view(1, scale_m, scale_k, 2, 256)
def triton_mxfp8_blockscaled_linear(
input: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
input_scale: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor] = None,
output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
if not (_is_cuda and is_sm100_supported()):
raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100+).")
input_2d = input.view(-1, input.shape[-1]).contiguous()
output_shape = [*input.shape[:-1], weight.shape[0]]
block_m = 128
block_n = 256 if weight.shape[0] % 256 == 0 else 128
block_k = 128
m, k = input_2d.shape
n, k_w = weight.shape
assert k == k_w, f"{k=} does not match {k_w=}"
assert k % 128 == 0, f"{k=} must be divisible by 128 for MXFP8"
assert n % block_n == 0, f"{n=} must be divisible by {block_n}"
assert weight.dtype == torch.float8_e4m3fn, "MXFP8 weight must be FP8 E4M3."
assert weight_scale.dtype == torch.uint8, "MXFP8 weight_scale must be UE8M0 uint8."
if input_scale is None:
q_input, x_scale_u8 = mxfp8_group_quantize(input_2d)
else:
q_input = input_2d
x_scale_u8 = input_scale
assert x_scale_u8.dtype == torch.uint8, "MXFP8 input_scale must be UE8M0 uint8."
assert x_scale_u8.shape == (m, k // 32)
if output_dtype is None:
if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32):
output_dtype = input_2d.dtype
else:
output_dtype = torch.bfloat16
if m % block_m != 0:
pad_rows = ceil_div(m, block_m) * block_m - m
q_input = torch.cat(
[
q_input,
torch.zeros((pad_rows, k), device=q_input.device, dtype=q_input.dtype),
],
dim=0,
)
pad_scale = torch.full(
(pad_rows, k // 32),
127,
device=x_scale_u8.device,
dtype=x_scale_u8.dtype,
)
x_scale_u8 = torch.cat([x_scale_u8, pad_scale], dim=0)
a_scale_packed = _pack_mxfp8_scales(x_scale_u8)
b_scale_packed = _pack_mxfp8_scales(weight_scale)
output = mxfp8_block_scaled_matmul_triton(
q_input,
a_scale_packed,
weight.contiguous(),
b_scale_packed,
output_dtype=output_dtype,
block_m=block_m,
block_n=block_n,
block_k=block_k,
)
output = output[:m, :]
if bias is not None:
output += bias
return output.to(dtype=output_dtype).view(*output_shape)
def dequant_mxfp4(
w_block: torch.Tensor,
w_scale: torch.Tensor,

3
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -36,6 +36,7 @@ from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.layers.dp_attention import get_attention_tp_rank
from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp8Config,
@@ -227,6 +228,8 @@ def get_quant_config(
# If the quantization config is not found, use the default config.
if not possible_config_filenames:
if model_config.quantization == "mxfp8":
return Fp8Config(use_mxfp8=True, is_checkpoint_fp8_serialized=False)
return quant_cls()
config_files = glob.glob(os.path.join(hf_folder, "*.json"))

23
python/sglang/srt/server_args.py
View File

@@ -93,6 +93,7 @@ LOAD_FORMAT_CHOICES = [
QUANTIZATION_CHOICES = [
"awq",
"fp8",
"mxfp8",
"gptq",
"marlin",
"gptq_marlin",
@@ -2013,6 +2014,14 @@ class ServerArgs:
), "Please enable dp attention when setting enable_dp_lm_head. "
def _handle_moe_kernel_config(self):
if self.quantization == "mxfp8":
if self.moe_runner_backend not in ["auto", "cutlass"]:
logger.warning(
"mxfp8 quantization forces --moe-runner-backend=cutlass. "
f"Overriding {self.moe_runner_backend!r}."
)
self.moe_runner_backend = "cutlass"
if self.moe_runner_backend == "flashinfer_cutlass":
assert self.quantization in [
"modelopt_fp4",
@@ -2041,14 +2050,18 @@ class ServerArgs:
logger.warning(
"SGLANG_CUTLASS_MOE is deprecated, use --moe-runner-backend=cutlass and/or --speculative-moe-runner-backend=cutlass instead"
)
assert (
self.quantization == "fp8"
), "cutlass MoE is only supported with fp8 quantization"
assert self.quantization in [
"fp8",
"mxfp8",
], "cutlass MoE is only supported with fp8/mxfp8 quantization"
self.moe_runner_backend = "cutlass"
if self.moe_runner_backend == "cutlass" and self.quantization == "fp8":
if self.moe_runner_backend == "cutlass" and self.quantization in [
"fp8",
"mxfp8",
]:
assert (
self.ep_size == 1
), "FP8 Cutlass MoE is only supported with ep_size == 1"
), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1"
def _handle_a2a_moe(self):
if self.moe_a2a_backend == "deepep":

102
python/sglang/test/test_block_fp8.py
View File

@@ -1,5 +1,6 @@
import itertools
import unittest
from functools import lru_cache
import torch
@@ -13,12 +14,29 @@ from sglang.srt.layers.quantization.fp8_kernel import (
static_quant_fp8,
w8a8_block_fp8_matmul,
)
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.srt.layers.quantization.fp8_utils import (
input_to_float8,
mxfp8_group_quantize,
triton_mxfp8_blockscaled_linear,
)
from sglang.srt.utils import is_sm100_supported
from sglang.test.test_utils import CustomTestCase
_is_cuda = torch.cuda.is_available() and torch.version.cuda
# For test
@lru_cache(maxsize=1)
def _get_triton_mxfp8_upcast():
try:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp_torch
except Exception as err:
raise RuntimeError(
"MXFP8 dequantization requires triton_kernels with MXFP8 support."
) from err
return upcast_from_mxfp_torch
# For test
def native_per_token_group_quant_fp8(
x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn
@@ -414,6 +432,88 @@ class TestW8A8BlockFP8Matmul(CustomTestCase):
self._w8a8_block_fp8_matmul(*params)
def _mxfp8_group_dequant(q: torch.Tensor, scale_u8: torch.Tensor) -> torch.Tensor:
upcast_from_mxfp_torch = _get_triton_mxfp8_upcast()
return upcast_from_mxfp_torch(q, scale_u8, torch.float32, axis=1)
class TestMXFP8DenseLinear(CustomTestCase):
DTYPES = [torch.bfloat16]
M = [1, 127, 128, 129, 255, 256]
NKs = [
(256, 512),
(384, 1024),
(512, 2048),
(768, 1024),
]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
if not is_sm100_supported():
raise unittest.SkipTest("MXFP8 requires Blackwell (SM100+)")
torch.set_default_device("cuda")
def _mxfp8_dense_linear(self, M, NK, dtype, seed):
N, K = NK
torch.manual_seed(seed)
input_fp32 = torch.randn((M, K), dtype=torch.float32) / 4
input_fp16 = input_fp32.to(dtype)
weight_fp32 = torch.randn((N, K), dtype=torch.float32) / 4
weight_q, weight_scale_u8 = mxfp8_group_quantize(weight_fp32)
with torch.inference_mode():
q_input, input_scale_u8 = mxfp8_group_quantize(input_fp16.to(torch.float32))
a_dq = _mxfp8_group_dequant(q_input, input_scale_u8)
b_dq = _mxfp8_group_dequant(weight_q, weight_scale_u8)
ref_out = torch.matmul(a_dq, b_dq.t()).to(dtype)
out = triton_mxfp8_blockscaled_linear(
input=input_fp16,
weight=weight_q,
weight_scale=weight_scale_u8,
)
out_prequant = triton_mxfp8_blockscaled_linear(
input=q_input,
weight=weight_q,
weight_scale=weight_scale_u8,
input_scale=input_scale_u8,
output_dtype=dtype,
)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
self.assertTrue(
torch.mean(
torch.abs(out_prequant.to(torch.float32) - ref_out.to(torch.float32))
)
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
def test_mxfp8_dense_linear(self):
for params in itertools.product(
self.M,
self.NKs,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
dtype=params[2],
seed=params[3],
):
self._mxfp8_dense_linear(*params)
# For test
def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
"""This function performs fused moe with block-wise quantization using native torch."""