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 1, 18:45 UTC
[XPU] Support cpu kv offloading and tiering offloading on XPU platform (#36423)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jun 7, 2026 at 09:59 UTC
810966453a2576f45880c899db2af2e604207ed4
7 files changed
+86
-32
.buildkite/intel_jobs/misc_intel.yaml
+23
@@ -53,3 +53,26 @@ steps:
53
pytest -v -s v1/test_oracle.py &&
54
pytest -v -s v1/test_request.py &&
55
pytest -v -s v1/test_outputs.py'
56
+
57
+- label: XPU CPU Offload
58
+ timeout_in_minutes: 60
59
+ device: intel_gpu
60
+ no_plugin: true
61
+ working_dir: "."
62
+ env:
63
+ REGISTRY: "public.ecr.aws/q9t5s3a7"
64
+ REPO: "vllm-ci-test-repo"
65
+ VLLM_TEST_DEVICE: "xpu"
66
+ source_file_dependencies:
67
+ - vllm/
68
+ - vllm/v1/kv_offload/
69
+ - vllm/v1/kv_connector/
70
+ - tests/v1/kv_offload/
71
+ - tests/v1/kv_connector/unit/test_offloading_connector.py
72
+ commands:
73
+ - >-
74
+ bash .buildkite/scripts/hardware_ci/run-intel-test.sh
75
+ 'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
76
+ cd tests &&
77
+ pytest -v -s v1/kv_offload &&
78
+ pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+5
-2
@@ -30,6 +30,7 @@ BLOCK_SIZE = 16
30
NUM_KV_HEADS = 4
31
HEAD_SIZE = 64
32
DTYPE = torch.float16
33
+DEVICE_TYPE = current_platform.device_type
34
35
# Attention backends to test
36
ATTN_BACKENDS: list[str] = []
@@ -42,6 +43,8 @@ if current_platform.is_cuda():
43
]
44
elif current_platform.is_rocm():
45
ATTN_BACKENDS = ["TRITON_ATTN"]
46
+elif current_platform.is_xpu():
47
+ ATTN_BACKENDS = ["TRITON_ATTN", "FLASH_ATTN"]
48
49
# ---------------------------------------------------------------------------
50
# Helpers
@@ -270,7 +273,7 @@ def test_register_kv_caches(backend):
273
kv_caches = _allocate_and_reshape_kv_caches(
274
kv_cache_config,
275
attn_groups,
273
- device=torch.device("cuda:0"),
276
+ device=torch.device(f"{DEVICE_TYPE}:0"),
277
)
278
279
worker, spec = _make_worker(kv_cache_config)
@@ -413,7 +416,7 @@ def test_register_kv_caches_uniform_type(backend):
416
kv_caches = _allocate_and_reshape_kv_caches(
417
kv_cache_config,
418
attn_groups,
416
- device=torch.device("cuda:0"),
419
+ device=torch.device(f"{DEVICE_TYPE}:0"),
420
)
421
422
worker, spec = _make_worker(kv_cache_config)
tests/v1/kv_connector/unit/test_offloading_connector.py
+8
-5
@@ -14,11 +14,14 @@ from vllm.config import KVEventsConfig, KVTransferConfig
14
from vllm.distributed.kv_events import BlockStored, KVEventBatch
15
from vllm.platforms import current_platform
16
17
+CPU_BLOCK_SIZES: int = 64 if current_platform.is_xpu() else 48
18
_ATTN_BACKENDS: list[str] = []
19
if current_platform.is_cuda():
20
_ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
21
elif current_platform.is_rocm():
22
_ATTN_BACKENDS = ["TRITON_ATTN"]
23
+elif current_platform.is_xpu():
24
+ _ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"]
25
26
# (model, attn_backend | None, block_size | None, uses_hma)
27
#
@@ -30,14 +33,14 @@ elif current_platform.is_rocm():
33
# After page-size unification the mamba and attention groups have
34
# different block sizes.
35
MODEL_PARAMS: list[tuple[str, str | None, int | None, bool]] = [
33
- ("meta-llama/Llama-3.2-1B-Instruct", backend, 48, False)
36
+ ("meta-llama/Llama-3.2-1B-Instruct", backend, CPU_BLOCK_SIZES, False)
37
for backend in _ATTN_BACKENDS
38
]
39
# HMA / Mamba models are only tested on CUDA (not ROCm).
40
if current_platform.is_cuda():
41
MODEL_PARAMS += [
39
- ("google/gemma-3-1b-it", None, 48, True),
40
- ("state-spaces/mamba-130m-hf", None, 48, True),
42
+ ("google/gemma-3-1b-it", None, CPU_BLOCK_SIZES, True),
43
+ ("state-spaces/mamba-130m-hf", None, CPU_BLOCK_SIZES, True),
44
# Falcon-H1: parallel hybrid (every layer has both attention and SSM).
45
# The mamba and attention groups end up with different GPU block sizes
46
# after page-size unification, so we leave cpu_block_size=None
@@ -304,7 +307,7 @@ def test_tiering_offloading() -> None:
307
"""Tests OffloadingConnector with TieringOffloadingSpec."""
308
extra_config: dict = {
309
"cpu_bytes_to_use": 500 << 20,
307
- "block_size": 48,
310
+ "block_size": CPU_BLOCK_SIZES,
311
"spec_name": "TieringOffloadingSpec",
312
"secondary_tiers": [{"type": "example"}],
313
}
@@ -350,7 +353,7 @@ def test_fs_tiering_offloading(tmp_path) -> None:
353
+ fs secondary tier."""
354
extra_config: dict = {
355
"cpu_bytes_to_use": 1 << 30,
353
- "block_size": 48,
356
+ "block_size": CPU_BLOCK_SIZES,
357
"spec_name": "TieringOffloadingSpec",
358
"secondary_tiers": [{"type": "fs", "root_dir": str(tmp_path)}],
359
}
vllm/_custom_ops.py
+9
-4
@@ -2812,7 +2812,9 @@ def swap_blocks_batch(
2812
Batch version of swap_blocks: submit all copies in a single driver call.
2813
2814
Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i]
2815
- of sizes[i] bytes. All three tensors must be int64 CPU tensors.
2815
+ of sizes[i] bytes. All three tensors must be CPU tensors with the
2816
+ platform-appropriate pointer dtype: int64 on CUDA/ROCm (required by
2817
+ cache_kernels.cu) and uint64 on XPU (required by the XPU DMA engine).
2818
On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission
2819
overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync.
2820
@@ -2822,9 +2824,12 @@ def swap_blocks_batch(
2824
writing to the source. Defaults to False (STREAM ordering), which
2825
is always safe.
2826
"""
2825
- torch.ops._C_cache_ops.swap_blocks_batch(
2826
- src_ptrs, dst_ptrs, sizes, is_src_access_order_any
2827
- )
2827
+ if current_platform.is_xpu():
2828
+ torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
2829
+ else:
2830
+ torch.ops._C_cache_ops.swap_blocks_batch(
2831
+ src_ptrs, dst_ptrs, sizes, is_src_access_order_any
2832
+ )
2833
2834
2835
def convert_fp8(
vllm/v1/kv_offload/cpu/gpu_worker.py
+28
-13
@@ -11,6 +11,7 @@ from typing_extensions import override
11
12
from vllm import _custom_ops as ops
13
from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
from vllm.triton_utils import HAS_TRITON, triton
16
from vllm.utils.math_utils import cdiv
17
from vllm.utils.platform_utils import is_pin_memory_available
@@ -43,8 +44,10 @@ def _select_swap_blocks_fn(
44
if gpu_to_cpu:
45
return ops.swap_blocks_batch
46
# Fall back to the C++ DMA path on platforms where Triton isn't usable
46
- # (e.g. ROCm builds without Triton).
47
- if not HAS_TRITON:
47
+ # (e.g. ROCm builds without Triton) or where GPU kernels cannot directly
48
+ # dereference CPU pointers (XPU lacks CUDA's unified virtual address space,
49
+ # so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU).
50
+ if not HAS_TRITON or current_platform.is_xpu():
51
return ops.swap_blocks_batch
52
page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g]
53
# Triton wins only on small, 8-byte-aligned payloads.
@@ -92,7 +95,7 @@ def compute_sub_block_ptrs(
95
Args:
96
block_ids: array of block IDs at the tensor's native granularity.
97
block_size_factor: number of sub-blocks per block.
95
- output: pre-allocated int64 array to write pointers into.
98
+ output: pre-allocated pointer array to write pointers into.
99
tensor: the source or destination tensor.
100
skip_count: sub-blocks to skip in the first block.
101
"""
@@ -104,16 +107,16 @@ def compute_sub_block_ptrs(
107
108
if block_size_factor == 1:
109
# Fast path: 1:1 mapping, no sub-block expansion needed.
107
- output[:] = base_ptr + block_ids[:num_sub_blocks] * row_stride
110
+ output[:] = base_ptr + block_ids.astype(np.uint64)[:num_sub_blocks] * row_stride
111
return
112
113
# Vectorized expansion for block_size_factor > 1.
114
assert tensor.shape[1] % block_size_factor == 0
115
sub_block_size = tensor.shape[1] // block_size_factor
113
- sub_offsets = np.arange(block_size_factor, dtype=np.int64) * sub_block_size
116
+ sub_offsets = np.arange(block_size_factor, dtype=np.uint64) * sub_block_size
117
# (num_blocks, 1) + (1, block_size_factor) -> (num_blocks, block_size_factor)
118
all_ptrs = (
116
- base_ptr + block_ids.astype(np.int64)[:, np.newaxis] * row_stride
119
+ base_ptr + block_ids.astype(np.uint64)[:, np.newaxis] * row_stride
120
) + sub_offsets[np.newaxis, :]
121
# Flatten and apply skip_count / truncation
122
flat = all_ptrs.ravel()
@@ -122,6 +125,14 @@ def compute_sub_block_ptrs(
125
126
def pin_mmap_region(region: SharedOffloadRegion) -> None:
127
"""Register the entire mmap as CUDA pinned memory via cudaHostRegister."""
128
+ if not current_platform.is_cuda_alike():
129
+ logger.info(
130
+ "Skipping mmap host registration on %s; cudaHostRegister is only "
131
+ "available on CUDA/ROCm.",
132
+ current_platform.device_name,
133
+ )
134
+ return
135
+
136
rank = region.rank
137
138
base_ptr = region._base.data_ptr()
@@ -146,10 +157,12 @@ def _new_descriptor_buffers(
157
num_copy_ops: int,
158
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
159
pin = is_pin_memory_available()
160
+ # CUDA cache_kernels.cu requires int64; XPU DMA engine requires uint64.
161
+ ptr_dtype = torch.uint64 if current_platform.is_xpu() else torch.int64
162
return (
150
- torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
151
- torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
152
- torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
163
+ torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin),
164
+ torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin),
165
+ torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin),
166
)
167
168
@@ -190,7 +203,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
203
for gpu_tensor, cpu_tensor in zip(gpu_tensors, cpu_tensors):
204
assert gpu_tensor.dtype == torch.int8
205
assert gpu_tensor.ndim == 2
193
- assert gpu_tensor.is_cuda
206
+ assert gpu_tensor.is_cuda or gpu_tensor.is_xpu
207
assert cpu_tensor.dtype == torch.int8
208
assert cpu_tensor.ndim == 2
209
assert cpu_tensor.device.type == "cpu"
@@ -355,7 +368,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
368
assert dst_offset == num_dst_blocks
369
assert op_idx == num_copy_ops
370
358
- stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream()
371
+ stream = (
372
+ self._stream_pool.pop() if self._stream_pool else current_platform.Stream()
373
+ )
374
start_event = (
375
self._event_pool.pop()
376
if self._event_pool
@@ -369,7 +384,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
384
385
if self.gpu_to_cpu:
386
# wait for model computation to finish before offloading
372
- stream.wait_stream(torch.cuda.current_stream())
387
+ stream.wait_stream(current_platform.current_stream())
388
if self._transfers:
389
last_transfer: Transfer = self._transfers[-1]
390
last_event = last_transfer.end_event
@@ -382,7 +397,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
397
# writing; we must keep STREAM ordering so source reads are gated
398
# by the transfer stream's wait_stream(compute) barrier.
399
is_src_access_order_any = not self.gpu_to_cpu
385
- with torch.cuda.stream(stream):
400
+ with current_platform.stream(stream):
401
start_event.record(stream)
402
if num_copy_ops > 0:
403
self._swap_blocks_batch(
vllm/v1/kv_offload/cpu/shared_offload_region.py
+10
-6
@@ -7,6 +7,7 @@ import time
7
import torch
8
9
from vllm.logger import init_logger
10
+from vllm.platforms import current_platform
11
12
logger = init_logger(__name__)
13
@@ -171,12 +172,15 @@ class SharedOffloadRegion:
172
173
def cleanup(self) -> None:
174
if self.is_pinned and self._base is not None:
174
- base_ptr = self._base.data_ptr()
175
- result = torch.cuda.cudart().cudaHostUnregister(base_ptr)
176
- if result.value != 0:
177
- logger.warning(
178
- "cudaHostUnregister failed for rank=%d (code=%d)", self.rank, result
179
- )
175
+ if current_platform.is_cuda_alike():
176
+ base_ptr = self._base.data_ptr()
177
+ result = torch.cuda.cudart().cudaHostUnregister(base_ptr)
178
+ if result.value != 0:
179
+ logger.warning(
180
+ "cudaHostUnregister failed for rank=%d (code=%d)",
181
+ self.rank,
182
+ result,
183
+ )
184
self.is_pinned = False
185
# Release views before _base: each view holds a _base reference and a
186
# direct StorageImpl reference. Freeing views first lets both refcounts
vllm/v1/kv_offload/cpu/spec.py
+3
-2
@@ -107,9 +107,10 @@ class CPUOffloadingSpec(OffloadingSpec):
107
self, kv_caches: CanonicalKVCaches
108
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
109
if not self._handlers:
110
- if not current_platform.is_cuda_alike():
110
+ if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
111
raise Exception(
112
- "CPU Offloading is currently only supported on CUDA-alike GPUs"
112
+ "CPU Offloading is currently only supported on CUDA-alike "
113
+ "and XPU GPUs"
114
)
115
self._handlers = self.create_handlers(kv_caches)
116