From a40058ed3a180f4812912fdcb9bef6ef06f24abf Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Mon, 14 Sep 2026 14:24:56 +0800 Subject: [PATCH 1/2] fix(npu): support W4A8 expert weights in CAM MLP Signed-off-by: bjf-frz --- .../models/npu/deepseek_v2_attention_gate.py | 46 +++++- .../models/test_forward_context.py | 1 + .../models/test_w4a8_attention_gate.py | 141 ++++++++++++++++++ 3 files changed, 185 insertions(+), 3 deletions(-) create mode 100644 tests/unit/model_executor/models/test_w4a8_attention_gate.py diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py index c0e78871..0bb9c7eb 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py @@ -165,10 +165,36 @@ def compute_attention_gate_moe_ffn( ], w2_scale=[experts.get_eplb_parameter("w2_weight_scale")], ) + # ### PATCH START: W4A8 CAM expert weights + # Mirror AscendW4A8DynamicFusedMoEMethod.apply's weight payload; CAM + # already dispatched and quantized the activations, so use only its MLP. + elif quant_type == QuantType.W4A8: + owner = experts.routed_experts + if experts.dynamic_eplb: + moe_weights = MoEWeights( + w1=[w.view(torch.int32) for w in owner.w13_weight_list], + w2=[w.view(torch.int32) for w in owner.w2_weight_list], + w1_scale=owner.w13_weight_scale_list, + w2_scale=owner.w2_weight_scale_list, + w1_scale_bias=owner.w13_scale_bias_list, + w2_scale_bias=owner.w2_scale_bias_list, + ) + else: + bias1 = owner._parameters.get("w13_scale_bias") + bias2 = owner._parameters.get("w2_scale_bias") + moe_weights = MoEWeights( + w1=[owner.w13_weight], + w2=[owner.w2_weight], + w1_scale=[owner.w13_weight_scale], + w2_scale=[owner.w2_weight_scale], + w1_scale_bias=[bias1.detach()] if bias1 is not None else None, + w2_scale_bias=[bias2.detach()] if bias2 is not None else None, + ) + # ### PATCH END: W4A8 CAM expert weights else: raise RuntimeError( - "compute_gate_on_attention currently supports only unquantized " - f"or W8A8 Ascend MoE experts, got {quant_type}", + "compute_gate_on_attention supports unquantized, W8A8 or W4A8 " + f"Ascend MoE experts, got {quant_type}", ) use_gmmswigluquant_fusion = ( quant_type in (QuantType.W8A8, getattr(QuantType, "MXFP8", None)) @@ -217,7 +243,21 @@ def compute_attention_gate_moe_ffn( dynamic_scale=dynamic_scales, topk_scales=topk_scales, weights=moe_weights, - quant=MoEQuantParams(quant_type=quant_type), + # ### PATCH START: W4A8 MLP quantization contract + quant=MoEQuantParams( + quant_type=quant_type, + is_per_channel_weight=( + experts.routed_experts.quant_method.quant_method.is_per_channel_weight + if quant_type == QuantType.W4A8 + else False + ), + ), + swiglu_limit=( + float(layer.mlp.swiglu_limit or 0.0) + if quant_type == QuantType.W4A8 + else 0.0 + ), + # ### PATCH END: W4A8 MLP quantization contract fusion=use_gmmswigluquant_fusion, activation=experts.activation, need_trans=False, diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index 86ff567c..133a8f07 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -635,6 +635,7 @@ def test_deepseek_afd_ffn_skips_empty_rank_local_moe_work( class FakeQuantType: NONE = "none" W8A8 = "w8a8" + W4A8 = "w4a8" class KeywordArguments: def __init__(self, **kwargs): diff --git a/tests/unit/model_executor/models/test_w4a8_attention_gate.py b/tests/unit/model_executor/models/test_w4a8_attention_gate.py new file mode 100644 index 00000000..224ba413 --- /dev/null +++ b/tests/unit/model_executor/models/test_w4a8_attention_gate.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Exercise the CAM W4A8 MLP contract without loading an NPU runtime.""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + + +@pytest.mark.parametrize("dynamic_eplb", [False, True]) +@pytest.mark.parametrize("per_channel", [False, True]) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("rows", [0, 2]) +def test_w4a8_cam_mlp_contract(monkeypatch, dynamic_eplb, per_channel, with_bias, rows): + # Compile the production function in isolation: importing the model module + # would initialize optional vLLM/Ascend dependencies on CPU test hosts. + source = Path( + "afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py" + ).read_text() + function = next( + node + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + and node.name == "compute_attention_gate_moe_ffn" + ) + namespace = { + "torch": torch, + "AFDF2ATransferPayload": SimpleNamespace, + "_gmmswigluquant_fusion_enabled": lambda: False, + } + code = ast.Module( + body=[ + ast.ImportFrom( + module="__future__", names=[ast.alias(name="annotations")], level=0 + ), + function, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(code), "", "exec"), namespace) + quant_type = SimpleNamespace(NONE="none", W8A8="w8a8", W4A8="w4a8") + calls = [] + + def apply_mlp(*, mlp_compute_input): + calls.append(mlp_compute_input) + return torch.ones((rows, 4), dtype=torch.bfloat16), None + + modules = { + "vllm_ascend.ops.fused_moe.moe_mlp": SimpleNamespace( + unified_apply_mlp=apply_mlp + ), + "vllm_ascend.ops.fused_moe.moe_stage_contracts": SimpleNamespace( + MoEMlpComputeInput=SimpleNamespace, MoEWeights=SimpleNamespace + ), + "vllm_ascend.ops.fused_moe.moe_stage_params": SimpleNamespace( + MoEQuantParams=SimpleNamespace + ), + "vllm_ascend.quantization.quant_type": SimpleNamespace(QuantType=quant_type), + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + owner = torch.nn.Module() + for name in ("w13_weight", "w2_weight"): + owner.register_parameter( + name, + torch.nn.Parameter( + torch.ones((2, 4, 4), dtype=torch.int32), requires_grad=False + ), + ) + for name in ("w13_weight_scale", "w2_weight_scale"): + owner.register_parameter( + name, torch.nn.Parameter(torch.ones((2, 4)), requires_grad=False) + ) + for name in ("w13_scale_bias", "w2_scale_bias"): + owner.register_parameter( + name, + torch.nn.Parameter(torch.ones((2, 4)), requires_grad=False) + if with_bias + else None, + ) + owner.w13_weight_list = [owner.w13_weight] + owner.w2_weight_list = [owner.w2_weight] + owner.w13_weight_scale_list = [owner.w13_weight_scale] + owner.w2_weight_scale_list = [owner.w2_weight_scale] + owner.w13_scale_bias_list = [owner.w13_scale_bias] if with_bias else None + owner.w2_scale_bias_list = [owner.w2_scale_bias] if with_bias else None + owner.quant_method = SimpleNamespace( + quant_method=SimpleNamespace(is_per_channel_weight=per_channel) + ) + # The wrapper owns the clamp; do not accidentally read it from the owner. + owner.swiglu_limit = None + experts = SimpleNamespace( + quant_type=quant_type.W4A8, + dynamic_eplb=dynamic_eplb, + routed_experts=owner, + _shared_experts=None, + activation="silu", + ) + layer = SimpleNamespace( + mlp=SimpleNamespace( + experts=experts, swiglu_limit=10.0, routed_scaling_factor=1.0 + ) + ) + hidden_states = torch.ones((rows, 4), dtype=torch.int8) + scales = torch.ones(rows) + output = namespace["compute_attention_gate_moe_ffn"]( + layer, + hidden_states=hidden_states, + group_list=torch.tensor([rows, rows]), + dynamic_scales=scales, + expand_x_shared=None, + dynamic_scales_shared=None, + topk_scales=None, + group_list_type=0, + ) + assert len(calls) == int(rows > 0) + assert output.routed_output.shape == (rows, 4) + assert output.routed_output.dtype == torch.bfloat16 + if not rows: + return + contract = calls[0] + assert contract.hidden_states is hidden_states + assert contract.dynamic_scale is scales + assert contract.quant.quant_type == quant_type.W4A8 + assert contract.quant.is_per_channel_weight == per_channel + assert contract.swiglu_limit == 10.0 + assert contract.weights.w1[0].data_ptr() == owner.w13_weight.data_ptr() + assert contract.weights.w2[0].data_ptr() == owner.w2_weight.data_ptr() + assert contract.weights.w1[0].dtype == torch.int32 + assert contract.weights.w1_scale[0] is owner.w13_weight_scale + assert contract.weights.w2_scale[0] is owner.w2_weight_scale + assert (contract.weights.w1_scale_bias is not None) == with_bias + assert (contract.weights.w2_scale_bias is not None) == with_bias + assert contract.fusion is False From 14e2c22c57b6461147eaab8be49a3a481713957d Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Mon, 14 Sep 2026 15:00:31 +0800 Subject: [PATCH 2/2] style(npu): remove W4A8 patch markers from AFD-owned code Signed-off-by: bjf-frz --- .../model_executor/models/npu/deepseek_v2_attention_gate.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py index 0bb9c7eb..f0210aae 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py @@ -165,7 +165,6 @@ def compute_attention_gate_moe_ffn( ], w2_scale=[experts.get_eplb_parameter("w2_weight_scale")], ) - # ### PATCH START: W4A8 CAM expert weights # Mirror AscendW4A8DynamicFusedMoEMethod.apply's weight payload; CAM # already dispatched and quantized the activations, so use only its MLP. elif quant_type == QuantType.W4A8: @@ -190,7 +189,6 @@ def compute_attention_gate_moe_ffn( w1_scale_bias=[bias1.detach()] if bias1 is not None else None, w2_scale_bias=[bias2.detach()] if bias2 is not None else None, ) - # ### PATCH END: W4A8 CAM expert weights else: raise RuntimeError( "compute_gate_on_attention supports unquantized, W8A8 or W4A8 " @@ -243,7 +241,6 @@ def compute_attention_gate_moe_ffn( dynamic_scale=dynamic_scales, topk_scales=topk_scales, weights=moe_weights, - # ### PATCH START: W4A8 MLP quantization contract quant=MoEQuantParams( quant_type=quant_type, is_per_channel_weight=( @@ -257,7 +254,6 @@ def compute_attention_gate_moe_ffn( if quant_type == QuantType.W4A8 else 0.0 ), - # ### PATCH END: W4A8 MLP quantization contract fusion=use_gmmswigluquant_fusion, activation=experts.activation, need_trans=False,