Prhub

#44570 [MoE Refactor] Combine CompressedTensorsWNA16MarlinMoEMethod with CompressedTensorsWNA16MoEMethod

原始 PR 作者 bnellnm 合并时间 2026-08-01 05:13 文件变更 14 提交数 114 评论 28 代码增减 +545 / -855

执行摘要

合并 WNA16 Marlin 与通用 MoE 量化方法,删除冗余实现

PR body 明确说明:'Now that the int_wna16.py handles backend selection, the CompressedTensorsWNA16MarlinMoEMethod can be combined with/replaced by CompressedTensorsWNA16MoEMethod.' 该 PR 基于 #44120 引入的 oracle 抽象,目的是消除两个高度重叠的 MoE 量化方法类,收敛后端选择与权重处理逻辑。

值得精读。该 PR 展示了如何用 oracle 模式收敛多个量化方法类,并处理了 Marlin tile padding、w2 scale 的 TP 分片等细节。review 中关于 group_size 显式传递的讨论是很好的设计教训;建议关注后续是否补充 dispatcher 层单测。

讨论亮点

review 中最有价值的讨论集中在三类:

  1. wjabbour 在 ROCm gfx1201(RDNA4)实测 Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4 启动崩溃,报 qzeros 形状不匹配(kernel 期望 [16, 768] 但收到 [768, 16]),作者承认可能漏了转置 zero 张量。
  2. bedeks 指出从 weight_key.scale.group_shape[1] 推导 group_size 不可靠(kInt4StaticGroupScale 的 group_shape 是 -1 标记),可能把 Marlin 加入不支持该 group_size 的层,建议显式传参;提交记录显示已按建议修复。
  3. wjabbour 发现 dispatcher 缩进问题导致 FlydslMoEMethod 路径不可达,作者解释为 merge 问题并修复;wjabbour 进一步指出现有 flydsl 测试只验证 kernel 与 Triton 参考等价、不覆盖路由逻辑,建议补充 dispatch 单元测试(该建议未在本 PR 落地)。

实现拆解

  1. 扩展 int_wna16.py 的 oracle:select_wna16_moe_backend 增加显式 group_size 与 allow_tile_padding 参数,将 Marlin 支持性检查(check_moe_marlin_supports_config)下沉到 _backend_incompatibility_reason 中,并对 MoeWNA16Config 禁止 Marlin。
  2. 合并方法类:compressed_tensors_moe_wna16.py 的 init 吸收原 Marlin 类的 QuantKey 构造、oracle 选择、is_marlin/is_transposed 判断与 input_dtype 选择;新增 get_weight_shape 与 _w2_scale_sharding,create_weights 支持 actorder=group 时的 full-K w2 scale;process_weights_after_loading 统一处理 Marlin 与其它后端。
  3. 简化 dispatcher:compressed_tensors_moe.py 的 get_moe_method 不再调用 check_moe_marlin_supports_layer,去掉 Marlin 预判分支,仅保留 ROCm(rocm_moe_rdna/flydsl/emulation)特例后统一返回 CompressedTensorsWNA16MoEMethod。
  4. 清理死代码:从 modelopt.py、mxfp4.py、auto_gptq.py、auto_awq.py 与合并后的 wna16 方法中删除已失效的 select_gemm_impl 覆盖方法,并调整 AutoGPTQ/AWQ 的 may_have_zp 与 allow_tile_padding 传参。
  5. 测试配套:test_compressed_tensors.py 的 _w2_scale_sharding 测试改为引用合并后的类,test_moe_wna16.py 与 test_marlin_tile_padding.py 做相应适配;作者本地跑了 Kimi、Qwen1.5-MoE、Mixtral 的 gsm8k 验证。
文件 模块 状态 重要度
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py 量化方法 modified 8.86
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py 量化方法 removed 9.08
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py 方法分发 modified 7.28
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py 后端选择 modified 6.7
vllm/model_executor/layers/quantization/utils/marlin_utils.py Marlin 支持 modified 7.09
vllm/model_executor/layers/quantization/auto_gptq.py 量化方法 modified 6.2
tests/quantization/test_compressed_tensors.py 量化测试 modified 5.02

