@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 40b40d1d39

[Bugfix][KV Offload] Register the offload region in chunks (#51081)

Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Colton Ottley <coltonottley@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

drakosha committed Sep 17, 2026 at 13:19 UTC 40b40d1d39bb82563dd699be3456abaa95d0722a

5 files changed +198 -15

tests/v1/kv_offload/cpu/test_canonical_layout.py

@@ -241,6 +241,7 @@ def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int):

241 # The Triton load path dereferences CPU pointers on the GPU, which is

242 # only legal on pinned memory; production pins via CPUOffloadingWorker

243 pin_mmap_region(region)

244 + assert region.is_pinned

245 return region.create_next_canonical_view(_CANONICAL_PAGE)

246

247 try:

tests/v1/kv_offload/cpu/test_gpu_worker.py

+16

@@ -49,6 +49,22 @@ def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:

49 )

52 +@pytest.mark.skipif(not gpu_worker.HAS_TRITON, reason="Requires Triton")

53 +def test_unpinned_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:

54 + """The Triton load path dereferences CPU pointers on the GPU, so pageable

55 + host memory takes the DMA path even for pages where Triton would win."""

56 + monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False)

57 + monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: False)

58 +

59 + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]]

60 + assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is not (

61 + ops.swap_blocks_batch

62 + )

63 + assert gpu_worker._select_swap_blocks_fn(

64 + refs, gpu_to_cpu=False, host_memory_is_pinned=False

65 + ) is (ops.swap_blocks_batch)

66 +

67 +

68 def test_worker_shutdown_releases_region_and_runs_both_handlers() -> None:

69 """Both directions drain before the worker releases its shared region."""

70 worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker)

tests/v1/kv_offload/cpu/test_shared_offload_region.py

+82

@@ -15,6 +15,7 @@ from unittest.mock import MagicMock, call

15 import pytest

17 from vllm.utils.system_utils import get_mp_context

18 +from vllm.v1.kv_offload.cpu import gpu_worker

19 from vllm.v1.kv_offload.cpu import shared_offload_region as region_module

20 from vllm.v1.kv_offload.cpu.shared_offload_region import (

21 SharedOffloadRegion,

@@ -788,6 +789,87 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):

789 ]

790

791

792 +@pytest.fixture

793 +def host_register(monkeypatch):

794 + """Pretend to run on CUDA with host registrations capped at seven pages.

795 +

796 + Registration goes to a mock CudaRTLibrary and the region's cleanup to a

797 + separate mock torch runtime. Both stay installed through the cleanup that

798 + ``_region`` runs on exit.

799 + """

800 + cudart = MagicMock()

801 + cudart.cudaHostRegister.return_value = 0

802 + cudart.cudaHostUnregister.return_value = 0

803 + cudart.cudaGetLastError.return_value = 2

804 + torch_cudart = MagicMock()

805 + torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0)

806 + monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True)

807 + monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE)

808 + monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart)

809 + monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart)

810 + return cudart, torch_cudart

811 +

812 +

813 +def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register):

814 + """Chunks end on row boundaries: 3-page rows under a 7-page cap register

815 + as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1."""

816 + cudart, torch_cudart = host_register

817 + with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:

818 + gpu_worker.pin_mmap_region(region)

819 + base = region._base.data_ptr()

820 + assert cudart.mock_calls == [

821 + call.cudaHostRegister(base, 6 * PAGE_SIZE),

822 + call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE),

823 + call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE),

824 + ]

825 + assert region.is_pinned

826 +

827 + assert torch_cudart.cudaHostUnregister.call_args_list == [

828 + call(base + 12 * PAGE_SIZE),

829 + call(base + 6 * PAGE_SIZE),

830 + call(base),

831 + ]

832 +

833 +

834 +@pytest.mark.parametrize("fail_at", [0, 1, 2])

835 +def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at):

836 + """A failed chunk drains its pending error and unregisters the chunks

837 + before it on the same handle, so the region is pinned whole or not at all."""

838 + cudart, torch_cudart = host_register

839 + cudart.cudaHostRegister.side_effect = [0] * fail_at + [2]

840 + with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:

841 + gpu_worker.pin_mmap_region(region)

842 + base = region._base.data_ptr()

843 + assert not region.is_pinned

844 + assert region.pinned_addresses == []

845 +

846 + registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)]

847 + failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE

848 + assert cudart.mock_calls == [

849 + *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered),

850 + call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size),

851 + call.cudaGetLastError(),

852 + *(call.cudaHostUnregister(address) for address in reversed(registered)),

853 + ]

854 + torch_cudart.cudaHostUnregister.assert_not_called()

855 +

856 +

857 +def test_pin_mmap_region_without_cuda_runtime_stays_pageable(

858 + iid, host_register, monkeypatch

859 +):

860 + """A CUDA runtime that cannot be loaded leaves the region pageable instead

861 + of failing worker startup."""

862 +

863 + def missing_runtime():

