Prhub

#30319 [NPU] Add mxfp4-w4a4 MOE Quantization Support for NPU

原始 PR 作者 LinyuanLi0046 合并时间 2026-08-19 00:06 文件变更 7 提交数 10 评论 20 代码增减 +199 / -6

执行摘要

NPU 新增 ModelSlim W4A4 MXFP4 MoE 量化路径

PR body 明确说明:SGLang 已在 Ascend NPU 上支持 dense 线性层的 ModelSlim W4A4_MXFP4 量化,但对应的 MoE 路径缺失,导致专家权重以 W4A4_MXFP4 方案导出的 MoE 模型无法通过 ModelSlim 量化后端加载和执行。同时 Ascend MoE 实现已重构为模块化 AscendRunner 架构,因此 MXFP4 MoE 支持需要集成新的 per-weight-group 量化方法和 runner 流程,而不是恢复旧的 fused MoE 实现。

值得 NPU 量化相关开发者精读。本 PR 展示了三个值得借鉴的设计决策:一是在 dispatcher 能力受限时选择 BF16 dispatch + 紧邻 GMM 前动态量化的延迟量化策略;二是通过 use_mx_quant 开关把 FP4 这类 torch dtype 对象无法直接识别的 MX 数据类型收敛进公共量化器;三是面对性能优化建议时以 profiling 数据为依据、明确拒绝并建议后续独立 PR,保持改动范围可控。

讨论亮点

Review 由 OrangeRedeng(NPU 维护者)主导,讨论聚焦在性能与正确性两点:

  • 性能建议被拒绝:OrangeRedeng 两次建议使用 GroupedMatmulSwigluQuant(fused SwiGLU 量化)以及在 init_moe_routing 中实现 MXFP4 量化(aclnnMoeInitRoutingV3)。作者 LinyuanLi0046 回复称 profiling 算子延迟数据显示两者均无明显性能收益,为保持改动精简不予采纳,并建议后续单独提交 PR。

  • contiguous() 正确性争论:OrangeRedeng 建议对权重调用 .continues()(实际为 .contiguous()),作者指出对 FP4 NZ 张量调用会导致错误;reviewer 随后自答说明在 w4a8 mxfp 场景下对 continues() 张量也有过问题,同意保持现状。

  • 量化逻辑收敛到公共类:OrangeRedeng 要求把激活量化的 reshape 逻辑移入 HiddenStatesDynamicQuant(与 mxfp8 的做法一致),作者以新增 use_mx_quant 参数的方式修复。

实现拆解

