Prhub

#53674 [Bugfix][ROCm] Fix MiniMax-M3 fused MXFP8 block scale

原始 PR 作者 maithilijoshi20 合并时间 2026-09-16 10:59 文件变更 2 提交数 6 评论 10 代码增减 +28 / -4

执行摘要

修复 MiniMax-M3 ROCm fused MXFP8 block scale 下溢,补回归测试

Issue #53668 指出 PR #53110 只修复了通用 ROCm MXFP8 Triton quantizer,而 MiniMax-M3 的 fused SwiGLU-OAI + MXFP8 kernel「仍然独立使用旧公式」:sb = floor(log2(amax)) + 127 会把每个 block 最大值映射到 [1, 2),而正确的 E4M3 有限范围高达 448。复现实验显示新旧公式的 scale 指数相差约 8 步,小值(如 2**-10)与最大值共享 block 时会被量化到 0.0,而 reference 能保留 0.25,直接影响两个 GEMM 之间的激活精度。该路径用于 mxfp8_native_moe.py 的 ROCm native MXFP8 MoE 实现,属于量化精度正确性问题。

值得精读,重点在于 Issue #53668 的复现方法论:作者在无 ROCm 硬件的环境下用 Triton CPU interpreter 精确复现了量化下溢,给出最小复现与新旧公式对照,并说明该问题属于 PR #53110 未覆盖的第二个实现。PR 本体改动极小且被 reviewer 快速确认,是「同源 bug 双实现遗漏」的典型修复样本;对维护者而言,可考虑后续将 fused kernel 的 scale 逻辑直接抽取复用共享 reference,从根本上杜绝再次分叉。

讨论亮点

Reviewer AndreasKaratzas 在审批时给出了 LGTM 但附带协作建议:PR 描述过度依赖 AI 生成,「the PR description is heavily AI generated and although you've declared it at the end, I would personally appreciate it if the PR description was friendlier to the reviewers... talking about how ruff checks are passing just wastes everybody's time」。作者 maithilijoshi20 致歉并承诺后续改进。

另一条关键问答是 bug 复现方式:AndreasKaratzas 询问「how to repro the bug? Cause this is not linked to any open issue」,作者澄清 Issue #53668 已关联,并说明 bug 发生在 ROCm 上 fused SwiGLU MXFP8 路径处理全零 32 值 block 时。该 bug 在无 ROCm 硬件环境下用 Triton CPU interpreter 复现,具备完整的可验证性。

实现拆解

  1. 修正 kernel 的 block scale 公式:在 vllm/models/minimax_m3/amd/ops/swiglu_oai.py 的 _swiglu_oai_quant_kernel 中新增 FP8_MAX、TINY 两个 tl.constexpr 参数,scale 计算由旧的 floor(log2(amax)) + 127 改为 ceil(log2(amax / FP8_MAX)) + 127,其中 FP8_MAX 取 MXFP8_VALUE_DTYPE(E4M3)的有限最大值 448,TINY 取 float32 的最小正规数用于空 block 兜底。
  2. 同步更新调用点:swiglu_oai_quantize_mxfp8 在 kernel launch 时通过 torch.finfo(MXFP8_VALUE_DTYPE).max 与 torch.finfo(torch.float32).tiny 计算并传入这两个常量,保证 fused kernel 与共享 reference(_mxfp8_e4m3_quantize_torch / _triton)使用完全同源的常量。
  3. 新增回归测试:tests/kernels/test_minimax_m3_amd_ops.py 增加 test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale,构造 alpha=0、beta=0 使激活恰好等于 up,前两个 MX block 内放置 1.0 最大值与 2**-10 小值、第三个 block 全零,断言 fused kernel 与 torch reference 的 scale、量化结果完全相等,同时覆盖空 block 的 TINY 兜底路径。
  4. 修订既有测试注释:test_mxfp8_quant_triton_matches_torch 的注释从「floor(log2(amax)) + 127 算法」更正为「E4M3-aware scale calculation」,使测试意图与实际实现保持一致。
  5. 验证与 CI:作者本地通过 ruff 与 compileall;CI 由 reviewer 触发 Buildkite #86303/#86316,期间两次 pre-commit 格式检查失败后修复通过,最终由 reviewer 批准合入。
