From f82addd4a80ac6c7a07f14e7504baad717c71a95 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Wed, 24 Jun 2026 16:58:48 -0700 Subject: [PATCH] Support online MXFP8 quantization for ungated MoE (#27939) Co-authored-by: Brayden Zhong --- .../docs/advanced_features/quantization.mdx | 4 +++- .../moe/moe_runner/flashinfer_trtllm.py | 7 +++++- python/sglang/srt/layers/quantization/fp8.py | 10 +++++---- .../srt/layers/quantization/fp8_utils.py | 22 ++++++++++--------- 4 files changed, 27 insertions(+), 16 deletions(-) diff --git a/docs_new/docs/advanced_features/quantization.mdx b/docs_new/docs/advanced_features/quantization.mdx index c3350f5c7..d17a5d823 100644 --- a/docs_new/docs/advanced_features/quantization.mdx +++ b/docs_new/docs/advanced_features/quantization.mdx @@ -200,7 +200,7 @@ On Ascend, various layers quantization configurations are supported, see [Ascend ## GEMM Backends for FP4/FP8 Quantization -Backend selection is supported only for **blockwise FP8** and **NVFP4** GEMM. When running FP8 or FP4 quantized models, you can select the GEMM backend via `--fp8-gemm-backend` and `--fp4-gemm-backend`. +Backend selection applies to **blockwise FP8**, **MXFP8** (dense linear), and **NVFP4** GEMM. When running offline or online FP8 or FP4 quantized models, you can select the GEMM backend via `--fp8-gemm-backend` and `--fp4-gemm-backend`. ### `--fp8-gemm-backend` (Blockwise FP8 GEMM) @@ -259,6 +259,8 @@ Backend selection is supported only for **blockwise FP8** and **NVFP4** GEMM. Wh **`auto` selection order:** 1) DeepGEMM (SM90/SM100, installed); 2) FlashInfer TRTLLM (SM100, FlashInfer available); 3) CUTLASS (SM90/SM100/120); 4) AITER (AMD); 5) Triton. **Exception:** SM120 always resolves to Triton. +**MXFP8 dense linear:** `auto` uses `flashinfer_cutlass` on SM100 (else `triton`). `flashinfer_cutlass` is fastest on most shapes; `flashinfer_trtllm` is faster only at small M. + ### `--fp4-gemm-backend` (NVFP4 GEMM) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 09b77a5bb..4d44dfe6d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -322,6 +322,11 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: assert w13_scale.dtype == torch.uint8 assert w2_scale.dtype == torch.uint8 + if not is_gated: + intermediate = w2_weight.shape[2] + w13_weight = w13_weight[:, :intermediate, :].contiguous() + w13_scale = w13_scale[:, :intermediate, :].contiguous() + # Pad for kernel alignment (non-gated needs 128, gated needs 16) min_alignment = 16 if is_gated else 128 w13_weight, w13_scale, w2_weight, w2_scale, _ = _align_mxfp8_moe_weights( @@ -675,7 +680,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( assert quant_info.weight_block_k == 32 from flashinfer import mxfp8_quantize - a_q, a_sf = mxfp8_quantize(hidden_states, False) + a_q, a_sf = mxfp8_quantize(hidden_states, False, backend="cute-dsl") # FlashInfer TRT-LLM MxFP8 expects token-major activation scales: # [num_tokens, hidden_size // 32] (no transpose). a_sf_t = a_sf.view(torch.uint8).reshape(hidden_states.shape[0], -1) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 6d33b9121..5250003e9 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -587,7 +587,8 @@ class Fp8LinearMethod(LinearMethodBase): if not self.use_mxfp8: return - if get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): + backend = get_fp8_gemm_runner_backend() + if backend.is_flashinfer_trtllm(): from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a weight = layer.weight.data @@ -628,7 +629,7 @@ class Fp8LinearMethod(LinearMethodBase): .reshape_as(scale_u8) .contiguous(), ) - elif get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): + elif backend.is_flashinfer_cutlass(): from flashinfer import block_scale_interleave scale_u8 = layer.weight_scale_inv.data @@ -791,9 +792,10 @@ class Fp8LinearMethod(LinearMethodBase): ) if self.use_mxfp8: - if get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): + backend = get_fp8_gemm_runner_backend() + if backend.is_flashinfer_cutlass(): weight_scale = layer.weight_scale_inv_swizzled - elif get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): + elif backend.is_flashinfer_trtllm(): weight_scale = layer.weight_scale_inv_shuffled else: weight_scale = layer.weight_scale_inv diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 70cff5edf..132ed9399 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -410,15 +410,8 @@ def dispatch_w8a8_block_fp8_linear() -> Callable: def dispatch_w8a8_mxfp8_linear() -> Callable: - """Dispatch MXFP8 linear kernel by --fp8-gemm-backend. - - For MXFP8, Triton remains the default path. We only route to FlashInfer - when backend is explicitly set to flashinfer_cutlass or flashinfer_trtllm. - """ backend = get_fp8_gemm_runner_backend() - if backend.is_flashinfer_trtllm(): - return flashinfer_mxfp8_blockscaled_linear - elif backend.is_flashinfer_cutlass(): + if backend.is_flashinfer_cutlass() or backend.is_flashinfer_trtllm(): return flashinfer_mxfp8_blockscaled_linear return triton_mxfp8_blockscaled_linear @@ -516,7 +509,17 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported backend = "triton" - FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend(backend) + backend = Fp8GemmRunnerBackend(backend) + + if ( + backend.is_auto() + and server_args.quantization == "mxfp8" + and _is_sm100_supported + and is_flashinfer_available() + ): + backend = Fp8GemmRunnerBackend.FLASHINFER_CUTLASS + + FP8_GEMM_RUNNER_BACKEND = backend def get_fp8_gemm_runner_backend() -> Fp8GemmRunnerBackend: @@ -1121,7 +1124,6 @@ def flashinfer_mxfp8_blockscaled_linear( weight_t = weight.contiguous().t() if get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): - weight_scale_t = weight_scale.contiguous().view(-1) output = flashinfer_mm_mxfp8( q_input,