Prhub

#50126 [ROCm] Enable pinned memory on supported WSL2 kernels

原始 PR 作者 fcui-amd 合并时间 2026-08-07 23:43 文件变更 2 提交数 5 评论 9 代码增减 +29 / -3

执行摘要

ROCm 在 WSL2 下按内核版本启用 pinned memory

PR body 明确指出:RocmPlatform 未覆盖 is_pin_memory_available(),ROCm 构建在 WSL 下运行时总是回退到保守的基类实现并无条件禁用 pinned memory,造成性能损失。同时基类中的 WSL 警告若使用 warning_once() 会在 WSL 下触发循环导入(issue #48397),需要对齐 #48444 中已采用的 logger.warning() 模式。

值得快速阅读:若是 ROCm/WSL 相关开发者,可重点看 is_pin_memory_available() 的门控逻辑与 warning_once() 循环导入注释;若要评审,建议提醒作者将 XPU CI 脚本改动拆到独立 PR,并补充 _get_wsl_kernel_version() 的单测。

讨论亮点

该 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()

实现拆解

本变更的核心是让 ROCm 平台在 WSL 下具备与 CUDA 平台一致的 pinned memory 判定能力,分为以下几步:

  1. 平台探测层新增内核版本解析:在 vllm/platforms/rocm.py 中新增 import platform,并实现 @cache 装饰的 _get_wsl_kernel_version(),将 platform.uname().release 解析为 (major, minor, patch) 元组;解析失败返回 None,供调用方走保守路径。
  2. 覆盖 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 下的行为一致。
  3. 混入的 XPU CI 脚本修复:.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh 中为 docker manifest rm/create 的镜像引用加上双引号,防止 shell 分词,属于与主题无关的独立改动。
  4. 测试与验证:本次未新增测试文件,PR body 声明通过运行 python -c "...is_pin_memory_available()..."、已有测试 tests/rocm/test_platform.py 以及 pre-commit 校验完成验证;CI 由 maintainer 在合并前触发 Buildkite 构建确认通过。
文件 模块 状态 重要度
vllm/platforms/rocm.py 平台层 modified 7.16
.buildkite/scripts/xpu/create-xpu-ecr-manifest.sh CI 脚本 modified 2.32

关键符号

is_pin_memory_available _get_wsl_kernel_version

关键源码片段

vllm/platforms/rocm.py core-logic

核心源码改动:新增 _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

评论区精华

WSL 下 warning_once 循环导入风险 正确性

PR 在 WSL 警告路径刻意使用 logger.warning() 而非 warning_once(),防止 is_pin_memory_available() 在 torch_utils 模块加载中途被调用时触发 parallel_state 惰性导入形成的循环导入(issue #48397 的根因)。

结论:沿用 #48444 的修复方案:仅本调用点退回 logger.warning(),不动共享的 warning_once 去重机制。 · 已解决

风险与影响

风险点集中在 rocm.py 的新逻辑:

  • 回归风险:is_pin_memory_available() 现在返回 True 的范围扩大(新内核 WSL),若某些 WSL2 环境下 4.19.121 门槛判定不准确,可能错误启用 pinned memory 导致异常或性能问题;但该门控与 CudaPlatformBase 完全一致,实际回归概率低。
  • 解析健壮性:_get_wsl_kernel_version() 对 release 字符串解析失败时返回 None 并走保守的 False 路径,安全方向正确。
  • 导入顺序:rocm.py 新增 from .interface import in_wsl,该符号在基类中已存在,不引入新的循环依赖风险。
  • 混入无关改动:XPU CI 脚本的引号修复与本 PR 主题无关,虽属低风险正确性修复,但增加了审查噪音。
  • 测试缺口:未新增针对新函数的单测文件,依赖已有 tests/rocm/test_platform.py,对新分支(解析失败、老内核、新内核)的覆盖不明确。

影响范围限定在 ROCm + WSL2 使用者:内核版本不低于 4.19.121 的环境将从禁用 pinned memory 变为启用,可消除访存性能损失;老内核或无法解析内核版本的环境维持原有保守行为。对 CUDA、CPU 等其他平台无影响,因为改动仅作用于 RocmPlatform。对团队而言风险低、改动量小,但为以后 ROCm 平台能力补齐(与 CudaPlatformBase 对齐)建立了模式。

平台核心路径变更 缺少新增测试覆盖 混入无关 CI 改动

关联 Issue

#48397 Circular import on WSL when importing vLLM v0.25.0 (warning_once in is_pin_memory_available)
#48444 [Bugfix] Fix WSL circular import from pin_memory warning_once

完整报告

参与讨论