iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[XPU] Support register KV offload mmap region as pinned host memory on XPU (#51956)

Browse Source 
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
This commit is contained in:
Chaojun Zhang
2026-10-09 18:58:07 +08:00
committed by GitHub
parent 5cadabfafd
commit b027ac8312
6 changed files with 151 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+35 -46
View File
@@ -772,17 +772,14 @@ def test_cleanup_idempotent(iid):
def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):
"""cleanup() must release every chunk from a split host registration."""
r = _make_region(iid)
cudart = MagicMock()
cudart.cudaHostUnregister.return_value = MagicMock(value=0)
monkeypatch.setattr(region_module, "current_platform", MagicMock())
region_module.current_platform.is_cuda_alike.return_value = True
monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart)
unregister = MagicMock()
monkeypatch.setattr(region_module, "host_unregister", unregister)
r.pinned_addresses = [0x100000, 0x200000, 0x300000]
r.is_pinned = True
r.cleanup()
assert cudart.cudaHostUnregister.call_args_list == [
assert unregister.call_args_list == [
call(0x300000),
call(0x200000),
call(0x100000),
@@ -791,40 +788,35 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):
@pytest.fixture
def host_register(monkeypatch):
"""Pretend to run on CUDA with host registrations capped at seven pages.
Registration goes to a mock CudaRTLibrary and the region's cleanup to a
separate mock torch runtime. Both stay installed through the cleanup that
``_region`` runs on exit.
"""Mock the cross-platform host_register/host_unregister helpers with
registrations capped at seven pages. The same unregister mock backs both
pin_mmap_region's rollback path and the region's cleanup() on exit from
``_region``.
"""
cudart = MagicMock()
cudart.cudaHostRegister.return_value = 0
cudart.cudaHostUnregister.return_value = 0
cudart.cudaGetLastError.return_value = 2
torch_cudart = MagicMock()
torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0)
monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True)
register = MagicMock(return_value=True)
unregister = MagicMock()
monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE)
monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart)
monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart)
return cudart, torch_cudart
monkeypatch.setattr(gpu_worker, "host_register", register)
monkeypatch.setattr(gpu_worker, "host_unregister", unregister)
monkeypatch.setattr(region_module, "host_unregister", unregister)
return register, unregister
def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register):
"""Chunks end on row boundaries: 3-page rows under a 7-page cap register
as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1."""
cudart, torch_cudart = host_register
register, unregister = host_register
with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
gpu_worker.pin_mmap_region(region)
base = region._base.data_ptr()
assert cudart.mock_calls == [
call.cudaHostRegister(base, 6 * PAGE_SIZE),
call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE),
call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE),
assert register.call_args_list == [
call(base, 6 * PAGE_SIZE),
call(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE),
call(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE),
]
assert region.is_pinned
assert torch_cudart.cudaHostUnregister.call_args_list == [
assert unregister.call_args_list == [
call(base + 12 * PAGE_SIZE),
call(base + 6 * PAGE_SIZE),
call(base),
@@ -833,10 +825,10 @@ def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register):
@pytest.mark.parametrize("fail_at", [0, 1, 2])
def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at):
"""A failed chunk drains its pending error and unregisters the chunks
before it on the same handle, so the region is pinned whole or not at all."""
cudart, torch_cudart = host_register
cudart.cudaHostRegister.side_effect = [0] * fail_at + [2]
"""A failed chunk unregisters the chunks before it on the same range, so
the region is pinned whole or not at all."""
register, unregister = host_register
register.side_effect = [True] * fail_at + [False]
with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
gpu_worker.pin_mmap_region(region)
base = region._base.data_ptr()
@@ -845,25 +837,22 @@ def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail
registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)]
failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE
assert cudart.mock_calls == [
*(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered),
call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size),
call.cudaGetLastError(),
*(call.cudaHostUnregister(address) for address in reversed(registered)),
assert register.call_args_list == [
*(call(address, 6 * PAGE_SIZE) for address in registered),
call(base + fail_at * 6 * PAGE_SIZE, failed_size),
]
assert unregister.call_args_list == [
call(address) for address in reversed(registered)
]
torch_cudart.cudaHostUnregister.assert_not_called()
def test_pin_mmap_region_without_cuda_runtime_stays_pageable(
iid, host_register, monkeypatch
def test_pin_mmap_region_when_host_register_unsupported_stays_pageable(
iid, monkeypatch
):
"""A CUDA runtime that cannot be loaded leaves the region pageable instead
of failing worker startup."""
def missing_runtime():
raise AssertionError("libcudart is not loaded in the current process")
monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime)
"""A platform without host registration support (host_register()
returning False) leaves the region pageable instead of failing worker
startup."""
monkeypatch.setattr(gpu_worker, "host_register", lambda ptr, num_bytes: False)
with _region(iid) as region:
gpu_worker.pin_mmap_region(region)
assert not region.is_pinned
tests/v1/kv_offload/test_factory.py
+3
View File
@@ -307,6 +307,7 @@ def test_cpu_spec_replicated_disabled_without_shared_region(
# those private buffers empty and corrupting subsequent loads.
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False)
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
@@ -341,6 +342,7 @@ def test_cpu_spec_replicated_layout_truth_matrix(
# actually allocates on the shared region.
import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False)
monkeypatch.setattr(
cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike
)
@@ -417,6 +419,7 @@ def test_cpu_spec_create_worker_uses_tensor_path_without_shared_region(
worker_calls.append(kwargs)
return MagicMock()
monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False)
monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
vllm/v1/kv_offload/cpu/gpu_worker.py
+7 -38
View File
@@ -11,7 +11,6 @@ import numpy as np
import torch
from vllm import _custom_ops as ops
from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, triton
@@ -27,6 +26,7 @@ from vllm.v1.kv_offload.base import (
OffloadingWorker,
TransferResult,
)
from vllm.v1.kv_offload.cpu.host_register import host_register, host_unregister
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.kv_offload.cpu.swap_blocks_triton import (
THRESHOLD_BYTES,
@@ -200,26 +200,7 @@ MAX_HOST_REGISTER_CHUNK_BYTES = 64 * 1024**3
def pin_mmap_region(region: SharedOffloadRegion) -> None:
"""Register row-aligned chunks, rolling back on failure."""
if not current_platform.is_cuda_alike():
logger.info(
"Skipping mmap host registration on %s; cudaHostRegister is only "
"available on CUDA/ROCm.",
current_platform.device_name,
)
return
rank = region.rank
try:
cudart = CudaRTLibrary()
except (AssertionError, AttributeError, OSError):
logger.warning(
"Could not load the CUDA runtime for host registration on rank=%d; "
"the offload region stays pageable",
rank,
exc_info=True,
)
return
base_ptr = region._base.data_ptr()
total_size = region.total_size_bytes
# Chunks end on block-row boundaries, which are page aligned, so neither the
@@ -228,42 +209,30 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:
rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1)
chunk_size = rows_per_chunk * region._row_stride
# Register, drain and roll back through the same runtime handle, so a
# failed chunk leaves neither a pending error nor a partly pinned region.
# Register, drain and roll back through the same helper, so a failed
# chunk leaves neither a pending error nor a partly pinned region.
addresses: list[int] = []
for offset in range(0, total_size, chunk_size):
address = base_ptr + offset
size = min(chunk_size, total_size - offset)
result = cudart.cudaHostRegister(address, size)
if result == 0:
if host_register(address, size):
addresses.append(address)
continue
cudart.cudaGetLastError()
logger.warning(
"cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); "
"host_register failed for rank=%d at %.2f of %.2f GB; "
"the offload region stays pageable",
rank,
offset / 1e9,
total_size / 1e9,
result,
)
for registered in reversed(addresses):
unregister_result = cudart.cudaHostUnregister(registered)
if unregister_result != 0:
cudart.cudaGetLastError()
logger.warning(
"cudaHostUnregister failed for rank=%d at %#x (code=%d); "
"that chunk stays registered until the process exits",
rank,
registered,
unregister_result,
)
host_unregister(registered)
return
region.pinned_addresses.extend(addresses)
region.is_pinned = True
logger.debug(
"cudaHostRegister rank=%d %.2f GB in %d chunk(s)",
"Host-registered mmap region rank=%d %.2f GB in %d chunk(s)",
rank,
total_size / 1e9,
len(addresses),
vllm/v1/kv_offload/cpu/host_register.py
+97
View File
@@ -0,0 +1,97 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Registration of existing host memory for pinned DMA transfers.
Host allocations that the offloading workers share across processes (the
`SharedOffloadRegion` mmap) cannot be allocated through the accelerator's
pinned allocator, so they must be registered after the fact. Pinning is a
pure bandwidth optimization: when registration is unavailable the transfers
still run correctly, just as pageable (slower) copies.
"""
import torch
from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
from vllm.logger import init_logger
from vllm.platforms import current_platform
logger = init_logger(__name__)
def _load_cudart() -> CudaRTLibrary | None:
try:
return CudaRTLibrary()
except (AssertionError, AttributeError, OSError):
logger.warning(
"Could not load the CUDA runtime for host registration; "
"the offload region stays pageable",
exc_info=True,
)
return None
def host_register(ptr: int, num_bytes: int) -> bool:
"""Register a host range as device-accessible pinned memory.
Args:
ptr: Base address of the host range.
num_bytes: Size of the host range in bytes.
Returns:
True if the range was registered and later needs `host_unregister`.
"""
if not ptr or num_bytes <= 0:
return False
if current_platform.is_cuda_alike():
cudart = _load_cudart()
if cudart is None:
return False
result = cudart.cudaHostRegister(ptr, num_bytes)
if result != 0:
cudart.cudaGetLastError()
logger.warning(
"cudaHostRegister failed (code=%d) — transfers will still work "
"but may be slower (unpinned DMA)",
result,
)
return False
return True
if current_platform.is_xpu():
if not torch.ops._C.xpu_host_register(ptr, num_bytes):
logger.warning(
"xpu_host_register failed — transfers will still work "
"but may be slower (unpinned DMA)"
)
return False
return True
logger.info_once(
"Host registration is not implemented on %s; KV offload transfers will "
"use unpinned host memory.",
current_platform.device_name,
)
return False
def host_unregister(ptr: int) -> None:
"""Release a host range previously registered by `host_register`."""
if not ptr:
return
if current_platform.is_cuda_alike():
cudart = _load_cudart()
if cudart is None:
return
result = cudart.cudaHostUnregister(ptr)
if result != 0:
cudart.cudaGetLastError()
logger.warning("cudaHostUnregister failed (code=%d)", result)
return
if current_platform.is_xpu():
if not torch.ops._C.xpu_host_unregister(ptr):
logger.warning("xpu_host_unregister failed")
return
vllm/v1/kv_offload/cpu/shared_offload_region.py
+6 -15
View File
@@ -14,7 +14,7 @@ from vllm.distributed.device_communicators.shm_broadcast import (
check_shm_free_space,
)
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.v1.kv_offload.cpu.host_register import host_unregister
logger = init_logger(__name__)
@@ -343,20 +343,11 @@ class SharedOffloadRegion:
def cleanup(self) -> None:
if self.is_pinned and self._base is not None:
if current_platform.is_cuda_alike():
base_ptr = self._base.data_ptr()
addresses = self.pinned_addresses or [base_ptr]
for address in reversed(addresses):
result = torch.cuda.cudart().cudaHostUnregister(address)
if result.value != 0:
logger.warning(
"cudaHostUnregister failed for rank=%d, "
"address=%#x (code=%d)",
self.rank,
address,
result.value,
)
self.pinned_addresses.clear()
base_ptr = self._base.data_ptr()
addresses = self.pinned_addresses or [base_ptr]
for address in reversed(addresses):
host_unregister(address)
self.pinned_addresses.clear()
self.is_pinned = False
# Release views before _base: each view holds a _base reference and a
# direct StorageImpl reference. Freeing views first lets both refcounts
vllm/v1/kv_offload/cpu/spec.py
+3 -1
View File
@@ -162,7 +162,9 @@ class CPUOffloadingSpec(OffloadingSpec):
def _uses_shared_region(self) -> bool:
"""Whether the worker CPU buffer is the shared mmap region (vs a private
per-rank tensor); replicated-layout dedup is gated on this being True."""
return current_platform.is_cuda_alike() and not current_platform.is_rocm()
return (
current_platform.is_cuda_alike() and not current_platform.is_rocm()
) or current_platform.is_xpu()
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
mmap_region: SharedOffloadRegion | None = None