iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Use UVA-backed contents for MRV2 apply_write (#55819)

Browse Source 
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
bjf-frz
2026-09-10 07:27:16 -07:00
committed by GitHub
co-authored by Nick Hill
parent 7d8d71e989
commit 2e0ee66cab
2 changed files with 166 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/core/test_uva.py
+147
View File
@@ -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)
vllm/v1/worker/gpu/buffer_utils.py
+19 -5
View File
@@ -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,)](