Prhub

#50516 [ROCm][CI] Fall back to lossless Kimi K3 MXFP4 emulation on gfx942

原始 PR 作者 AndreasKaratzas 合并时间 2026-07-31 23:33 文件变更 4 提交数 1 评论 0 代码增减 +34 / -5

执行摘要

gfx942 回退到无损 Kimi K3 MXFP4 仿真

PR #50000 将 Kimi K3 默认改为 MXFP4 RoutedExperts,但 gfx942 上 AITER 后端拒绝 SiTU 激活,导致模型构造时无可选后端并抛 NotImplementedError。AMD CI 构建 11504 出现 Kimi K3 模型初始化失败。

值得精读,展示了后端 fallback 设计模式:当首选后端不支持特定激活或设备时,如何无损回退到仿真路径,并正确处理量化语义。

讨论亮点

本 PR 没有实质性的 review 讨论。claude[bot] 自动提示“此 PR 来自 fork,自动审查已禁用”,维护者 mgoin 直接批准。没有公开的评论线程涉及设计权衡。

实现拆解

  1. 后端选择回退:在 vllm/model_executor/layers/fused_moe/oracle/mxfp4.py_get_priority_backends 中,ROCm 平台候选后端从仅 AITER_MXFP4_BF16 扩展为 [AITER_MXFP4_BF16, EMULATION],当 AITER 因设备或激活不匹配被拒绝时,oracle 可继续选择 EMULATION 后端。
  2. 权重布局保持:同一文件的权重准备逻辑将 EMULATIONXPU 归为同一分支,不对 MXFP4 权重做 swizzle 或转换,直接使用原始 checkpoint 布局,由仿真后端在运行时反量化。
  3. 量化配置修正:在 vllm/model_executor/layers/quantization/mxfp4.pyMxfp4Config.get_fused_moe_quant_config 中,当后端为 EMULATION 时改用 mxfp4_w4a16_moe_quant_config,显式保留 BF16 激活,而非通用 EMULATION 的 W4A4 配置。
  4. 仿真后端激活语义:在 vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py__init__ 中,为 w_mxfp4w_mxfp6_e3m2w_mxfp6_e2m3 等 weight-only 方案设置 quant_dtype = None,与 W4A16 语义一致。
  5. 激活函数支持:在 vllm/model_executor/layers/fused_moe/experts/triton_moe.pyTritonExperts._supports_activation 中新增 MoEActivation.SITU,使 Triton(EMULATION 的后端载体)能够接受 Kimi K3 的激活。

测试与 CI:本 PR 没有新增测试文件,依赖 AMD CI 上已有的 Kimi K3 模型初始化测试来覆盖该回退路径。

文件 模块 状态 重要度
vllm/model_executor/layers/quantization/mxfp4.py 量化层 modified 6.62
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py 后端路由 modified 6.12
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py 仿真专家 modified 6.07
vllm/model_executor/layers/fused_moe/experts/triton_moe.py 专家内核 modified 4.65

关键符号

Mxfp4Config.get_fused_moe_quant_config _get_priority_backends OCP_MXQuantizationEmulationTritonExperts.__init__ TritonExperts._supports_activation

关键源码片段

vllm/model_executor/layers/quantization/mxfp4.py data-contract

核心逻辑:为 EMULATION 后端选择 W4A16 量化配置,确保仅权重量化、激活保留 BF16

# vllm/model_executor/layers/quantization/mxfp4.py
def get_fused_moe_quant_config(
    self,
    layer: RoutedExperts,
) -> FusedMoEQuantConfig | None:
    w1_bias = getattr(layer, "w13_bias", None)
    w2_bias = getattr(layer, "w2_bias", None)
    swiglu_limit = getattr(layer, "swiglu_limit", None)
​
    from vllm.platforms.rocm import on_gfx1250
​
    if self.mxfp4_backend in TRITON_BACKENDS or (
        self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
    ):
        # TRITON backends 在 swizzle 后会释放原始 scale 参数,
        # swizzled scale 存放在 precision config 中,这里直接引用。
        assert self.w13_precision_config is not None
        assert self.w2_precision_config is not None
        w1_scale = self.w13_precision_config
        w2_scale = self.w2_precision_config
    else:
        w1_scale = layer.w13_weight_scale
        w2_scale = layer.w2_weight_scale
​
    if self.mxfp4_backend == Mxfp4MoeBackend.EMULATION:
        # Kimi K3 等 checkpoint 是 weight-only W4A16,
        # 通用 EMULATION 配置是 W4A4,这里保留 BF16 激活,
        # 仅对权重做反量化,保证无损回退。
        return mxfp4_w4a16_moe_quant_config(
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            w1_bias=w1_bias,
            w2_bias=w2_bias,
            gemm1_clamp_limit=swiglu_limit,
        )
​
    return make_mxfp4_moe_quant_config(
        mxfp4_backend=self.mxfp4_backend,
        w1_scale=w1_scale,
        w2_scale=w2_scale,
        w1_bias=w1_bias,
        w2_bias=w2_bias,
        swiglu_limit=swiglu_limit,
        layer=layer,
    )
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py data-contract

修复 weight-only 方案被当作 W4A4 处理的问题,将 quant_dtype 置为 None

# vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
    def __init__(self, moe_config, quant_config):
        super().__init__(moe_config, quant_config)
        # 省略 warning 和 scale 设置 ...
        self.quantization_emulation = True
​
        if self.ocp_mx_scheme in {
            OCP_MX_Scheme.w_mxfp4,
            OCP_MX_Scheme.w_mxfp6_e3m2,
            OCP_MX_Scheme.w_mxfp6_e2m3,
        }:
            # Weight-only 方案只量化权重,激活保持 BF16。
            self._quant_dtype = None
        elif self.ocp_mx_scheme in {
            OCP_MX_Scheme.w_mxfp4_a_mxfp4,
        }:
            # 权重和激活都量化时,激活按 mxfp4 处理。
            self._quant_dtype = "mxfp4"
        elif self.ocp_mx_scheme in [
            OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
            OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
            OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2,
            OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3,
        ]:
            self._quant_dtype = "mxfp6"
        elif self.ocp_mx_scheme in [
            OCP_MX_Scheme.w_mxfp4_a_fp8,
            OCP_MX_Scheme.w_mxfp6_e3m2_a_fp8,
        ]:
            self._quant_dtype = current_platform.fp8_dtype()

评论区精华

自动化审查因 fork 被禁用 other

claude[bot] 评论:此 PR 来自 fork,自动审查已禁用。维护者可以评论 @claude review 来运行一次审查。

结论:没有进一步评论,维护者 mgoin 直接批准并合并。 · closed

风险与影响

性能风险:EMULATION 后端在运行时反量化权重,比原生 AITER/TRTLLM 后端慢,但仅影响 gfx942 上 AITER 不支持的场景。兼容性风险:ocp_mx_emulation_moe.py 中 quant_dtype 的语义变更可能影响其他使用同等 scheme 的模型。测试覆盖缺口:未新增单元测试,完全依赖 AMD CI 的端到端测试。后端优先级变化可能掩盖 AITER 的兼容性问题。

主要影响 AMD gfx942 上使用 Kimi K3 的用户,使其能正常加载并运行模型;同时让 ROCm 上 MXFP4 通用后端选择更健壮。影响范围限于 ROCm 平台和 MXFP4 量化路径,对其他平台无影响。

平台特定回退 缺少测试覆盖 量化路径变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论