@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f485081e8b

[XPU] Support EC connector KV Offloading on XPU (#49532)

Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>

Chaojun Zhang committed Aug 19, 2026 at 10:25 UTC f485081e8b76135091398818ea63040ffe7e1c5b

5 files changed +93 -58

.buildkite/intel_jobs/misc_intel.yaml

+4 -2

@@ -8,7 +8,7 @@ steps:

8 agent_tags:

9 label: production

10 gpu: 1+

11 - mem: 16+

11 + mem: 24+

12 no_plugin: true

13 working_dir: "."

14 env:

@@ -22,6 +22,7 @@ steps:

22 - tests/v1/kv_offload

23 - tests/v1/worker

24 - tests/v1/kv_connector/unit

25 + - tests/v1/ec_connector/unit

26 - tests/v1/metrics

27 - tests/entrypoints/openai/correctness/test_lmeval.py

28 commands:

@@ -30,7 +31,8 @@ steps:

31 'pip install -r requirements/kv_connectors.txt &&

32 export VLLM_WORKER_MULTIPROC_METHOD=spawn &&

33 cd tests &&

33 - pytest -v -s v1/executor'

34 + pytest -v -s v1/executor &&

35 + VLLM_BATCH_INVARIANT=1 pytest -v -s -m 'not cpu_test' v1/ec_connector/unit'

37 - label: V1 Sample + Logits

38 timeout_in_minutes: 90

tests/v1/ec_connector/unit/cpu/worker/test_worker.py

+42 -37

@@ -2,10 +2,10 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Tests for ECCPUWorker.

5 -The byte-level tests exercise real CUDA stream/event coordination

6 -against a real ``ECSharedRegion`` mmap and are skipped on hosts without

7 -CUDA. The lifecycle tests don't fire any ``torch.cuda.*`` primitives and

8 -run anywhere.

5 +The byte-level tests exercise real accelerator (CUDA/XPU/...) stream/event

6 +coordination against a real ``ECSharedRegion`` mmap and are skipped on hosts

7 +without an accelerator. The lifecycle tests don't fire any

8 +``current_platform.*`` device primitives and run anywhere.

10 Mocking policy

11 --------------

@@ -37,6 +37,7 @@ from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (

37 ECSharedRegion,

38 )

39 from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ECCPUWorker

40 +from vllm.platforms import current_platform

42 # ── shape constants ──────────────────────────────────────────────────────────

@@ -47,9 +48,11 @@ _DTYPE = torch.float16

48 _BLOCK_SIZE_BYTES = _HIDDEN_DIM * _DTYPE.itemsize

49 _NUM_BLOCKS = 8

50 -_requires_cuda = pytest.mark.skipif(

51 - not torch.cuda.is_available(),

52 - reason="exercises real CUDA stream/event coordination in ECCPUWorker",

51 +DEVICE_TYPE = current_platform.device_type

52 +

53 +_requires_accelerator = pytest.mark.skipif(

54 + not (current_platform.is_cuda() or current_platform.is_xpu()),

55 + reason="exercises real accelerator stream/event coordination in ECCPUWorker",

56 )

@@ -136,7 +139,7 @@ def make_worker():

139 # ── save_caches ──────────────────────────────────────────────────────────────

140

141

139 -@_requires_cuda

142 +@_requires_accelerator

143 @pytest.mark.parametrize(

144 "n_elements,n_blocks",

145 [

@@ -158,7 +161,7 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks

161 sentinel = 0x5A

162 worker._region.blocks.fill_(sentinel)

163

161 - src = torch.arange(n_elements, dtype=_DTYPE, device="cuda")

164 + src = torch.arange(n_elements, dtype=_DTYPE, device=DEVICE_TYPE)

165 expected_bytes = src.cpu().reshape(-1).view(torch.uint8)

166 total_bytes = n_elements * _DTYPE.itemsize

167

@@ -228,15 +231,15 @@ def test_save_caches_raises_when_allocated_blocks_too_small(make_worker):

231 worker.save_caches({"h": src}, "h", _meta(saves={"h": [0, 1]}))

232

233

231 -@_requires_cuda

234 +@_requires_accelerator

235 def test_save_caches_batches_multiple_hashes(make_worker):

236 """Multiple save_caches calls are batched into a single flush."""

237 worker = make_worker()

238 sentinel = 0x5A

239 worker._region.blocks.fill_(sentinel)

240

238 - src_a = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda")

239 - src_b = torch.arange(_HIDDEN_DIM, 2 * _HIDDEN_DIM, dtype=_DTYPE, device="cuda")

241 + src_a = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

242 + src_b = torch.arange(_HIDDEN_DIM, 2 * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

243

244 cache = {"a": src_a, "b": src_b}

245 worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]}))

@@ -258,7 +261,7 @@ def test_save_caches_batches_multiple_hashes(make_worker):

261 # ── start_load_caches ────────────────────────────────────────────────────────

262

263

261 -@_requires_cuda

264 +@_requires_accelerator

265 def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker):

