From b027ac831253e83ff605a2ff5d4aa7c00619802c Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Fri, 9 Oct 2026 18:58:07 +0800 Subject: [PATCH] [XPU] Support register KV offload mmap region as pinned host memory on XPU (#51956) Signed-off-by: Chaojun Zhang --- .../cpu/test_shared_offload_region.py | 81 +++++++--------- tests/v1/kv_offload/test_factory.py | 3 + vllm/v1/kv_offload/cpu/gpu_worker.py | 45 ++------- vllm/v1/kv_offload/cpu/host_register.py | 97 +++++++++++++++++++ .../kv_offload/cpu/shared_offload_region.py | 21 ++-- vllm/v1/kv_offload/cpu/spec.py | 4 +- 6 files changed, 151 insertions(+), 100 deletions(-) create mode 100644 vllm/v1/kv_offload/cpu/host_register.py diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index f84e1bd729..b7259e6fd6 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -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 diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 01ef756cae..84b9da36e0 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -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) diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index a56febf702..f6053480a7 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -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), diff --git a/vllm/v1/kv_offload/cpu/host_register.py b/vllm/v1/kv_offload/cpu/host_register.py new file mode 100644 index 0000000000..3744885740 --- /dev/null +++ b/vllm/v1/kv_offload/cpu/host_register.py @@ -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 diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index af183fbc4f..7e0e97aa12 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -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 diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index fddf566085..bf289d37a6 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -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