文件 模块 状态 重要度
vllm/models/minimax_m3/amd/ops/swiglu_oai.py 量化算子 modified 4.74
tests/kernels/test_minimax_m3_amd_ops.py 量化测试 modified 5.35

关键符号

_swiglu_oai_quant_kernel swiglu_oai_quantize_mxfp8 test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale

关键源码片段

vllm/models/minimax_m3/amd/ops/swiglu_oai.py core-logic

核心修复文件:fused SwiGLU-OAI + MXFP8 量化 kernel 的 E8M0 block scale 公式从旧的 floor(log2(amax)) + 127 改为与共享 MXFP8 reference 一致的 ceil(log2(amax / FP8_MAX)) + 127,并新增 FP8_MAX、TINY 两个 tl.constexpr 常量。

# MiniMax-M3 的 fused SwiGLU-OAI + MXFP8 量化 kernel。
# 修复点:把 E8M0 block scale 从旧的 floor(log2(amax)) + 127 换成与共享
# MXFP8 reference 一致的 E4M3 感知公式,避免小激活值随 block 最大值下溢。
​
def _swiglu_oai_quant_kernel(
    aq_ptr,
    asc_ptr,
    gate_up_ptr,
    alpha,
    beta,
    limit,
    HAS_LIMIT: tl.constexpr,
    BLOCK_M: tl.constexpr,
    FP8_MAX: tl.constexpr, # E4M3 有限最大值(448.0)
    TINY: tl.constexpr, # float32 最小正规数,空 block 兜底
):
    """SwiGLU-OAI(split 布局)融合 per-32-block MXFP8(E4M3 + E8M0)量化。"""
    pid = tl.program_id(0)
    offs = pid * 32 + tl.arange(0, 32)
​
    # 每个 program 读入 [BLOCK_M, 32] 的 gate 与 up,计算 SwiGLU 激活
    gate = tl.load(gate_up_ptr + pid * 64 + offs)
    up = tl.load(gate_up_ptr + pid * 64 + 32 + offs)
    if HAS_LIMIT:
        gate = tl.minimum(gate, limit)
        up = tl.minimum(tl.maximum(up, -limit), limit)
    act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32
​
    # 旧公式 floor(log2(amax)) + 127 把最大值映射到 [1, 2),而对 E4M3 来说
    # 有限范围高达 448,小值会因此被压到 0;除以 FP8_MAX 后取 ceil 得到的
    # scale 才能覆盖整个 E4M3 动态范围。空 block 由 TINY 兜底再 clamp 到 0。
    amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M]
    sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0)
    descale = tl.exp2(sb - 127.0)
    aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty)
​
    # 按 [BLOCK_M, 32] 写回量化结果与 E8M0 指数
    tl.store(
        aq_ptr + pid * BLOCK_M * 32 + tl.arange(0, BLOCK_M)[:, None] * 32 + offs[None, :],
        aq,
    )
    tl.store(asc_ptr + pid * BLOCK_M + tl.arange(0, BLOCK_M),
             sb.to(asc_ptr.dtype.element_ty))
# 调用点:常量与 mxfp8_utils 的 reference 同源,保证 fused kernel 与
# _mxfp8_e4m3_quantize_torch / _triton 行为完全一致。
_swiglu_oai_quant_kernel[grid](
    aq, asc, gate_up, alpha, beta,
    0.0 if limit is None else float(limit),
    HAS_LIMIT=limit is not None,
    BLOCK_M=block_m,
    FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
    TINY=float(torch.finfo(torch.float32).tiny),
    num_warps=4,
)
tests/kernels/test_minimax_m3_amd_ops.py test-coverage