864 + raise AssertionError("libcudart is not loaded in the current process")

865 +

866 + monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime)

867 + with _region(iid) as region:

868 + gpu_worker.pin_mmap_region(region)

869 + assert not region.is_pinned

870 + assert region.pinned_addresses == []

871 +

872 +

873 def test_cleanup_after_create_next_worker_view_releases_mmap(iid):

874 """cleanup() must close the mmap even after create_next_worker_view was called.

875 create_next_worker_view returns a view that shares storage with _base; both must be

vllm/distributed/device_communicators/cuda_wrapper.py

+26

@@ -78,6 +78,16 @@ class CudaRTLibrary:

78 cudaError_t,

79 [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],

80 ),

81 + # cudaError_t cudaHostRegister ( void* ptr, size_t size, unsigned int flags )

82 + Function(

83 + "cudaHostRegister",

84 + cudaError_t,

85 + [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint],

86 + ),

87 + # cudaError_t cudaHostUnregister ( void* ptr )

88 + Function("cudaHostUnregister", cudaError_t, [ctypes.c_void_p]),

89 + # cudaError_t cudaGetLastError ( void )

90 + Function("cudaGetLastError", cudaError_t, []),

91 ]

93 # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa

@@ -92,6 +102,9 @@ class CudaRTLibrary:

102 "cudaMemcpy": "hipMemcpy",

103 "cudaIpcGetMemHandle": "hipIpcGetMemHandle",

104 "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",

105 + "cudaHostRegister": "hipHostRegister",

106 + "cudaHostUnregister": "hipHostUnregister",

107 + "cudaGetLastError": "hipGetLastError",

108 }

109

110 # class attribute to store the mapping from the path to the library

@@ -185,3 +198,16 @@ class CudaRTLibrary:

198 )

199 )

200 return devPtr

201 +

202 + def cudaHostRegister(self, ptr: int, size: int, flags: int = 0) -> int:

203 + """Return the raw error code instead of raising, since callers can

204 + fall back to pageable memory. Drain a failure with cudaGetLastError."""

205 + return self.funcs["cudaHostRegister"](ctypes.c_void_p(ptr), size, flags)

206 +

207 + def cudaHostUnregister(self, ptr: int) -> int:

208 + """Return the raw error code instead of raising."""

209 + return self.funcs["cudaHostUnregister"](ctypes.c_void_p(ptr))

210 +

211 + def cudaGetLastError(self) -> int:

212 + """Return and clear the error pending on this thread."""

213 + return self.funcs["cudaGetLastError"]()

vllm/v1/kv_offload/cpu/gpu_worker.py

+73 -15

@@ -11,6 +11,7 @@ import numpy as np

11 import torch

13 from vllm import _custom_ops as ops

14 +from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary

15 from vllm.logger import init_logger

16 from vllm.platforms import current_platform

17 from vllm.triton_utils import HAS_TRITON, triton

@@ -38,10 +39,13 @@ logger = init_logger(__name__)

39 def _select_swap_blocks_fn(

40 layer_refs_per_group: list[list[CanonicalKVCacheRef]],

41 gpu_to_cpu: bool,

42 + host_memory_is_pinned: bool = True,

43 ):

44 """Resolve the swap_blocks function for a handler at init time."""

45 # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton.

44 - if gpu_to_cpu:

46 + # The Triton kernel dereferences CPU pointers on the GPU, which is only

47 + # valid for pinned host memory.

48 + if gpu_to_cpu or not host_memory_is_pinned:

49 return ops.swap_blocks_batch

50 # Fall back to the C++ DMA path on platforms where Triton isn't usable

51 # (e.g. ROCm host mappings) or where GPU kernels cannot directly

@@ -190,8 +194,12 @@ def _canonical_block_sizes(

194 return canonical_bytes_per_block

195

196

197 +# Bound registration size to avoid driver limits on large host allocations.

198 +MAX_HOST_REGISTER_CHUNK_BYTES = 64 * 1024**3

199 +

200 +

201 def pin_mmap_region(region: SharedOffloadRegion) -> None:

194 - """Register the entire mmap as CUDA pinned memory via cudaHostRegister."""

202 + """Register row-aligned chunks, rolling back on failure."""

203 if not current_platform.is_cuda_alike():

204 logger.info(

205 "Skipping mmap host registration on %s; cudaHostRegister is only "

@@ -201,23 +209,65 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:

209 return

210

211 rank = region.rank

212 + try:

213 + cudart = CudaRTLibrary()

214 + except (AssertionError, AttributeError, OSError):

215 + logger.warning(

216 + "Could not load the CUDA runtime for host registration on rank=%d; "

217 + "the offload region stays pageable",

218 + rank,

219 + exc_info=True,

220 + )

221 + return

222

223 base_ptr = region._base.data_ptr()

206 - result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0)

207 - if result.value != 0:

224 + total_size = region.total_size_bytes

225 + # Chunks end on block-row boundaries, which are page aligned, so neither the

226 + # driver's page rounding nor any single block transfer straddles two

227 + # registrations.

228 + rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1)

229 + chunk_size = rows_per_chunk * region._row_stride

230 +

231 + # Register, drain and roll back through the same runtime handle, so a

232 + # failed chunk leaves neither a pending error nor a partly pinned region.

233 + addresses: list[int] = []

234 + for offset in range(0, total_size, chunk_size):

235 + address = base_ptr + offset

236 + size = min(chunk_size, total_size - offset)

237 + result = cudart.cudaHostRegister(address, size)

238 + if result == 0:

239 + addresses.append(address)

240 + continue

241 + cudart.cudaGetLastError()

242 logger.warning(

209 - "cudaHostRegister failed for rank=%d (code=%d) — "

210 - "transfers will still work but may be slower (unpinned DMA)",

243 + "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); "

244 + "the offload region stays pageable",

245 rank,

246 + offset / 1e9,

247 + total_size / 1e9,

248 result,

249 )

214 - else:

215 - logger.debug(

216 - "cudaHostRegister rank=%d %.2f GB",

217 - rank,

218 - region.total_size_bytes / 1e9,

219 - )

220 - region.is_pinned = True

250 + for registered in reversed(addresses):

251 + unregister_result = cudart.cudaHostUnregister(registered)

252 + if unregister_result != 0:

253 + cudart.cudaGetLastError()

254 + logger.warning(

255 + "cudaHostUnregister failed for rank=%d at %#x (code=%d); "

256 + "that chunk stays registered until the process exits",

257 + rank,

258 + registered,

259 + unregister_result,

260 + )

261 + return

262 +

263 + region.pinned_addresses.extend(addresses)

264 + region.is_pinned = True

265 + logger.debug(

266 + "cudaHostRegister rank=%d %.2f GB in %d chunk(s)",

267 + rank,

268 + total_size / 1e9,

269 + len(addresses),

270 + )

271

272

273 def _new_descriptor_buffers(

@@ -248,6 +298,7 @@ class SingleDirectionOffloadingHandler:

298 layer_refs_per_group: list[list[CanonicalKVCacheRef]],

299 gpu_to_cpu: bool,

300 canonical_layout: bool = False,

301 + host_memory_is_pinned: bool = True,

302 ):

303 """Initialize a SingleDirectionOffloadingHandler.

