Add mxfp8 support for online quantization, Triton dense linear, and CUTLASS MoE

docs/advanced_features/server_arguments.md

@@ -112,7 +112,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) |

python/sglang/srt/layers/moe/cutlass_moe.py

@@ -5,14 +5,16 @@
 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:
     from sgl_kernel import (
         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,

python/sglang/srt/layers/quantization/__init__.py

@@ -52,6 +52,7 @@ def override_quantization_method(self, *args, **kwargs):
 # 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,