From b4808d44da2cfec3850a7666d67369fbe9c1425c Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Thu, 28 May 2026 00:01:31 -0700 Subject: [PATCH] Use Cute-DSL MXFP8 quantize kernels (#25486) Co-authored-by: b8zhong --- .../layers/quantization/mxfp4_flashinfer_trtllm_moe.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 6ec432394..7a2e34d47 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -21,7 +21,9 @@ from sglang.srt.utils import ( log_info_on_rank0, set_weight_attrs, ) -from sglang.srt.utils.common import next_power_of_2 +from sglang.srt.utils.common import is_sm100_supported, next_power_of_2 + +_MXFP8_QUANTIZE_BACKEND = "cute-dsl" if is_sm100_supported() else "cuda" if is_flashinfer_available(): from flashinfer import mxfp8_quantize, shuffle_matrix_a, shuffle_matrix_sf_a @@ -379,7 +381,10 @@ class Mxfp4FlashinferTrtllmMoEMethod: ) elif precision == "default": x_quant, x_scale = mxfp8_quantize( - hidden_states, False, alignment=hidden_size + hidden_states, + False, + alignment=hidden_size, + backend=_MXFP8_QUANTIZE_BACKEND, ) x_scale = x_scale.view(torch.float8_e4m3fn).reshape( *hidden_states.shape[:-1], -1