新增回归测试 test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale,直接对比 fused kernel 与 torch reference 的 scale 和量化结果,覆盖混合量级 block 与全零 block 两种关键场景。

@torch.inference_mode()
def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale():
    # 回归场景来自 Issue #53668:前两个 MX block 在最大值 1.0 旁放 2**-10
    # 的小值(旧公式会把它下溢为 0.0),第三个 block 全零验证 TINY 兜底。
    m, inter = 3, 96
    gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16)
    up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16)
    up[:, :64:32] = 1.0 # block 0/1 放入最大值
    up[:, 64:] = 0.0 # block 2 全零
    gate_up = torch.cat((gate, up), dim=-1)
​
    # alpha=0 使 sigmoid(alpha * gate)=0.5,配合 beta=0 时激活恒等于 up,
    # 因此可直接与 _mxfp8_e4m3_quantize_torch 的输出做全等比较。
    got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None)
    ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False)
​
    # scale 与量化值都必须与共享 reference 完全一致,不允许任何偏差
    assert torch.equal(got_s, ref_s)
    assert torch.equal(got_q, ref_q)

评论区精华

PR 描述过度依赖 AI 生成 style

Reviewer AndreasKaratzas 指出描述中诸如 ruff 检查通过之类的信息对 review 没有帮助,希望描述更便于 reviewer 阅读。

结论:作者 maithilijoshi20 致歉并承诺后续避免使用 AI 生成描述;reviewer 随后正式批准。 · 已解决

Bug 复现方式与 Issue 关联 question

AndreasKaratzas 询问「how to repro the bug? Cause this is not linked to any open issue」,因为 PR body 未直接点明关联 Issue。

结论:作者澄清已关联 Issue #53668,并说明 bug 发生在 ROCm 上 fused SwiGLU MXFP8 路径处理全零 32 值 block 时;复现脚本在 Issue 中给出,可用 Triton CPU interpreter 在无 ROCm 硬件环境运行。 · 已解决

CI 触发与 pre-commit 校验 测试

作者因 PR 数不足无法本地跑 pre-commit,mergify 两次报告格式检查失败;reviewer 两次触发 Buildkite CI(#86303、#86316)。

结论:格式问题修复后 CI 通过,PR 合入 main。 · 已解决

风险与影响

  1. 数值核心路径变更:改动直接改变 ROCm native MXFP8 MoE(mxfp8_native_moe.py)中 fused 激活量化的 scale 计算,若新公式在极端输入下与权重方预量化 scale 失配,理论上可能产生精度变化;但新公式与共享 reference 完全同源,回归风险低。
  2. 空 block 语义依赖:新公式依赖 TINY 兜底再 clamp 到 0 的路径,测试覆盖了全零 block,但 kernel 对 NaN/Inf 输入仍无显式防护(旧版同样如此),不构成本 PR 引入的新风险。
  3. 测试平台限制:测试模块在非 ROCm 平台整体 skip(模块级 skip),新增回归保护仅对 ROCm CI 生效;作者未在真实 ROCm 硬件上验证,依赖 reviewer 触发的 AMD CI。
  4. 性能影响:ceil/log2/divide 仅在每个 32 元素 block 计算一次,相对 GEMM 开销可忽略。

用户侧:提升 MiniMax-M3 在 ROCm(gfx942/gfx950)上 native MXFP8 MoE 路径的激活量化精度,修复小激活值下溢导致的输出质量退化。系统侧:改动仅 2 文件、28 行,无接口、配置、schema 或部署变化,不影响其他模型和平台。团队侧:为 fused kernel 与共享 MXFP8 reference 的一致性确立了测试契约,避免同类「只修一个实现」的遗漏再次发生,对后续 ROCm 量化内核维护有示范意义。

ROCm 特定路径 量化精度回归风险 缺少 ROCm 硬件验证 新测试仅 ROCm CI 生效

关联 Issue

#53668 [Bug]: MiniMax-M3 fused SwiGLU MXFP8 kernel still uses the pre-#53110 block scale formula

完整报告

参与讨论