A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 19:00 UTC
support view_from_cpu_tensor on XPU (#33868)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Feb 6, 2026 at 16:34 UTC
e969a169ef3842fd18fdf922dd23a9f67bea98a4
4 files changed
+13
-9
tests/kernels/core/test_uva.py
+3
-3
@@ -4,7 +4,7 @@ import pytest
4
import torch
5
6
from vllm.utils.platform_utils import is_uva_available
7
-from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
7
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
8
9
CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
10
@@ -14,7 +14,7 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e
14
def test_cpu_write(device):
15
torch.set_default_device(device)
16
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
17
- cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor)
17
+ cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
18
assert cuda_view.device.type == "cuda"
19
20
assert cuda_view[0, 0] == 0
@@ -36,7 +36,7 @@ def test_cpu_write(device):
36
def test_gpu_write(device):
37
torch.set_default_device(device)
38
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
39
- cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor)
39
+ cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
40
assert cuda_view.device.type == "cuda"
41
42
assert cuda_view[0, 0] == 0
vllm/model_executor/models/utils.py
+2
-2
@@ -36,7 +36,7 @@ from vllm.utils.platform_utils import (
36
)
37
from vllm.utils.torch_utils import (
38
direct_register_custom_op,
39
- get_cuda_view_from_cpu_tensor,
39
+ get_accelerator_view_from_cpu_tensor,
40
)
41
42
logger = init_logger(__name__)
@@ -663,7 +663,7 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
663
else:
664
# keep the cpu data alive
665
p._vllm_offloaded_cpu_data = cpu_data
666
- p.data = get_cuda_view_from_cpu_tensor(cpu_data)
666
+ p.data = get_accelerator_view_from_cpu_tensor(cpu_data)
667
_CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
668
offloaded_parameters = True
669
vllm/utils/torch_utils.py
+6
-2
@@ -674,11 +674,15 @@ def weak_ref_tensors(
674
raise ValueError("Invalid type for tensors")
675
676
677
-def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
677
+def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
678
"""
679
- Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
679
+ Get an accelerator view of a CPU tensor using Unified Virtual Addressing (UVA).
680
"""
681
assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
682
+ from vllm.platforms import current_platform
683
+
684
+ if current_platform.is_xpu():
685
+ return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor)
686
return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
687
688
vllm/v1/worker/gpu/buffer_utils.py
+2
-2
@@ -9,7 +9,7 @@ import torch
9
from vllm.triton_utils import tl, triton
10
from vllm.utils.math_utils import next_power_of_2
11
from vllm.utils.platform_utils import is_uva_available
12
-from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
12
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
13
14
15
def async_copy_to_gpu(
@@ -38,7 +38,7 @@ class UvaBuffer:
38
raise RuntimeError("UVA is not available")
39
self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
40
self.np = self.cpu.numpy()
41
- self.uva = get_cuda_view_from_cpu_tensor(self.cpu)
41
+ self.uva = get_accelerator_view_from_cpu_tensor(self.cpu)
42
43
44
class UvaBufferPool: