diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 2ad4681a9..195178689 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -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) | diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 135211282..9fbc2764d 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 734b7f037..9174c08ff 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 573f69a3c..679e62fdb 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 7701f9757..c3b6c89d4 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 61219f6b0..6268066fe 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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, diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 1bfe0facd..a46b8ffc6 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -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")) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 59098f7fe..b12400c75 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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": diff --git a/python/sglang/test/test_block_fp8.py b/python/sglang/test/test_block_fp8.py index 2390489ca..edfd98e42 100644 --- a/python/sglang/test/test_block_fp8.py +++ b/python/sglang/test/test_block_fp8.py @@ -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."""