变更从量化方案注册开始,到权重创建、权重后处理、动态激活量化和文档配套共五步:

  1. 方案注册与导出:在 python/sglang/srt/layers/quantization/modelslim/modelslim.pyget_moe_schememoe_quant_schemes 列表中新增 ("W4A4_MXFP4", ModelSlimW4A4MXFP4MoE),并在 schemes/__init__.py 中导入和加入 __all__,使 MoE 量化检测能识别 W4A4_MXFP4 权重组。这是整个功能的检测入口。

  2. 权重创建(数据契约):新增 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py,按 w13/w2 两个权重组分别创建参数:权重张量用 torch.uint8 容器(对应 checkpoint 中 FP4 值装在 float8_e4m3fn 容器),权重 scale 用 uint8 的 UE8M0 block scale(block size 32),并通过 FusedMoeWeightScaleSupported.BLOCK 标记量化是 block 级。

  3. 权重后处理与布局转换NPUW4A4MXFP4MoEMethod.process_weights_after_loading 将 FP4 权重通过 npu_format_cast 转成 NPU 需要的 NZ 布局并转置,将 block scale reshape 为 [E, K//32, N, 2] 供 grouped matmul 使用;由于重构后的 Ascend dispatcher 只支持 BF16/INT8 输出,w13 的 dispatcher 输出显式固定为 BF16。

  4. 动态激活量化python/sglang/srt/hardware_backend/npu/moe/quant.pyHiddenStatesDynamicQuant 新增 use_mx_quant 参数,使 FP4 等 MX dtype 也能选择 torch.ops.npu.npu_dynamic_mx_quantapply 中在每次 expert GMM 之前把 BF16 激活动态量化成 MXFP4,并组装 scale/per_token_scale/x_dtype/weight_dtype 参数调用 GroupedMatmul.forward

  5. 文档配套:在 docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx 增加支持矩阵行、示例命令与实现说明,并在 docs/docs/advanced_features/quantization.mdx 更新支持范围和 ModelSlim 支持清单。测试方面:该 PR 未新增任何单元测试文件,精度验证以 end-to-end eval 形式在 PR body 中给出。

文件 模块 状态 重要度
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py 量化方案 added 8.72
python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py MoE 算子 modified 8.26
python/sglang/srt/hardware_backend/npu/moe/quant.py 激活量化 modified 5.77
python/sglang/srt/layers/quantization/modelslim/modelslim.py 方案注册 modified 4.53
python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py 量化导出 modified 4.7
docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx 文档 modified 3.46
docs/docs/advanced_features/quantization.mdx 文档 modified 2.24

关键符号

ModelSlimW4A4MXFP4MoE.__init__ ModelSlimW4A4MXFP4MoE.create_weights ModelSlimW4A4MXFP4MoE.process_weights_after_loading NPUW4A4MXFP4MoEMethod.__init__ NPUW4A4MXFP4MoEMethod.process_weights_after_loading NPUW4A4MXFP4MoEMethod.apply HiddenStatesDynamicQuant.__init__ ModelSlimConfig.get_moe_scheme

关键源码片段

python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py data-contract

新增的 ModelSlim W4A4_MXFP4 MoE 方案类,是权重创建的入口,定义 w13/w2 权重组的参数形状与 block scale 契约。

# python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py
# ModelSlim W4A4_MXFP4 MoE scheme:负责按 w13 / w2 分组创建 expert 权重参数
class ModelSlimW4A4MXFP4MoE(ModelSlimMoEScheme):
    """Create one ModelSlim MXFP4 expert-weight group (w13 or w2)."""
​
    def __init__(self, quant_config, weight_prefix):
        if weight_prefix not in ("w13", "w2"):
            raise ValueError(
                f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'"
            )
        self.quant_config = quant_config
        self.weight_prefix = weight_prefix
        # 直接绑定 NPU 侧的量化执行方法,后续 process_weights 会转给它
        self.kernel = NPUW4A4MXFP4MoEMethod()
​
    def create_weights(self, layer, num_experts, hidden_size,
                       intermediate_size_per_partition, **extra_weight_attrs):
        # 标记为 block 级量化,供上层 MoE 权重处理逻辑识别
        extra_weight_attrs.update(
            {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
        )
        if self.weight_prefix == "w13":
            output_size = 2 * intermediate_size_per_partition
            input_size = hidden_size
        else:
            output_size = hidden_size
            input_size = intermediate_size_per_partition
​
        # checkpoint 中 FP4 值实际存放在 float8_e4m3fn 容器(uint8 存储),
        # 因此先用 uint8 创建空参数,加载后再 repack 成 FP4 布局
        weight = torch.nn.Parameter(
            torch.empty(num_experts, output_size, input_size // 2, dtype=torch.uint8),
            requires_grad=False,
        )
        layer.register_parameter(f"{self.weight_prefix}_weight", weight)
        set_weight_attrs(weight, extra_weight_attrs)
​
        # UE8M0 block scale:block size 为 32,先按 K 方向 ceil 分配
        weight_scale = torch.nn.Parameter(
            torch.zeros(
                num_experts,
                output_size,
                (input_size + MXFP4_BLOCK_SIZE - 1) // MXFP4_BLOCK_SIZE,
                dtype=torch.uint8,
            ),
            requires_grad=False,
        )
        layer.register_parameter(f"{self.weight_prefix}_weight_scale", weight_scale)
        set_weight_attrs(weight_scale, extra_weight_attrs)
​
    def process_weights_after_loading(self, layer):
        # 布局转换与 repack 全部下沉到 NPU kernel method 里完成
        self.kernel.process_weights_after_loading(layer, self.weight_prefix)
python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py core-logic

NPUW4A4MXFP4MoEMethod 是本功能的核心执行逻辑:权重布局转换、block scale reshape、dispatcher 输出控制与 GMM 前动态激活量化都在这里完成。

# python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py
# NPUW4A4MXFP4MoEMethod:ModelSlim W4A4 MXFP4 MoE 在 Ascend NPU 上的核心执行方法
class NPUW4A4MXFP4MoEMethod(_NPUMoEMethodBase):
    """ModelSlim W4A4 MXFP4 MoE with single-level FP4 weights and activations."""
​
    def __init__(self):
        super().__init__(quant_config=None)
        self.matmul = GroupedMatmul()
        # FP4 权重 dtype 是 torch_npu 的 float4_e2m1fn_x2(int 枚举),
        # 不能直接用 torch.float4_e2m1fn_x2,需要从枚举中解析
        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
        if fp4_dtype is None:
            raise RuntimeError("NPU W4A4 MXFP4 MoE requires float4 support.")
        # use_mx_quant=True 强制走 npu_dynamic_mx_quant(MX 块缩放路径)
        self.hidden_states_quantizer = HiddenStatesDynamicQuant(
            quant_dtype=fp4_dtype,
            use_mx_quant=True,
        )
​
    def process_weights_after_loading(self, layer, weight_prefix):
        self._validate_weight_prefix(layer, weight_prefix)
​
        # offline checkpoint 的 FP4 值装在 float8_e4m3fn 容器里,
        # 这里转成 NPU 的 NZ 布局并转置,供 grouped matmul 使用
        weight = getattr(layer, f"{weight_prefix}_weight")
        weight.data = npu_format_cast(weight.data).transpose(-1, -2)
​
        # UE8M0 block scale 从 [E, N, K // 32] reshape 成
        # [E, K // 32, N, 2],匹配 grouped matmul 的 scale 布局
        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
​
        # 重构后的 Ascend dispatcher 目前只输出 BF16 或 INT8,
        # 所以 w13 的 routing 输出保持 BF16,激活量化放到 apply 里做
        if weight_prefix == "w13":
            self._set_dispatcher_output_dtype(layer, "bf16")
​
    def apply(self, quant_info, hidden_states, expert_tokens, pertoken_scale,
              output_dtype, weight_prefix, group_list_type):
        fp4_dtype = self.hidden_states_quantizer.quant_dtype
        e8m0_dtype = _require_e8m0_dtype()
​
        # BF16 激活在两次 expert GMM 之前动态量化成 MXFP4
        if pertoken_scale is None:
            hidden_states, pertoken_scale = self.hidden_states_quantizer(hidden_states)
        elif pertoken_scale is not None:
            pertoken_scale = pertoken_scale.reshape(
                hidden_states.shape[0], hidden_states.shape[1] // 32, 2
            )
​
        scale_args = {
            "scale": [getattr(quant_info, f"{weight_prefix}_weight_scale", None)],
            "scale_dtype": e8m0_dtype,
            "per_token_scale": [pertoken_scale],
            "per_token_scale_dtype": e8m0_dtype,
            "x_dtype": fp4_dtype,
            "weight_dtype": fp4_dtype,
        }
        scale_args.update(self._get_bias_args(quant_info, weight_prefix))
        return self.matmul.forward(
            quant_info,
            weight_prefix,
            hidden_states,
            expert_tokens.to(torch.int64),
            output_dtype,
            group_list_type=group_list_type,
            transposed=True,
            **scale_args,
        )
python/sglang/srt/hardware_backend/npu/moe/quant.py core-logic

HiddenStatesDynamicQuant 新增 use_mx_quant 参数,使 FP4 等 MX dtype 也能走 npu_dynamic_mx_quant,是本功能激活量化能力的公共基础设施。

# python/sglang/srt/hardware_backend/npu/moe/quant.py
# HiddenStatesDynamicQuant:NPU 动态量化激活的封装
class HiddenStatesDynamicQuant(BaseHiddenStatesQuant):
    """Dynamic per-token quantisation of hidden states."""
​
    def __init__(self, quant_dtype, use_mx_quant=False):
        super().__init__(quant_dtype)
        # 新增 use_mx_quant 开关:FP4 等 MX dtype 的 torch 对象无法直接
        # 与 float8_e4m3fn 比较,需要显式指定走 npu_dynamic_mx_quant
        if use_mx_quant or quant_dtype == torch.float8_e4m3fn:
            self._op = torch.ops.npu.npu_dynamic_mx_quant
        elif quant_dtype in (torch.int8, torch.quint4x2):
            self._op = torch.ops.npu.npu_dynamic_quant
        else:
            raise ValueError(f"Unsupported dynamic quant dtype: {quant_dtype}")
​
    def __call__(self, hidden_states):
        quantized, scale = self._op(hidden_states, dst_type=self.quant_dtype)
        return quantized, scale

评论区精华

是否改用 GroupedMatmulSwigluQuant 提升性能 性能

OrangeRedeng 在 moe_methods.py 第 281 行建议复用 w4a8 场景的 GroupedMatmulSwigluQuant 以提升性能;作者回复:基于性能数据与 profiling 算子延迟,该方法无明显收益,为保持改动精简不采纳。

结论:维持 GroupedMatmul 实现,作者建议后续如有收益场景再单独提交 PR。 · 已解决

是否在 init_moe_routing 中实现 MXFP4 量化 性能

OrangeRedeng 建议在 init_moe_routing 中做 mxfp4 量化(kernel 已支持 aclnnMoeInitRoutingV3);作者回复:测试显示无明显性能提升,为保持改动精简不加入。

结论:不实现,作者建议后续单独 PR 探索。 · 已解决

FP4 NZ 张量调用 .contiguous() 的正确性 正确性

OrangeRedeng 建议在权重后处理中使用 .contiguous();作者说明 FP4 NZ 张量上调用会导致错误;reviewer 随后自答:w4a8 mxfp 场景对 contiguous 张量也有过问题,同意保持现状。

结论:不调用 .contiguous(),维持 npu_format_cast + transpose 的现状。 · 已解决

激活量化逻辑是否收敛到 HiddenStatesDynamicQuant 设计

OrangeRedeng 要求像 mxfp8 一样把激活量化的 reshape 逻辑移入 HiddenStatesDynamicQuant;作者通过新增 use_mx_quant 参数完成收敛。

结论:已修复,HiddenStatesDynamicQuant 增加 use_mx_quant 开关。 · 已解决

风险与影响

主要风险集中在:

  1. 缺少单元测试覆盖:本 PR 未新增任何 test 文件,实现正确性仅依赖端到端精度验证;FP4 布局转换、block scale reshape 这类极易出错的逻辑缺乏回归保护。

  2. FP4 布局转换敏感moe_methods.pynpu_format_cast + transpose(-1, -2) 以及 scale 的三段 reshape/transpose 强依赖 checkpoint 布局;review 中作者明确说明 FP4 NZ 张量上调用 contiguous() 会报错,说明该路径对内存布局非常敏感。

  3. 潜在的量化开销:dispatcher 输出 BF16 后再在每次 GMM 前动态量化,相比直接 MXFP4 dispatch 多一次量化算子;PR 未附速度测试数据,无法确认性能影响。

  4. 硬件范围限制:文档明确该路径要求 Ascend 950+ 产品,低端 NPU 上会直接不适用;CI 分析显示各平台失败均与本改动无关,但 NPU 专项回归仍有限。

对用户:Qwen3 系列等 MoE 模型的 W4A4_MXFP4 离线量化 checkpoint 现在可以在 Ascend NPU 上自动检测并加载,无需显式传 --quantization 参数,扩大了 NPU 上的低比特 MoE 推理能力。对系统:复用 refactored AscendRunner 的 routing、activation 与 finalization 流程,仅新增 per-weight-group 量化方法与 GMM 前置量化,不影响现有 dense W4A4_MXFP4 线性层实现。对团队:确立了 ModelSlim MoE 量化的扩展模式(scheme 类 + kernel method + 公共量化器),后续 W4A8/W4A4 系列方案可以复用 HiddenStatesDynamicQuantuse_mx_quant 机制。

缺少独立单元测试 FP4 布局转换敏感 GMM 前动态量化额外开销 仅支持 Ascend 950+ 无速度基准数据

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论