执行摘要
- 一句话:为ZenCPU平台增加运行时日志和文档
- 推荐动作:此PR值得精读,特别是日志添加模式(debug_once统一为所有路径添加)和测试设计(硬件无关的mock测试),可作为平台可观测性增强的参考范本。
功能与动机
为了增加AMD Zen CPU平台的运行时可见性,方便用户通过实时日志确认ZenCpuPlatform是否激活、未量化的GEMM是否走zentorch路径(PR body原文)。
实现拆解
- 提升平台激活日志级别:在
vllm/platforms/__init__.py的cpu_platform_plugin函数中,将Zen平台匹配成功时的logger.debug改为logger.info,使得启动时直接输出“AMD Zen CPU detected with zentorch installed, using ZenCpuPlatform.”。
- 添加GEMM分发路径日志:在
vllm/model_executor/layers/utils.py的dispatch_cpu_unquantized_gemm函数中,对每个分发分支(zentorch、sgl-kernel、oneDNN、fallback)分别添加logger.debug_once调用,输出当前使用的内核名称和关键参数(如prepack状态)。
- 新增单元测试验证日志:在
tests/model_executor/test_cpu_unquantized_gemm_dispatch.py中新增test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch,通过mock zentorch op和平台检测,验证dispatch_cpu_unquantized_gemm在Zen路径下输出预期的日志字符串。
- 更新安装与验证文档:在
docs/getting_started/installation/cpu.x86.inc.md中新增“AMD Zen optimizations”章节(包含Docker构建、检测规则、支持dtype、环境变量等);在docs/getting_started/installation/cpu.md中引用该章节并新增FAQ;在docs/models/hardware_supported_models/cpu.md中添加AMD Zen说明提示。
关键文件:
tests/model_executor/test_cpu_unquantized_gemm_dispatch.py(模块 分发逻辑;类别 test;类型 test-coverage;符号 test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch): 新增测试用例验证zentorch路径日志输出,确保日志行为正确且可被CI覆盖。
vllm/model_executor/layers/utils.py(模块 GEMM分发;类别 source;类型 core-logic): 在 dispatch_cpu_unquantized_gemm 的所有分支添加 debug_once 日志,增强分发可见性。
vllm/platforms/__init__.py(模块 平台检测;类别 source;类型 core-logic;符号 cpu_platform_plugin): 将 Zen 平台激活日志从 debug 提升为 info,使用户启动时即可确认平台类型。
docs/getting_started/installation/cpu.x86.inc.md(模块 安装文档;类别 docs;类型 documentation): 新增 AMD Zen 优化章节,包括检测规则、Docker 构建、环境变量等。
docs/getting_started/installation/cpu.md(模块 安装文档;类别 docs;类型 documentation): 添加 AMD Zen 优化章节链接和环境变量说明。
docs/models/hardware_supported_models/cpu.md(模块 模型文档;类别 docs;类型 documentation): 添加 AMD Zen CPU 支持说明,提醒用户模型兼容性不变。
关键符号:cpu_platform_plugin, dispatch_cpu_unquantized_gemm, test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch
关键源码片段
tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
新增测试用例验证zentorch路径日志输出,确保日志行为正确且可被CI覆盖。
# 测试 zentorch 路径的日志输出(硬件无关的单元测试)
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
def test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch(monkeypatch):
# 模拟当前平台为 Zen CPU
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
# 期望的 prepacked 状态由环境变量和 zentorch op 能力决定
expected_prepacked = bool(utils.envs.VLLM_ZENTORCH_WEIGHT_PREPACK) and hasattr(
torch.ops.zentorch, "zentorch_weight_prepack_for_linear"
)
# 捕获 logger.debug_once 调用
log_calls = []
monkeypatch.setattr(utils.logger, "debug_once", lambda *args: log_calls.append(args))
layer = torch.nn.Linear(16, 8, bias=True)
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
# 断言日志内容与预期一致
assert log_calls == [
(
"CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
expected_prepacked,
)
]
vllm/model_executor/layers/utils.py
在 dispatch_cpu_unquantized_gemm 的所有分支添加 debug_once 日志,增强分发可见性。
def dispatch_cpu_unquantized_gemm(
layer: torch.nn.Module,
remove_weight: bool = False,
) -> None:
# ... 前面的代码 ...
# Zen CPU 路径:zentorch_linear_unary
if current_platform.is_zen_cpu() and hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
zen_weight = layer.weight.detach()
is_prepacked = False
if envs.VLLM_ZENTORCH_WEIGHT_PREPACK and hasattr(
torch.ops.zentorch, "zentorch_weight_prepack_for_linear"
):
zen_weight = torch.ops.zentorch.zentorch_weight_prepack_for_linear(zen_weight)
is_prepacked = True
layer.cpu_linear = lambda x, weight, bias, _p=is_prepacked: (
torch.ops.zentorch.zentorch_linear_unary(x, zen_weight, bias, is_weight_prepacked=_p)
)
if remove_weight:
layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
# 添加日志:记录当前使用的内核及 prepack 状态
logger.debug_once(
"CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
is_prepacked,
)
return
# sgl-kernel 路径
if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
# ... 原有代码 ...
logger.debug_once("CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear")
return
# oneDNN 路径
elif (
ops._supports_onednn
and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC
):
# ... 原有代码 ...
logger.debug_once("CPU unquantized GEMM dispatch: using oneDNN onednn_mm")
return
# fallback 路径
# ... 原有代码 ...
logger.debug_once(
"CPU unquantized GEMM dispatch: using torch.nn.functional.linear (fallback)"
)
评论区精华
- 测试是否应跳过非Zen平台:AndreasKaratzas 询问新测试是否应该跳过非Zen CPU硬件;amd-lalithnc 解释这是一个硬件无关的单元测试,通过mock模拟Zen环境,目的是在CI中也能验证zentorch路径的日志输出,因此不应跳过。结论:保持当前设计。
- 日志级别与统一后端日志:tlrmchlsmth 建议将zentorch路径的info_once改为debug_once,并为所有后端添加类似日志,以便统一分发可见性。amd-lalithnc 最初只添加了zentorch日志,后根据建议为所有路径添加了debug_once日志。结论:所有分发路径都使用logger.debug_once。
- 安装文档简化:tlrmchlsmth 指出安装命令过于复杂、版本过时,并建议使用wheels.vllm.ai索引简化安装步骤。amd-lalithnc 根据反馈改为使用索引和GitHub API获取最新版本。结论:安装命令简化。
- 测试是否应跳过非Zen平台 (testing): 保持当前设计,测试不跳过,在所有CI平台上执行。
- 日志级别与统一后端日志 (design): 所有分发路径都添加了logger.debug_once日志,级别统一为debug。
- 安装文档简化 (documentation): 安装命令简化为使用wheels.vllm.ai索引,并动态获取最新release版本。
风险与影响
关联脉络
参与讨论