266 """Single batched load across all hashes with correct byte→dtype→shape."""

267 worker = make_worker()

@@ -276,20 +279,22 @@ def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker

279 worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))

280

281 out = encoder_cache["h"]

279 - assert out.is_cuda, "consumer worker must place the tensor on the GPU"

282 + assert out.device.type == DEVICE_TYPE, (

283 + "consumer worker must place the tensor on the accelerator"

284 + )

285 assert out.shape == (n_blocks, _HIDDEN_DIM)

286 assert out.dtype == _DTYPE

287 assert torch.equal(out.cpu(), src_orig)

288

289

285 -@_requires_cuda

290 +@_requires_accelerator

291 def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker):

292 """If ``encoder_cache`` already holds the ``mm_hash``, the worker must

293 not overwrite it."""

294 worker = make_worker()

295 worker._region.blocks[0].fill_(0x42)

296

292 - sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device="cuda")

297 + sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device=DEVICE_TYPE)

298 encoder_cache = {"h": sentinel}

299 worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]}))

300

@@ -298,7 +303,7 @@ def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker):

303 )

304

305

301 -@_requires_cuda

306 +@_requires_accelerator

307 def test_start_load_caches_noop_when_loads_is_empty(make_worker):

308 """When ``meta.loads`` is empty the early-return must fire."""

309 worker = make_worker()

@@ -308,7 +313,7 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker):

313 assert encoder_cache == {}

314

315

311 -@_requires_cuda

316 +@_requires_accelerator

317 def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):

318 """Cached entries are preserved while new ones are loaded."""

319 worker = make_worker()

@@ -321,7 +326,7 @@ def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):

326 for i, idx in enumerate(new_block_ids):

327 worker._region.blocks[idx].copy_(src_int8[i])

328

324 - cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device="cuda")

329 + cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device=DEVICE_TYPE)

330 encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor}

331 worker.start_load_caches(

332 encoder_cache,

@@ -336,7 +341,7 @@ def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker):

341 assert torch.equal(out.cpu(), src_orig)

342

343

339 -@_requires_cuda

344 +@_requires_accelerator

345 @pytest.mark.parametrize(

346 "tp_rank,pcp_rank",

347 [(0, 0), (1, 0), (0, 1), (1, 1)],

@@ -358,23 +363,23 @@ def test_start_load_caches_works_on_all_ranks(make_worker, tp_rank, pcp_rank):

363 worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids}))

364

365 out = encoder_cache["h"]

361 - assert out.is_cuda

366 + assert out.device.type == DEVICE_TYPE

367 assert torch.equal(out.cpu(), src_orig)

368

369

370 # ── round-trip ───────────────────────────────────────────────────────────────

371

372

368 -@_requires_cuda

373 +@_requires_accelerator

374 def test_save_then_load_round_trips_bytes(make_worker):

375 """Full producer→mmap→consumer byte path in one shot."""

376 worker = make_worker()

377 n_blocks = 3

378 block_ids = [5, 1, 6]

379

375 - src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape(

376 - n_blocks, _HIDDEN_DIM

377 - )

380 + src = torch.arange(

381 + n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE

382 + ).reshape(n_blocks, _HIDDEN_DIM)

383 worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids}))

384 worker.flush_saves()

385

@@ -390,12 +395,12 @@ def test_save_then_load_round_trips_bytes(make_worker):

395 # ── buffer recycling ────────────────────────────────────────────────────────

396

397

393 -@_requires_cuda

398 +@_requires_accelerator

399 def test_buffer_pool_is_reused_across_save_steps(make_worker):

400 """After flush_saves, descriptor buffers are returned to the pool and

401 reused on the next flush — no reallocation."""

402 worker = make_worker()

398 - src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device="cuda")

403 + src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE)

404

405 worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]}))

