From 12b643bb7945ff0f36e479cd281c3691c99886f2 Mon Sep 17 00:00:00 2001 From: iridiumine Date: Wed, 23 Sep 2026 08:40:22 +0000 Subject: [PATCH] [NPU] Enable MXFP8 low-latency dispatch for Fp8MoEMethod FP4 experts Fp8MoEMethod with is_fp4_expert on NPU delegates execution to NPUW4A8MXFP4MoEMethod kernels but never configured the DeepEP dispatcher quant dtype, leaving dispatch in BF16 and forcing a separate DynamicMxQuant before each GMM. Declare the required dispatch wire format as a DISPATCHER_QUANT_CONFIG class attribute on NPUW4A8MXFP4MoEMethod (low_latency=mxfp8, normal stays BF16 since A5 MXFP8 normal dispatch is intranode-only), have Fp8MoEMethod pick it up from the attached w13_kernel instead of platform if/else patches, and reuse the same declaration in NPUW4A8MXFP4FusedMoEMethod, replacing its hand-written copy. --- .../npu/quantization/fp4_moe_methods.py | 5 +---- .../npu/quantization/moe_methods.py | 20 +++++++++---------- python/sglang/srt/layers/quantization/fp8.py | 11 +++++++++- 3 files changed, 21 insertions(+), 15 deletions(-) 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 b0b174f44d63..d66bdbce9300 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -2492,8 +2492,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, }