diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py index e8df34a525e4..06c52f2fa058 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py @@ -171,10 +171,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if hasattr(layer, "dispatcher"): layer.dispatcher.set_quant_config( - { - "normal_dispatcher_output_dtype": "bf16", - "low_latency_dispatcher_output_dtype": "mxfp8", - } + dict(self.w13_kernel.DISPATCHER_QUANT_CONFIG) ) def apply(self, layer: torch.nn.Module, dispatch_output: "DispatchOutput"): diff --git a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py index 8c5534984ea9..0189d580b8f8 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py @@ -279,6 +279,15 @@ def _get_bias_args( class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): """ModelSlim W4A8 MoE with packed MXFP4 weights and MXFP8 activations.""" + # Dispatch wire format required by the W4A8 MXFP GMMs: low-latency + # dispatch quantizes activations to MXFP8 (E8M0 block scales); normal + # dispatch stays BF16 because A5 MXFP8 normal dispatch is intranode-only, + # which also preserves multi-node prefill when ``deepep-mode=auto``. + DISPATCHER_QUANT_CONFIG = { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "mxfp8", + } + def __init__(self, dynamic_quant_kwargs=_DEFAULT_DYNAMIC_QUANT): super().__init__(quant_config=None) self.dynamic_quant_kwargs = dynamic_quant_kwargs @@ -308,17 +317,8 @@ def process_weights_after_loading( weight.data, weight_scale.data ) - # A5 DeepEP low-latency dispatch quantizes the valid received rows to - # MXFP8 and returns the matching E8M0 block scales. Keep normal mode - # in BF16 because A5 MXFP8 normal dispatch is intranode-only; this also - # preserves multi-node prefill when ``deepep-mode=auto``. if weight_prefix == "w13" and hasattr(layer, "dispatcher"): - layer.dispatcher.set_quant_config( - { - "normal_dispatcher_output_dtype": "bf16", - "low_latency_dispatcher_output_dtype": "mxfp8", - } - ) + layer.dispatcher.set_quant_config(dict(self.DISPATCHER_QUANT_CONFIG)) def apply( self, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 9d0c3e0116fa..f08fdbf9a94d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -2671,8 +2671,17 @@ def process_weights_after_loading(self, layer: Module) -> None: self._prepare_hpc_ops_weights(layer) if hasattr(layer, "dispatcher"): + # Execution kernels may declare the dispatch wire format they + # need (DISPATCHER_QUANT_CONFIG, e.g. the NPU W4A8 MXFP GMMs + # consume MXFP8 activations); otherwise report the checkpoint + # dtype and let the dispatcher pick. + kernel_quant_config = getattr( + getattr(layer, "w13_kernel", None), "DISPATCHER_QUANT_CONFIG", None + ) layer.dispatcher.set_quant_config( - { + kernel_quant_config + if kernel_quant_config is not None + else { "weight_dtype": layer.w13_weight.dtype, "use_mxfp8": self.use_mxfp8, }