| 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.platforms import current_platform |
| 8 | from vllm.utils.platform_utils import is_uva_available |
| 9 | from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor |
| 10 | from vllm.v1.worker.gpu import buffer_utils |
| 11 | from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor |
| 12 | |
| 13 | DEVICE_TYPE = current_platform.device_type |
| 14 | DEVICES = [ |
| 15 | f"{DEVICE_TYPE}:{i}" |
| 16 | for i in range(1 if torch.accelerator.device_count() == 1 else 2) |
| 17 | ] |
| 18 | |
| 19 | |
| 20 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 21 | @pytest.mark.parametrize("device", DEVICES) |
| 22 | def test_cpu_write(device): |
| 23 | torch.set_default_device(device) |
| 24 | cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) |
| 25 | gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) |
| 26 | assert gpu_view.device.type == DEVICE_TYPE |
| 27 | |
| 28 | assert gpu_view[0, 0] == 0 |
| 29 | assert gpu_view[2, 3] == 0 |
| 30 | assert gpu_view[4, 5] == 0 |
| 31 | |
| 32 | cpu_tensor[0, 0] = 1 |
| 33 | cpu_tensor[2, 3] = 2 |
| 34 | cpu_tensor[4, 5] = -1 |
| 35 | |
| 36 | gpu_view.mul_(2) |
| 37 | assert gpu_view[0, 0] == 2 |
| 38 | assert gpu_view[2, 3] == 4 |
| 39 | assert gpu_view[4, 5] == -2 |
| 40 | |
| 41 | |
| 42 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 43 | @pytest.mark.parametrize("device", DEVICES) |
| 44 | def test_gpu_write(device): |
| 45 | torch.set_default_device(device) |
| 46 | cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) |
| 47 | gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) |
| 48 | assert gpu_view.device.type == DEVICE_TYPE |
| 49 | |
| 50 | assert gpu_view[0, 0] == 0 |
| 51 | assert gpu_view[2, 3] == 0 |
| 52 | assert gpu_view[4, 5] == 0 |
| 53 | |
| 54 | gpu_view[0, 0] = 1 |
| 55 | gpu_view[2, 3] = 2 |
| 56 | gpu_view[4, 5] = -1 |
| 57 | gpu_view.mul_(2) |
| 58 | |
| 59 | assert cpu_tensor[0, 0] == 2 |
| 60 | assert cpu_tensor[2, 3] == 4 |
| 61 | assert cpu_tensor[4, 5] == -2 |
| 62 | |
| 63 | |
| 64 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 65 | @pytest.mark.parametrize("device", DEVICES) |
| 66 | def test_staged_write_uses_uva_contents_for_uva_target(device, monkeypatch): |
| 67 | def fail_async_tensor_h2d(*args, **kwargs): |
| 68 | pytest.fail("UVA-backed targets should not copy write contents to the GPU") |
| 69 | |
| 70 | monkeypatch.setattr(buffer_utils, "async_tensor_h2d", fail_async_tensor_h2d) |
| 71 | staged = StagedWriteTensor( |
| 72 | (3, 4096), |
| 73 | dtype=torch.int32, |
| 74 | device=torch.device(device), |
| 75 | max_concurrency=2, |
| 76 | uva_instead_of_gpu=True, |
| 77 | ) |
| 78 | |
| 79 | staged.stage_write(2, 3, [11, 12, 13]) |
| 80 | staged.apply_write() |
| 81 | torch.accelerator.synchronize() |
| 82 | staged.stage_write(1, 7, [21, 22]) |
| 83 | staged.apply_write() |
| 84 | torch.accelerator.synchronize() |
| 85 | staged.stage_write(0, 1020, range(1500)) |
| 86 | staged.apply_write() |
| 87 | torch.accelerator.synchronize() |
| 88 | |
| 89 | assert staged.gpu[2, 3:6].tolist() == [11, 12, 13] |
| 90 | assert staged.gpu[1, 7:9].tolist() == [21, 22] |
| 91 | assert staged.gpu[0, 1020:2520].tolist() == list(range(1500)) |
| 92 | |
| 93 | |
| 94 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 95 | @pytest.mark.parametrize("input_type", ["list", "numpy", "tensor"]) |
| 96 | @pytest.mark.parametrize("size", [(4,), (4, 3)]) |
| 97 | def test_uva_pool_overwrites_exposed_prefix(input_type, size): |
| 98 | """Both slots expose only current contents across growth and shorter reuse.""" |
| 99 | pool = buffer_utils.UvaBufferPool(size, torch.int32, max_concurrency=2) |
| 100 | for buf in pool._uva_bufs: |
| 101 | assert tuple(buf.cpu.shape) == size |
| 102 | assert buf.cpu.is_pinned() |
| 103 | assert torch.count_nonzero(buf.cpu).item() == 0 |
| 104 | lengths = [0, 0, 3, 3, 4, 4, 5, 3, 3, 5, 1024, 1024, 1025, 1025, 2, 2] |
| 105 | for step, length in enumerate(lengths): |
| 106 | if input_type == "list" and len(size) > 1 and length == 0: |
| 107 | # An empty list has no trailing shape; preserve NumPy's rejection. |
| 108 | with pytest.raises(ValueError, match="could not broadcast"): |
| 109 | pool.copy_to_uva([]) |
| 110 | continue |
| 111 | shape = (length, *size[1:]) |
| 112 | expected = ( |
| 113 | torch.arange(int(np.prod(shape)), dtype=torch.int32, device="cpu").reshape( |
| 114 | shape |
| 115 | ) |
| 116 | - step |
| 117 | ) |
| 118 | values = expected.tolist() |
| 119 | if input_type == "numpy": |
| 120 | values = expected.numpy() |
| 121 | elif input_type == "tensor": |
| 122 | values = expected |
| 123 | before = list(pool._uva_bufs) |
| 124 | slot = (pool._curr + 1) % pool.max_concurrency |
| 125 | result = pool.copy_to_uva(values) |
| 126 | assert tuple(result.shape) == shape |
| 127 | assert pool._curr == slot |
| 128 | assert pool._uva_bufs[1 - slot] is before[1 - slot] |
| 129 | if length <= before[slot].cpu.shape[0]: |
| 130 | assert pool._uva_bufs[slot] is before[slot] |
| 131 | else: |
| 132 | assert tuple(pool._uva_bufs[slot].cpu.shape) == ( |
| 133 | 1 << (length - 1).bit_length(), |
| 134 | *size[1:], |
| 135 | ) |
| 136 | assert pool.size == size |
| 137 | # Blocking read also retires GPU readers before the next slot reuse. |
| 138 | torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0) |
| 139 | |
| 140 | |
| 141 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 142 | @pytest.mark.parametrize("use_out", [False, True]) |
| 143 | @pytest.mark.parametrize("input_type", ["numpy", "tensor"]) |
| 144 | def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type): |
| 145 | pool = buffer_utils.UvaBufferPool((2, 3), torch.int32, max_concurrency=2) |
| 146 | for length in (2, 5, 3, 6): |
| 147 | expected = torch.arange(length * 3, dtype=torch.int32, device="cpu").reshape( |
| 148 | length, 3 |
| 149 | ) |
| 150 | values = expected.numpy() if input_type == "numpy" else expected |
| 151 | out = ( |
| 152 | torch.empty(expected.shape, dtype=torch.int32, device=DEVICE_TYPE) |
| 153 | if use_out |
| 154 | else None |
| 155 | ) |
| 156 | result = pool.copy_to_gpu(values, out=out) |
| 157 | if use_out: |
| 158 | assert result is out |
| 159 | torch.testing.assert_close(result.cpu(), expected, rtol=0, atol=0) |
| 160 | |
| 161 | |
| 162 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 163 | @pytest.mark.parametrize("uva_target", [False, True]) |
| 164 | @pytest.mark.parametrize("dtype", [torch.int32, torch.int64, torch.float32]) |
| 165 | def test_staged_write_inflight(uva_target, dtype): |
| 166 | """Preserve every generation until its consumer finishes before slot reuse.""" |
| 167 | device = torch.device(f"{DEVICE_TYPE}:0") |
| 168 | with torch.accelerator.device_index(device.index): |
| 169 | state = StagedWriteTensor( |
| 170 | (4, 4096), |
| 171 | dtype, |
| 172 | device, |
| 173 | max_concurrency=2, |
| 174 | uva_instead_of_gpu=uva_target, |
| 175 | ) |
| 176 | assert (state.write_contents is not None) == uva_target |
| 177 | stream = torch.Stream(device=device) |
| 178 | stream.wait_stream(torch.accelerator.current_stream()) |
| 179 | pending: list[tuple[torch.Event, torch.Tensor, torch.Tensor]] = [] |
| 180 | expected = torch.zeros((4, 4096), dtype=dtype, device="cpu") |
| 181 | for step in range(24): |
| 182 | if len(pending) == 2: |
| 183 | event, snapshot, reference = pending.pop(0) |
| 184 | event.synchronize() |
| 185 | torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0) |
| 186 | # Growing and shrinking lengths exercise reallocation and reuse. |
| 187 | length = [3, 17, 1500, 4090][step % 4] |
| 188 | row = step % 3 |
| 189 | values = torch.arange(length, dtype=dtype, device="cpu") + step * 8192 |
| 190 | if dtype == torch.float32: |
| 191 | values += 0.25 |
| 192 | expected[row, 2 : 2 + length] = values |
| 193 | expected[3, 1:4] = step |
| 194 | with stream: |
| 195 | state.stage_write(row, 2, values.tolist()) |
| 196 | state.stage_write(3, 1, [step] * 3) |
| 197 | state.apply_write() |
| 198 | # A GPU consumer observes this generation before the next update. |
| 199 | snapshot = state.gpu.clone() |
| 200 | event = torch.Event() |
| 201 | event.record(stream) |
| 202 | pending.append((event, snapshot, expected.clone())) |
| 203 | for event, snapshot, reference in pending: |
| 204 | event.synchronize() |
| 205 | torch.testing.assert_close(snapshot.cpu(), reference, rtol=0, atol=0) |
| 206 | |
| 207 | |
| 208 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 209 | @pytest.mark.parametrize("device", DEVICES) |
| 210 | def test_non_pinned_cpu_tensor(device): |
| 211 | # Non-pinned CPU tensors are internally copied into a pinned buffer, |
| 212 | # so the resulting gpu view reflects the values at creation time but |
| 213 | # is decoupled from further writes to the original `cpu_tensor`. |
| 214 | torch.set_default_device(device) |
| 215 | cpu_tensor = torch.arange(100, dtype=torch.int32, device="cpu").view(10, 10) |
| 216 | assert not cpu_tensor.is_pinned() |
| 217 | gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) |
| 218 | assert gpu_view.device.type == DEVICE_TYPE |
| 219 | |
| 220 | assert gpu_view[0, 0] == 0 |
| 221 | assert gpu_view[2, 3] == 23 |
| 222 | assert gpu_view[9, 9] == 99 |
| 223 | |
| 224 | # Writes to the original (unpinned) CPU tensor must not affect the view, |
| 225 | # since a private pinned copy was made. |
| 226 | cpu_tensor[0, 0] = -1 |
| 227 | assert gpu_view[0, 0] == 0 |
| 228 | |
| 229 | # The view itself remains writable and independently usable. |
| 230 | gpu_view.mul_(2) |
| 231 | assert gpu_view[2, 3] == 46 |
| 232 | assert gpu_view[9, 9] == 198 |
| 233 | |
| 234 | |
| 235 | @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") |
| 236 | @pytest.mark.skipif( |
| 237 | not current_platform.is_xpu(), reason="XPU non-contiguous UVA test." |
| 238 | ) |
| 239 | @pytest.mark.parametrize("pinned", [False, True]) |
| 240 | @pytest.mark.parametrize("device", DEVICES) |
| 241 | def test_non_contiguous_strided_view(device, pinned): |
| 242 | torch.set_default_device(device) |
| 243 | # Simulate scale_kn: a [32, 16] contiguous tensor |
| 244 | scale_kn = torch.arange( |
| 245 | 512, dtype=torch.float32, device="cpu", pin_memory=pinned |
| 246 | ).view(32, 16) |
| 247 | # Transposed view: shape [16, 32], non-contiguous stride (1, 16) |
| 248 | cpu_view = scale_kn.t() |
| 249 | assert cpu_view.shape == (16, 32) |
| 250 | assert cpu_view.stride() == (1, 16) |
| 251 | assert not cpu_view.is_contiguous() |
| 252 | |
| 253 | gpu_view = get_accelerator_view_from_cpu_tensor(cpu_view) |
| 254 | assert gpu_view.device.type == DEVICE_TYPE |
| 255 | assert gpu_view.shape == (16, 32) |
| 256 | assert gpu_view.stride() == (1, 16) |
| 257 | assert not gpu_view.is_contiguous() |
| 258 | |
| 259 | # Transposing tensor_view back should yield a contiguous [32, 16] tensor |
| 260 | gpu_view_t = gpu_view.t() |
| 261 | assert gpu_view_t.shape == (32, 16) |
| 262 | assert gpu_view_t.is_contiguous() |
| 263 | |
| 264 | # Correctness check: values match original transposed CPU tensor |
| 265 | assert torch.equal(gpu_view.cpu(), cpu_view) |