406 worker.flush_saves()

@@ -411,7 +416,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker):

416 assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id

417

418

414 -@_requires_cuda

419 +@_requires_accelerator

420 def test_buffer_pool_is_reused_across_load_steps(make_worker):

421 """After start_load_caches, descriptor buffers are returned to the pool

422 and reused on the next call."""

@@ -435,11 +440,11 @@ def test_buffer_pool_is_reused_across_load_steps(make_worker):

440 # ── stream management ────────────────────────────────────────────────────────

441

442

438 -@_requires_cuda

443 +@_requires_accelerator

444 def test_stream_initialized_at_construction(make_worker):

440 - """``_load_stream`` must be a fully initialized CUDA stream."""

445 + """``_load_stream`` must be a fully initialized accelerator stream."""

446 worker = make_worker()

442 - assert isinstance(worker._load_stream, torch.cuda.Stream)

447 + assert isinstance(worker._load_stream, current_platform.Stream)

448

449

450 # ── lifecycle ────────────────────────────────────────────────────────────────

@@ -518,14 +523,14 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm):

523 # ── e2e: scheduler + worker pipeline ────────────────────────────────────────

524

525

521 -@_requires_cuda

526 +@_requires_accelerator

527 def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

528 """Full pipeline: scheduler allocates blocks, worker saves GPU tensor to

529 mmap via flush_saves, scheduler marks ready after step delay, worker

530 loads from mmap back to GPU, and the result matches the original.

531

532 Exercises the real scheduler + worker cooperation through a shared

528 - ECSharedRegion, with real CUDA transfers and stream coordination.

533 + ECSharedRegion, with real accelerator transfers and stream coordination.

534 """

535 import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod

536 from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import (

@@ -551,9 +556,9 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

556

557 # -- Step 1: scheduler allocates, worker saves --

558 n_blocks = 3

554 - src = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device="cuda").reshape(

555 - n_blocks, _HIDDEN_DIM

556 - )

559 + src = torch.arange(

560 + n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE

561 + ).reshape(n_blocks, _HIDDEN_DIM)

562

563 class _Pos:

564 offset = 0

@@ -590,7 +595,7 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch):

595 worker.start_load_caches(load_cache, meta_load)

596

597 out = load_cache["img_001"]

593 - assert out.is_cuda

598 + assert out.device.type == DEVICE_TYPE

599 assert out.shape == src.shape

600 assert out.dtype == src.dtype

601 assert torch.equal(out.cpu(), src.cpu())

tests/v1/ec_connector/unit/test_ec_cpu_connector.py

+7 -4

@@ -7,7 +7,7 @@ Verifies:

7 - Accuracy: outputs from EC CPU cache match fresh encoder computation.

8 - Latency: loading from EC CPU cache is faster than a cold encoder run.

10 -Requires a CUDA GPU and the Qwen2-VL-2B-Instruct model.

10 +Requires a CUDA or XPU GPU and the Qwen2-VL-2B-Instruct model.

