[Model Runner V2] Tolerate lack of pinned memory support (#56908)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Sep 14, 2026 at 21:28 UTC
f84b0c4bceee120936030dfb0c90af0a818263ed
2 files changed
+40
-11
vllm/v1/worker/cpu/buffer_utils.py
+3
-5
@@ -4,13 +4,11 @@ from collections.abc import Sequence
4
5
import torch
6
7
-from vllm.utils.platform_utils import is_uva_available
8
-
7
8
class UvaBuffer:
9
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
12
- if not is_uva_available():
13
- raise RuntimeError("UVA is not available")
10
self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
11
self.np = self.cpu.numpy()
16
- self.uva = self.cpu
12
+
13
+ def uva(self, n: int | None = None):
14
+ return self.cpu[:n] if n is not None else self.cpu
vllm/v1/worker/gpu/buffer_utils.py
+37
-6
@@ -6,13 +6,17 @@ from functools import partial
6
import numpy as np
7
import torch
8
9
+from vllm.logger import init_logger
10
from vllm.triton_utils import tl, triton
11
from vllm.utils.platform_utils import is_uva_available
12
from vllm.utils.torch_utils import (
13
+ PIN_MEMORY,
14
async_tensor_h2d,
15
get_accelerator_view_from_cpu_tensor,
16
)
17
18
+logger = init_logger(__name__)
19
+
20
# Default round-robin depth for the UVA buffer pools. Must be >= the number of
21
# concurrent in-flight steps (engine batch_queue_size).
22
_DEFAULT_MAX_CONCURRENCY = 2
@@ -37,7 +41,7 @@ def async_copy_to_gpu(
41
out = torch.empty_like(x, device=device)
42
43
# pin_memory() is no-op if the memory is already pinned.
40
- pinned = x.pin_memory()
44
+ pinned = x.pin_memory() if PIN_MEMORY else x
45
return out.copy_(pinned, non_blocking=True)
46
47
@@ -47,7 +51,30 @@ class UvaBuffer:
51
raise RuntimeError("UVA is not available")
52
self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
53
self.np = self.cpu.numpy()
50
- self.uva = get_accelerator_view_from_cpu_tensor(self.cpu)
54
+ self._uva = get_accelerator_view_from_cpu_tensor(self.cpu)
55
+
56
+ def uva(self, n: int | None = None) -> torch.Tensor:
57
+ return self._uva[:n] if n is not None else self._uva
58
+
59
+
60
+class NonUvaBuffer:
61
+ """Explicit-copy fallback for platforms without pinned memory."""
62
+
63
+ def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
64
+ from vllm.platforms import current_platform
65
+
66
+ logger.warning_once(
67
+ "Pinned memory is not available on this platform; falling back "
68
+ "to device memory for UVA buffers."
69
+ )
70
+ self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
71
+ self.np = self.cpu.numpy()
72
+ self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type)
73
+
74
+ def uva(self, n: int | None = None) -> torch.Tensor:
75
+ if n is None:
76
+ return self._uva.copy_(self.cpu)
77
+ return self._uva[:n].copy_(self.cpu[:n])
78
79
80
class UvaBufferPool:
@@ -69,7 +96,8 @@ class UvaBufferPool:
96
self.max_concurrency = max_concurrency
97
98
# UVA buffers for concurrency
72
- self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)]
99
+ self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer
100
+ self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)]
101
# Current buffer index
102
self._curr = 0
103
@@ -80,12 +108,12 @@ class UvaBufferPool:
108
n = len(x)
109
if n > buf.cpu.shape[0]:
110
capacity = 1 << (n - 1).bit_length()
83
- buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype)
111
+ buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype)
112
self._uva_bufs[self._curr] = buf
113
# CPU-to-CPU copy
114
dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np
115
dst[:n] = x
88
- return buf.uva[:n]
116
+ return buf.uva(n)
117
118
def copy_to_gpu(
119
self,
@@ -141,6 +169,9 @@ class StagedWriteTensor:
169
self.device = device
170
self.max_concurrency = max_concurrency
171
172
+ # Without UVA, large-but-cold tensors live on the GPU directly.
173
+ uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available()
174
+
175
if not uva_instead_of_gpu:
176
# Create a GPU tensor (default)
177
self.gpu = torch.zeros(size, dtype=dtype, device=device)
@@ -148,7 +179,7 @@ class StagedWriteTensor:
179
# For a large but not-frequently-accessed tensor, we can use UVA instead of
180
# GPU to save GPU memory
181
self._uva_buf = UvaBuffer(size, dtype)
151
- self.gpu = self._uva_buf.uva
182
+ self.gpu = self._uva_buf.uva()
183
184
self._staged_write_indices: list[int] = []
185
self._staged_write_starts: list[int] = []