Prhub

#51411 [Bugfix][Quantization] Fix INT8 W8A8 MoE crash in TritonExperts

原始 PR 作者 djramic 合并时间 2026-08-08 11:06 文件变更 1 提交数 3 评论 6 代码增减 +2 / -0

执行摘要

修复 Triton MoE 路径 INT8 W8A8 激活崩溃

PR body 说明:TritonExperts did not recognize torch.int8 as a valid input dtype, causing INT8 W8A8 MoE models to crash when processing quantized activations. 在 review 中 fxmarty-amd 补充指出该失败由 PR #50833 引入,属于回归修复而非新功能。

这是一个 2 行的迷你 bugfix,值得快速阅读。它的价值在于:一是展示了量化类型扩展时内核入口 dtype 契约容易遗漏的问题;二是作为回归案例提醒,在跟进 #50833 等上游改动时需同步检查所有入口校验。代码逻辑简单,无需精读。

讨论亮点

核心讨论围绕回归根因:

fxmarty-amd(CHANGES_REQUESTED):cc @ILikeIneine @DarkLight1337 @NickLucche as this failure was caused by https://github.com/vllm-project/vllm/pull/50833

后续 fxmarty-amd 再次审核时转为 APPROVED,说明经过确认该修复合理。

AndreasKaratzas(APPROVED):LGTM

claude[bot]:来自 fork 的 PR 自动 review 被禁用。

实现拆解

  1. 定位问题:vllm/model_executor/layers/fused_moe/experts/triton_moe.pyTritonExperts.applyhidden_states.dtype 校验列表缺少 torch.int8,当 INT8 W8A8 量化模型传入量化激活时触发 AssertionError,导致 EngineCore 初始化失败。
  2. 放开 dtype 校验:在允许的 dtype 列表中加入 torch.int8,使 INT8 激活通过前置断言。
  3. 映射计算类型:在 compute_type 分支中将 torch.int8float8_e4m3fn/float8_e4m3fnuz 并列,统一映射到 tl.bfloat16,确保 Triton 内核按 bf16 计算,避免 int8 乘累加溢出或精度问题。
  4. 验证:运行 pytest tests/quantization/test_quark.py::test_quark_int8_w8a8_moe,结果从 FAILED 变为 PASSED;本次未新增测试文件,属于修复已有测试的可运行性。
文件 模块 状态 重要度
vllm/model_executor/layers/fused_moe/experts/triton_moe.py MoE 专家 modified 4.93

关键符号

apply

关键源码片段

vllm/model_executor/layers/fused_moe/experts/triton_moe.py data-contract

这是 Triton MoE 专家前向路径的核心文件,apply 方法中的 dtype 断言直接导致 INT8 W8A8 崩溃;本次两处修改均在此方法内。

# vllm/model_executor/layers/fused_moe/experts/triton_moe.py
# TritonExperts.apply 中的输入 dtype 校验与计算类型映射# 校验激活 dtype:int8 是 W8A8 量化场景下的合法输入,
# 此前缺少该类型导致 INT8 激活直接触发断言崩溃
assert hidden_states.dtype in [
    torch.float32,
    torch.float16,
    torch.bfloat16,
    torch.float8_e4m3fn,
    torch.float8_e4m3fnuz,
    torch.int8,
]# 将输入 dtype 映射为 Triton 计算类型:
# int8 与 float8 一样统一使用 bf16 计算,
# 既避免 int8 乘累加溢出,也保持与现有效率路径一致
if hidden_states.dtype == torch.bfloat16:
    compute_type = tl.bfloat16
elif hidden_states.dtype == torch.float16:
    compute_type = tl.float16
elif hidden_states.dtype == torch.float32:
    compute_type = tl.float32
elif (
    hidden_states.dtype == torch.float8_e4m3fn
    or hidden_states.dtype == torch.float8_e4m3fnuz
    or hidden_states.dtype == torch.int8
):
    compute_type = tl.bfloat16
else:
    raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

评论区精华

INT8 崩溃是否为 #50833 引入的回归 question

fxmarty-amd 在 CHANGES_REQUESTED 中指出该失败由 https://github.com/vllm-project/vllm/pull/50833 引起,并 @ 了 ILikeIneine、DarkLight1337、NickLucche 等维护者。

结论:确认是回归修复,后续 review 中 fxmarty-amd 与 AndreasKaratzas 均批准。 · 已解决

风险与影响

  1. compute_type 统一到 bf16:与 float8 路径一致,但需确认 Triton kernel 内部是否对 int8 输入有额外反量化处理;当前改动仅放开入口,若内核未适配 int8 数据布局,后续计算可能出错(上下文不足,无法验证内核内部实现)。
  2. 回归面:放宽断言可能让非预期路径(如未量化的 int8 输入)通过校验,需确认调用方只会在 W8A8 量化下传入 int8。
  3. 测试覆盖有限:仅覆盖 Quark 单测一种配置,其他 INT8 W8A8 组合(不同 block_shape、per_act_token_quant、EP/DP)未新增测试。

用户侧:INT8 W8A8 MoE 模型可正常启动和运行 Triton 专家路径,解除启动崩溃。系统侧:无接口或行为变化,compute_type 与原 float8 路径保持一致。团队侧:改动面极小,风险可控;但该修复依赖已有测试覆盖,建议后续补充配置矩阵回归测试。

回归源来自 #50833 测试覆盖有限 int8 计算依赖 bf16

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论