304

@@ -262,6 +313,8 @@ class SingleDirectionOffloadingHandler:

313 gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU.

314 canonical_layout: if True, CPU pages use the canonical layout

315 described by the refs' mappings.

316 + host_memory_is_pinned: whether the CPU tensors are pinned, so GPU

317 + kernels may dereference them directly.

318

319 """

320 assert len(gpu_tensors) == len(cpu_tensors)

@@ -299,7 +352,7 @@ class SingleDirectionOffloadingHandler:

352 self.gpu_to_cpu: bool = gpu_to_cpu

353 self.layer_refs_per_group = layer_refs_per_group

354 self._swap_blocks_batch = _select_swap_blocks_fn(

302 - layer_refs_per_group, gpu_to_cpu

355 + layer_refs_per_group, gpu_to_cpu, host_memory_is_pinned

356 )

357

358 # GPU blocks may be smaller

@@ -646,7 +699,7 @@ class SingleDirectionOffloadingHandler:

699 last_event = last_transfer.end_event

700 # assure job will start only after the previous one completes

701 stream.wait_event(last_event)

649 - # CPU->GPU reads from host pinned memory, which is never written

702 + # CPU->GPU reads from host memory, which is never written

703 # by a concurrent GPU stream, so CU_MEMCPY_SRC_ACCESS_ORDER_ANY is

704 # safe and lets the driver pipeline source reads. GPU->CPU reads

705 # from the live GPU KV cache, which the compute stream keeps

@@ -764,6 +817,9 @@ class CPUOffloadingWorker(OffloadingWorker):

817 logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))

818 if mmap_region is not None and pin_memory:

819 pin_mmap_region(mmap_region)

820 + host_memory_is_pinned = pin_memory and (

821 + mmap_region is None or mmap_region.is_pinned

822 + )

823

824 canonical_bytes_per_block = (

825 _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors))

@@ -813,6 +869,7 @@ class CPUOffloadingWorker(OffloadingWorker):

869 layer_refs_per_group=kv_caches.group_data_refs,

870 gpu_to_cpu=True,

871 canonical_layout=canonical_layout,

872 + host_memory_is_pinned=host_memory_is_pinned,

873 )

874

875 self._load_handler = SingleDirectionOffloadingHandler(

@@ -822,6 +879,7 @@ class CPUOffloadingWorker(OffloadingWorker):

879 layer_refs_per_group=kv_caches.group_data_refs,

880 gpu_to_cpu=False,

881 canonical_layout=canonical_layout,

882 + host_memory_is_pinned=host_memory_is_pinned,

883 )

884

885 def submit_store(