From 7230dfea501b232ce92d057f1c5ee9cb009c4cb4 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Tue, 29 Sep 2026 09:43:07 +0800 Subject: [PATCH] [XPU] Preserve non-contiguous strides when pinning CPU tensors for UVA view (#54874) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/kernels_intel.yaml | 23 +++++ tests/kernels/core/test_uva.py | 123 +++++++++++++++++------ vllm/utils/torch_utils.py | 13 ++- 3 files changed, 125 insertions(+), 34 deletions(-) diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index da6086c1dc..3650a729d7 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -25,3 +25,26 @@ steps: 'cd tests && pytest -v -s ir && pytest -v -s kernels/ir' + +- label: UVA Kernel Tests + key: uva-kernel-tests + timeout_in_minutes: 20 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - tests/kernels/core/test_uva.py + - vllm/utils/torch_utils.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s kernels/core/test_uva.py' diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index f8c297afc6..b7c2984d6a 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -4,54 +4,57 @@ import numpy as np import pytest import torch +from vllm.platforms import current_platform 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) +DEVICE_TYPE = current_platform.device_type +DEVICES = [ + f"{DEVICE_TYPE}:{i}" + for i in range(1 if torch.accelerator.device_count() == 1 else 2) ] @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) def test_cpu_write(device): torch.set_default_device(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) - cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) - assert cuda_view.device.type == "cuda" + gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) + assert gpu_view.device.type == DEVICE_TYPE - assert cuda_view[0, 0] == 0 - assert cuda_view[2, 3] == 0 - assert cuda_view[4, 5] == 0 + assert gpu_view[0, 0] == 0 + assert gpu_view[2, 3] == 0 + assert gpu_view[4, 5] == 0 cpu_tensor[0, 0] = 1 cpu_tensor[2, 3] = 2 cpu_tensor[4, 5] = -1 - cuda_view.mul_(2) - assert cuda_view[0, 0] == 2 - assert cuda_view[2, 3] == 4 - assert cuda_view[4, 5] == -2 + gpu_view.mul_(2) + assert gpu_view[0, 0] == 2 + assert gpu_view[2, 3] == 4 + assert gpu_view[4, 5] == -2 @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) def test_gpu_write(device): torch.set_default_device(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) - cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) - assert cuda_view.device.type == "cuda" + gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) + assert gpu_view.device.type == DEVICE_TYPE - assert cuda_view[0, 0] == 0 - assert cuda_view[2, 3] == 0 - assert cuda_view[4, 5] == 0 + assert gpu_view[0, 0] == 0 + assert gpu_view[2, 3] == 0 + assert gpu_view[4, 5] == 0 - cuda_view[0, 0] = 1 - cuda_view[2, 3] = 2 - cuda_view[4, 5] = -1 - cuda_view.mul_(2) + gpu_view[0, 0] = 1 + gpu_view[2, 3] = 2 + gpu_view[4, 5] = -1 + gpu_view.mul_(2) assert cpu_tensor[0, 0] == 2 assert cpu_tensor[2, 3] == 4 @@ -59,7 +62,7 @@ def test_gpu_write(device): @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", 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") @@ -146,7 +149,7 @@ def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type): ) values = expected.numpy() if input_type == "numpy" else expected out = ( - torch.empty(expected.shape, dtype=torch.int32, device="cuda") + torch.empty(expected.shape, dtype=torch.int32, device=DEVICE_TYPE) if use_out else None ) @@ -161,7 +164,7 @@ def test_uva_pool_copy_to_gpu_preserves_shape_and_out(use_out, input_type): @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") + device = torch.device(f"{DEVICE_TYPE}:0") with torch.accelerator.device_index(device.index): state = StagedWriteTensor( (4, 4096), @@ -171,9 +174,9 @@ def test_staged_write_inflight(uva_target, dtype): 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]] = [] + stream = torch.Stream(device=device) + stream.wait_stream(torch.accelerator.current_stream()) + pending: list[tuple[torch.Event, torch.Tensor, torch.Tensor]] = [] expected = torch.zeros((4, 4096), dtype=dtype, device="cpu") for step in range(24): if len(pending) == 2: @@ -188,15 +191,75 @@ def test_staged_write_inflight(uva_target, dtype): values += 0.25 expected[row, 2 : 2 + length] = values expected[3, 1:4] = step - with torch.cuda.stream(stream): + with 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 = torch.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) + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("device", DEVICES) +def test_non_pinned_cpu_tensor(device): + # Non-pinned CPU tensors are internally copied into a pinned buffer, + # so the resulting gpu view reflects the values at creation time but + # is decoupled from further writes to the original `cpu_tensor`. + torch.set_default_device(device) + cpu_tensor = torch.arange(100, dtype=torch.int32, device="cpu").view(10, 10) + assert not cpu_tensor.is_pinned() + gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) + assert gpu_view.device.type == DEVICE_TYPE + + assert gpu_view[0, 0] == 0 + assert gpu_view[2, 3] == 23 + assert gpu_view[9, 9] == 99 + + # Writes to the original (unpinned) CPU tensor must not affect the view, + # since a private pinned copy was made. + cpu_tensor[0, 0] = -1 + assert gpu_view[0, 0] == 0 + + # The view itself remains writable and independently usable. + gpu_view.mul_(2) + assert gpu_view[2, 3] == 46 + assert gpu_view[9, 9] == 198 + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.skipif( + not current_platform.is_xpu(), reason="XPU non-contiguous UVA test." +) +@pytest.mark.parametrize("pinned", [False, True]) +@pytest.mark.parametrize("device", DEVICES) +def test_non_contiguous_strided_view(device, pinned): + torch.set_default_device(device) + # Simulate scale_kn: a [32, 16] contiguous tensor + scale_kn = torch.arange( + 512, dtype=torch.float32, device="cpu", pin_memory=pinned + ).view(32, 16) + # Transposed view: shape [16, 32], non-contiguous stride (1, 16) + cpu_view = scale_kn.t() + assert cpu_view.shape == (16, 32) + assert cpu_view.stride() == (1, 16) + assert not cpu_view.is_contiguous() + + gpu_view = get_accelerator_view_from_cpu_tensor(cpu_view) + assert gpu_view.device.type == DEVICE_TYPE + assert gpu_view.shape == (16, 32) + assert gpu_view.stride() == (1, 16) + assert not gpu_view.is_contiguous() + + # Transposing tensor_view back should yield a contiguous [32, 16] tensor + gpu_view_t = gpu_view.t() + assert gpu_view_t.shape == (32, 16) + assert gpu_view_t.is_contiguous() + + # Correctness check: values match original transposed CPU tensor + assert torch.equal(gpu_view.cpu(), cpu_view) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 849d0a38f5..b4a13419d5 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -912,10 +912,15 @@ def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tens if cpu_tensor.numel() == 0: return torch.empty(cpu_tensor.shape, dtype=cpu_tensor.dtype, device="xpu") if not cpu_tensor.is_pinned(): - contiguous_cpu = cpu_tensor.contiguous() - pinned = torch.empty_like(contiguous_cpu, pin_memory=True) - pinned.copy_(contiguous_cpu) - cpu_tensor = pinned + pinned = torch.empty_strided( + cpu_tensor.size(), + cpu_tensor.stride(), + dtype=cpu_tensor.dtype, + device="cpu", + pin_memory=True, + ) + pinned.copy_(cpu_tensor) + return torch.ops._C.get_xpu_view_from_cpu_tensor(pinned) return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor) elif current_platform.is_cuda_alike(): return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)