diff --git a/python/sglang/srt/hardware_backend/npu/moe/activation.py b/python/sglang/srt/hardware_backend/npu/moe/activation.py index 3ec8819de..2f43ea016 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/activation.py +++ b/python/sglang/srt/hardware_backend/npu/moe/activation.py @@ -201,6 +201,35 @@ class NPUSwigluStepAndMul(BaseActivation): return gate * up +class NPUSwigluMxfp8Quant(BaseActivation): + """DeepSeek-V4 grouped SwiGLU with MXFP8 requantization for GMM2.""" + + def __init__(self, limit: float): + self._limit = float(limit) + + def _apply_activation( + self, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int, + ): + # The op sums the group list as per-expert counts and has no cumulative layout; + # a cusum list passed through would silently process the wrong rows. + if group_list_type != 1: + raise ValueError( + "swiglu_group_quant takes a per-expert count group list, got " + f"group_list_type={group_list_type}" + ) + out, scale, _ = torch.ops.npu.swiglu_group_quant( + x=hidden_states, + group_index=group_list, + quant_mode=2, # MX: one e8m0 scale per 32-element block + group_list_type=0, # sglang numbers the count layout 1, the op numbers it 0 + clamp_value=self._limit, + ) + return out, scale + + # ============================================================================= # Generic TP all‑gather wrapper – used by the runner when needed # ============================================================================= 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 ab5e433b6..e8df34a52 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 @@ -1,69 +1,28 @@ """MXFP4 routed-expert MoE method for Ascend A5 (Ascend 950). DeepSeek-V4's FP4 expert checkpoint stores block-32 MXFP4 weights with E8M0 -scales. This module wires those weights to the A5 grouped-matmul kernels, both -for the plain (init-routing) path and for the DeepEP dispatch path. +scales. This module adapts those checkpoint weights to the shared Ascend MoE +runner and A5 grouped-matmul kernels. """ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING import torch -from sgl_kernel_npu.activation.swiglu_mxfp8_quant import swiglu_quant -from sglang.srt.environ import envs -from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( - _get_float4_e2m1fn_x2_dtype, - _get_float8_e8m0fnu_dtype, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A8MXFP4MoEMethod, + prepare_w4a8_mxfp_weight, ) -from sglang.srt.hardware_backend.npu.utils import is_npu_arch35 from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase from sglang.srt.utils import set_weight_attrs if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput + from sglang.srt.layers.moe.token_dispatcher import DispatchOutput # MXFP4 group size, fixed at 32 by the msmodelslim export format. MXFP4_BLOCK_SIZE = 32 -def _configure_dsv4_deepep_dispatcher(layer: torch.nn.Module) -> None: - """Select the DSV4 FP4 DeepEP wire format without changing other MoEs.""" - dispatcher = getattr(layer, "dispatcher", None) - if dispatcher is None: - return - - # This method is only instantiated for DSV4 FP4 experts on A5 today, but - # retain the former BF16 setting if that selection changes in the future. - if not is_npu_arch35(): - dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"}) - return - - # Import lazily to avoid importing the MoE backend during quant method - # module initialization. - from sglang.srt.layers.moe import get_moe_a2a_backend - - if not get_moe_a2a_backend().is_deepep(): - dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"}) - return - - low_latency_dtype = envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.get() - if low_latency_dtype not in {"mxfp8", "bf16"}: - raise ValueError( - "SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE must be one of " - "'mxfp8' or 'bf16' for A5 DSV4 DeepEP low-latency dispatch; " - f"got {low_latency_dtype!r}." - ) - - # The concrete dispatcher selects one mode-specific value. Normal (prefill) - # remains BF16, while low-latency (decode) defaults to MXFP8. - dispatcher.set_quant_config( - { - "normal_dispatcher_output_dtype": "bf16", - "low_latency_dispatcher_output_dtype": low_latency_dtype, - } - ) - - def _wrap_mxfp4_scale_weight_loader(weight_loader): def load_scale(param, loaded_weight, *args, **kwargs): if param.dtype == torch.uint8 and loaded_weight.dtype == torch.float8_e8m0fnu: @@ -73,17 +32,25 @@ def _wrap_mxfp4_scale_weight_loader(weight_loader): return load_scale -class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase): +class NPUW4A8MXFP4FusedMoEMethod(FusedMoEMethodBase): """DeepSeek-V4 routed experts on Ascend A5: W4A8 MXFP weights. - Delegates nothing to ``fp8_method`` except the shared runner config; it is - held so the FP8 method sees the same ``moe_runner_config`` the layer built. + The checkpoint-specific loading remains here while execution is delegated + to the shared Ascend MoE runner. """ - def __init__(self, fp8_method, prefix: str = ""): - self._fp8 = fp8_method + def __init__(self, prefix: str = ""): self.prefix = prefix - self.moe_runner_config = None + # ``None`` selects the full MX dynamic-quant defaults used by the + # original DeepSeek-V4 path; the shared ModelSlim path keeps its + # historical explicit ``dst_type`` behavior. + # TODO: Fuse DeepSeek-V4 W13 GMM + SwiGLU + MXFP8 requant with + # npu_grouped_matmul_swiglu_quant_v2 once it accepts swiglu_limit. + # V4 sets swiglu_limit=10.0; the current fused op implements only + # standard SwiGLU and would skip the required gate/up clamps. + self.w13_kernel = NPUW4A8MXFP4MoEMethod(dynamic_quant_kwargs=None) + self.w2_kernel = NPUW4A8MXFP4MoEMethod(dynamic_quant_kwargs=None) + self.runner = None def create_weights( self, @@ -153,11 +120,27 @@ class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase): set_weight_attrs(w2_weight_scale, scale_attrs) def create_moe_runner(self, layer: torch.nn.Module, moe_runner_config): - self.moe_runner_config = moe_runner_config - self._fp8.moe_runner_config = moe_runner_config + from sglang.srt.layers.moe.moe_runner.runner import MoeRunner + from sglang.srt.layers.moe.utils import ( + MoeRunnerBackend, + get_moe_runner_backend, + ) + + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + if not backend.is_ascend(): + raise ValueError( + f"NPU W4A8 MXFP4 requires the Ascend MoE runner, got {backend.value}" + ) + + layer.w13_kernel = self.w13_kernel + layer.w2_kernel = self.w2_kernel + moe_runner_config.layer = layer + self.runner = MoeRunner(backend, moe_runner_config) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast + from sglang.srt.hardware_backend.npu.utils import NPUACLFormat if layer.w13_weight_scale_inv.data.max() == 0: raise RuntimeError( @@ -172,459 +155,38 @@ class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase): "not match w2_weight_scale_inv." ) - nz_kwargs = { - "customize_dtype": torch.float8_e4m3fn, - "input_dtype": _get_float4_e2m1fn_x2_dtype(), - } nz_format = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ - layer.w13_weight.data = npu_format_cast( - layer.w13_weight.data.view(torch.uint8), nz_format, **nz_kwargs - ).transpose(1, 2) - layer.w2_weight.data = npu_format_cast( - layer.w2_weight.data.view(torch.uint8), nz_format, **nz_kwargs - ).transpose(1, 2) - - layer.w13_weight_scale_inv = torch.nn.Parameter( - _reshape_mxfp4_scale_for_npu(layer.w13_weight_scale_inv.data), - requires_grad=False, + layer.w13_weight.data, w13_scale = prepare_w4a8_mxfp_weight( + layer.w13_weight.data.view(torch.uint8), + layer.w13_weight_scale_inv.data, + npu_format=nz_format, ) - layer.w2_weight_scale_inv = torch.nn.Parameter( - _reshape_mxfp4_scale_for_npu(layer.w2_weight_scale_inv.data), - requires_grad=False, + layer.w2_weight.data, w2_scale = prepare_w4a8_mxfp_weight( + layer.w2_weight.data.view(torch.uint8), + layer.w2_weight_scale_inv.data, + npu_format=nz_format, ) + layer.w13_weight_scale_inv = torch.nn.Parameter(w13_scale, requires_grad=False) + layer.w2_weight_scale_inv = torch.nn.Parameter(w2_scale, requires_grad=False) - _configure_dsv4_deepep_dispatcher(layer) - - def apply( - self, - layer: torch.nn.Module, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": - combine_input = npu_apply_w4a8_mxfp_moe_deepep(layer, dispatch_output) - if combine_input is not None: - return combine_input - - combine_input = npu_apply_w4a4_mxfp_moe_ascend_tp(layer, dispatch_output) - if combine_input is not None: - return combine_input - - # Standard dispatch. Unreachable on NPU today — create_moe_dispatcher - # picks AscendTPDispatcher whenever is_npu() and no a2a backend is set — - # but kept so this method is not silently wrong if that changes. - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - hidden_states = dispatch_output.hidden_states - topk_weights, topk_ids, _ = dispatch_output.topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(hidden_states.dtype) - moe_runner_config = layer.moe_runner_config - - output = npu_fused_experts_w4a4_mxfp( - hidden_states, - layer.w13_weight, - layer.w13_weight_scale_inv, - layer.w2_weight, - layer.w2_weight_scale_inv, - topk_weights, - topk_ids, - moe_runner_config.top_k, - swiglu_limit=moe_runner_config.swiglu_limit, - ) - return StandardCombineInput(hidden_states=output) - - -def _reshape_mxfp4_scale_for_npu(scale: torch.Tensor) -> torch.Tensor: - """``[E, N, K/32] -> [E, K/64, N, 2]``, the packed-pair layout the GMM wants.""" - if scale.dim() != 3: - return scale - num_experts, n, k32 = scale.shape - if k32 % 2 != 0: - raise ValueError( - "MXFP4 scale K dimension must be divisible by 2 for the " - f"[E, K/64, N, 2] layout, got {tuple(scale.shape)}." - ) - return scale.view(num_experts, n, k32 // 2, 2).transpose(1, 2) - - -def _apply_swiglu_limit_npu( - gate_up: torch.Tensor, swiglu_limit: Optional[float] -) -> None: - """Clamp the SwiGLU input in place before ``npu_swiglu`` (DeepSeek-V4). - - gate (first half) <= limit; up (second half) in - [-limit, limit]. ``chunk`` returns views, so the in-place clamps mutate - ``gate_up`` directly. No-op when ``swiglu_limit`` is unset or <= 0. - """ - if swiglu_limit is None or swiglu_limit <= 0: - return - gate, up = gate_up.chunk(2, dim=-1) - gate.clamp_(max=swiglu_limit) - up.clamp_(min=-swiglu_limit, max=swiglu_limit) - - -def npu_fused_experts_w4a4_mxfp( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - swiglu_limit: Optional[float] = None, - **kwargs, -): - if torch.npu.is_current_stream_capturing(): - return npu_fused_experts_w4a4_mxfp_decode( - hidden_states=hidden_states, - w13=w13, - w13_weight_scale_inv=w13_weight_scale_inv, - w2=w2, - w2_weight_scale_inv=w2_weight_scale_inv, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=top_k, - swiglu_limit=swiglu_limit, - **kwargs, - ) - - original_shape = hidden_states.shape - original_dtype = hidden_states.dtype - if len(original_shape) == 3: - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - num_tokens = hidden_states.shape[0] - num_experts = w13.shape[0] - row_idx = ( - torch.arange( - 0, num_tokens * top_k, dtype=torch.int32, device=topk_weights.device - ) - .view(top_k, -1) - .permute(1, 0) - .contiguous() - ) - hidden_states, expanded_row_idx, expanded_expert_idx = ( - torch.ops.npu.npu_moe_init_routing( - hidden_states, - row_idx=row_idx, - expert_idx=topk_ids, - active_num=num_tokens, - ) - ) - expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( - expanded_expert_idx, num_experts - ).to(torch.int64) - - # npu_moe_init_routing pads its output to the worst case; rows past the last - # expert boundary hold garbage and must not reach finalize_routing. - row_ids = torch.arange( - hidden_states.shape[0], device=hidden_states.device, dtype=torch.int64 - ) - valid_mask_2d = (row_ids < expert_tokens[-1]).unsqueeze(1) - - hidden_states = w4a8_mxfp_gmm( - input=hidden_states, - input_scale=None, - weight=w13, - weight_scale=w13_weight_scale_inv, - group_list_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - ) - assert swiglu_limit is not None - hidden_states, hidden_states_scale = swiglu_quant( - hidden_states, - group_list=expert_tokens, - group_list_type=0, - need_quant=True, - do_limit=True, - limit=swiglu_limit, - ) - hidden_states = w4a8_mxfp_gmm( - input=hidden_states, - input_scale=hidden_states_scale, - weight=w2, - weight_scale=w2_weight_scale_inv, - group_list_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - ) - hidden_states = hidden_states * valid_mask_2d.to(hidden_states.dtype) - - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - return final_hidden_states - - -def npu_fused_experts_w4a4_mxfp_decode( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_weight_scale_inv: torch.Tensor, - w2: torch.Tensor, - w2_weight_scale_inv: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - swiglu_limit: Optional[float] = None, - **kwargs, -): - """Graph-capturable variant: routing v2 + token_unpermute, no host syncs.""" - num_tokens = hidden_states.shape[:-1].numel() - global_num_experts = w13.shape[0] - original_shape = hidden_states.shape - original_dtype = hidden_states.dtype - group_list_type = 1 - - hidden_states, expanded_row_idx, expert_tokens, _ = ( - torch.ops.npu.npu_moe_init_routing_v2( - hidden_states, - topk_ids, - active_num=num_tokens * top_k, - expert_num=global_num_experts, - expert_tokens_num_type=group_list_type, - expert_tokens_num_flag=True, - active_expert_range=[0, global_num_experts], - quant_mode=-1, - ) - ) - expert_tokens = expert_tokens.to(torch.int64) - hidden_states = w4a8_mxfp_gmm( - input=hidden_states, - input_scale=None, - weight=w13, - weight_scale=w13_weight_scale_inv, - group_list_type=group_list_type, - group_list=expert_tokens, - output_dtype=original_dtype, - ) - assert swiglu_limit is not None - hidden_states, hidden_states_scale = swiglu_quant( - hidden_states, - group_list=expert_tokens, - group_list_type=group_list_type, - need_quant=True, - do_limit=True, - limit=swiglu_limit, - ) - hidden_states = w4a8_mxfp_gmm( - input=hidden_states, - input_scale=hidden_states_scale, - weight=w2, - weight_scale=w2_weight_scale_inv, - group_list_type=group_list_type, - group_list=expert_tokens, - output_dtype=original_dtype, - ) - - final_hidden_states = torch.ops.npu.npu_moe_token_unpermute( - permuted_tokens=hidden_states, - sorted_indices=torch.abs(expanded_row_idx), - probs=topk_weights, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - return final_hidden_states - - -def npu_apply_w4a4_mxfp_moe_ascend_tp( - layer: torch.nn.Module, - dispatch_output: "DispatchOutput", -) -> Optional["CombineInput"]: - """Ascend TP path. Returns ``None`` when the dispatch is not an Ascend TP one. - - AscendTPDispatcher already ran npu_moe_init_routing_v2 on dispatch and runs - npu_moe_finalize_routing (with topk_weights) on combine, so this only owns - the grouped-matmul chain in between — no permute, no routing-weight apply. - """ - from sglang.srt.layers.moe.token_dispatcher import AscendTPCombineInput - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker - - if not DispatchOutputChecker.format_is_ascend_tp(dispatch_output): - return None - - hidden_states = npu_apply_without_routing_weights_w4a4_mxfp( - layer, - dispatch_output.hidden_states, - dispatch_output.hidden_states_scale, - group_list_type=dispatch_output.group_list_type, - group_list=dispatch_output.expert_tokens, - output_dtype=torch.bfloat16, - ) - return AscendTPCombineInput(hidden_states=hidden_states) - - -def npu_apply_w4a8_mxfp_moe_deepep( - layer: torch.nn.Module, - dispatch_output: "DispatchOutput", -) -> Optional["CombineInput"]: - """DeepEP path. Returns ``None`` when the dispatch is not a DeepEP one.""" - from sglang.srt.layers.moe.token_dispatcher import ( - DeepEPLLCombineInput, - DeepEPNormalCombineInput, - ) - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker - - if not dispatch_output.format.is_deepep(): - return None - - if DispatchOutputChecker.format_is_deepep_normal(dispatch_output): - hidden_states, hidden_states_scale, _, _, num_recv_tokens_per_expert = ( - dispatch_output - ) - group_list = torch.tensor( - num_recv_tokens_per_expert, dtype=torch.int64, device=hidden_states.device - ) - combine_cls = DeepEPNormalCombineInput - else: - hidden_states, hidden_states_scale, _, _, group_list, _ = dispatch_output - group_list = group_list.to(torch.int64) - combine_cls = DeepEPLLCombineInput - - hidden_states = npu_apply_without_routing_weights_w4a4_mxfp( - layer, - hidden_states, - hidden_states_scale, - group_list_type=1, - group_list=group_list, - output_dtype=torch.bfloat16, - ) - return combine_cls( - hidden_states=hidden_states, - topk_ids=dispatch_output.topk_ids, - topk_weights=dispatch_output.topk_weights, - ) - - -def npu_apply_without_routing_weights_w4a4_mxfp( - layer, - hidden_states, - hidden_states_scale, - *, - group_list_type, - group_list, - output_dtype, -): - hidden_states = w4a8_mxfp_gmm( - input=hidden_states, - input_scale=hidden_states_scale, - weight=layer.w13_weight, - weight_scale=layer.w13_weight_scale_inv, - group_list_type=group_list_type, - group_list=group_list, - output_dtype=output_dtype, - ) - assert layer.moe_runner_config.swiglu_limit is not None - hidden_states, hidden_states_scale = swiglu_quant( - hidden_states, - group_list=group_list, - group_list_type=group_list_type, - need_quant=True, - do_limit=True, - limit=layer.moe_runner_config.swiglu_limit, - ) - return w4a8_mxfp_gmm( - input=hidden_states, - input_scale=hidden_states_scale, - weight=layer.w2_weight, - weight_scale=layer.w2_weight_scale_inv, - group_list_type=group_list_type, - group_list=group_list, - output_dtype=output_dtype, - ) - - -def _pair_pack_mxfp_act_scale( - scale: torch.Tensor, input_shape: Optional[tuple[int, int]] = None -) -> torch.Tensor: - """Adapt MXFP activation scales to the A5 GMM ``[M, K/64, 2]`` layout. - - Low-latency DeepEP MXFP8 returns a flat E8M0 scale buffer, one byte for - every 32 activation elements. The grouped-matmul kernel expects those - bytes paired on the final dimension instead. - """ - if scale.ndim == 1: - if input_shape is None or len(input_shape) != 2: - raise ValueError( - "A flat MXFP activation scale requires its two-dimensional " - "activation input shape." + if hasattr(layer, "dispatcher"): + layer.dispatcher.set_quant_config( + { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "mxfp8", + } ) - num_tokens, hidden_size = input_shape - if hidden_size % (2 * MXFP4_BLOCK_SIZE) != 0: - raise ValueError( - "MXFP activation hidden size must be divisible by " - f"{2 * MXFP4_BLOCK_SIZE}; got {hidden_size}." - ) - expected_num_scales = num_tokens * (hidden_size // MXFP4_BLOCK_SIZE) - if scale.numel() != expected_num_scales: - raise ValueError( - "Invalid flat MXFP activation scale length: expected " - f"{expected_num_scales} for input shape {input_shape}, got " - f"{scale.numel()}." - ) - scale = scale.reshape(num_tokens, hidden_size // MXFP4_BLOCK_SIZE) - # ``[M, K/32] -> [M, K/64, 2]`` MX per-token scale layout for the A5 GMM. - if scale.ndim != 2: - return scale - if scale.shape[-1] % 2 != 0: - raise ValueError(f"Invalid MXFP per-token scale shape: {tuple(scale.shape)}") - return scale.reshape(scale.shape[0], scale.shape[1] // 2, 2) + def apply(self, layer: torch.nn.Module, dispatch_output: "DispatchOutput"): + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + if self.runner is None: + raise RuntimeError("The NPU FP4 MoE runner has not been initialized") -def w4a8_mxfp_gmm( - *, - input: torch.Tensor, - input_scale: Optional[torch.Tensor], - weight: torch.Tensor, - weight_scale: torch.Tensor, - group_list_type: int, - group_list: torch.Tensor, - output_dtype: torch.dtype, - scale_alg=None, -) -> torch.Tensor: - """FP4 weight x FP8-e4m3 activation (the checkpoint's W4A8_MXFP scheme). - - W4A8MXFP GMM call: FP8 ``x_dtype``, FP4 - ``weight_dtype``, and the weight block scales fed through ``antiquant_scale`` - with ``scale=None`` — the ``scale=`` + ``scale_dtype=`` form belongs to - W4A4_MXFP4 and dequantizes differently. - """ - group_list = group_list.to(torch.int64) - if input_scale is None: - x, x_scale = torch.ops.npu.npu_dynamic_mx_quant( - input, - axis=1, - round_mode="rint", - dst_type=torch.float8_e4m3fn, - block_size=MXFP4_BLOCK_SIZE, - scale_alg=scale_alg, + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale_inv, + w2_weight_scale=layer.w2_weight_scale_inv, ) - else: - x, x_scale = input, input_scale - - return torch.ops.npu.npu_grouped_matmul( - [x], - [weight], - scale=None, - antiquant_scale=[weight_scale], - scale_dtype=None, - per_token_scale=[ - _pair_pack_mxfp_act_scale(x_scale, input_shape=tuple(x.shape)) - ], - split_item=2, - group_type=0, - group_list=group_list, - group_list_type=group_list_type, - output_dtype=output_dtype, - x_dtype=torch.float8_e4m3fn, - weight_dtype=_get_float4_e2m1fn_x2_dtype(), - per_token_scale_dtype=_get_float8_e8m0fnu_dtype(), - )[0] + return self.runner.run(dispatch_output, quant_info) 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 df21796f8..188d8e5fd 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py @@ -9,8 +9,8 @@ from sglang.srt.hardware_backend.npu.utils import npu_format_cast from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase if TYPE_CHECKING: - from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + from sglang.srt.layers.quantization.base_config import QuantizationConfig import logging @@ -28,6 +28,7 @@ logger = logging.getLogger(__name__) _E8M0_DTYPE = None +_DEFAULT_DYNAMIC_QUANT = object() def _require_e8m0_dtype(): @@ -65,6 +66,93 @@ def _require_e8m0_dtype(): return _E8M0_DTYPE +def reshape_w4a8_mxfp_weight_scale_for_npu(scale: torch.Tensor) -> torch.Tensor: + """Pack MXFP4 scales from ``[E, N, K/32]`` to the A5 GMM layout.""" + if scale.dim() != 3: + return scale + num_experts, n, k32 = scale.shape + if k32 % 2 != 0: + raise ValueError( + "MXFP4 scale K dimension must be divisible by 2 for the " + f"[E, K/64, N, 2] layout, got {tuple(scale.shape)}." + ) + return scale.view(num_experts, n, k32 // 2, 2).transpose(1, 2) + + +def prepare_w4a8_mxfp_weight( + weight: torch.Tensor, + weight_scale: torch.Tensor, + *, + npu_format=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Convert W4A8 MXFP weights and scales to the shared A5 GMM layout.""" + cast_args = { + "customize_dtype": torch.float8_e4m3fn, + "input_dtype": _get_float4_e2m1fn_x2_dtype(), + } + if npu_format is None: + weight = npu_format_cast(weight, **cast_args) + else: + weight = npu_format_cast(weight, npu_format, **cast_args) + return weight.transpose(-1, -2), reshape_w4a8_mxfp_weight_scale_for_npu( + weight_scale + ) + + +def _pair_pack_mxfp_act_scale(scale: torch.Tensor) -> torch.Tensor: + """Pack MXFP activation scales from ``[M, K/32]`` to A5 GMM layout.""" + if scale.ndim != 2: + return scale + if scale.shape[-1] % 2 != 0: + raise ValueError(f"Invalid MXFP per-token scale shape: {tuple(scale.shape)}") + return scale.reshape(scale.shape[0], scale.shape[1] // 2, 2) + + +def w4a8_mxfp_gmm( + *, + input: torch.Tensor, + input_scale: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + group_list_type: int, + group_list: torch.Tensor, + output_dtype: torch.dtype, + scale_alg=None, + dynamic_quant_kwargs: Optional[Dict[str, Any]] = None, +) -> torch.Tensor: + """Run the shared A5 FP4-weight × FP8-activation grouped matmul.""" + group_list = group_list.to(torch.int64) + if input_scale is None: + if dynamic_quant_kwargs is None: + dynamic_quant_kwargs = { + "axis": 1, + "round_mode": "rint", + "dst_type": torch.float8_e4m3fn, + "block_size": 32, + "scale_alg": scale_alg, + } + x, x_scale = torch.ops.npu.npu_dynamic_mx_quant(input, **dynamic_quant_kwargs) + else: + x, x_scale = input, input_scale + + return torch.ops.npu.npu_grouped_matmul( + [x], + [weight], + scale=None, + antiquant_scale=[weight_scale], + scale_dtype=None, + per_token_scale=[_pair_pack_mxfp_act_scale(x_scale)], + split_item=2, + group_type=0, + group_list=group_list, + group_list_type=group_list_type, + output_dtype=output_dtype, + x_dtype=torch.float8_e4m3fn, + weight_dtype=_get_float4_e2m1fn_x2_dtype(), + per_token_scale_dtype=_require_e8m0_dtype(), + )[0] + + # DEPRECATED METHOD # TODO: Remove in future realeses def fused_moe_npu( @@ -191,12 +279,9 @@ class _NPUMoEMethodBase(FusedMoEMethodBase): class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): """ModelSlim W4A8 MoE with packed MXFP4 weights and MXFP8 activations.""" - def __init__(self): + def __init__(self, dynamic_quant_kwargs=_DEFAULT_DYNAMIC_QUANT): super().__init__(quant_config=None) - self.matmul = GroupedMatmul() - self.hidden_states_quantizer = HiddenStatesDynamicQuant( - quant_dtype=torch.float8_e4m3fn - ) + self.dynamic_quant_kwargs = dynamic_quant_kwargs def process_weights_after_loading( self, layer: torch.nn.Module, weight_prefix: str @@ -208,20 +293,10 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.") weight = getattr(layer, f"{weight_prefix}_weight") - weight.data = npu_format_cast( - weight.data, - customize_dtype=torch.float8_e4m3fn, - input_dtype=fp4_dtype, - ).transpose(-1, -2) - weight_scale = getattr(layer, f"{weight_prefix}_weight_scale") - scale = weight_scale.data.reshape( - weight_scale.shape[0], - weight_scale.shape[1], - weight_scale.shape[2] // 2, - 2, - ).transpose(1, 2) - weight_scale.data = scale + weight.data, weight_scale.data = prepare_w4a8_mxfp_weight( + weight.data, weight_scale.data + ) # The refactored NPU dispatchers currently support BF16 and INT8. # Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM. @@ -238,35 +313,24 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): weight_prefix: str, group_list_type: int, ) -> torch.Tensor: - fp4_dtype = _get_float4_e2m1fn_x2_dtype() - if fp4_dtype is None: - raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.") - e8m0_dtype = _require_e8m0_dtype() - - if pertoken_scale is None: - hidden_states, pertoken_scale = self.hidden_states_quantizer(hidden_states) - elif pertoken_scale is not None: + if pertoken_scale is not None: pertoken_scale = pertoken_scale.reshape( hidden_states.shape[0], hidden_states.shape[1] // 64, 2 ) - return self.matmul.forward( - quant_info, - weight_prefix, - hidden_states, - expert_tokens.to(torch.int64), - output_dtype, + dynamic_quant_kwargs = self.dynamic_quant_kwargs + if dynamic_quant_kwargs is _DEFAULT_DYNAMIC_QUANT: + dynamic_quant_kwargs = {"dst_type": torch.float8_e4m3fn} + + return w4a8_mxfp_gmm( + input=hidden_states, + input_scale=pertoken_scale, + weight=getattr(quant_info, f"{weight_prefix}_weight"), + weight_scale=getattr(quant_info, f"{weight_prefix}_weight_scale"), group_list_type=group_list_type, - transposed=True, - scale=None, - scale_dtype=None, - per_token_scale=[pertoken_scale], - antiquant_scale=[ - getattr(quant_info, f"{weight_prefix}_weight_scale", None) - ], - x_dtype=torch.float8_e4m3fn, - weight_dtype=fp4_dtype, - per_token_scale_dtype=e8m0_dtype, + group_list=expert_tokens, + output_dtype=output_dtype, + dynamic_quant_kwargs=dynamic_quant_kwargs, ) diff --git a/python/sglang/srt/layers/moe/moe_runner/ascend.py b/python/sglang/srt/layers/moe/moe_runner/ascend.py index 7bc97d17f..0015218cf 100644 --- a/python/sglang/srt/layers/moe/moe_runner/ascend.py +++ b/python/sglang/srt/layers/moe/moe_runner/ascend.py @@ -13,6 +13,7 @@ from sglang.srt.hardware_backend.npu.moe.activation import ( NPUSitu, NPUSwiglu, NPUSwigluDeepEPKernel, + NPUSwigluMxfp8Quant, NPUSwigluOAI, NPUSwigluQuant, NPUSwigluStepAndMul, @@ -20,6 +21,7 @@ from sglang.srt.hardware_backend.npu.moe.activation import ( from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( NPUMXFP8MoEMethod, NPUW4A8Int8MoEMethod, + NPUW4A8MXFP4MoEMethod, NPUW8A8Int8MoEMethod, ) from sglang.srt.layers.moe.moe_runner.base import ( @@ -97,6 +99,12 @@ class AscendRunnerCore(MoeRunnerCore): # both dispatchers: ascend_tp gets its activation quant fused into # routing, DeepEP dispatches bf16 and gmm1 quantises it itself. self.activation = None + elif ( + isinstance(kernel, NPUW4A8MXFP4MoEMethod) + and config.swiglu_limit is not None + and config.swiglu_limit > 0 + ): + self.activation = NPUSwigluMxfp8Quant(config.swiglu_limit) elif get_moe_a2a_backend().is_deepep(): # DeepEP path: use a unified kernel that decides quantisation is_quant_kernel = isinstance( @@ -186,7 +194,7 @@ class AscendRunnerCore(MoeRunnerCore): # Grouped-row activations require dispatch metadata. if isinstance( self.activation, - (NPUSwigluDeepEPKernel, NPUSitu), + (NPUSwigluDeepEPKernel, NPUSitu, NPUSwigluMxfp8Quant), ): hidden_states, pertoken_scale = self.activation._apply_activation( hidden_states, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 290155641..b8f65a0cb 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -421,10 +421,10 @@ class Fp8Config(QuantizationConfig): and self.is_dsv4_fp4_experts ): from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import ( - NPUW4A4Fp4MoEMethod, + NPUW4A8MXFP4FusedMoEMethod, ) - return NPUW4A4Fp4MoEMethod(fp8_method, prefix=prefix) + return NPUW4A8MXFP4FusedMoEMethod(prefix=prefix) if self.is_fp4_experts and get_moe_runner_backend().is_marlin(): from sglang.srt.layers.quantization.mxfp4_marlin_moe import ( diff --git a/test/registered/unit/npu/quantization/test_fp4_moe_methods.py b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py index a21ac940f..1594ae8b8 100644 --- a/test/registered/unit/npu/quantization/test_fp4_moe_methods.py +++ b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py @@ -1,5 +1,3 @@ -import inspect -import os import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -18,23 +16,19 @@ register_npu_ci(est_time=1, suite="stage-a-unit-test-npu") # initialized, `_get_float8_e8m0fnu_dtype` not yet defined). Initializing the # package first mirrors how the engine loads quantization at model-config time. import sglang.srt.layers.quantization # noqa: F401 -from sglang.srt.environ import envs -from sglang.srt.hardware_backend.npu.quantization import fp4_moe_methods +from sglang.srt.hardware_backend.npu.moe.activation import NPUSwigluMxfp8Quant from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import ( - NPUW4A4Fp4MoEMethod, - _apply_swiglu_limit_npu, - _configure_dsv4_deepep_dispatcher, + NPUW4A8MXFP4FusedMoEMethod, +) +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( _pair_pack_mxfp_act_scale, - _reshape_mxfp4_scale_for_npu, - npu_apply_without_routing_weights_w4a4_mxfp, + prepare_w4a8_mxfp_weight, + reshape_w4a8_mxfp_weight_scale_for_npu, 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 -_NOT_PASSED = object() - class TestFP4MethodGate(unittest.TestCase): def test_pre_arch35_keeps_fp8_moe_method(self): @@ -52,25 +46,63 @@ class TestFP4MethodGate(unittest.TestCase): self.assertIsInstance(method, Fp8MoEMethod) + def test_arch35_uses_ascend_runner_method(self): + config = Fp8Config(is_fp4_experts=True) + config.is_dsv4_fp4_experts = True + layer = FusedMoE.__new__(FusedMoE) -class TestApplySwiGLULimitNpu(unittest.TestCase): - def test_clamps_gate_and_up_asymmetrically(self): - # DeepSeek-V4 clamps gate (first half) to <= limit but only the upper - # bound, while up (second half) is clamped symmetrically to [-limit, limit]. - # A regression that swapped these would silently change expert activations. - gate_up = torch.tensor([[8.0, -9.0, 9.0, -9.0]]) - _apply_swiglu_limit_npu(gate_up, 7.0) - self.assertTrue(torch.equal(gate_up, torch.tensor([[7.0, -9.0, 7.0, -7.0]]))) + with ( + patch("sglang.srt.layers.quantization.fp8.is_npu", return_value=True), + patch( + "sglang.srt.layers.quantization.fp8.is_npu_arch35", + return_value=True, + ), + ): + method = config.get_quant_method(layer, "model.layers.0.experts") - def test_noop_when_limit_none(self): - gate_up = torch.tensor([[8.0, -9.0]]) - _apply_swiglu_limit_npu(gate_up, None) - self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]]))) + self.assertIsInstance(method, NPUW4A8MXFP4FusedMoEMethod) - def test_noop_when_limit_nonpositive(self): - gate_up = torch.tensor([[8.0, -9.0]]) - _apply_swiglu_limit_npu(gate_up, 0.0) - self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]]))) + +class TestNPUSwigluMxfp8Quant(unittest.TestCase): + def test_translates_runner_conventions_onto_the_ascend_c_op(self): + # The sgl-kernel-npu op numbers the group-list layouts the other way round + # from sglang, so count layout arrives as 1 and must leave as 0. The clamp is + # inert at 0.0, so a limit that missed clamp_value would silently disable it. + activation = NPUSwigluMxfp8Quant(7.0) + output = torch.empty(2, 4, dtype=torch.float8_e4m3fn) + scale = torch.empty(2, 1, 2, dtype=torch.float8_e8m0fnu) + group_list = torch.tensor([1, 1], dtype=torch.int64) + hidden_states = torch.empty(2, 8) + + with patch.object( + torch.ops.npu, + "swiglu_group_quant", + return_value=(output, scale, None), + create=True, + ) as kernel: + actual_output, actual_scale = activation._apply_activation( + hidden_states, group_list, group_list_type=1 + ) + + self.assertIs(actual_output, output) + self.assertIs(actual_scale, scale) + kwargs = kernel.call_args.kwargs + self.assertIs(kwargs["x"], hidden_states) + self.assertIs(kwargs["group_index"], group_list) + self.assertEqual(kwargs["quant_mode"], 2) + self.assertEqual(kwargs["group_list_type"], 0) + self.assertEqual(kwargs["clamp_value"], 7.0) + + def test_rejects_a_cumulative_group_list(self): + # The op has no cusum layout and sums its group list as counts, so a cusum list + # must fail here rather than derive a row count from the sum of prefix sums. + activation = NPUSwigluMxfp8Quant(7.0) + with self.assertRaises(ValueError): + activation._apply_activation( + torch.empty(2, 8), + torch.tensor([1, 2], dtype=torch.int64), + group_list_type=0, + ) class TestReshapeMxfp4ScaleForNpu(unittest.TestCase): @@ -78,13 +110,35 @@ class TestReshapeMxfp4ScaleForNpu(unittest.TestCase): # [E, N, K/32] -> [E, K/64, N, 2] is the packed-pair layout the GMM reads; # getting the transpose axis wrong silently dequantizes with the wrong scale. scale = torch.arange(8, dtype=torch.uint8).view(1, 2, 4) - out = _reshape_mxfp4_scale_for_npu(scale) + out = reshape_w4a8_mxfp_weight_scale_for_npu(scale) self.assertEqual(tuple(out.shape), (1, 2, 2, 2)) self.assertTrue(torch.equal(out, scale.view(1, 2, 2, 2).transpose(1, 2))) def test_rejects_odd_k_dim(self): with self.assertRaises(ValueError): - _reshape_mxfp4_scale_for_npu(torch.zeros(1, 2, 3, dtype=torch.uint8)) + reshape_w4a8_mxfp_weight_scale_for_npu( + torch.zeros(1, 2, 3, dtype=torch.uint8) + ) + + +class TestPrepareW4A8MxfpWeight(unittest.TestCase): + def test_uses_shared_weight_and_scale_layout(self): + # A wrong transpose or scale packing makes both ModelSlim W4A8 and + # DeepSeek-V4 W4A8 read different blocks from the same checkpoint. + weight = torch.arange(16, dtype=torch.uint8).view(1, 2, 8) + scale = torch.arange(8, dtype=torch.uint8).view(1, 2, 4) + formatted_weight = torch.arange(16, dtype=torch.uint8).view(1, 2, 8) + + with patch( + "sglang.srt.hardware_backend.npu.quantization.moe_methods.npu_format_cast", + return_value=formatted_weight, + ): + prepared_weight, prepared_scale = prepare_w4a8_mxfp_weight(weight, scale) + + self.assertTrue(torch.equal(prepared_weight, formatted_weight.transpose(1, 2))) + self.assertTrue( + torch.equal(prepared_scale, scale.view(1, 2, 2, 2).transpose(1, 2)) + ) class TestMxfp4ScaleWeightLoader(unittest.TestCase): @@ -95,7 +149,7 @@ class TestMxfp4ScaleWeightLoader(unittest.TestCase): loaded.append(loaded_weight.clone()) layer = torch.nn.Module() - method = NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test") + method = NPUW4A8MXFP4FusedMoEMethod(prefix="test") method.create_weights( layer, num_experts=1, @@ -131,264 +185,6 @@ class TestPairPackMxfpActScale(unittest.TestCase): with self.assertRaises(ValueError): _pair_pack_mxfp_act_scale(torch.zeros(2, 3)) - def test_unflattens_low_latency_deepep_scale_as_view(self): - # DeepEP returns one flat E8M0 scale per 32-element block. Passing - # that flat buffer to GMM would use the wrong scale layout and either - # fail or dequantize activations incorrectly. - flat = torch.arange(4, dtype=torch.uint8) - packed = _pair_pack_mxfp_act_scale(flat, input_shape=(2, 64)) - - self.assertEqual(tuple(packed.shape), (2, 1, 2)) - self.assertEqual(packed.data_ptr(), flat.data_ptr()) - self.assertTrue(torch.equal(packed, torch.tensor([[[0, 1]], [[2, 3]]]))) - - def test_rejects_low_latency_deepep_scale_with_wrong_length(self): - with self.assertRaises(ValueError): - _pair_pack_mxfp_act_scale( - torch.zeros(3, dtype=torch.uint8), input_shape=(2, 64) - ) - - -class TestDsv4DeepEPMxfp8DispatcherConfig(unittest.TestCase): - @staticmethod - def _deepep_backend(): - return SimpleNamespace(is_deepep=lambda: True) - - def test_a5_deepep_defaults_low_latency_dispatch_to_mxfp8(self): - dispatcher = MagicMock() - layer = SimpleNamespace(dispatcher=dispatcher) - - with ( - patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), - patch( - "sglang.srt.layers.moe.get_moe_a2a_backend", - return_value=self._deepep_backend(), - ), - patch.dict(os.environ, {}, clear=True), - ): - _configure_dsv4_deepep_dispatcher(layer) - - dispatcher.set_quant_config.assert_called_once_with( - { - "normal_dispatcher_output_dtype": "bf16", - "low_latency_dispatcher_output_dtype": "mxfp8", - } - ) - - def test_non_deepep_ignores_the_low_latency_quant_environment(self): - dispatcher = MagicMock() - layer = SimpleNamespace(dispatcher=dispatcher) - - with ( - patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), - patch( - "sglang.srt.layers.moe.get_moe_a2a_backend", - return_value=SimpleNamespace(is_deepep=lambda: False), - ), - envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), - ): - _configure_dsv4_deepep_dispatcher(layer) - - dispatcher.set_quant_config.assert_called_once_with( - {"dispatcher_output_dtype": "bf16"} - ) - - def test_a5_deepep_allows_bf16_low_latency_fallback(self): - dispatcher = MagicMock() - layer = SimpleNamespace(dispatcher=dispatcher) - - with ( - patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), - patch( - "sglang.srt.layers.moe.get_moe_a2a_backend", - return_value=self._deepep_backend(), - ), - envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("bf16"), - ): - _configure_dsv4_deepep_dispatcher(layer) - - dispatcher.set_quant_config.assert_called_once_with( - { - "normal_dispatcher_output_dtype": "bf16", - "low_latency_dispatcher_output_dtype": "bf16", - } - ) - - def test_a5_deepep_rejects_an_invalid_low_latency_quant_mode(self): - layer = SimpleNamespace(dispatcher=MagicMock()) - - with ( - patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), - patch( - "sglang.srt.layers.moe.get_moe_a2a_backend", - return_value=self._deepep_backend(), - ), - envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), - self.assertRaisesRegex(ValueError, "SGLANG_NPU_DSV4"), - ): - _configure_dsv4_deepep_dispatcher(layer) - - def test_non_a5_ignores_the_low_latency_quant_environment(self): - dispatcher = MagicMock() - layer = SimpleNamespace(dispatcher=dispatcher) - - with ( - patch.object(fp4_moe_methods, "is_npu_arch35", return_value=False), - patch( - "sglang.srt.layers.moe.get_moe_a2a_backend", - return_value=self._deepep_backend(), - ), - envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), - ): - _configure_dsv4_deepep_dispatcher(layer) - - dispatcher.set_quant_config.assert_called_once_with( - {"dispatcher_output_dtype": "bf16"} - ) - - -class _LowLatencyBuffer: - 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, - quant_mode=_NOT_PASSED, - **kwargs, - ): - self.kwargs = {"use_fp8": use_fp8, "quant_mode": quant_mode, **kwargs} - return torch.empty(0), torch.empty(0), object(), object(), object() - - -class _LegacyLowLatencyBuffer: - def low_latency_dispatch( - self, - hidden_states, - topk_ids, - num_max_dispatch_tokens_per_rank, - num_experts, - *, - use_fp8, - **kwargs, - ): - return torch.empty(0), torch.empty(0), object(), object(), object() - - -class TestDeepEPLowLatencyMxfp8Dispatch(unittest.TestCase): - @staticmethod - def _dispatcher(quant_mode, buffer): - dispatcher = object.__new__(deepep._DeepEPDispatcherImplLowLatency) - dispatcher.quant_config = {} - dispatcher.use_fp8 = False - dispatcher.use_nvfp4 = False - dispatcher.low_latency_quant_mode = quant_mode - dispatcher._low_latency_quant_mode_runtime_checked = 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_kernel_quant_mode(self): - buffer = _LowLatencyBuffer() - dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) - - with ( - patch.dict(os.environ, {}, clear=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.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3") - - def test_mxfp8_ops_strategy_uses_legacy_mxfp8_flags(self): - buffer = _LowLatencyBuffer() - dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) - - with ( - patch.dict(os.environ, {"DEEP_USE_MODE": "ops"}, clear=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.assertTrue(buffer.kwargs["use_fp8"]) - self.assertTrue(buffer.kwargs["use_ue8m0"]) - self.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3") - - def test_mxfp8_rejects_an_unsupported_low_latency_strategy(self): - dispatcher = self._dispatcher("mx_fp8_e4m3", _LowLatencyBuffer()) - - with ( - patch.dict(os.environ, {"DEEP_USE_MODE": "alltoall"}, clear=True), - self.assertRaisesRegex(RuntimeError, "DEEP_USE_MODE"), - ): - dispatcher._dispatch_core( - torch.zeros(1, 64), - torch.zeros(1, 1, dtype=torch.int64), - torch.ones(1, 1), - ) - - def test_mxfp8_checks_runtime_interface_once_per_dispatcher(self): - buffer = _LowLatencyBuffer() - dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) - - with ( - patch.dict(os.environ, {}, clear=True), - patch.object(deepep, "_deepep_precompile_tp_barrier"), - patch.object( - deepep.inspect, "signature", wraps=inspect.signature - ) as signature, - ): - dispatcher._dispatch_core( - torch.zeros(1, 64), - torch.zeros(1, 1, dtype=torch.int64), - torch.ones(1, 1), - ) - dispatcher._dispatch_core( - torch.zeros(1, 64), - torch.zeros(1, 1, dtype=torch.int64), - torch.ones(1, 1), - ) - - self.assertEqual(signature.call_count, 1) - - def test_bf16_does_not_pass_a_quant_mode(self): - buffer = _LowLatencyBuffer() - dispatcher = self._dispatcher(None, buffer) - - with 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.assertIs(buffer.kwargs["quant_mode"], _NOT_PASSED) - - def test_mxfp8_rejects_legacy_runtime_without_quant_mode(self): - dispatcher = self._dispatcher("mx_fp8_e4m3", _LegacyLowLatencyBuffer()) - - with self.assertRaisesRegex(RuntimeError, "quant_mode"): - dispatcher._dispatch_core( - torch.zeros(1, 64), - torch.zeros(1, 1, dtype=torch.int64), - torch.ones(1, 1), - ) - class TestW4A8MxfpGmmInputScale(unittest.TestCase): def setUp(self): @@ -431,28 +227,6 @@ class TestW4A8MxfpGmmInputScale(unittest.TestCase): self.assertEqual(call_kwargs["group_list"].dtype, torch.int64) self.assertTrue(torch.equal(call_kwargs["group_list"], self.group_list)) - def test_flat_deepep_scale_skips_dynamic_quant_after_layout_adaptation(self): - flat_scale = torch.arange(4, dtype=torch.uint8) - expected = torch.randn(2, 32) - with ( - patch.object( - torch.ops.npu, "npu_dynamic_mx_quant", create=True - ) as dynamic_quant, - patch.object( - torch.ops.npu, - "npu_grouped_matmul", - return_value=[expected], - create=True, - ) as grouped_matmul, - ): - output = self._call_gmm(flat_scale) - - dynamic_quant.assert_not_called() - self.assertIs(output, expected) - packed_scale = grouped_matmul.call_args.kwargs["per_token_scale"][0] - self.assertEqual(tuple(packed_scale.shape), (2, 1, 2)) - self.assertEqual(packed_scale.data_ptr(), flat_scale.data_ptr()) - def test_missing_scale_uses_dynamic_quant(self): quantized = torch.empty(2, 64, dtype=torch.float8_e4m3fn) quantized_scale = torch.ones(2, 1, 2) @@ -480,51 +254,55 @@ class TestW4A8MxfpGmmInputScale(unittest.TestCase): ) -class TestW4A8MxfpGmmChain(unittest.TestCase): - def test_passes_swiglu_limit_to_quant(self): - gate_up = torch.randn(1, 64) - activated = torch.randn(1, 32) - activated_scale = torch.randn(1, 1) - expected = torch.randn(1, 32) +class TestRunnerDelegation(unittest.TestCase): + def test_apply_delegates_with_dsv4_scale_names(self): + method = NPUW4A8MXFP4FusedMoEMethod(prefix="test") + expected = object() + method.runner = MagicMock() + method.runner.run.return_value = expected layer = SimpleNamespace( w13_weight=MagicMock(), w13_weight_scale_inv=MagicMock(), w2_weight=MagicMock(), w2_weight_scale_inv=MagicMock(), - moe_runner_config=SimpleNamespace(swiglu_limit=7.0), ) + dispatch_output = object() - with ( - patch.object( - fp4_moe_methods, "w4a8_mxfp_gmm", side_effect=[gate_up, expected] - ) as gmm, - patch.object( - fp4_moe_methods, - "swiglu_quant", - return_value=(activated, activated_scale), - ) as swiglu, - ): - output = npu_apply_without_routing_weights_w4a4_mxfp( - layer, - torch.randn(1, 4), - torch.ones(1, 1, 2), - group_list_type=1, - group_list=torch.tensor([1], dtype=torch.int64), - output_dtype=torch.bfloat16, - ) + output = method.apply(layer, dispatch_output) self.assertIs(output, expected) - self.assertTrue(torch.equal(swiglu.call_args.args[0], gate_up)) - self.assertTrue(swiglu.call_args.kwargs["do_limit"]) - self.assertEqual(swiglu.call_args.kwargs["limit"], 7.0) - self.assertIs(gmm.call_args_list[1].kwargs["input"], activated) - self.assertIs(gmm.call_args_list[1].kwargs["input_scale"], activated_scale) + method.runner.run.assert_called_once() + self.assertIs(method.runner.run.call_args.args[0], dispatch_output) + quant_info = method.runner.run.call_args.args[1] + self.assertIs(quant_info.w13_weight_scale, layer.w13_weight_scale_inv) + self.assertIs(quant_info.w2_weight_scale, layer.w2_weight_scale_inv) + + def test_create_runner_installs_internal_kernels_before_runner(self): + method = NPUW4A8MXFP4FusedMoEMethod(prefix="test") + layer = SimpleNamespace() + config = SimpleNamespace(layer=None) + backend = MagicMock() + backend.is_auto.return_value = True + + with ( + patch("sglang.srt.layers.moe.moe_runner.runner.MoeRunner") as moe_runner, + patch( + "sglang.srt.layers.moe.utils.get_moe_runner_backend", + return_value=backend, + ), + ): + method.create_moe_runner(layer, config) + + self.assertIs(layer.w13_kernel, method.w13_kernel) + self.assertIs(layer.w2_kernel, method.w2_kernel) + self.assertIs(config.layer, layer) + moe_runner.assert_called_once() class TestProcessWeightsAfterLoadingZeroScale(unittest.TestCase): @staticmethod def _method(): - return NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test") + return NPUW4A8MXFP4FusedMoEMethod(prefix="test") def test_raises_when_w13_scales_never_loaded(self): # An all-zero scale is the signature of a checkpoint whose scale names @@ -553,6 +331,38 @@ class TestProcessWeightsAfterLoadingZeroScale(unittest.TestCase): with self.assertRaises(RuntimeError): self._method().process_weights_after_loading(layer) + def test_keeps_low_latency_dispatch_in_mxfp8(self): + layer = SimpleNamespace( + w13_weight=torch.nn.Parameter( + torch.ones(1, 2, 32, dtype=torch.uint8), requires_grad=False + ), + w13_weight_scale_inv=torch.nn.Parameter( + torch.ones(1, 2, 1, dtype=torch.uint8), requires_grad=False + ), + w2_weight=torch.nn.Parameter( + torch.ones(1, 32, 1, dtype=torch.uint8), requires_grad=False + ), + w2_weight_scale_inv=torch.nn.Parameter( + torch.ones(1, 32, 1, dtype=torch.uint8), requires_grad=False + ), + dispatcher=MagicMock(), + ) + prepared_weight = torch.ones(1, 32, 2, dtype=torch.uint8) + prepared_scale = torch.ones(1, 1, 2, 2, dtype=torch.uint8) + + with patch( + "sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods.prepare_w4a8_mxfp_weight", + return_value=(prepared_weight, prepared_scale), + ): + self._method().process_weights_after_loading(layer) + + layer.dispatcher.set_quant_config.assert_called_once_with( + { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "mxfp8", + } + ) + if __name__ == "__main__": unittest.main()