From 44db041700d2f8e1aa8de7d1efb58b42fd4e1c64 Mon Sep 17 00:00:00 2001 From: Sahithi Chigurupati <58884509+csahithi@users.noreply.github.com> Date: Sun, 23 Aug 2026 03:17:03 -0700 Subject: [PATCH] [NVIDIA] Fix SM107 MXFP8 activation prep (#35405) Signed-off-by: Sahithi Chigurupati Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> --- .../srt/layers/quantization/fp8_utils.py | 35 ++--- .../sglang/srt/layers/quantization/mxfp4.py | 81 +++++++----- .../quantization/test_fp8_utils_mxfp4.py | 25 +++- .../test_mxfp4_flashinfer_activation_prep.py | 125 ++++++++++++++++++ 4 files changed, 216 insertions(+), 50 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 93e1811bc..49920c0cb 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -418,6 +418,26 @@ if flashinfer_per_tensor_fp8_supported(): ).view(m, n) +def _fake_flashinfer_mxfp8_quantize( + input: torch.Tensor, + _is_sf_swizzled_layout: bool = True, + alignment: int = 32, + backend: str = "cute-dsl", +) -> Tuple[torch.Tensor, torch.Tensor]: + m = input.numel() // input.shape[-1] + k_aligned = ((input.shape[-1] + alignment - 1) // alignment) * alignment + q_input = input.new_empty((m, k_aligned), dtype=torch.float8_e4m3fn) + sf_columns = k_aligned // 32 + if _is_sf_swizzled_layout: + padded_rows = ((m + 127) // 128) * 128 + padded_sf_columns = ((sf_columns + 3) // 4) * 4 + scale_size = padded_rows * padded_sf_columns + else: + scale_size = m * sf_columns + scale = input.new_empty((scale_size,), dtype=torch.uint8) + return q_input, scale + + if is_blackwell_supported() and is_flashinfer_available(): from flashinfer import SfLayout from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 @@ -479,21 +499,6 @@ if is_blackwell_supported() and is_flashinfer_available(): # Wrap MXFP8 ops as custom ops so torch.compile does not trace into # flashinfer's JIT compilation path (filesystem checks/cubin loader). - def _fake_flashinfer_mxfp8_quantize( - input: torch.Tensor, - _is_sf_swizzled_layout: bool = True, - alignment: int = 32, - backend: str = "cute-dsl", - ) -> Tuple[torch.Tensor, torch.Tensor]: - # Fake mode only needs dtypes and output rank to propagate compile graph. - # The scale tensor shape is not consumed before the following fake mm op. - k_aligned = ((input.shape[1] + alignment - 1) // alignment) * alignment - q_input = input.new_empty( - (input.shape[0], k_aligned), dtype=torch.float8_e4m3fn - ) - scale = input.new_empty((1,), dtype=torch.uint8) - return q_input, scale - @register_custom_op( op_name="flashinfer_mxfp8_quantize", mutates_args=[], diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 12ed0cf27..953ba7b03 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -19,6 +19,7 @@ from __future__ import annotations import os from dataclasses import replace +from functools import lru_cache from typing import TYPE_CHECKING, List, Optional import torch @@ -51,6 +52,7 @@ from sglang.srt.layers.quantization.utils import is_layer_skipped from sglang.srt.runtime_context import get_exec from sglang.srt.utils import ( cpu_has_amx_support, + get_device_capability, is_cpu, is_flashinfer_available, is_gfx95_supported, @@ -74,6 +76,48 @@ has_triton_kernels = is_triton_kernels_available() _UE8M0_ONE = 127 +@lru_cache(maxsize=1) +def _is_sm107_supported() -> bool: + return get_device_capability() == (10, 7) + + +def _prepare_flashinfer_mxfp8_activations( + x: torch.Tensor, hidden_size: int +) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: + prepared = None + if x.shape[-1] == hidden_size: + if x.dim() > 2: + x = x.view(-1, x.shape[-1]) + # K3's routing dispatch may already have quantized these rows and + # packed the topk ids. Other models use FlashInfer's own activation + # preparation so the producer matches the fused-MoE input contract. + from sglang.srt.layers.moe import route_quant_handoff + + prepared = route_quant_handoff.take(x) + + 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(): + from sglang.srt.layers.quantization.fp8_utils import ( + flashinfer_mxfp8_quantize, + ) + + prepared_packed_topk = None + x_quant, x_scale = flashinfer_mxfp8_quantize(x, False, alignment=hidden_size) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1) + else: + from sglang.kernels.ops.quantization.per_token_group_quant import ( + per_token_group_quant, + ) + + prepared_packed_topk = None + x_quant, x_scale = per_token_group_quant(x, group_size=32, scale_ue8m0=True) + x_scale = x_scale.view(torch.float8_e4m3fn) + + return x, prepared_packed_topk, x_quant, x_scale + + if is_flashinfer_available(): from flashinfer import ( nvfp4_block_scale_interleave, @@ -1464,40 +1508,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): value=0.0, ) elif self.flashinfer_mxfp4_moe_precision == "default": - if x.shape[-1] == self.hidden_size: - if x.dim() > 2: - x = x.view(-1, x.shape[-1]) - # K3 staged fusion (route_quant_handoff): the routing - # dispatch already quantized these rows and packed the - # topk ids in the fused route launch — consume both and - # skip the two standalone kernels. Identity-verified; - # a miss runs the unfused chain below. - from sglang.srt.layers.moe import route_quant_handoff - - prepared = route_quant_handoff.take(x) - if prepared is not None: - prepared_packed_topk, x_quant, x_scale = prepared - x_scale = x_scale.view(torch.float8_e4m3fn) - else: - from sglang.kernels.ops.quantization.per_token_group_quant import ( - per_token_group_quant, - ) - - x_quant, x_scale = per_token_group_quant( - x, group_size=32, scale_ue8m0=True - ) - x_scale = x_scale.view(torch.float8_e4m3fn) - else: - from sglang.srt.layers.quantization.fp8_utils import ( - flashinfer_mxfp8_quantize, - ) - - x_quant, x_scale = flashinfer_mxfp8_quantize( - x, False, alignment=self.hidden_size - ) - x_scale = x_scale.view(torch.float8_e4m3fn).reshape( - *x.shape[:-1], -1 - ) + x, prepared_packed_topk, x_quant, x_scale = ( + _prepare_flashinfer_mxfp8_activations(x, self.hidden_size) + ) else: raise NotImplementedError() diff --git a/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py b/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py index dca230142..2275f7de2 100644 --- a/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py +++ b/test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py @@ -1,16 +1,39 @@ +"""CPU unit tests for MXFP4 conversion and MXFP8 fake-output metadata.""" + import unittest import torch from sglang.srt.layers.quantization.fp8_utils import ( + _fake_flashinfer_mxfp8_quantize, quantize_block_fp8_weight_to_mxfp4, ) from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=4, suite="base-a-test-cpu") -class TestFp8UtilsMxfp4(unittest.TestCase): +class TestFp8UtilsMxfp4(CustomTestCase): + def test_fake_flashinfer_mxfp8_quantize_linear_scale_shape(self): + """The fake op must flatten leading dimensions and preserve scale groups.""" + input = torch.empty((2, 3, 96), dtype=torch.bfloat16) + + quantized, scale = _fake_flashinfer_mxfp8_quantize(input, False, alignment=128) + + self.assertEqual(quantized.shape, torch.Size([6, 128])) + self.assertEqual(quantized.dtype, torch.float8_e4m3fn) + self.assertEqual(scale.shape, torch.Size([24])) + self.assertEqual(scale.dtype, torch.uint8) + + def test_fake_flashinfer_mxfp8_quantize_swizzled_scale_shape(self): + input = torch.empty((3, 64), dtype=torch.bfloat16) + + quantized, scale = _fake_flashinfer_mxfp8_quantize(input, True, alignment=64) + + self.assertEqual(quantized.shape, torch.Size([3, 64])) + self.assertEqual(scale.shape, torch.Size([512])) + def test_quantize_block_fp8_weight_to_mxfp4_shapes_and_dtype(self): fp8_weight = ( torch.linspace(-2.0, 2.0, 32 * 32, dtype=torch.float32) 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 new file mode 100644 index 000000000..a444648ed --- /dev/null +++ b/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py @@ -0,0 +1,125 @@ +"""CPU unit tests for MXFP8 activation-preparation dispatch.""" + +import importlib +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.layers.quantization.mxfp4 import ( + _prepare_flashinfer_mxfp8_activations, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + +per_token_group_quant_module = importlib.import_module( + "sglang.kernels.ops.quantization.per_token_group_quant" +) + + +class TestMxfp4FlashinferActivationPrep(CustomTestCase): + def test_sm107_handoff_miss_uses_flashinfer_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) + + 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( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + return_value=(x_quant, x_scale), + create=True, + ) as 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, False, alignment=64) + 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_triton_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) + + with patch( + "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None + ), patch( + "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", + return_value=False, + ), 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) + ) + + 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_padded_input_keeps_flashinfer_quantizer(self): + """A group-aligned input must use hidden-size-aligned quantization.""" + x = torch.randn(3, 96, dtype=torch.bfloat16) + x_quant = torch.empty(3, 128, dtype=torch.float8_e4m3fn) + x_scale = torch.arange(12, dtype=torch.uint8) + + with patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + return_value=(x_quant, x_scale), + create=True, + ) as quantize, patch("sglang.srt.layers.moe.route_quant_handoff.take") as take: + actual_x, packed_topk, actual_quant, actual_scale = ( + _prepare_flashinfer_mxfp8_activations(x, 128) + ) + + take.assert_not_called() + quantize.assert_called_once_with(x, False, alignment=128) + self.assertIs(actual_x, x) + self.assertIsNone(packed_topk) + self.assertIs(actual_quant, x_quant) + self.assertEqual(actual_scale.shape, torch.Size([3, 4])) + + def test_kimi_handoff_skips_flashinfer_quantizer(self): + x = torch.randn(2, 64, dtype=torch.bfloat16) + packed_topk = torch.zeros(2, 4, dtype=torch.int32) + x_quant = torch.empty(2, 64, dtype=torch.float8_e4m3fn) + x_scale = torch.arange(4, dtype=torch.uint8).reshape(2, 2) + + with patch( + "sglang.srt.layers.moe.route_quant_handoff.take", + return_value=(packed_topk, x_quant, x_scale), + ), patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + create=True, + ) as quantize: + actual_x, actual_packed, actual_quant, actual_scale = ( + _prepare_flashinfer_mxfp8_activations(x, 64) + ) + + quantize.assert_not_called() + self.assertIs(actual_x, x) + self.assertIs(actual_packed, packed_topk) + self.assertIs(actual_quant, x_quant) + self.assertTrue(torch.equal(actual_scale.view(torch.uint8), x_scale)) + + +if __name__ == "__main__": + unittest.main()