执行摘要
- 一句话:合并 WNA16 Marlin 与通用 MoE 量化方法,删除冗余实现
- 推荐动作:值得精读。该 PR 展示了如何用 oracle 模式收敛多个量化方法类,并处理了 Marlin tile padding、w2 scale 的 TP 分片等细节。review 中关于 group_size 显式传递的讨论是很好的设计教训;建议关注后续是否补充 dispatcher 层单测。
功能与动机
PR body 明确说明:'Now that the int_wna16.py handles backend selection, the CompressedTensorsWNA16MarlinMoEMethod can be combined with/replaced by CompressedTensorsWNA16MoEMethod.' 该 PR 基于 #44120 引入的 oracle 抽象,目的是消除两个高度重叠的 MoE 量化方法类,收敛后端选择与权重处理逻辑。
实现拆解
- 扩展 int_wna16.py 的 oracle:select_wna16_moe_backend 增加显式 group_size 与 allow_tile_padding 参数,将 Marlin 支持性检查(check_moe_marlin_supports_config)下沉到 _backend_incompatibility_reason 中,并对 MoeWNA16Config 禁止 Marlin。
- 合并方法类: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 与其它后端。
- 简化 dispatcher:compressed_tensors_moe.py 的 get_moe_method 不再调用 check_moe_marlin_supports_layer,去掉 Marlin 预判分支,仅保留 ROCm(rocm_moe_rdna/flydsl/emulation)特例后统一返回 CompressedTensorsWNA16MoEMethod。
- 清理死代码:从 modelopt.py、mxfp4.py、auto_gptq.py、auto_awq.py 与合并后的 wna16 方法中删除已失效的 select_gemm_impl 覆盖方法,并调整 AutoGPTQ/AWQ 的 may_have_zp 与 allow_tile_padding 传参。
- 测试配套: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(模块 量化方法;类别 source;类型 data-contract;符号 CompressedTensorsWNA16MoEMethod, init, get_weight_shape, _w2_scale_sharding): 合并后的核心方法类,吸收了 Marlin 专用逻辑并统一走 oracle 选择后端
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py(模块 量化方法;类别 source;类型 deletion;符号 CompressedTensorsWNA16MarlinMoEMethod, init, get_weight_shape, _w2_scale_sharding): 被整个删除的 Marlin 专用方法类,601 行冗余代码移除
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py(模块 方法分发;类别 source;类型 data-contract;符号 get_moe_method): dispatcher 简化,不再预判 Marlin,直接返回统一方法类,保留 ROCm 特例
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py(模块 后端选择;类别 source;类型 data-contract;符号 select_wna16_moe_backend, _backend_incompatibility_reason): oracle 增加显式 group_size 与 allow_tile_padding 参数,Marlin 支持性检查下沉,是本合并的关键支撑
vllm/model_executor/layers/quantization/utils/marlin_utils.py(模块 Marlin 支持;类别 source;类型 data-contract;符号 check_moe_marlin_supports_config, check_moe_marlin_supports_layer): 新增 check_moe_marlin_supports_config,原 layer 版变为包装,支持在 oracle 侧用 FusedMoEConfig 做形状检查
vllm/model_executor/layers/quantization/auto_gptq.py(模块 量化方法;类别 source;类型 data-contract;符号 select_gemm_impl): may_have_zp 改为依据 is_sym,allow_tile_padding 按 desc_act 传递,并删除 select_gemm_impl
tests/quantization/test_compressed_tensors.py(模块 量化测试;类别 test;类型 test-coverage;符号 test_wna16_moe_w2_scale_sharding): w2 scale sharding 测试改为引用合并后的 CompressedTensorsWNA16MoEMethod
关键符号: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
合并后的核心方法类,吸收了 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
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
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)的兼容性检查 ...
评论区精华
review 中最有价值的讨论集中在三类:
- wjabbour 在 ROCm gfx1201(RDNA4)实测 Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4 启动崩溃,报 qzeros 形状不匹配(kernel 期望 [16, 768] 但收到 [768, 16]),作者承认可能漏了转置 zero 张量。
- bedeks 指出从 weight_key.scale.group_shape[1] 推导 group_size 不可靠(kInt4StaticGroupScale 的 group_shape 是 -1 标记),可能把 Marlin 加入不支持该 group_size 的层,建议显式传参;提交记录显示已按建议修复。
- wjabbour 发现 dispatcher 缩进问题导致 FlydslMoEMethod 路径不可达,作者解释为 merge 问题并修复;wjabbour 进一步指出现有 flydsl 测试只验证 kernel 与 Triton 参考等价、不覆盖路由逻辑,建议补充 dispatch 单元测试(该建议未在本 PR 落地)。
- GPTQ qzeros 转置导致 ROCm gfx1201 启动崩溃 (correctness): bnellnm 承认可能漏了转置 zero 张量,后续提交中修复。
- group_size 应从调用方显式传入 oracle (design): 作者按建议修改,提交信息为 'pass group_size to select. remove select_gemm_impl methods'。
- dispatcher 缩进问题导致 FlydslMoEMethod 不可达 (design): 缩进已修复,Flydsl 路径恢复。
- 缺少 dispatcher 层单元测试 (testing): 本 PR 未补充该测试,遗留为后续工作。
风险与影响
- 风险:
- 回归风险:合并后所有 CompressedTensors WNA16 MoE 模型(Qwen1.5-MoE、Mixtral、Kimi 等)统一走新路径,oracle 分支或权重处理任何错误都会影响全量用户;gsm8k 覆盖有限,未覆盖所有 group_size/TP 组合。
- 平台兼容:ROCm gfx1201 上 GPTQ qzeros 转置 bug 曾在 CI 中暴露,虽经修复但 BuildKite AMD 流水线缺少 RDNA4 常驻 agent,回归可能再次漏网。
- Marlin 选择逻辑:group_size 由 weight_key 推导改为显式传入后,依赖调用方正确传参;若后续新增量化类型忘记传 group_size,可能静默选择错误后端。
- 测试缺口: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+ 行代码
关联脉络
- PR #44120 MoE kernel abstraction oracle for WNA16: PR body 明确说明本合并基于 #44120 引入的 oracle 抽象,int_wna16.py 是其后端选择核心。
- PR #43389 Unrelated MoE PR waiting for this merge: wjabbour 在评论中引用:'TJ would like this PR merged before #43389',说明本 PR 是后续 MoE 重构的前置条件。
- PR #44941 [MoE Refactor] Rename FusedMoE to FusedMoEFactory: 同属 MoE 重构系列,且都触及 fused_moe 目录与量化方法类命名,属于同一演进主线。
参与讨论