Prhub

#49444 [Misc] Enable test_silu_mul_fp8_quant_deep_gemm on XPU

原始 PR 作者 pmanczak 合并时间 2026-08-11 23:19 文件变更 2 提交数 2 评论 6 代码增减 +19 / -17

执行摘要

修复 XPU 上 FP8 SiLU 内核崩溃并启用对应测试

PR body 明确指出:persistent_masked_m_silu_mul_quant() 调用 current_platform.get_device_capability().to_int() 并断言其非 None,而 XpuPlatform 返回 None,导致该 wrapper 在 XPU 上直接崩溃而非回退到 Triton 路径。此次变更的目标是让 deep_gemm 的 SiLU + Mul + FP8 量化融合内核在 XPU 上可回退、可测试,使 Intel GPU 用户获得与 ROCm 相同的 Triton 路径覆盖。

值得快速精读,作为平台抽象(current_platform)用法的范例:"能力缺失时优雅回退而非断言失败"。对维护多后端(CUDA/ROCm/XPU)的团队有参考价值,尤其是 has_device_capability() 替换 get_device_capability()+assert 的写法值得在同类平台分支中推广。

讨论亮点
  1. oonyshch 批准时提示 please also check the path this kernel goes thru,作者 pmanczak 回复确认:XPU 上 Triton 是唯一路径而非慢路径,因为该 C++ 内核是 CUDA-only(activation_kernels.cu 主体在 #ifndef USE_ROCM 下),仅在 CUDA dispatch key 下注册,且未编译进 XPU wheel,XPU 上不存在等价 CUDA 内核可选。
  2. jikunshang 在测试中针对 device=gate.device 写法提问 why not use DEVICE?,作者认可并统一为模块级 DEVICE 常量,最终保持代码风格一致。

实现拆解

  1. 平台判断重构(源码):在 vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.pypersistent_masked_m_silu_mul_quant() 中,删除 get_device_capability() + assert + to_int() 三段式判断,改为 if current_platform.is_cuda() and current_platform.has_device_capability(80)has_device_capability() 在平台无能力概念时返回 False,因此 CUDA sm_80+ 仍走 C++ 原生内核,ROCm/XPU 自动落入 Triton 回退路径;删除 assert 同时消除了 XPU 上的崩溃窗口。
  2. 测试去除 CUDA 硬编码tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py 引入 DEVICE = current_platform.device_type,替换 pack_scalesref_with_scale_fmttoken_random 以及 tokens_per_expert 分配中的 device="cuda",参考实现张量跟随输入设备,测试可在任意受支持平台上执行。
  3. 平台门控与容差扩展:新增模块级 pytestmark = pytest.mark.skipif(not (is_cuda_alike() or is_xpu()), ...) 整文件门控;UE8M0 打包 scale 用例仍仅在 is_cuda() 时追加(依赖 C++ 内核);由于 XPU 与 ROCm 共用 Triton 回退内核,断言处以 is_rocm() or is_xpu() 统一放宽为 atol=32.0 / rtol=0.2 的 1-ULP FP8 容差。
  4. 验证:作者在 Intel Arc Pro B70 与 H200 上跑通 23/23 用例;jikunshang 触发 Buildkite CI #83358 全量验证。无配置、schema 或部署配套改动。
文件 模块 状态 重要度
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py MoE 内核 modified 5.29
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py 单元测试 modified 4.36

关键符号

persistent_masked_m_silu_mul_quant test_silu_mul_fp8_quant_deep_gemm pack_scales ref_with_scale_fmt token_random

关键源码片段

vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py core-logic

核心修复文件:`persistent_masked_m_silu_mul_quant()` 的平台判断从 `get_device_capability()+assert` 改为 `is_cuda() and has_device_capability(80)`,消除 XPU 崩溃并让 ROCm/XPU 统一走 Triton 回退路径。

def persistent_masked_m_silu_mul_quant(
    y: torch.Tensor,
    tokens_per_expert: torch.Tensor,
    group_size: int = 128,
    quant_scale_fmt: DeepGemmQuantScaleFMT = DeepGemmQuantScaleFMT.FLOAT32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """SiLU + Mul + FP8 量化融合内核的统一入口(deep_gemm 通道)。    CUDA sm_80+ 走原生 C++ 内核,ROCm / XPU 走 Triton 回退。
    """
    E, T, H2 = y.shape # (E, T, 2*H):gate 与 up 拼接后的输入
    H = H2 // 2
    G = (H + group_size - 1) // group_size
​
    tokens_per_expert = tokens_per_expert.to(device=y.device, dtype=torch.int32)
    fp8_dtype = current_platform.fp8_dtype()
    y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
    # ... 中间省略 y_s(量化 scale)按 quant_scale_fmt 布局分配并清零 ...
    ceil_ue8m0 = quant_scale_fmt in (
        DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
        DeepGemmQuantScaleFMT.UE8M0,
    )
​
    # 关键变更:旧代码先取 get_device_capability() 再断言非 None。
    # XpuPlatform 没有 " 设备计算能力 " 概念,返回 None 直接触发断言崩溃。
    # has_device_capability(80) 在平台无能力概念时返回 False,
    # 因此 CUDA sm_80+ 走 C++ 内核,ROCm / XPU 自动落入 Triton 路径。
    if current_platform.is_cuda() and current_platform.has_device_capability(80):
        torch.ops._C.persistent_masked_m_silu_mul_quant(
            y, tokens_per_expert, y_q, y_s, ceil_ue8m0
        )
    else:
        # Triton 回退(ROCm 与 XPU 共用):C++ 内核受
        # activation_kernels.cu 的 #ifndef USE_ROCM 保护,且未打进 XPU wheel,
        # 此处以静态 grid 覆盖 (E * G) 个专家分组,token 维在核内循环。
        grid = (E * G,)
        stride_i_e, stride_i_t, stride_i_h = y.stride()
        stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
        _silu_mul_fp8_quant_deep_gemm[grid](
            y,
            y_q,
            y_s,
            tokens_per_expert,
            H,
            group_size,
            stride_i_e,
            stride_i_t,
            stride_i_h,
            stride_yq_e,
            stride_yq_t,
            stride_yq_h,
            ys_strides[0],
            # ... 其余 stride / 形状参数 ...
        )
    return y_q, y_s
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py test-coverage

测试平台化:引入 DEVICE 常量替换硬编码 "cuda",新增 CUDA-alike/XPU 的 skipif 门控,并将 1-ULP FP8 容差扩展到 XPU,使测试可在 Intel GPU 上运行并通过。

# 平台无关的设备常量:所有测试张量分配统一使用当前平台设备
DEVICE = current_platform.device_type# 融合内核仅对 CUDA 系(含 ROCm)与 XPU 有意义,其余平台整文件跳过
pytestmark = pytest.mark.skipif(
    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
    reason="persistent_masked_m_silu_mul_quant requires a CUDA-alike or XPU device.",
)# 对比断言处:ROCm 与 XPU 共用 Triton 回退内核
if current_platform.is_rocm() or current_platform.is_xpu():
    # Triton 的 f32 数学内建函数(如 tl.exp)与 torch.exp 最多差 1 ULP,
    # 在 FP8 量化边界可能翻转一个可表示值,允许 1 个 FP8 量子的误差。
    torch.testing.assert_close(
        y_q[e, :nt].to(torch.float32),
        ref_y_q[e, :nt].to(torch.float32),
        atol=32.0,
        rtol=0.2,
    )
else:
    torch.testing.assert_close(
        y_q[e, :nt].to(torch.float32),
        ref_y_q[e, :nt].to(torch.float32),
    )

评论区精华

确认 XPU 上内核走 Triton 路径而非慢路径 question

oonyshch 批准时提示 `please also check the path this kernel goes thru`,作者 pmanczak 回复确认:XPU 上 Triton 是唯一路径;C++ 内核是 CUDA-only(`activation_kernels.cu` 主体在 `#ifndef USE_ROCM` 下),仅在 CUDA dispatch key 下注册,且未编译进 XPU wheel,不存在更慢路径回退的问题。

结论:确认 XPU 上无 CUDA 等价内核可选,Triton 即最优路径,不存在性能回退风险。 · 已解决

为何用 gate.device 而非 DEVICE 常量 style

jikunshang 在测试的 `ref_with_scale_fmt` 中看到 `device=gate.device` 的写法,提问 `why not use DEVICE?`;作者 pmanczak 回复 "Fair point, functionally identical, but device is simpler, switched both spots.",最终统一为模块级 `DEVICE = current_platform.device_type`。

结论:作者接受建议,测试内所有张量分配统一改用 DEVICE 常量,保持风格一致。 · 已解决

风险与影响

低风险。主要关注点:其一,CUDA 分支条件从 cuda_arch >= 80 改为 has_device_capability(80),行为等价性依赖 has_device_capability 内部实现与原先 to_int() 比较语义一致,建议在 CUDA CI 上回归确认;其二,测试将 atol=32.0 / rtol=0.2 容差扩展到 XPU,理论上可能掩盖真实数值偏差,但其依据与 ROCm 相同(Triton 回退内核的 tl.exptorch.exp 最多差 1 ULP,FP8 量化边界可能翻转一个可表示值),理由成立;其三,skipif 依赖 is_cuda_alike() 对 ROCm 返回 True 的语义,若未来平台分类调整需同步更新门控。

对 XPU(Intel GPU)用户而言,该融合内核路径从必然崩溃变为可正常回退 Triton 并通过测试,运行期稳定性提升;对 CUDA/ROCm 无运行时行为变化。对团队而言,Intel GPU CI 新增一条 FP8 MoE 量化内核路径的测试覆盖,且测试代码平台化后更易移植到其他后端。影响范围小、程度中等偏轻。

平台判断行为依赖 has_device_capability 实现 CUDA 路径条件等价性需回归验证 测试容差放宽可能掩盖数值偏差

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论