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

[NPU][bugfix] update low latency quantization input and update MXFP8 tests (#38831)

Browse Source 
Co-authored-by: AndyLi429 <AndyLi429@noreply.gitcode.com>
This commit is contained in:
AndyLi429
2026-09-20 09:55:30 +08:00
committed by GitHub
co-authored by AndyLi429
parent f9c2791460
commit d903351a66
3 changed files with 518 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+86 -69
View File
@@ -441,10 +441,14 @@ class _DeepEPDispatcherImplBase:
config_map = {
DispatcherOutputDtype.BF16: {
"use_fp8": False,
"use_mxfp4": False,
"use_mxfp8": False,
"use_nvfp4": False,
},
DispatcherOutputDtype.FP8: {
"use_fp8": True,
"use_mxfp4": False,
"use_mxfp8": False,
"use_nvfp4": False,
},
# Needed for Ascend A2/A3 NPU case,
@@ -452,14 +456,26 @@ class _DeepEPDispatcherImplBase:
# quantization will be performed in int8
DispatcherOutputDtype.INT8: {
"use_fp8": True,
"use_mxfp4": False,
"use_mxfp8": False,
"use_nvfp4": False,
},
DispatcherOutputDtype.NVFP4: {
"use_fp8": False,
"use_mxfp4": False,
"use_mxfp8": False,
"use_nvfp4": True,
},
DispatcherOutputDtype.MXFP4: {
"use_fp8": False,
"use_mxfp4": True,
"use_mxfp8": False,
"use_nvfp4": False,
},
DispatcherOutputDtype.MXFP8: {
"use_fp8": False,
"use_mxfp4": False,
"use_mxfp8": True,
"use_nvfp4": False,
},
}
@@ -470,6 +486,8 @@ class _DeepEPDispatcherImplBase:
# Apply configuration
config = config_map[self.deepep_output_dtype]
self.use_fp8 = config["use_fp8"]
self.use_mxfp4 = config["use_mxfp4"]
self.use_mxfp8 = config["use_mxfp8"]
self.use_nvfp4 = config["use_nvfp4"]
# Handle environment variables
@@ -478,33 +496,17 @@ class _DeepEPDispatcherImplBase:
def _validate_and_adjust_dtype(self) -> None:
"""Validate dtype against hardware and adjust if necessary."""
self.low_latency_quant_mode = None
self._low_latency_quant_mode_runtime_checked = False
if self.deepep_output_dtype == DispatcherOutputDtype.MXFP8:
if not _is_npu or self.dispatch_mode != DeepEPMode.LOW_LATENCY:
raise RuntimeError(
"MXFP8 DeepEP dispatch is supported only for A5 "
"low-latency dispatch."
)
if _is_npu and self.deepep_output_dtype == DispatcherOutputDtype.FP8:
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
if not is_npu_arch35():
raise RuntimeError(
"MXFP8 DeepEP dispatch is supported only on Ascend A5 "
"in low-latency mode."
)
self.low_latency_quant_mode = "mx_fp8_e4m3"
return
if _is_npu:
if self.deepep_output_dtype == DispatcherOutputDtype.FP8:
logger.warning_once(
"Ascend A2/A3 NPU does not support fp8 "
"deepep_dispatcher_output_dtype, switching to int8..."
"deepep_dispatcher_output_dtype; DeepEP will use int8."
)
self.deepep_output_dtype = DispatcherOutputDtype.INT8
elif self.deepep_output_dtype == DispatcherOutputDtype.NVFP4:
if _is_npu:
if self.deepep_output_dtype == DispatcherOutputDtype.NVFP4:
raise RuntimeError(
"Ascend A2/A3 NPU does not support nvfp4 deepep_dispatcher_output_dtype."
)
@@ -561,6 +563,40 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
previous_event = Buffer.capture() if self.async_finish else None
return hidden_states, topk_ids, topk_weights, previous_event
def _get_quantization_kwargs(self, buffer: Buffer) -> dict:
if not _is_npu:
return {}
dispatch_params = inspect.signature(buffer.dispatch).parameters
flag_kwargs = {
"use_fp8": self.use_fp8,
"use_mxfp4": self.use_mxfp4,
"use_mxfp8": self.use_mxfp8,
}
if all(name in dispatch_params for name in flag_kwargs):
return flag_kwargs
if "quant_mode" in dispatch_params:
if self.use_mxfp4:
quant_mode = "mx_fp4_e2m1"
elif self.use_mxfp8:
quant_mode = "mx_fp8_e4m3"
elif self.use_fp8:
quant_mode = "int8"
else:
quant_mode = "bf16"
return {"quant_mode": quant_mode}
if not self.use_mxfp4 and not self.use_mxfp8:
# A3's legacy pybind Buffer does not expose its dispatch signature.
# It selects BF16/INT8 dispatch through the DeepEP runtime instead.
return {}
raise RuntimeError(
"Installed DeepEP normal dispatch does not support either "
"use_fp8/use_mxfp4/use_mxfp8 or quant_mode."
)
def dispatch_b(self, hidden_states, topk_ids, topk_weights, previous_event):
(
hidden_states,
@@ -610,6 +646,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
# `handle` as a member variable works.
_deepep_precompile_tp_barrier()
npu_quantization_opts = self._get_quantization_kwargs(buffer)
(
recv_x,
recv_topk_ids,
@@ -630,6 +667,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
allocate_on_comm_stream=(previous_event is not None) and self.async_finish,
expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1,
config=DeepEPConfig.get_instance().normal_dispatch_config,
**npu_quantization_opts,
)
get_global_expert_distribution_recorder().on_deepep_dispatch_normal(
num_recv_tokens_per_expert,
@@ -774,9 +812,8 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
):
input_global_scale = self.quant_config.get("input_global_scale", None)
# round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP
# to return int32-packed UE8M0 scales that don't feed the flashinfer
# cutedsl kernel.
# round_scale / use_ue8m0 are FP8-DeepGEMM specific. Dropping use_ue8m0
# makes DeepEP return fp32 column-major scales the e8m0 cast cannot view.
fp8_deepgemm_scale_opts = (
dict(
round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
@@ -789,58 +826,16 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
)
buffer = self._get_buffer()
if (
self.low_latency_quant_mode is not None
and not self._low_latency_quant_mode_runtime_checked
):
try:
dispatch_signature = inspect.signature(buffer.low_latency_dispatch)
except (TypeError, ValueError) as exc:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch requires a recent "
"sgl-kernel-npu/DeepEP runtime exposing "
"low_latency_dispatch(..., quant_mode=...)."
) from exc
if "quant_mode" not in dispatch_signature.parameters:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch requires a recent "
"sgl-kernel-npu/DeepEP runtime exposing "
"low_latency_dispatch(..., quant_mode=...)."
)
self._low_latency_quant_mode_runtime_checked = True
use_fp8 = self.use_fp8
low_latency_quant_kwargs = {}
if self.low_latency_quant_mode is not None:
deep_use_mode = os.environ.get("DEEP_USE_MODE", "default")
if deep_use_mode == "default":
low_latency_quant_kwargs = {
"quant_mode": self.low_latency_quant_mode,
}
elif deep_use_mode == "ops":
# The ops strategy ignores quant_mode and uses the legacy
# flags. Pass both forms so the request is explicit and the
# strategy still produces E4M3 + E8M0 MXFP8 tensors.
use_fp8 = True
low_latency_quant_kwargs = {
"quant_mode": self.low_latency_quant_mode,
"use_ue8m0": True,
}
else:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch supports only "
"DEEP_USE_MODE=default or DEEP_USE_MODE=ops; got "
f"{deep_use_mode!r}."
)
_deepep_precompile_tp_barrier()
npu_mxfp_quantization_opts = self._get_npu_mxfp_quantization_kwargs(buffer)
packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
buffer.low_latency_dispatch(
hidden_states,
topk_ids,
self.num_max_dispatch_tokens_per_rank,
self.num_experts,
use_fp8=use_fp8,
**low_latency_quant_kwargs,
use_fp8=self.use_fp8,
**npu_mxfp_quantization_opts,
**(
dict(topk_weights=topk_weights)
if _is_npu and not _use_zbal
@@ -859,6 +854,28 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
)
return packed_recv_hidden, self.packed_recv_count, event, hook
def _get_npu_mxfp_quantization_kwargs(self, buffer: Buffer) -> dict:
if not _is_npu:
return {}
parameters = inspect.signature(buffer.low_latency_dispatch).parameters
if any(
parameter.kind is inspect.Parameter.VAR_KEYWORD
for parameter in parameters.values()
):
return {
"use_mxfp4": self.use_mxfp4,
"use_mxfp8": self.use_mxfp8,
}
return {
name: value
for name, value in {
"use_mxfp4": self.use_mxfp4,
"use_mxfp8": self.use_mxfp8,
}.items()
if name in parameters
}
def combine_a(
self,
hidden_states: torch.Tensor,
python/sglang/srt/layers/moe/utils.py
+2
View File
@@ -284,6 +284,7 @@ class DispatcherOutputDtype(Enum):
- FP8: dispatch hidden states in fp8
- INT8: dispatch hidden states in int8
- NVFP4: dispatch hidden states in nvfp4
- MXFP4: dispatch hidden states in mxfp4 (fp4_e2m1 + e8m0 block scale)
- MXFP8: dispatch hidden states in mxfp8 (fp8_e4m3 + e8m0 block scale)
"""
@@ -291,6 +292,7 @@ class DispatcherOutputDtype(Enum):
FP8 = "fp8"
INT8 = "int8"
NVFP4 = "nvfp4"
MXFP4 = "mxfp4"
MXFP8 = "mxfp8"
test/registered/unit/npu/quantization/test_fp4_moe_methods.py
+430
View File
@@ -27,6 +27,7 @@ from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
w4a8_mxfp_gmm,
)
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
from sglang.srt.layers.moe.token_dispatcher import deepep
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
@@ -186,6 +187,435 @@ class TestPairPackMxfpActScale(unittest.TestCase):
_pair_pack_mxfp_act_scale(torch.zeros(2, 3))
class _LowLatencyBuffer:
"""The MXFP8-era Buffer: bool flags, no quant_mode."""
def __init__(self):
self.kwargs = None
def low_latency_dispatch(
self,
hidden_states,
topk_ids,
num_max_dispatch_tokens_per_rank,
num_experts,
*,
use_fp8,
use_mxfp4=False,
use_mxfp8=False,
**kwargs,
):
self.kwargs = {
"use_fp8": use_fp8,
"use_mxfp4": use_mxfp4,
"use_mxfp8": use_mxfp8,
**kwargs,
}
return torch.empty(0), torch.empty(0), object(), object(), object()
class _LegacyLowLatencyBuffer:
"""A DeepEP API version that predates the use_mxfp8 flag."""
def __init__(self):
self.use_mxfp4 = None
def low_latency_dispatch(
self,
hidden_states,
topk_ids,
num_max_dispatch_tokens_per_rank,
num_experts,
*,
use_fp8,
use_mxfp4=False,
topk_weights,
async_finish,
return_recv_hook,
):
self.use_mxfp4 = use_mxfp4
return torch.empty(0), torch.empty(0), object(), object(), object()
class _CudaLowLatencyBuffer:
"""CUDA's Buffer API does not accept the NPU-only MXFP flags."""
def __init__(self):
self.use_fp8 = None
def low_latency_dispatch(
self,
hidden_states,
topk_ids,
num_max_dispatch_tokens_per_rank,
num_experts,
*,
use_fp8,
round_scale=False,
use_ue8m0=False,
async_finish=False,
return_recv_hook=False,
):
self.use_fp8 = use_fp8
return torch.empty(0), torch.empty(0), object(), object(), object()
class _CudaNormalBuffer:
"""CUDA's normal Buffer API receives an already-quantized input tuple."""
def __init__(self):
self.dispatched = False
def get_dispatch_layout(self, *args, **kwargs):
return (
torch.ones(1, dtype=torch.int32),
None,
torch.ones(2, dtype=torch.int32),
torch.ones(1, 1, dtype=torch.bool),
None,
)
def dispatch(
self,
x,
*,
topk_idx,
topk_weights,
num_tokens_per_rank,
num_tokens_per_rdma_rank,
is_token_in_rank,
num_tokens_per_expert,
previous_event,
async_finish,
allocate_on_comm_stream,
expert_alignment,
config,
):
self.dispatched = True
return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object()
class _FlagNormalBuffer(_CudaNormalBuffer):
def __init__(self):
super().__init__()
self.quantization_kwargs = None
def dispatch(
self,
x,
*,
topk_idx,
topk_weights,
num_tokens_per_rank,
num_tokens_per_rdma_rank,
is_token_in_rank,
num_tokens_per_expert,
previous_event,
async_finish,
allocate_on_comm_stream,
expert_alignment,
config,
use_fp8,
use_mxfp4,
use_mxfp8,
):
self.dispatched = True
self.quantization_kwargs = {
"use_fp8": use_fp8,
"use_mxfp4": use_mxfp4,
"use_mxfp8": use_mxfp8,
}
return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object()
class _LegacyNormalBuffer:
"""The pre-bool-flags DeepEP normal-dispatch API used by CI."""
def __init__(self):
self.quant_mode = None
def get_dispatch_layout(self, *args, **kwargs):
return (
torch.ones(1, dtype=torch.int32),
None,
torch.ones(2, dtype=torch.int32),
torch.ones(1, 1, dtype=torch.bool),
None,
)
def dispatch(
self,
x,
*,
topk_idx,
topk_weights,
num_tokens_per_rank,
num_tokens_per_rdma_rank,
is_token_in_rank,
num_tokens_per_expert,
previous_event,
async_finish,
allocate_on_comm_stream,
expert_alignment,
config,
quant_mode,
):
self.quant_mode = quant_mode
return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object()
class _OpaqueNormalBuffer(_CudaNormalBuffer):
"""The A3 pybind Buffer API whose dispatch signature hides quantization args."""
def __init__(self):
super().__init__()
self.dispatch_kwargs = None
def dispatch(self, *args, **kwargs):
self.dispatched = True
self.dispatch_kwargs = kwargs
return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object()
class TestDeepEPLowLatencyMxfp8Dispatch(unittest.TestCase):
def test_mxfp4_output_dtype_enables_only_mxfp4(self):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplBase)
with patch.object(
deepep,
"get_deepep_output_dtype",
return_value=deepep.DispatcherOutputDtype.MXFP4,
):
dispatcher.set_deepep_dispatcher_dtype()
self.assertFalse(dispatcher.use_fp8)
self.assertTrue(dispatcher.use_mxfp4)
self.assertFalse(dispatcher.use_mxfp8)
@staticmethod
def _dispatcher(quant_mode, buffer):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplLowLatency)
dispatcher.quant_config = {}
dispatcher.use_fp8 = False
dispatcher.use_mxfp4 = False
dispatcher.use_mxfp8 = quant_mode == "mxfp8"
dispatcher.use_nvfp4 = False
dispatcher.num_max_dispatch_tokens_per_rank = 2
dispatcher.num_experts = 2
dispatcher.return_recv_hook = False
dispatcher._get_buffer = lambda: buffer
return dispatcher
def test_mxfp8_passes_the_mxfp8_flag_without_ue8m0(self):
buffer = _LowLatencyBuffer()
dispatcher = self._dispatcher("mxfp8", buffer)
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
)
self.assertFalse(buffer.kwargs["use_fp8"])
self.assertTrue(buffer.kwargs["use_mxfp8"])
self.assertNotIn("use_ue8m0", buffer.kwargs)
def test_mxfp4_passes_the_mxfp4_flag(self):
buffer = _LowLatencyBuffer()
dispatcher = self._dispatcher("mxfp4", buffer)
dispatcher.use_mxfp4 = True
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
)
self.assertFalse(buffer.kwargs["use_fp8"])
self.assertTrue(buffer.kwargs["use_mxfp4"])
self.assertFalse(buffer.kwargs["use_mxfp8"])
def test_bf16_omits_unsupported_mxfp8_flag_for_legacy_buffer(self):
buffer = _LegacyLowLatencyBuffer()
dispatcher = self._dispatcher("bf16", buffer)
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
)
self.assertFalse(buffer.use_mxfp4)
def test_normal_dispatch_passes_quantization_flags(self):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal)
dispatcher.num_experts = 2
dispatcher.async_finish = False
dispatcher.use_fp8 = False
dispatcher.use_mxfp4 = False
dispatcher.use_mxfp8 = True
buffer = _FlagNormalBuffer()
dispatcher._get_buffer = lambda: buffer
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
patch.object(
deepep.DeepEPConfig,
"get_instance",
return_value=SimpleNamespace(normal_dispatch_config=None),
),
patch.object(
deepep,
"get_global_expert_distribution_recorder",
return_value=MagicMock(),
),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
None,
)
self.assertTrue(buffer.quantization_kwargs["use_mxfp8"])
self.assertFalse(buffer.quantization_kwargs["use_fp8"])
self.assertFalse(buffer.quantization_kwargs["use_mxfp4"])
def test_normal_dispatch_uses_legacy_quant_mode_when_flags_are_unsupported(self):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal)
dispatcher.num_experts = 2
dispatcher.async_finish = False
dispatcher.use_fp8 = False
dispatcher.use_mxfp4 = False
dispatcher.use_mxfp8 = False
buffer = _LegacyNormalBuffer()
dispatcher._get_buffer = lambda: buffer
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
patch.object(
deepep.DeepEPConfig,
"get_instance",
return_value=SimpleNamespace(normal_dispatch_config=None),
),
patch.object(
deepep,
"get_global_expert_distribution_recorder",
return_value=MagicMock(),
),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
None,
)
self.assertEqual(buffer.quant_mode, "bf16")
def test_normal_dispatch_keeps_a3_legacy_path_for_opaque_signature(self):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal)
dispatcher.num_experts = 2
dispatcher.async_finish = False
dispatcher.use_fp8 = True
dispatcher.use_mxfp4 = False
dispatcher.use_mxfp8 = False
buffer = _OpaqueNormalBuffer()
dispatcher._get_buffer = lambda: buffer
with (
patch.object(deepep, "_is_npu", True),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
patch.object(
deepep.DeepEPConfig,
"get_instance",
return_value=SimpleNamespace(normal_dispatch_config=None),
),
patch.object(
deepep,
"get_global_expert_distribution_recorder",
return_value=MagicMock(),
),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
None,
)
self.assertTrue(buffer.dispatched)
self.assertNotIn("use_fp8", buffer.dispatch_kwargs)
self.assertNotIn("use_mxfp4", buffer.dispatch_kwargs)
self.assertNotIn("use_mxfp8", buffer.dispatch_kwargs)
self.assertNotIn("quant_mode", buffer.dispatch_kwargs)
def test_cuda_normal_dispatch_omits_npu_quantization_flags(self):
dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal)
dispatcher.num_experts = 2
dispatcher.async_finish = False
dispatcher.use_fp8 = True
dispatcher.use_mxfp4 = True
dispatcher.use_mxfp8 = True
buffer = _CudaNormalBuffer()
dispatcher._get_buffer = lambda: buffer
with (
patch.object(deepep, "_is_npu", False),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
patch.object(
deepep.DeepEPConfig,
"get_instance",
return_value=SimpleNamespace(normal_dispatch_config=None),
),
patch.object(
deepep,
"get_global_expert_distribution_recorder",
return_value=MagicMock(),
),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
None,
)
self.assertTrue(buffer.dispatched)
def test_cuda_low_latency_dispatch_omits_npu_mxfp_flags(self):
buffer = _CudaLowLatencyBuffer()
dispatcher = self._dispatcher("mxfp8", buffer)
dispatcher.use_fp8 = True
dispatcher.use_mxfp4 = True
with (
patch.object(deepep, "_is_npu", False),
patch.object(deepep, "_deepep_precompile_tp_barrier"),
):
dispatcher._dispatch_core(
torch.zeros(1, 64),
torch.zeros(1, 1, dtype=torch.int64),
torch.ones(1, 1),
)
self.assertTrue(buffer.use_fp8)
class TestW4A8MxfpGmmInputScale(unittest.TestCase):
def setUp(self):
self.input = torch.randn(2, 64)