执行摘要
- 一句话:ROCm 在 WSL2 下按内核版本启用 pinned memory
- 推荐动作:值得快速阅读:若是 ROCm/WSL 相关开发者,可重点看 is_pin_memory_available() 的门控逻辑与 warning_once() 循环导入注释;若要评审,建议提醒作者将 XPU CI 脚本改动拆到独立 PR,并补充 _get_wsl_kernel_version() 的单测。
功能与动机
PR body 明确指出:RocmPlatform 未覆盖 is_pin_memory_available(),ROCm 构建在 WSL 下运行时总是回退到保守的基类实现并无条件禁用 pinned memory,造成性能损失。同时基类中的 WSL 警告若使用 warning_once() 会在 WSL 下触发循环导入(issue #48397),需要对齐 #48444 中已采用的 logger.warning() 模式。
实现拆解
本变更的核心是让 ROCm 平台在 WSL 下具备与 CUDA 平台一致的 pinned memory 判定能力,分为以下几步:
- 平台探测层新增内核版本解析:在 vllm/platforms/rocm.py 中新增 import platform,并实现 @cache 装饰的 _get_wsl_kernel_version(),将
platform.uname().release 解析为 (major, minor, patch) 元组;解析失败返回 None,供调用方走保守路径。
- 覆盖 is_pin_memory_available():在 RocmPlatform 上新增类方法,从 .interface 导入 in_wsl 判断是否为 WSL 环境;若为 WSL 且内核版本低于 4.19.121(或解析失败),记录 logger.warning() 并返回 False,否则返回 True。日志特意使用 warning() 而非 warning_once(),避免在模块加载中途触发 parallel_state 惰性导入导致循环导入。该逻辑与 CudaPlatformBase 完全对齐,保证 ROCm 与 CUDA 在 WSL 下的行为一致。
- 混入的 XPU CI 脚本修复:.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh 中为 docker manifest rm/create 的镜像引用加上双引号,防止 shell 分词,属于与主题无关的独立改动。
- 测试与验证:本次未新增测试文件,PR body 声明通过运行
python -c "...is_pin_memory_available()..."、已有测试 tests/rocm/test_platform.py 以及 pre-commit 校验完成验证;CI 由 maintainer 在合并前触发 Buildkite 构建确认通过。
关键文件:
vllm/platforms/rocm.py(模块 平台层;类别 source;类型 core-logic;符号 _get_wsl_kernel_version, is_pin_memory_available): 核心源码改动:新增 _get_wsl_kernel_version() 与覆盖 is_pin_memory_available(),使 ROCm 平台在 WSL2 新内核下启用 pinned memory,并与 CudaPlatformBase 行为对齐,同时规避 WSL 循环导入。
.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh(模块 CI 脚本;类别 infra;类型 bugfix): 混入的无关 CI 脚本改动:为 docker manifest rm/create 的镜像引用加双引号,防止 shell 分词,属于低风险的 shell 正确性修复。
关键符号:is_pin_memory_available, _get_wsl_kernel_version
关键源码片段
vllm/platforms/rocm.py
核心源码改动:新增 _get_wsl_kernel_version() 与覆盖 is_pin_memory_available(),使 ROCm 平台在 WSL2 新内核下启用 pinned memory,并与 CudaPlatformBase 行为对齐,同时规避 WSL 循环导入。
# rocm.py 中新增的 WSL 内核版本解析辅助函数
@cache
def _get_wsl_kernel_version() -> tuple[int, ...] | None:
# 将 uname 的 release 字段解析为 (major, minor, patch) 元组,
# 例如 "5.15.90.1-microsoft-standard-WSL2" → (5, 15, 90)。
# 解析失败返回 None,调用方会走保守路径禁用 pinned memory。
try:
release = platform.uname().release
parts = release.split("-")[0].split(".")
return tuple(int(part) for part in parts[:3])
except (TypeError, ValueError):
return None
@classmethod
def is_pin_memory_available(cls) -> bool:
# 覆盖基类 Platform.is_pin_memory_available():此前 ROCm 构建在 WSL
# 下始终回退到基类的无条件 False,导致 pinned memory 被永久禁用。
# 这里复用 CudaPlatformBase 的内核版本门控(4.19.121 起 WSL2 内核
# 才具备稳定支持 pin_memory 的语义)。
if in_wsl():
version = _get_wsl_kernel_version()
if version is None or version < (4, 19, 121):
# 必须使用 logger.warning() 而不是 warning_once():后者会触发
# vllm.distributed.parallel_state 的惰性导入,而本函数可能在
# vllm.utils.torch_utils 模块加载中途被调用,造成循环导入
# (见 issue #48397,修复见 PR #48444)。
logger.warning(
"Using 'pin_memory=False' as WSL is detected and the "
"WSL2 kernel version is below 4.19.121. This may slow "
"down performance. Please run `wsl --update`."
)
return False
return True
评论区精华
该 PR 的评审过程非常简短,没有代码级评论:
- claude[bot] 指出 fork 来源的 PR 默认禁用自动审查,可手动触发;
- dllehr-amd 直接批准("This looks good to me!")。
技术层面的关键讨论实际来自关联 issue #48397/#48444:is_pin_memory_available() 可能在 vllm.utils.torch_utils 模块加载中途被调用(PIN_MEMORY 模块级变量),而 warning_once() 会经 _should_log_with_scope() 惰性导入 vllm.distributed.parallel_state,后者又反向导入仍在初始化中的 torch_utils,形成循环导入。该 PR 沿用 #48444 的结论,在 WSL 警告路径使用普通 logger.warning()。
- WSL 下 warning_once 循环导入风险 (correctness): 沿用 #48444 的修复方案:仅本调用点退回 logger.warning(),不动共享的 warning_once 去重机制。
风险与影响
关联脉络
- PR #48444 [Bugfix] Fix WSL circular import from pin_memory warning_once: 本 PR 的 WSL 警告日志正是对齐 #48444 的修复结论(warning_once 改回 warning),且明确引用了该 PR 对应 issue #48397。
参与讨论