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

Signed-off-by: Nick Hill <nickhill123@gmail.com>

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

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] = []