From 2e0ee66cab1e7a0fd2ccfb0992a0e4b5e940196d Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Thu, 10 Sep 2026 22:27:16 +0800 Subject: [PATCH] [Perf] Use UVA-backed contents for MRV2 apply_write (#55819) Signed-off-by: bjf-frz Co-authored-by: Nick Hill --- tests/kernels/core/test_uva.py | 147 +++++++++++++++++++++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 24 ++++- 2 files changed, 166 insertions(+), 5 deletions(-) diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index 7c25612500..f8c297afc6 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor +from vllm.v1.worker.gpu import buffer_utils +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor CUDA_DEVICES = [ 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): assert cpu_tensor[0, 0] == 2 assert cpu_tensor[2, 3] == 4 assert cpu_tensor[4, 5] == -2 + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("device", CUDA_DEVICES) +def test_staged_write_uses_uva_contents_for_uva_target(device, monkeypatch): + def fail_async_tensor_h2d(*args, **kwargs): + pytest.fail("UVA-backed targets should not copy write contents to the GPU") + + monkeypatch.setattr(buffer_utils, "async_tensor_h2d", fail_async_tensor_h2d) + staged = StagedWriteTensor( + (3, 4096), + dtype=torch.int32, + device=torch.device(device), + max_concurrency=2, + uva_instead_of_gpu=True, + ) + + staged.stage_write(2, 3, [11, 12, 13]) + staged.apply_write() + torch.accelerator.synchronize() + staged.stage_write(1, 7, [21, 22]) + staged.apply_write() + torch.accelerator.synchronize() + staged.stage_write(0, 1020, range(1500)) + staged.apply_write() + torch.accelerator.synchronize() + + assert staged.gpu[2, 3:6].tolist() == [11, 12, 13] + assert staged.gpu[1, 7:9].tolist() == [21, 22] + assert staged.gpu[0, 1020:2520].tolist() == list(range(1500)) + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("input_type", ["list", "numpy", "tensor"]) +@pytest.mark.parametrize("size", [(4,), (4, 3)]) +def test_uva_pool_overwrites_exposed_prefix(input_type, size): + """Both slots expose only current contents across growth and shorter reuse.""" + pool = buffer_utils.UvaBufferPool(size, torch.int32, max_concurrency=2) + for buf in pool._uva_bufs: + assert tuple(buf.cpu.shape) == size + assert buf.cpu.is_pinned() + assert torch.count_nonzero(buf.cpu).item() == 0 + lengths = [0, 0, 3, 3, 4, 4, 5, 3, 3, 5, 1024, 1024, 1025, 1025, 2, 2] + for step, length in enumerate(lengths): + if input_type == "list" and len(size) > 1 and length == 0: + # An empty list has no trailing shape; preserve NumPy's rejection. + with pytest.raises(ValueError, match="could not broadcast"): + pool.copy_to_uva([]) + continue + shape = (length, *size[1:]) + expected = ( + torch.arange(int(np.prod(shape)), dtype=torch.int32, device="cpu").reshape( + shape + ) + - step + ) + values = expected.tolist() + if input_type == "numpy": + values = expected.numpy() + elif input_type == "tensor": + values = expected + before = list(pool._uva_bufs) + slot = (pool._curr + 1) % pool.max_concurrency + result = pool.copy_to_uva(values) + assert tuple(result.shape) == shape + assert pool._curr == slot + assert pool._uva_bufs[1 - slot] is before[1 - slot] + if length <= before[slot].cpu.shape[0]: + assert pool._uva_bufs[slot] is before[slot] + else: + assert tuple(pool._uva_bufs[slot].cpu.shape) == ( + 1 << (length - 1).bit_length(), + *size[1:], + ) + assert pool.size == size + # Blocking read also retires GPU readers before the next slot reuse. + torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("use_out", [False, True]) +@pytest.mark.parametrize("input_type", ["numpy", "tensor"]) +def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type): + pool = buffer_utils.UvaBufferPool((2, 3), torch.int32, max_concurrency=2) + for length in (2, 5, 3, 6): + expected = torch.arange(length * 3, dtype=torch.int32, device="cpu").reshape( + length, 3 + ) + values = expected.numpy() if input_type == "numpy" else expected + out = ( + torch.empty(expected.shape, dtype=torch.int32, device="cuda") + if use_out + else None + ) + result = pool.copy_to_gpu(values, out=out) + if use_out: + assert result is out + torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("uva_target", [False, True]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64, torch.float32]) +def test_staged_write_inflight(uva_target, dtype): + """Preserve every generation until its consumer finishes before slot reuse.""" + device = torch.device("cuda:0") + with torch.accelerator.device_index(device.index): + state = StagedWriteTensor( + (4, 4096), + dtype, + device, + max_concurrency=2, + uva_instead_of_gpu=uva_target, + ) + assert (state.write_contents is not None) == uva_target + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + pending: list[tuple[torch.cuda.Event, torch.Tensor, torch.Tensor]] = [] + expected = torch.zeros((4, 4096), dtype=dtype, device="cpu") + for step in range(24): + if len(pending) == 2: + event, snapshot, reference = pending.pop(0) + event.synchronize() + torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0) + # Growing and shrinking lengths exercise reallocation and reuse. + length = [3, 17, 1500, 4090][step % 4] + row = step % 3 + values = torch.arange(length, dtype=dtype, device="cpu") + step * 8192 + if dtype == torch.float32: + values += 0.25 + expected[row, 2 : 2 + length] = values + expected[3, 1:4] = step + with torch.cuda.stream(stream): + state.stage_write(row, 2, values.tolist()) + state.stage_write(3, 1, [step] * 3) + state.apply_write() + # A GPU consumer observes this generation before the next update. + snapshot = state.gpu.clone() + event = torch.cuda.Event() + event.record(stream) + pending.append((event, snapshot, expected.clone())) + for event, snapshot, reference in pending: + event.synchronize() + torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index f8336fa074..2c1dac2d07 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -51,6 +51,11 @@ class UvaBuffer: 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], @@ -72,9 +77,13 @@ class UvaBufferPool: # Round robin to the next buffer. self._curr = (self._curr + 1) % self.max_concurrency buf = self._uva_bufs[self._curr] + n = len(x) + if n > buf.cpu.shape[0]: + capacity = 1 << (n - 1).bit_length() + buf = UvaBuffer((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 - n = len(x) dst[:n] = x return buf.uva[:n] @@ -151,6 +160,7 @@ class StagedWriteTensor: self.write_indices = new_buffer(self.num_rows, dtype=torch.int32) self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) + self.write_contents = new_buffer(1, dtype=dtype) if uva_instead_of_gpu else None def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] @@ -180,10 +190,14 @@ class StagedWriteTensor: starts_uva = self.write_starts.copy_to_uva(self._staged_write_starts) cu_lens_uva = self.write_cu_lens.copy_to_uva(self._staged_write_cu_lens) - # Special handling for write_contents - write_contents = async_tensor_h2d( - self._staged_write_contents, device=self.device, dtype=self.dtype - ) + if self.write_contents is None: + write_contents = async_tensor_h2d( + self._staged_write_contents, device=self.device, dtype=self.dtype + ) + else: + write_contents = self.write_contents.copy_to_uva( + self._staged_write_contents + ) # Write diffs to the GPU buffer _apply_write_kernel[(n,)](