From 23ccd3ea533c6b4068f03c3d878ddaa65603ce74 Mon Sep 17 00:00:00 2001 From: jooho-xcena Date: Wed, 15 Jul 2026 03:44:00 +0000 Subject: [PATCH] feat: chunked CUDA host registration with eager/lazy pin modes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on the rc-check fix (#65) to make >= 512GiB pools actually pin. A single cudaHostRegister call of >= 2^39 bytes (512GiB) fails on NVIDIA r580+ drivers (per-registration page-table kvzalloc > INT_MAX, "NVRM: failed to allocate page table"); the cap is per call, not cumulative, so large regions are registered in chunks. - Register regions in MARU_PIN_CHUNK_GB chunks (default 128GiB, clamped to 256GiB; 0 = single call), recording (addr, size) per pinned chunk on MappedRegion and unregistering per chunk on unmap/close (replaces the single _cuda_pinned flag from #65). - MARU_PIN_MODE=eager (default) pins synchronously inside map_region(), preserving the fully-pinned-on-return contract; MARU_PIN_MODE=lazy pins from a background thread. A per-region pin lock + stop event serialize the pin loop against unmap-time unpinning, bounding teardown to one in-flight chunk. Verified: unit suite 766 passed; E2E on gaia (/dev/dax0.0) with pool_size=520 — previously crashed at startup, now pins 5 chunks in 4.45s and the benchmark passes 20/20 with cache-hit TTFT 1.8s. --- maru_handler/memory/mapper.py | 308 +++++++++++++++++++++++++--------- maru_handler/memory/types.py | 22 ++- maru_vllm/connector.py | 10 +- tests/unit/test_mapper.py | 183 ++++++++++++++++++-- 4 files changed, 422 insertions(+), 101 deletions(-) diff --git a/maru_handler/memory/mapper.py b/maru_handler/memory/mapper.py index 2228f51..259c6bc 100644 --- a/maru_handler/memory/mapper.py +++ b/maru_handler/memory/mapper.py @@ -24,6 +24,34 @@ _PREFAULT_ENABLED = os.environ.get("MARU_PREFAULT", "1") != "0" _PAGE_SIZE = os.sysconf("SC_PAGESIZE") if hasattr(os, "sysconf") else 4096 +# CUDA pinning strategy. +# +# A single cudaHostRegister call fails with cudaErrorMemoryAllocation for +# ranges >= 2^39 bytes (512 GiB) — the NVIDIA kernel driver cannot build a +# page table for one mapping that large ("NVRM: failed to allocate page +# table"). The limit is per call, not cumulative, so large regions are +# registered in chunks. +# +# MARU_PIN_MODE: "eager" (default) — register all chunks synchronously +# inside map_region(); the region is fully pinned when +# map_region() returns. +# "lazy" — return immediately and register chunks in a +# background thread (LMCache LazyMemoryAllocator style). +# Chunks not yet pinned degrade GPU DMA to pageable +# copies but remain fully usable. +# MARU_PIN_CHUNK_GB: chunk size in GiB (default 128, clamped to 256 so a +# chunk can never reach the 512 GiB wall; 0 = single +# call, i.e. the pre-chunking behavior). +_PIN_MODE = os.environ.get("MARU_PIN_MODE", "eager") +if _PIN_MODE not in ("eager", "lazy"): + _PIN_MODE = "eager" +try: + _PIN_CHUNK_BYTES = int(os.environ.get("MARU_PIN_CHUNK_GB", "128")) << 30 +except ValueError: + _PIN_CHUNK_BYTES = 128 << 30 +if _PIN_CHUNK_BYTES > 256 << 30: + _PIN_CHUNK_BYTES = 256 << 30 + def _clear_cuda_sticky_error() -> None: """Consume CUDA's per-thread sticky error after a failed cudaHostRegister. @@ -49,6 +77,67 @@ def _clear_cuda_sticky_error() -> None: logger.debug("could not clear CUDA sticky error", exc_info=True) +def _cuda_pin_chunks( + cudart, + base_addr: int, + length: int, + region_id: int, + records: list[tuple[int, int]], + stop_event: threading.Event | None = None, +) -> None: + """Register [base_addr, base_addr+length) with CUDA in chunks. + + Appends an (addr, size) record to ``records`` for every successfully + registered chunk (records are what cudaHostUnregister must be called + with later). A failed chunk is logged and its sticky CUDA error + cleared; pinning continues so the remaining chunks still get DMA. + """ + chunk = _PIN_CHUNK_BYTES if _PIN_CHUNK_BYTES > 0 else length + off = 0 + while off < length: + if stop_event is not None and stop_event.is_set(): + return + n = min(chunk, length - off) + rc = cudart.cudaHostRegister(base_addr + off, n, 0) + rc = int(rc[0]) if isinstance(rc, tuple) else int(rc) + if rc == 0: + records.append((base_addr + off, n)) + else: + _clear_cuda_sticky_error() + logger.warning( + "cudaHostRegister failed for region %d chunk [+%d, +%d): " + "rc=%d; chunk stays unpinned (GPU DMA degraded)", + region_id, + off, + off + n, + rc, + ) + off += n + + +def _cuda_unpin_records( + cudart, records: list[tuple[int, int]], region_id: int +) -> None: + """Unregister every previously pinned chunk of a region.""" + for addr, _size in records: + try: + rc = cudart.cudaHostUnregister(addr) + rc = int(rc[0]) if isinstance(rc, tuple) else int(rc) + if rc != 0: + _clear_cuda_sticky_error() + logger.warning( + "cudaHostUnregister failed for region %d addr=%#x: rc=%d", + region_id, + addr, + rc, + ) + except (RuntimeError, OSError) as e: + logger.warning( + "cudaHostUnregister failed for region %d: %s", region_id, e + ) + records.clear() + + class DaxMapper: """Maps shared memory regions via MaruShmClient. @@ -127,6 +216,15 @@ def map_region( ) self._regions[region_id] = region + # Base address for CUDA pinning — resolved while the lock still + # guarantees the buffer view is alive (a concurrent unmap after + # publication would otherwise race the from_buffer export). + pin_addr = ( + ctypes.addressof(ctypes.c_char.from_buffer(region._buffer_view)) + if region._buffer_view is not None + else None + ) + logger.debug( "Mapped region %d: length=%d", region_id, @@ -140,43 +238,55 @@ def map_region( prefault_ms = (time.monotonic() - t0) * 1000 # Outside lock: CUDA pin is idempotent - if region._buffer_view is not None: + if pin_addr is not None: try: import torch if torch.cuda.is_available(): - addr = ctypes.addressof( - ctypes.c_char.from_buffer(region._buffer_view) - ) - t0 = time.monotonic() - rc = torch.cuda.cudart().cudaHostRegister( - addr, handle.length, 0 - ) - cuda_pin_ms = (time.monotonic() - t0) * 1000 - rc = int(rc[0]) if isinstance(rc, tuple) else int(rc) - if rc == 0: - region._cuda_pinned = True + cudart = torch.cuda.cudart() + if _PIN_MODE == "lazy": + thread = threading.Thread( + target=self._pin_worker, + args=(cudart, region, pin_addr, handle.length), + daemon=True, + name=f"maru-pin-{region_id}", + ) + thread.start() + # Published only after a successful start() so unmap + # never joins a never-started thread. + region._pin_thread = thread logger.info( - "CUDA pinned region %d (%d bytes)", + "CUDA pinning region %d in background (%d bytes)", region_id, handle.length, ) else: - # NVIDIA r580+ drivers reject a single registration - # of >= 512 GiB (per-registration page table hits - # the kernel's kvmalloc INT_MAX cap: "NVRM: failed - # to allocate page table"). Clear the per-thread - # sticky error so an unrelated later CUDA call does - # not crash with a misattributed OOM. - _clear_cuda_sticky_error() - logger.warning( - "cudaHostRegister failed for region %d " - "(%d bytes): rc=%d — region stays unpinned, " - "GPU transfers fall back to pageable copies", - region_id, - handle.length, - rc, - ) + t0 = time.monotonic() + with region._pin_lock: + _cuda_pin_chunks( + cudart, pin_addr, handle.length, region_id, + region._pin_records, + stop_event=region._pin_stop, + ) + pinned = sum(n for _, n in region._pin_records) + n_chunks = len(region._pin_records) + cuda_pin_ms = (time.monotonic() - t0) * 1000 + if pinned == handle.length: + logger.info( + "CUDA pinned region %d (%d bytes, %d chunks)", + region_id, + pinned, + n_chunks, + ) + else: + logger.warning( + "CUDA pinned region %d PARTIALLY: %d/%d bytes" + " — unpinned ranges fall back to pageable" + " copies", + region_id, + pinned, + handle.length, + ) except (ImportError, RuntimeError, OSError) as e: if not isinstance(e, ImportError): logger.warning( @@ -199,6 +309,91 @@ def map_region( return region + @staticmethod + def _cuda_unpin_region(region) -> None: + """Stop/join any in-flight pinning, then unregister pinned chunks. + + Setting _pin_stop makes a concurrent pin loop (eager in another + thread, or the lazy worker) exit at its next chunk boundary, so + acquiring _pin_lock below waits for at most one chunk. + """ + region._pin_stop.set() + if region._pin_thread is not None: + region._pin_thread.join() + region._pin_thread = None + with region._pin_lock: + if not region._pin_records: + return + try: + import torch + + if torch.cuda.is_available(): + _cuda_unpin_records( + torch.cuda.cudart(), + region._pin_records, + region.region_id, + ) + except (ImportError, RuntimeError, OSError) as e: + if not isinstance(e, ImportError): + logger.warning( + "cudaHostUnregister failed for region %d: %s", + region.region_id, + e, + ) + + def _pin_worker( + self, cudart, region, base_addr: int, length: int + ) -> None: + """Background chunk-pinning worker (MARU_PIN_MODE=lazy). + + Appends to region._pin_records as chunks are registered so that + unmap_region()/close() — which join this thread first — always + unregister exactly what was pinned. + """ + t0 = time.monotonic() + try: + with region._pin_lock: + _cuda_pin_chunks( + cudart, + base_addr, + length, + region.region_id, + region._pin_records, + stop_event=region._pin_stop, + ) + pinned = sum(n for _, n in region._pin_records) + except (RuntimeError, OSError) as e: + logger.warning( + "cudaHostRegister failed for region %d: %s", + region.region_id, + e, + ) + return + if region._pin_stop.is_set() and pinned < length: + logger.info( + "CUDA background pinning region %d cancelled at %d/%d bytes", + region.region_id, + pinned, + length, + ) + elif pinned < length: + logger.warning( + "CUDA background pinning region %d PARTIAL: %d/%d bytes — " + "unpinned ranges fall back to pageable copies", + region.region_id, + pinned, + length, + ) + else: + logger.info( + "CUDA background pinning region %d finished: %d/%d bytes " + "in %.1f ms", + region.region_id, + pinned, + length, + (time.monotonic() - t0) * 1000, + ) + @staticmethod def _prefault_region(mmap_obj: mmap_module.mmap, region_id: int, size: int) -> None: """Pre-fault all pages in a mapped region. @@ -248,33 +443,9 @@ def unmap_region(self, region_id: int) -> bool: return False try: - # CUDA unpin before munmap (order matters — needs buffer_view for addr) - if region._cuda_pinned and region._buffer_view is not None: - try: - import torch - - if torch.cuda.is_available(): - addr = ctypes.addressof( - ctypes.c_char.from_buffer(region._buffer_view) - ) - rc = torch.cuda.cudart().cudaHostUnregister(addr) - rc = int(rc[0]) if isinstance(rc, tuple) else int(rc) - if rc != 0: - _clear_cuda_sticky_error() - logger.warning( - "cudaHostUnregister failed for region " - "%d: rc=%d", - region_id, - rc, - ) - region._cuda_pinned = False - except (ImportError, RuntimeError, OSError) as e: - if not isinstance(e, ImportError): - logger.warning( - "cudaHostUnregister failed for region %d: %s", - region_id, - e, - ) + # CUDA unpin before munmap (join lazy pin thread first so + # _pin_records is final, then unregister each pinned chunk) + self._cuda_unpin_region(region) if region.is_mapped: region.release() @@ -330,32 +501,7 @@ def close(self) -> None: continue try: # CUDA unpin before munmap (order matters!) - if region._cuda_pinned and region._buffer_view is not None: - try: - import torch - - if torch.cuda.is_available(): - addr = ctypes.addressof( - ctypes.c_char.from_buffer(region._buffer_view) - ) - rc = torch.cuda.cudart().cudaHostUnregister(addr) - rc = int(rc[0]) if isinstance(rc, tuple) else int(rc) - if rc != 0: - _clear_cuda_sticky_error() - logger.warning( - "cudaHostUnregister failed for region " - "%d: rc=%d", - rid, - rc, - ) - region._cuda_pinned = False - except (ImportError, RuntimeError, OSError) as e: - if not isinstance(e, ImportError): - logger.warning( - "cudaHostUnregister failed for region %d: %s", - rid, - e, - ) + self._cuda_unpin_region(region) if region.is_mapped: region.release() diff --git a/maru_handler/memory/types.py b/maru_handler/memory/types.py index 96dbb40..113378e 100644 --- a/maru_handler/memory/types.py +++ b/maru_handler/memory/types.py @@ -4,6 +4,7 @@ import logging import mmap as mmap_module +import threading from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -31,9 +32,24 @@ class MappedRegion: _mmap_obj: mmap_module.mmap | None = field(default=None, repr=False) # memoryview of the entire mmap — created eagerly in __post_init__ _buffer_view: memoryview | None = field(default=None, repr=False, init=False) - # True only when cudaHostRegister actually succeeded (rc == 0), so - # unmap/close know whether cudaHostUnregister is needed - _cuda_pinned: bool = field(default=False, repr=False, init=False) + # (addr, size) per successfully cudaHostRegister'ed chunk — the exact + # base addresses cudaHostUnregister must be called with on unmap + _pin_records: list[tuple[int, int]] = field( + default_factory=list, repr=False, init=False + ) + # Serializes the chunk-pin loop against unmap-time unpinning; _pin_stop + # makes an in-flight pin loop exit at the next chunk boundary so unpin + # never waits for more than one chunk. + _pin_lock: threading.Lock = field( + default_factory=threading.Lock, repr=False, init=False + ) + _pin_stop: threading.Event = field( + default_factory=threading.Event, repr=False, init=False + ) + # Background pinning thread (MARU_PIN_MODE=lazy only) + _pin_thread: threading.Thread | None = field( + default=None, repr=False, init=False + ) def __post_init__(self): if self._mmap_obj is not None: diff --git a/maru_vllm/connector.py b/maru_vllm/connector.py index 75270d1..a7e21e4 100644 --- a/maru_vllm/connector.py +++ b/maru_vllm/connector.py @@ -710,9 +710,13 @@ def start_load_kv( ) success = False break - # CXL mmap is already cudaHostRegistered by - # DaxMapper — no clone() needed, .cuda() uses - # DMA directly from pinned CXL memory + # CXL mmap is cudaHostRegistered by DaxMapper + # (fully, before map_region returns, with the + # default MARU_PIN_MODE=eager) — no clone() needed, + # .cuda() uses DMA directly from pinned CXL memory. + # With MARU_PIN_MODE=lazy or after a chunk pin + # failure, unpinned ranges still work but fall back + # to pageable copies. chunk_tensor = torch.frombuffer( info.view, dtype=kv_cache_layer.dtype ) diff --git a/tests/unit/test_mapper.py b/tests/unit/test_mapper.py index 808b7d0..914267f 100644 --- a/tests/unit/test_mapper.py +++ b/tests/unit/test_mapper.py @@ -189,23 +189,73 @@ def test_cuda_unavailable_graceful(self): assert region.is_mapped -class TestDaxMapperPinRcCheck: - """cudaHostRegister return-code checking (previously ignored).""" +class TestDaxMapperChunkedPin: + """Test chunked cudaHostRegister (512GiB-per-call driver limit workaround).""" - def test_register_rc_failure_warns_and_skips_unregister(self, monkeypatch): - """rc != 0 → warning, region unpinned, no unregister on unmap.""" + def test_chunked_register_calls(self, monkeypatch): + """Region larger than the chunk size is registered in multiple calls.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 1024) + mock_torch, mock_cudart = _mock_torch_cuda() + + mapper = DaxMapper() + handle = _make_handle(1, 4096) + + with patch.dict("sys.modules", {"torch": mock_torch}): + region = mapper.map_region(handle) + + base = ctypes.addressof(ctypes.c_char.from_buffer(region._buffer_view)) + calls = mock_cudart.cudaHostRegister.call_args_list + assert [c.args for c in calls] == [ + (base, 1024, 0), + (base + 1024, 1024, 0), + (base + 2048, 1024, 0), + (base + 3072, 1024, 0), + ] + assert region._pin_records == [ + (base, 1024), + (base + 1024, 1024), + (base + 2048, 1024), + (base + 3072, 1024), + ] + + def test_chunked_unregister_per_record(self, monkeypatch): + """Unmap unregisters each pinned chunk's base address.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 2048) + mock_torch, mock_cudart = _mock_torch_cuda() + + mapper = DaxMapper() + handle = _make_handle(1, 4096) + + with patch.dict("sys.modules", {"torch": mock_torch}): + region = mapper.map_region(handle) + base = ctypes.addressof( + ctypes.c_char.from_buffer(region._buffer_view) + ) + mapper.unmap_region(1) + + assert [c.args for c in mock_cudart.cudaHostUnregister.call_args_list] == [ + (base,), + (base + 2048,), + ] + + def test_register_failure_rc_checked(self, monkeypatch): + """Non-zero cudaHostRegister rc → warning, no record, error cleared.""" import logging import maru_handler.memory.mapper as mapper_mod + mock_torch, mock_cudart = _mock_torch_cuda() + mock_cudart.cudaHostRegister.return_value = (2,) # cudaErrorMemoryAllocation + cleared = [] monkeypatch.setattr( mapper_mod, "_clear_cuda_sticky_error", lambda: cleared.append(1) ) - mock_torch, mock_cudart = _mock_torch_cuda() - mock_cudart.cudaHostRegister.return_value = (2,) # cudaErrorMemoryAllocation - mapper = DaxMapper() handle = _make_handle(1, 4096) @@ -216,14 +266,115 @@ def test_register_rc_failure_warns_and_skips_unregister(self, monkeypatch): region = mapper.map_region(handle) mapper.unmap_region(1) - assert region._cuda_pinned is False + assert region._pin_records == [] assert cleared == [1] - assert mock_warning.called - assert "cudaHostRegister failed" in mock_warning.call_args[0][0] + warnings = [c.args[0] for c in mock_warning.call_args_list] + assert any("cudaHostRegister failed" in w for w in warnings) + assert any("PARTIALLY" in w for w in warnings) mock_cudart.cudaHostUnregister.assert_not_called() - def test_register_rc_success_sets_pinned_and_unregisters(self): - """rc == 0 → pinned flag set, unmap unregisters once.""" + def test_partial_register_failure(self, monkeypatch): + """One failed chunk is skipped; the others are pinned and unpinned.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 1024) + monkeypatch.setattr(mapper_mod, "_clear_cuda_sticky_error", lambda: None) + mock_torch, mock_cudart = _mock_torch_cuda() + mock_cudart.cudaHostRegister.side_effect = [(0,), (2,), (0,), (0,)] + + mapper = DaxMapper() + handle = _make_handle(1, 4096) + + with patch.dict("sys.modules", {"torch": mock_torch}): + region = mapper.map_region(handle) + base = ctypes.addressof( + ctypes.c_char.from_buffer(region._buffer_view) + ) + assert region._pin_records == [ + (base, 1024), + (base + 2048, 1024), + (base + 3072, 1024), + ] + mapper.unmap_region(1) + + assert mock_cudart.cudaHostUnregister.call_count == 3 + + def test_single_call_mode_chunk_zero(self, monkeypatch): + """MARU_PIN_CHUNK_GB=0 registers the whole region in one call.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 0) + mock_torch, mock_cudart = _mock_torch_cuda() + + mapper = DaxMapper() + handle = _make_handle(1, 4096) + + with patch.dict("sys.modules", {"torch": mock_torch}): + region = mapper.map_region(handle) + + base = ctypes.addressof(ctypes.c_char.from_buffer(region._buffer_view)) + mock_cudart.cudaHostRegister.assert_called_once_with(base, 4096, 0) + + def test_stop_event_truncates_pinning(self, monkeypatch): + """A pre-set stop event makes the pin loop register nothing.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 1024) + mock_torch, mock_cudart = _mock_torch_cuda() + + mapper = DaxMapper() + handle = _make_handle(1, 4096) + + # Pre-set the stop flag via the dataclass default the region will + # get — simulate by patching _cuda_pin_chunks call path: map, then + # verify unmap-after-stop unregisters exactly what was pinned. + with patch.dict("sys.modules", {"torch": mock_torch}): + region = mapper.map_region(handle) + region._pin_stop.set() + # a second pin loop over the same region would now no-op + mapper_mod._cuda_pin_chunks( + mock_cudart, + 0x1000, + 4096, + 1, + [], + stop_event=region._pin_stop, + ) + mapper.unmap_region(1) + + # 4 chunks from map_region only — the stopped loop added none + assert mock_cudart.cudaHostRegister.call_count == 4 + assert mock_cudart.cudaHostUnregister.call_count == 4 + + def test_lazy_close_joins_and_unregisters(self, monkeypatch): + """close() in lazy mode joins the worker and unregisters chunks.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_MODE", "lazy") + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 2048) + mock_torch, mock_cudart = _mock_torch_cuda() + + mapper = DaxMapper() + + with patch.dict("sys.modules", {"torch": mock_torch}): + r1 = mapper.map_region(_make_handle(1, 4096)) + r2 = mapper.map_region(_make_handle(2, 4096)) + for r in (r1, r2): + if r._pin_thread is not None: + r._pin_thread.join(timeout=5) + mapper.close() + + assert mock_cudart.cudaHostRegister.call_count == 4 + assert mock_cudart.cudaHostUnregister.call_count == 4 + assert r1._pin_thread is None and r2._pin_thread is None + assert mapper.get_region(1) is None and mapper.get_region(2) is None + + def test_lazy_mode_pins_in_background(self, monkeypatch): + """MARU_PIN_MODE=lazy pins on a background thread; unmap joins it.""" + import maru_handler.memory.mapper as mapper_mod + + monkeypatch.setattr(mapper_mod, "_PIN_MODE", "lazy") + monkeypatch.setattr(mapper_mod, "_PIN_CHUNK_BYTES", 1024) mock_torch, mock_cudart = _mock_torch_cuda() mapper = DaxMapper() @@ -231,10 +382,14 @@ def test_register_rc_success_sets_pinned_and_unregisters(self): with patch.dict("sys.modules", {"torch": mock_torch}): region = mapper.map_region(handle) - assert region._cuda_pinned is True + assert region._pin_thread is not None + region._pin_thread.join(timeout=5) + assert len(region._pin_records) == 4 mapper.unmap_region(1) - mock_cudart.cudaHostUnregister.assert_called_once() + assert mock_cudart.cudaHostRegister.call_count == 4 + assert mock_cudart.cudaHostUnregister.call_count == 4 + assert region._pin_thread is None class TestDaxMapperErrorPaths: