Prhub

#56908 [Model Runner V2] Tolerate lack of pinned memory support

原始 PR 作者 njhill 合并时间 2026-09-15 12:28 文件变更 2 提交数 1 评论 7 代码增减 +40 / -11

执行摘要

MRV2 增加无 pinned memory 平台回退缓冲

MRV2 使用 UVA 与静态指针,前提是系统支持 pinned memory。PR body 说明:'In MRV2 we use UVA with static pointers which requires pinned memory. This change allows it to work on platforms pinned memory isn't supported/enabled, in particular WSL.' 因此本次变更的核心动机是消除 MRV2 对 pinned memory 的硬依赖,扩展平台覆盖范围。

值得精读。改动虽小,但展示了如何用最小的接口调整(uva 属性改方法)来包装一个平台相关的回退实现,保持上层调用方无感。建议重点关注 NonUvaBuffer 的同步拷贝开销以及 StagedWriteTensor 在无 UVA 时的显存回退策略,这两处决定了新平台的性能与稳定性。若团队计划支持更多无 pinned memory 平台,可在此基础上补充针对回退路径的基准测试。

讨论亮点

该 PR 没有实质性的技术 review 线程:WoosukKwon 直接 approve(空 body)。PR 评论区的 7 条消息全部来自 CI 机器人:两次 /ci run 成功触发 Buildkite 构建;一次因分支落后 main 2 commits 被拒绝,需先 rebase;一次 pre-commit 检查失败,提示运行 pre-commit 并提交修复。最终在 rebase 到最新 main 后 CI 通过。整体流程体现了 vLLM 对分支新鲜度和代码格式的硬性要求,而设计本身没有引发讨论。

实现拆解

  1. 引入平台能力开关:在 vllm/v1/worker/gpu/buffer_utils.py 顶部新增 init_logger 导入,并从 vllm.utils.torch_utils 引入 PIN_MEMORY 全局常量;async_copy_to_gpu 在 PIN_MEMORY 为假时跳过 pin_memory(),避免在无 pinned memory 平台走不可用路径。
  2. 统一 buffer 访问接口:UvaBuffer.uva 由公开属性改为带可选参数 n 的方法,支持按长度截取;UvaBufferPool.copy_to_uva 相应改为 buf.uva(n)。这是为了让真实 UVA 与回退实现暴露相同的方法签名。
  3. 新增 NonUvaBuffer 回退实现:CPU 侧仍用普通内存,uva 则创建在 current_platform.device_type 对应的加速器上;uva() 通过同步 copy 把 CPU 切片拷到设备。构造时通过 logger.warning_once 提示一次回退原因。
  4. 池化与落盘逻辑按平台分支:UvaBufferPool.init 用 is_uva_available() 决定 _buffer_cls,槽位创建与扩容路径都使用该类型;StagedWriteTensor.init 中 uva_instead_of_gpu 与 is_uva_available() 做与运算,保证无 UVA 时“大而冷”的张量直接放 GPU 内存,而不是走不存在的 UVA 路径。
  5. CPU 侧配套简化:vllm/v1/worker/cpu/buffer_utils.py 删除 is_uva_available 检查与 RuntimeError,UvaBuffer.uva 变为返回 CPU 切片的方法;由于 CPU 平台本身没有 UVA,此文件不再需要动态回退逻辑。
  6. 测试与 CI 配套:本次改动没有新增或修改测试文件。PR 经过一次 rebase 使分支包含最新 main 后才通过 Buildkite CI;pre-commit 曾失败一次并在后续提交中修复。
文件 模块 状态 重要度
vllm/v1/worker/gpu/buffer_utils.py 缓冲层 modified 7.24
vllm/v1/worker/cpu/buffer_utils.py 缓冲层 modified 5.43

关键符号

NonUvaBuffer.__init__ NonUvaBuffer.uva UvaBuffer.uva UvaBufferPool.__init__ UvaBufferPool.copy_to_uva StagedWriteTensor.__init__ async_copy_to_gpu

关键源码片段

vllm/v1/worker/gpu/buffer_utils.py core-logic

核心逻辑改动:新增 NonUvaBuffer 显式拷贝回退,UvaBufferPool 按平台选择缓冲类型,uva 由属性改为方法,StagedWriteTensor 增加 is_uva_available 保护。

# vllm/v1/worker/gpu/buffer_utils.py(整理后的关键实现)
​
class NonUvaBuffer:
    """无 pinned memory 平台使用的显式拷贝回退实现。"""
​
    def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
        from vllm.platforms import current_platform
