minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NVIDIA] Fix SM107 MXFP8 activation prep (#35405)

Browse Source 
Signed-off-by: Sahithi Chigurupati <chigurupati.sahithi@gmail.com>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Sahithi Chigurupati
2026-08-23 18:17:03 +08:00
committed by GitHub
co-authored by Mohammad Miadh Angkad
parent 27aa48bca1
commit 44db041700
4 changed files with 216 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8_utils.py
+20 -15
View File
@@ -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=[],
python/sglang/srt/layers/quantization/mxfp4.py
+47 -34
View File
@@ -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()
test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py
+24 -1
View File
@@ -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)
test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py
+125
View File
@@ -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()