[Perf] Use UVA-backed contents for MRV2 apply_write (#55819)
Signed-off-by: bjf-frz <frz123db@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
bjf-frz committed
Sep 10, 2026 at 22:27 UTC
2e0ee66cab1e7a0fd2ccfb0992a0e4b5e940196d
2 files changed
+166
-5
tests/kernels/core/test_uva.py
+147
@@ -1,10 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
import pytest
5
import torch
6
7
from vllm.utils.platform_utils import is_uva_available
8
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
9
+from vllm.v1.worker.gpu import buffer_utils
10
+from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor
11
12
CUDA_DEVICES = [
13
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
@@ -53,3 +56,147 @@ def test_gpu_write(device):
56
assert cpu_tensor[0, 0] == 2
57
assert cpu_tensor[2, 3] == 4
58
assert cpu_tensor[4, 5] == -2
59
+
60
+
61
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
62
+@pytest.mark.parametrize("device", CUDA_DEVICES)
63
+def test_staged_write_uses_uva_contents_for_uva_target(device, monkeypatch):
64
+ def fail_async_tensor_h2d(*args, **kwargs):
65
+ pytest.fail("UVA-backed targets should not copy write contents to the GPU")
66
+
67
+ monkeypatch.setattr(buffer_utils, "async_tensor_h2d", fail_async_tensor_h2d)
68
+ staged = StagedWriteTensor(
69
+ (3, 4096),
70
+ dtype=torch.int32,
71
+ device=torch.device(device),
72
+ max_concurrency=2,
73
+ uva_instead_of_gpu=True,
74
+ )
75
+
76
+ staged.stage_write(2, 3, [11, 12, 13])
77
+ staged.apply_write()
78
+ torch.accelerator.synchronize()
79
+ staged.stage_write(1, 7, [21, 22])
80
+ staged.apply_write()
81
+ torch.accelerator.synchronize()
82
+ staged.stage_write(0, 1020, range(1500))
83
+ staged.apply_write()
84
+ torch.accelerator.synchronize()
85
+
86
+ assert staged.gpu[2, 3:6].tolist() == [11, 12, 13]
87
+ assert staged.gpu[1, 7:9].tolist() == [21, 22]
88
+ assert staged.gpu[0, 1020:2520].tolist() == list(range(1500))
89
+
90
+
91
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
92
+@pytest.mark.parametrize("input_type", ["list", "numpy", "tensor"])
93
+@pytest.mark.parametrize("size", [(4,), (4, 3)])
94
+def test_uva_pool_overwrites_exposed_prefix(input_type, size):
95
+ """Both slots expose only current contents across growth and shorter reuse."""
96
+ pool = buffer_utils.UvaBufferPool(size, torch.int32, max_concurrency=2)
97
+ for buf in pool._uva_bufs:
98
+ assert tuple(buf.cpu.shape) == size
99
+ assert buf.cpu.is_pinned()
100
+ assert torch.count_nonzero(buf.cpu).item() == 0
101
+ lengths = [0, 0, 3, 3, 4, 4, 5, 3, 3, 5, 1024, 1024, 1025, 1025, 2, 2]
102
+ for step, length in enumerate(lengths):
103
+ if input_type == "list" and len(size) > 1 and length == 0:
104
+ # An empty list has no trailing shape; preserve NumPy's rejection.
105
+ with pytest.raises(ValueError, match="could not broadcast"):
106
+ pool.copy_to_uva([])
107
+ continue
108
+ shape = (length, *size[1:])
109
+ expected = (
110
+ torch.arange(int(np.prod(shape)), dtype=torch.int32, device="cpu").reshape(
111
+ shape
112
+ )
113
+ - step
114
+ )
115
+ values = expected.tolist()
116
+ if input_type == "numpy":
117
+ values = expected.numpy()
118
+ elif input_type == "tensor":
119
+ values = expected
120
+ before = list(pool._uva_bufs)
121
+ slot = (pool._curr + 1) % pool.max_concurrency
122
+ result = pool.copy_to_uva(values)
123
+ assert tuple(result.shape) == shape
124
+ assert pool._curr == slot
125
+ assert pool._uva_bufs[1 - slot] is before[1 - slot]
126
+ if length <= before[slot].cpu.shape[0]:
127
+ assert pool._uva_bufs[slot] is before[slot]
128
+ else:
129
+ assert tuple(pool._uva_bufs[slot].cpu.shape) == (
130
+ 1 << (length - 1).bit_length(),
131
+ *size[1:],
132
+ )
133
+ assert pool.size == size
134
+ # Blocking read also retires GPU readers before the next slot reuse.
135
+ torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0)
136
+
137
+
138
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
139
+@pytest.mark.parametrize("use_out", [False, True])
140
+@pytest.mark.parametrize("input_type", ["numpy", "tensor"])
141
+def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type):
142
+ pool = buffer_utils.UvaBufferPool((2, 3), torch.int32, max_concurrency=2)
143
+ for length in (2, 5, 3, 6):
144
+ expected = torch.arange(length * 3, dtype=torch.int32, device="cpu").reshape(
145
+ length, 3
146
+ )
147
+ values = expected.numpy() if input_type == "numpy" else expected
148
+ out = (
149
+ torch.empty(expected.shape, dtype=torch.int32, device="cuda")
150
+ if use_out
151
+ else None
152
+ )
153
+ result = pool.copy_to_gpu(values, out=out)
154
+ if use_out:
155
+ assert result is out
156
+ torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0)
157
+
158
+
159
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
160
+@pytest.mark.parametrize("uva_target", [False, True])
161
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64, torch.float32])
162
+def test_staged_write_inflight(uva_target, dtype):
163
+ """Preserve every generation until its consumer finishes before slot reuse."""
164
+ device = torch.device("cuda:0")
165
+ with torch.accelerator.device_index(device.index):
166
+ state = StagedWriteTensor(
167
+ (4, 4096),
168
+ dtype,
169
+ device,
170
+ max_concurrency=2,
171
+ uva_instead_of_gpu=uva_target,
172
+ )
173
+ assert (state.write_contents is not None) == uva_target
174
+ stream = torch.cuda.Stream()
175
+ stream.wait_stream(torch.cuda.current_stream())
176
+ pending: list[tuple[torch.cuda.Event, torch.Tensor, torch.Tensor]] = []
177
+ expected = torch.zeros((4, 4096), dtype=dtype, device="cpu")
178
+ for step in range(24):
179
+ if len(pending) == 2:
180
+ event, snapshot, reference = pending.pop(0)
181
+ event.synchronize()
182
+ torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0)
183
+ # Growing and shrinking lengths exercise reallocation and reuse.
184
+ length = [3, 17, 1500, 4090][step % 4]
185
+ row = step % 3
186
+ values = torch.arange(length, dtype=dtype, device="cpu") + step * 8192
187
+ if dtype == torch.float32:
188
+ values += 0.25
189
+ expected[row, 2 : 2 + length] = values
190
+ expected[3, 1:4] = step
191
+ with torch.cuda.stream(stream):
192
+ state.stage_write(row, 2, values.tolist())
193
+ state.stage_write(3, 1, [step] * 3)
194
+ state.apply_write()
195
+ # A GPU consumer observes this generation before the next update.
196
+ snapshot = state.gpu.clone()
197
+ event = torch.cuda.Event()
198
+ event.record(stream)
199
+ pending.append((event, snapshot, expected.clone()))
200
+ for event, snapshot, reference in pending:
201
+ event.synchronize()
202
+ torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0)
vllm/v1/worker/gpu/buffer_utils.py
+19
-5
@@ -51,6 +51,11 @@ class UvaBuffer:
51
52
53
class UvaBufferPool:
54
+ """Preallocate each slot at size, growing its first dimension as needed.
55
+
56
+ Callers must retire a slot's GPU readers before reuse, including growth.
57
+ """
58
+
59
def __init__(
60
self,
61
size: int | Sequence[int],
@@ -72,9 +77,13 @@ class UvaBufferPool:
77
# Round robin to the next buffer.
78
self._curr = (self._curr + 1) % self.max_concurrency
79
buf = self._uva_bufs[self._curr]
80
+ n = len(x)
81
+ if n > buf.cpu.shape[0]:
82
+ capacity = 1 << (n - 1).bit_length()
83
+ buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype)
84
+ self._uva_bufs[self._curr] = buf
85
# CPU-to-CPU copy
86
dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np
77
- n = len(x)
87
dst[:n] = x
88
return buf.uva[:n]
89
@@ -151,6 +160,7 @@ class StagedWriteTensor:
160
self.write_indices = new_buffer(self.num_rows, dtype=torch.int32)
161
self.write_starts = new_buffer(self.num_rows, dtype=torch.int32)
162
self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32)
163
+ self.write_contents = new_buffer(1, dtype=dtype) if uva_instead_of_gpu else None
164
165
def stage_write(
166
self, index: int, start: int, x: Iterable[int] | Iterable[float]
@@ -180,10 +190,14 @@ class StagedWriteTensor:
190
starts_uva = self.write_starts.copy_to_uva(self._staged_write_starts)
191
cu_lens_uva = self.write_cu_lens.copy_to_uva(self._staged_write_cu_lens)
192
183
- # Special handling for write_contents
184
- write_contents = async_tensor_h2d(
185
- self._staged_write_contents, device=self.device, dtype=self.dtype
186
- )
193
+ if self.write_contents is None:
194
+ write_contents = async_tensor_h2d(
195
+ self._staged_write_contents, device=self.device, dtype=self.dtype
196
+ )
197
+ else:
198
+ write_contents = self.write_contents.copy_to_uva(
199
+ self._staged_write_contents
200
+ )
201
202
# Write diffs to the GPU buffer
203
_apply_write_kernel[(n,)](