From f84b0c4bceee120936030dfb0c90af0a818263ed Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 21:28:32 -0700 Subject: [PATCH] [Model Runner V2] Tolerate lack of pinned memory support (#56908) Signed-off-by: Nick Hill --- vllm/v1/worker/cpu/buffer_utils.py | 8 +++--- vllm/v1/worker/gpu/buffer_utils.py | 43 +++++++++++++++++++++++++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/cpu/buffer_utils.py b/vllm/v1/worker/cpu/buffer_utils.py index 8ee257a28f..cdc76ab8b4 100644 --- a/vllm/v1/worker/cpu/buffer_utils.py +++ b/vllm/v1/worker/cpu/buffer_utils.py @@ -4,13 +4,11 @@ from collections.abc import Sequence import torch -from vllm.utils.platform_utils import is_uva_available - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): - if not is_uva_available(): - raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu") self.np = self.cpu.numpy() - self.uva = self.cpu + + def uva(self, n: int | None = None): + return self.cpu[:n] if n is not None else self.cpu diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2c1dac2d07..0066f80907 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -6,13 +6,17 @@ from functools import partial import numpy as np import torch +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( + PIN_MEMORY, async_tensor_h2d, get_accelerator_view_from_cpu_tensor, ) +logger = init_logger(__name__) + # Default round-robin depth for the UVA buffer pools. Must be >= the number of # concurrent in-flight steps (engine batch_queue_size). _DEFAULT_MAX_CONCURRENCY = 2 @@ -37,7 +41,7 @@ def async_copy_to_gpu( out = torch.empty_like(x, device=device) # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() + pinned = x.pin_memory() if PIN_MEMORY else x return out.copy_(pinned, non_blocking=True) @@ -47,7 +51,30 @@ class UvaBuffer: raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True) self.np = self.cpu.numpy() - self.uva = get_accelerator_view_from_cpu_tensor(self.cpu) + self._uva = get_accelerator_view_from_cpu_tensor(self.cpu) + + def uva(self, n: int | None = None) -> torch.Tensor: + return self._uva[:n] if n is not None else self._uva + + +class NonUvaBuffer: + """Explicit-copy fallback for platforms without 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." + ) + self.cpu = torch.zeros(size, dtype=dtype, device="cpu") + self.np = self.cpu.numpy() + self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type) + + def uva(self, n: int | None = None) -> torch.Tensor: + if n is None: + return self._uva.copy_(self.cpu) + return self._uva[:n].copy_(self.cpu[:n]) class UvaBufferPool: @@ -69,7 +96,8 @@ class UvaBufferPool: self.max_concurrency = max_concurrency # UVA buffers for concurrency - self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)] + self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer + self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)] # Current buffer index self._curr = 0 @@ -80,12 +108,12 @@ class UvaBufferPool: n = len(x) if n > buf.cpu.shape[0]: capacity = 1 << (n - 1).bit_length() - buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype) + buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype) self._uva_bufs[self._curr] = buf # CPU-to-CPU copy dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np dst[:n] = x - return buf.uva[:n] + return buf.uva(n) def copy_to_gpu( self, @@ -141,6 +169,9 @@ class StagedWriteTensor: self.device = device self.max_concurrency = max_concurrency + # Without UVA, large-but-cold tensors live on the GPU directly. + uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available() + if not uva_instead_of_gpu: # Create a GPU tensor (default) self.gpu = torch.zeros(size, dtype=dtype, device=device) @@ -148,7 +179,7 @@ class StagedWriteTensor: # For a large but not-frequently-accessed tensor, we can use UVA instead of # GPU to save GPU memory self._uva_buf = UvaBuffer(size, dtype) - self.gpu = self._uva_buf.uva + self.gpu = self._uva_buf.uva() self._staged_write_indices: list[int] = [] self._staged_write_starts: list[int] = []