关键符号

CompressedTensorsWNA16MoEMethod.__init__ get_weight_shape _w2_scale_sharding select_wna16_moe_backend check_moe_marlin_supports_config get_moe_method process_weights_after_loading

关键源码片段

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py data-contract

合并后的核心方法类,吸收了 Marlin 专用逻辑并统一走 oracle 选择后端

# 合并后的 CompressedTensorsWNA16MoEMethod 构造逻辑
def __init__(self, weight_quant, input_quant, moe, layer_name=None):
    super().__init__(moe)
    self.weight_quant = weight_quant
    self.input_quant = input_quant
    # 统一提取量化属性,Marlin 与通用路径共用
    self.symmetric = weight_quant.symmetric
    self.num_bits = weight_quant.num_bits
    self.packed_factor = 32 // weight_quant.num_bits
    self.strategy = weight_quant.strategy
    self.group_size = weight_quant.group_size
    self.actorder = weight_quant.actorder
​
    # 依据位宽与对称性构造 QuantKey,供 oracle 查询后端
    self.quant_type = (
        WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
        if self.symmetric
        else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]
    )
    if self.num_bits == 4:
        scale = kInt4Static32GroupScale if self.group_size == 32 else kInt4StaticGroupScale
    elif self.num_bits == 8:
        assert self.group_size == -1
        scale = kInt8StaticGroupScale
    else:
        raise ValueError('CompressedTensorsWNA16MoEMethod only supports int4 and int8 now.')
​
    weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
​
    # actorder=group/dynamic 时 Marlin 不允许 tile padding,需限制后端选择
    is_actorder = (
        self.strategy == QuantizationStrategy.GROUP
        and self.actorder in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
    )
​
    # 统一由 oracle 选择 WNA16 MoE 后端(Marlin / Flashinfer / Triton / emulation)
    self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
        config=self.moe,
        weight_key=weight_key,
        quant_config=self.weight_quant,
        may_have_zp=not self.symmetric,
        may_have_bias=False,
        allow_tile_padding=not is_actorder,
    )
    self.is_marlin = self.wna16_backend in (
        WNA16MoEBackend.MARLIN,
        WNA16MoEBackend.BATCHED_MARLIN,
    )
    self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
​
    if self.is_marlin:
        assert check_moe_marlin_supports_config(
            self.moe, self.group_size, allow_tile_padding=not is_actorder
        )
        self.input_dtype = get_marlin_input_dtype(layer_name)
    else:
        # 非 Marlin 内核仍保留原有限制:仅支持 group 策略、无 group actorder、必须对称
        assert weight_quant.strategy == 'group'
        assert weight_quant.actorder != 'group'
        assert self.symmetric, 'Only symmetric quantization is supported for MoE'
        self.input_dtype = torch.bfloat16
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py data-contract

dispatcher 简化,不再预判 Marlin,直接返回统一方法类,保留 ROCm 特例

# get_moe_method 中 WNA16 group/channel 分支(合并后)
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
    # ... 格式与位宽校验省略 ...
​
    # 原生 ROCm HIP 内核(RDNA3 等)优先
    if current_platform.is_rocm():
        from . import rocm_moe_rdna
        if rocm_moe_rdna.is_supported(weight_quant):
            return rocm_moe_rdna.make_method(weight_quant, input_quant, layer.moe_config)
        # gfx950 特定 flydsl 后端
        if (moe_backend == 'flydsl' and <gfx950 等条件>):
            return CompressedTensorsW4A16FlydslMoEMethod(weight_quant, input_quant, layer.moe_config)
        elif moe_backend == 'emulation':
            logger.info_once('Using CompressedTensorsWNA16MoEMethod (emulation backend requested)')
​
    # 统一走 CompressedTensorsWNA16MoEMethod,后端选择交给 oracle
    from .compressed_tensors_moe_wna16 import CompressedTensorsWNA16MoEMethod
    logger.info_once('Using CompressedTensorsWNA16MoEMethod')
    return CompressedTensorsWNA16MoEMethod(weight_quant, input_quant, layer.moe_config)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py data-contract

