[XPU] Preserve non-contiguous strides when pinning CPU tensors for UVA view (#54874)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Sep 29, 2026 at 09:43 UTC
7230dfea501b232ce92d057f1c5ee9cb009c4cb4
3 files changed
+125
-34
.buildkite/intel_jobs/kernels_intel.yaml
+23
@@ -25,3 +25,26 @@ steps:
25
'cd tests &&
26
pytest -v -s ir &&
27
pytest -v -s kernels/ir'
28
+
29
+- label: UVA Kernel Tests
30
+ key: uva-kernel-tests
31
+ timeout_in_minutes: 20
32
+ device: intel_gpu
33
+ agent_tags:
34
+ label: production
35
+ gpu: 1+
36
+ mem: 16+
37
+ no_plugin: true
38
+ working_dir: "."
39
+ env:
40
+ REGISTRY: "public.ecr.aws/q9t5s3a7"
41
+ REPO: "vllm-ci-test-repo"
42
+ VLLM_TEST_DEVICE: "xpu"
43
+ source_file_dependencies:
44
+ - tests/kernels/core/test_uva.py
45
+ - vllm/utils/torch_utils.py
46
+ commands:
47
+ - >-
48
+ bash .buildkite/scripts/hardware_ci/run-intel-test.sh
49
+ 'cd tests &&
50
+ pytest -v -s kernels/core/test_uva.py'
tests/kernels/core/test_uva.py
+93
-30
@@ -4,54 +4,57 @@ 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
12
-CUDA_DEVICES = [
13
- f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
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.")
18
-@pytest.mark.parametrize("device", CUDA_DEVICES)
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)
22
- cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
23
- assert cuda_view.device.type == "cuda"
25
+ gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
26
+ assert gpu_view.device.type == DEVICE_TYPE
27
25
- assert cuda_view[0, 0] == 0
26
- assert cuda_view[2, 3] == 0
27
- assert cuda_view[4, 5] == 0
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
33
- cuda_view.mul_(2)
34
- assert cuda_view[0, 0] == 2
35
- assert cuda_view[2, 3] == 4
36
- assert cuda_view[4, 5] == -2
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.")
40
-@pytest.mark.parametrize("device", CUDA_DEVICES)
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)
44
- cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
45
- assert cuda_view.device.type == "cuda"
47
+ gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
48
+ assert gpu_view.device.type == DEVICE_TYPE
49
47
- assert cuda_view[0, 0] == 0
48
- assert cuda_view[2, 3] == 0
49
- assert cuda_view[4, 5] == 0
50
+ assert gpu_view[0, 0] == 0
51
+ assert gpu_view[2, 3] == 0
52
+ assert gpu_view[4, 5] == 0
53
51
- cuda_view[0, 0] = 1
52
- cuda_view[2, 3] = 2
53
- cuda_view[4, 5] = -1
54
- cuda_view.mul_(2)
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
@@ -59,7 +62,7 @@ def test_gpu_write(device):
62
63
64
@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
62
-@pytest.mark.parametrize("device", CUDA_DEVICES)
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")
@@ -146,7 +149,7 @@ def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type):
149
)
150
values = expected.numpy() if input_type == "numpy" else expected
151
out = (
149
- torch.empty(expected.shape, dtype=torch.int32, device="cuda")
152
+ torch.empty(expected.shape, dtype=torch.int32, device=DEVICE_TYPE)
153
if use_out
154
else None
155
)
@@ -161,7 +164,7 @@ def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type):
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."""
164
- device = torch.device("cuda:0")
167
+ device = torch.device(f"{DEVICE_TYPE}:0")
168
with torch.accelerator.device_index(device.index):
169
state = StagedWriteTensor(
170
(4, 4096),
@@ -171,9 +174,9 @@ def test_staged_write_inflight(uva_target, dtype):
174
uva_instead_of_gpu=uva_target,
175
)
176
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
+ 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:
@@ -188,15 +191,75 @@ def test_staged_write_inflight(uva_target, dtype):
191
values += 0.25
192
expected[row, 2 : 2 + length] = values
193
expected[3, 1:4] = step
191
- with torch.cuda.stream(stream):
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()
197
- event = torch.cuda.Event()
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)
vllm/utils/torch_utils.py
+9
-4
@@ -912,10 +912,15 @@ def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tens
912
if cpu_tensor.numel() == 0:
913
return torch.empty(cpu_tensor.shape, dtype=cpu_tensor.dtype, device="xpu")
914
if not cpu_tensor.is_pinned():
915
- contiguous_cpu = cpu_tensor.contiguous()
916
- pinned = torch.empty_like(contiguous_cpu, pin_memory=True)
917
- pinned.copy_(contiguous_cpu)
918
- cpu_tensor = pinned
915
+ pinned = torch.empty_strided(
916
+ cpu_tensor.size(),
917
+ cpu_tensor.stride(),
918
+ dtype=cpu_tensor.dtype,
919
+ device="cpu",
920
+ pin_memory=True,
921
+ )
922
+ pinned.copy_(cpu_tensor)
923
+ return torch.ops._C.get_xpu_view_from_cpu_tensor(pinned)
924
return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor)
925
elif current_platform.is_cuda_alike():
926
return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)