​
        logger.warning_once(
            "Pinned memory is not available on this platform; falling back "
            "to device memory for UVA buffers."
        )
        # 源数据仍是普通 CPU 内存,与 UvaBuffer 保持一致
        self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
        self.np = self.cpu.numpy()
        # 设备侧缓冲放在 accelerator 上,等待显式拷贝
        self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type)
​
    def uva(self, n: int | None = None) -> torch.Tensor:
        # 没有 UVA 时只能用同步 copy_ 把 CPU 数据搬到设备
        if n is None:
            return self._uva.copy_(self.cpu)
        return self._uva[:n].copy_(self.cpu[:n])
​
​
class UvaBufferPool:
    """Preallocate each slot at size, growing its first dimension as needed.
​
    Callers must retire a slot's GPU readers before reuse, including growth.
    """
​
    def __init__(
        self,
        size: int | Sequence[int],
        dtype: torch.dtype,
        max_concurrency: int | None = None,
    ):
        if max_concurrency is None:
            max_concurrency = _DEFAULT_MAX_CONCURRENCY
        self.size = size
        self.dtype = dtype
        self.max_concurrency = max_concurrency
​
        # 依据平台能力选择真实 UVA 或显式拷贝回退,对上层调用方透明
        self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer
        self._uva_bufs = [
            self._buffer_cls(size, dtype) for _ in range(max_concurrency)
        ]
        self._curr = 0
​
    def copy_to_uva(self, x: torch.Tensor | np.ndarray | list) -> torch.Tensor:
        # 轮转使用下一个槽位
        self._curr = (self._curr + 1) % self.max_concurrency
        buf = self._uva_bufs[self._curr]
        n = len(x)
        if n > buf.cpu.shape[0]:
            # 按 2 的幂扩容,扩容时仍使用相同的 buffer 类型
            capacity = 1 << (n - 1).bit_length()
            buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype)
            self._uva_bufs[self._curr] = buf
        # 先把数据写入 CPU 侧缓冲区
        dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np
        dst[:n] = x
        # 返回设备视图(UVA 零拷贝)或显式拷贝结果(NonUvaBuffer)
        return buf.uva(n)
vllm/v1/worker/cpu/buffer_utils.py dependency-wiring

CPU 平台配套简化:移除 is_uva_available 检查与 RuntimeError,UvaBuffer 不再需要平台判断,uva 改为方法保持接口一致。

# vllm/v1/worker/cpu/buffer_utils.py(整理后的关键实现)
​
class UvaBuffer:
    def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
        # CPU 平台没有真实 UVA,这里 UvaBuffer 就是普通 CPU 张量
        self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
        self.np = self.cpu.numpy()
​
    def uva(self, n: int | None = None):
        # 保持与 GPU 版一致的接口,支持按行数截取
        return self.cpu[:n] if n is not None else self.cpu

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

1) NonUvaBuffer 的 uva() 每次执行同步 copy_,会把 CPU 到设备拷贝串行化;在需要频繁访问缓冲的调度路径上(如 write_indices、write_starts)可能成为性能瓶颈,但该路径只在无 UVA 平台启用,WSL 场景下可接受。
2) StagedWriteTensor 在无 UVA 时将 uva_instead_of_gpu 显式置 False,大而冷的张量直接占用 GPU 显存,可能增大显存压力,属于与同步拷贝之间的权衡。
3) UvaBuffer.uva 由属性改为方法,所有调用方都必须改为 buf.uva(n) 形式;若外部代码仍以属性方式访问会报错,需确保仓库内无遗漏调用点。
4) 没有配套单元测试,后续若有平台行为变化,回退路径容易被回归。
5) CPU buffer_utils 删除了 is_uva_available 检查,语义变为“CPU 平台没有真实 UVA”,与 GPU 版接口保持一致,但若未来引入其他平台直接复用 CPU 版,需重新评估。

对用户:WSL、未启用 pinned memory 的 CPU 环境(如部分容器)可以运行 MRV2,但会付出显式拷贝的同步开销。对系统:gpu/buffer_utils.py 的缓冲层从单一 UVA 实现变为“UVA / NonUva 双实现”,池化与扩容逻辑通过 _buffer_cls 统一,接口收敛为 uva(n) 方法。对团队:MRV2 的平台兼容性进一步扩展,与近期 MRV2 系列 PR(PP pooling、acceptance estimation、Proton profiler)形成连续的演进脉络。

无配套测试 同步拷贝性能开销 平台行为分叉

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论