11 """

13 import time

@@ -57,8 +57,8 @@ def _wait_for_ec_ready(llm: LLM) -> None:

58 def _latency_test(llm: LLM) -> None:

59 """Verify EC CPU cache hit is faster than cold encoder computation."""

60 - if not current_platform.is_cuda():

61 - pytest.skip("Latency test requires CUDA")

60 + if not (current_platform.is_cuda() or current_platform.is_xpu()):

61 + pytest.skip("Latency test requires an accelerator (CUDA or XPU)")

63 sampling_params = SamplingParams(max_tokens=1, temperature=0)

64 base_image = ImageAsset("stop_sign").pil_image

@@ -134,7 +134,10 @@ def _accuracy_test(llm: LLM) -> None:

134 )

135

136

137 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")

137 +@pytest.mark.skipif(

138 + not (current_platform.is_cuda() or current_platform.is_xpu()),

139 + reason="Requires an accelerator (CUDA or XPU)",

140 +)

141 def test_ec_cpu_offloading() -> None:

142 """Tests ECCPUConnector accuracy and latency with a VLM model."""

143 ec_transfer_config = ECTransferConfig(

vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py

+9 -7

@@ -102,16 +102,15 @@ class ECCPUWorker:

102

103 assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel()

104

105 - src_ptrs, dst_ptrs, sizes = self._save_bufs

105 + bufs = self._save_bufs

106 src_base = src.view(-1).view(torch.uint8).data_ptr()

107 dst_base = self._region.blocks.data_ptr()

108 idx = self._save_count

109

110 for i, block_idx in enumerate(block_ids):

111 start = i * block_size

112 - src_ptrs[idx] = src_base + start

113 - dst_ptrs[idx] = dst_base + block_idx * block_size

114 - sizes[idx] = min(block_size, total_bytes - start)

112 + bufs.set_ptrs(idx, src_base + start, dst_base + block_idx * block_size)

113 + bufs.sizes[idx] = min(block_size, total_bytes - start)

114 idx += 1

115

116 self._save_count = idx

@@ -123,7 +122,7 @@ class ECCPUWorker:

122

123 bufs = self._save_bufs

124 assert bufs is not None

126 - src_ptrs, dst_ptrs, sizes = bufs

125 + src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes

126 n = self._save_count

127 swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n])

128

@@ -173,8 +172,11 @@ class ECCPUWorker:

172 op_idx = 0

173 for block_ids in load_items.values():

174 for block_idx in block_ids:

176 - src_ptrs[op_idx] = src_base + block_idx * block_size

177 - dst_ptrs[op_idx] = dst_buf_base + op_idx * block_size

175 + bufs.set_ptrs(

176 + op_idx,

177 + src_base + block_idx * block_size,

178 + dst_buf_base + op_idx * block_size,

179 + )

180 op_idx += 1

181

182 swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True)

vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py

+31 -8

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Reusable pool of (src_ptrs, dst_ptrs, sizes) int64 tensor triples.

3 +"""Reusable pool of (src_ptrs, dst_ptrs, sizes) tensor triples.

5 Used by ECCPUWorker to batch swap_blocks_batch descriptors without

6 per-step allocation overhead.

@@ -8,20 +8,46 @@ per-step allocation overhead.

9 from typing import NamedTuple

11 +import numpy as np

12 import torch

14 +from vllm.platforms import current_platform

15 +

16 +# CUDA/ROCm cache_kernels.cu requires int64 pointers; the XPU DMA engine

17 +# requires uint64 (see vllm._custom_ops.swap_blocks_batch).

18 +_PTR_DTYPE = torch.uint64 if current_platform.is_xpu() else torch.int64

19 +

21 class DescriptorBuffers(NamedTuple):

22 src_ptrs: torch.Tensor

23 dst_ptrs: torch.Tensor

24 sizes: torch.Tensor

25 + # Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs().

26 + src_np: np.ndarray

27 + dst_np: np.ndarray

28 +

29 + def set_ptrs(self, idx: int, src: int, dst: int) -> None:

30 + """Record the source and destination address of descriptor *idx*.

31 +

32 + TODO(torch>=2.14): drop this indirection and assign the tensors

33 + directly once the minimum supported torch is 2.14. The numpy detour

34 + exists only because torch's setitem unpacks the value as a signed

35 + long long before 2.14 (pytorch#191458), rejecting XPU USM pointers

36 + >= 2**63, while the two's-complement rewrite those versions accept is

37 + in turn rejected by uint64 from 2.14 on. Numpy casts against the

38 + array dtype and so works on either. `sizes` holds byte counts and

39 + needs no such care.

40 + """

41 + self.src_np[idx] = src

42 + self.dst_np[idx] = dst

45 class DescriptorBufferPool:

46 """Pool of descriptor buffer triples for swap_blocks_batch.

23 - Each buffer is a `DescriptorBuffers` namedtuple of three 1-D int64

24 - tensors of equal length. Buffers are recycled across steps; if a

48 + Each buffer is a `DescriptorBuffers` namedtuple of three 1-D tensors

49 + (dtype `_PTR_DTYPE`, platform-dependent) of equal length, paired with

50 + numpy aliases used to fill them. Buffers are recycled across steps; if a

51 returned buffer is too small it is discarded and a fresh one allocated.

52 """

@@ -35,11 +61,8 @@ class DescriptorBufferPool:

61 bufs = self._pool.pop()

62 if bufs.src_ptrs.numel() >= n:

63 return bufs

38 - return DescriptorBuffers(

39 - torch.empty(n, dtype=torch.int64),

40 - torch.empty(n, dtype=torch.int64),

41 - torch.empty(n, dtype=torch.int64),

42 - )

64 + src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3))

65 + return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy())

67 def release(self, bufs: DescriptorBuffers) -> None:

68 """Return a buffer triple to the pool for reuse."""