From 6e41f1ad298a5d14ed3ba3f6539d31b10b38b68d Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Wed, 2 Sep 2026 20:55:14 -0400 Subject: [PATCH] [Fix] Preserve FP32 in SM107 MXFP8 fallback (#37489) Co-authored-by: Yangmin Li --- .../sglang/srt/layers/quantization/mxfp4.py | 7 +++- .../test_mxfp4_flashinfer_activation_prep.py | 32 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 33a8b6194..2f9155527 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -98,7 +98,12 @@ def _prepare_flashinfer_mxfp8_activations( if prepared is not None: prepared_packed_topk, x_quant, x_scale = prepared x_scale = x_scale.view(torch.float8_e4m3fn) - elif x.shape[-1] != hidden_size or _is_sm107_supported(): + # FlashInfer handles SM107 inputs unless K3 reaches this fallback with an + # exact-width FP32 tensor, which its quantizer rejects. Use SGLang's compatible + # MXFP8/UE8M0 quantizer for that case; padded inputs still need alignment. + elif x.shape[-1] != hidden_size or ( + _is_sm107_supported() and x.dtype != torch.float32 + ): from sglang.srt.layers.quantization.fp8_utils import ( flashinfer_mxfp8_quantize, ) diff --git a/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py b/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py index a444648ed..f87b3eb20 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py @@ -46,7 +46,37 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): self.assertIs(actual_quant, x_quant) self.assertTrue(torch.equal(actual_scale.view(torch.uint8), x_scale)) - def test_other_sm10x_handoff_miss_keeps_triton_quantizer(self): + def test_sm107_fp32_handoff_miss_uses_sglang_quantizer(self): + x = torch.randn(3, 64, dtype=torch.float32) + x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn) + x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2) + + with patch( + "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None + ) as take, patch( + "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", + return_value=True, + ), patch.object( + per_token_group_quant_module, + "per_token_group_quant", + return_value=(x_quant, x_scale), + ) as quantize, patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + create=True, + ) as flashinfer_quantize: + actual_x, packed_topk, actual_quant, actual_scale = ( + _prepare_flashinfer_mxfp8_activations(x, 64) + ) + + take.assert_called_once_with(x) + quantize.assert_called_once_with(x, group_size=32, scale_ue8m0=True) + flashinfer_quantize.assert_not_called() + self.assertIs(actual_x, x) + self.assertIsNone(packed_topk) + self.assertIs(actual_quant, x_quant) + self.assertTrue(torch.equal(actual_scale.view(torch.uint8), x_scale)) + + def test_other_sm10x_handoff_miss_keeps_sglang_quantizer(self): x = torch.randn(3, 64, dtype=torch.bfloat16) x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn) x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2)