oracle 增加显式 group_size 与 allow_tile_padding 参数,Marlin 支持性检查下沉,是本合并的关键支撑

# oracle 中 Marlin 支持性检查(显式 group_size 版本)
def _backend_incompatibility_reason(backend, moe_config, quant_config, may_have_zp, may_have_bias, allow_tile_padding, group_size):
    if backend in (WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN):
        # MoeWNA16Config 是纯 Triton/Flashinfer 配置,不允许 Marlin
        if isinstance(quant_config, MoeWNA16Config):
            return 'Marlin not supported for this layer'
        # 其余 QuantizationArgs / AWQ / GPTQ 配置用显式 group_size 校验形状支持
        if not check_moe_marlin_supports_config(moe_config, group_size, allow_tile_padding):
            return 'Marlin not supported for this layer'
    # ... 其余后端(Flashinfer、Triton、emulation)的兼容性检查 ...

评论区精华

GPTQ qzeros 转置导致 ROCm gfx1201 启动崩溃 正确性

wjabbour 实测 Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4 启动崩溃:AssertionError: qzeros shape mismatch,kernel 期望 [16, 768] 但收到 [768, 16];同命令在 main 上可正常加载。

结论:bnellnm 承认可能漏了转置 zero 张量,后续提交中修复。 · 已解决

group_size 应从调用方显式传入 oracle 设计

bedeks 指出从 weight_key.scale.group_shape[1] 推导 group_size 不可靠(kInt4StaticGroupScale 的 group_shape 是 -1),可能误将 Marlin 加入优先级列表;建议 select_wna16_moe_backend 显式接收 group_size。

结论:作者按建议修改,提交信息为 'pass group_size to select. remove select_gemm_impl methods'。 · 已解决

dispatcher 缩进问题导致 FlydslMoEMethod 不可达 设计

wjabbour 指出 dedentation 后 CompressedTensorsW4A16FlydslMoEMethod 分支不可达,询问是否故意;作者回应是 merge 问题,已修复缩进。

结论:缩进已修复,Flydsl 路径恢复。 · 已解决

缺少 dispatcher 层单元测试 测试

wjabbour 指出现有 flydsl 测试套件只验证 kernel 与 Triton 参考等价,不覆盖 get_moe_method 的路由逻辑;建议补充 dispatch 路径的单元测试。

结论:本 PR 未补充该测试,遗留为后续工作。 · 待处理

风险与影响

  1. 回归风险:合并后所有 CompressedTensors WNA16 MoE 模型(Qwen1.5-MoE、Mixtral、Kimi 等)统一走新路径,oracle 分支或权重处理任何错误都会影响全量用户;gsm8k 覆盖有限,未覆盖所有 group_size/TP 组合。
  2. 平台兼容:ROCm gfx1201 上 GPTQ qzeros 转置 bug 曾在 CI 中暴露,虽经修复但 BuildKite AMD 流水线缺少 RDNA4 常驻 agent,回归可能再次漏网。
  3. Marlin 选择逻辑:group_size 由 weight_key 推导改为显式传入后,依赖调用方正确传参;若后续新增量化类型忘记传 group_size,可能静默选择错误后端。
  4. 测试缺口:dispatcher 层没有单元测试,缩进导致不可达代码的 bug 未被 CI 捕获,说明路由逻辑缺少防护网。

对最终用户:量化 MoE 模型加载与推理行为基本不变,gsm8k 指标略有提升(Qwen 0.4640→0.4663,Mixtral 0.6929→0.6945)。对系统:删除约 600 行重复实现,后端选择逻辑收敛到 oracle,新增 WNA16 MoE 后端时只需扩展 int_wna16.py 与 experts 类,不再需要维护两个方法类。对团队:MoE 量化扩展点更清晰,但需关注 ROCm 平台验证与 dispatcher 测试补充。

核心路径重构 平台兼容风险(ROCm) 缺少 dispatch 层单元测试 删除 600+ 行代码

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论