diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md index 965ab81d8d..53a3093717 100644 --- a/docs/design/hisparse.md +++ b/docs/design/hisparse.md @@ -60,6 +60,16 @@ sees only device pools. The source group has `block_pool_id=None`; device-pool consumers must narrow it before indexing, so host ownership cannot masquerade as a numeric GPU pool. +For single-node MP tensor parallelism, every TP worker maps the same pinned host +pool and uses the same block and layer offsets. MLA source KV is replicated +across TP ranks, so this stores one physical copy instead of one copy per rank. +TP rank 0 writes the shared host pool; peers wait on its IPC events before +reading it. Other executor and parallel layouts retain private per-rank pools. + +The shared layout backs the per-replica logical capacity with one physical pool; +the private layout allocates one physical pool per rank. Physical pool size +includes block-stride alignment. + ## Code boundary ```text diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index e7eee03e7b..3b29799d3f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,6 +3,7 @@ import copy import hashlib import importlib +import mmap import subprocess import sys from collections.abc import Callable @@ -15,6 +16,7 @@ import pytest import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module from vllm.config import ( CacheConfig, KVTransferConfig, @@ -90,7 +92,13 @@ pytestmark = pytest.mark.cpu_test @pytest.mark.parametrize("gpu_block_size", [32, 64]) -def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( block_size=gpu_block_size, @@ -116,7 +124,15 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): hf_config=SimpleNamespace(index_topk=128), max_model_len=gpu_block_size, ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), cache_config=SimpleNamespace( num_gpu_blocks_override=7, prefix_cache_retention_interval=None, @@ -137,6 +153,12 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size): assert cache_config.hisparse_host_num_blocks > 7 host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride assert host_group.host_resident assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) host_layers = set(host_group.layer_names) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 10464d950d..f597c8fa27 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -367,6 +367,42 @@ def test_hisparse_offloads_only_indexer_group(): assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1] +def test_hisparse_partial_group_size_survives_scheduler_flattening(): + """Worker and scheduler representations must produce identical mmap rows.""" + layer_specs = {name: _full_attention_spec() for name in ("indexer.0", "indexer.1")} + wrapped = UniformTypeKVCacheSpecs.from_specs(layer_specs) + assert wrapped is not None + source = KVCacheGroupSpec( + ["source"], + _full_attention_spec(), + host_resident=True, + role=KVCacheGroupRole.HISPARSE_SOURCE, + ) + worker_indexer = KVCacheGroupSpec( + list(layer_specs), wrapped, role=KVCacheGroupRole.HISPARSE_INDEXER + ) + scheduler_indexer = KVCacheGroupSpec( + list(layer_specs), + next(iter(layer_specs.values())), + role=KVCacheGroupRole.HISPARSE_INDEXER, + ) + + def make_config(indexer: KVCacheGroupSpec) -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[source, indexer], + hisparse_host_num_blocks=4, + ) + + config = _make_vllm_config() + worker = build_offloading_config(config, make_config(worker_indexer)) + scheduler = build_offloading_config(config, make_config(scheduler_indexer)) + + assert worker.worker_kv_bytes_per_block == scheduler.worker_kv_bytes_per_block + assert worker.worker_kv_bytes_per_block == wrapped.page_size_bytes + + def test_zero_blocks_skips_tensor_layout_validation(): kv_cache_config = _make_sizing_kv_cache_config(packed=False) kv_cache_config.num_blocks = 0 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 c7ddac4fef..c448e76a9c 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -10,11 +10,12 @@ import os import threading import time import uuid -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest from vllm.utils.system_utils import get_mp_context +from vllm.v1.kv_offload.cpu import shared_offload_region as region_module from vllm.v1.kv_offload.cpu.shared_offload_region import ( SharedOffloadRegion, _wait_for_file_size, @@ -192,28 +193,43 @@ def _mp_barrier_construct_and_hold( barrier, fill_value: int, done_queue, + replicated: bool = False, ) -> None: """Construct with a real cross-process barrier, write, then hold the mapping (no cleanup) until the parent SIGKILLs this process.""" try: + populate_calls = 0 + get_populate_write_fn = region_module._get_populate_write_fn + + def track_population(mmap_obj): + nonlocal populate_calls + populate_calls += 1 + return get_populate_write_fn(mmap_obj) + + region_module._get_populate_write_fn = track_population region = SharedOffloadRegion( engine_id=engine_id, num_chunks=2, - rank=rank, - kv_bytes_per_chunk=num_workers * PAGE_SIZE, + rank=0 if replicated else rank, + kv_bytes_per_chunk=PAGE_SIZE if replicated else num_workers * PAGE_SIZE, cpu_page_size=PAGE_SIZE, barrier=lambda: barrier.wait(30), + populate_only_on_creator=replicated, ) # The constructor's barrier precedes the creator's unlink. barrier.wait(30) t = region.create_next_worker_view(PAGE_SIZE) - t[:, :] = fill_value + if not replicated or rank == 0: + t[:, :] = fill_value + barrier.wait(30) done_queue.put( { "rank": rank, "error": None, "inode": os.fstat(region.fd).st_ino, "path_exists": os.path.exists(region.mmap_path), + "populate_calls": populate_calls, + "contents_match": bool((t == (1 if replicated else fill_value)).all()), } ) time.sleep(60) # hold the mapping until killed @@ -752,6 +768,26 @@ def test_cleanup_idempotent(iid): r.cleanup() # must be a no-op +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) + r.pinned_addresses = [0x100000, 0x200000, 0x300000] + r.is_pinned = True + + r.cleanup() + + assert cudart.cudaHostUnregister.call_args_list == [ + call(0x300000), + call(0x200000), + call(0x100000), + ] + + def test_cleanup_after_create_next_worker_view_releases_mmap(iid): """cleanup() must close the mmap even after create_next_worker_view was called. create_next_worker_view returns a view that shares storage with _base; both must be @@ -811,11 +847,53 @@ def test_wait_for_file_size_timeout(tmp_path): os.close(fd) +@pytest.mark.skip_global_cleanup +def test_wait_for_file_size_rejects_unlinked_file(tmp_path): + path = tmp_path / "unlinked.mmap" + fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600) + try: + os.unlink(path) + with pytest.raises(RuntimeError, match="creator failed"): + _wait_for_file_size(fd, PAGE_SIZE, timeout=1.0) + finally: + os.close(fd) + + # --------------------------------------------------------------------------- # Constructor — capacity validation # --------------------------------------------------------------------------- +@pytest.mark.skip_global_cleanup +@pytest.mark.skipif(not os.path.isdir("/dev/shm"), reason="requires /dev/shm") +def test_creator_memory_check_runs_only_for_creator(iid): + checked_sizes: list[int] = [] + creator = SharedOffloadRegion( + engine_id=iid, + num_chunks=4, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + creator_memory_check=checked_sizes.append, + ) + joiner: SharedOffloadRegion | None = None + try: + joiner = SharedOffloadRegion( + engine_id=iid, + num_chunks=4, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + creator_memory_check=checked_sizes.append, + ) + assert checked_sizes == [4 * PAGE_SIZE] + finally: + if joiner is not None: + joiner.cleanup() + creator.cleanup() + _cleanup_file(creator.mmap_path) + + def test_insufficient_space_raises_clear_error(monkeypatch): """A failed creator capacity check must clean up and give a clear error.""" import vllm.v1.kv_offload.cpu.shared_offload_region as region @@ -916,7 +994,8 @@ def test_barrier_failure_unlinks_creator_and_raises(iid): assert not os.path.exists(path) -def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): +@pytest.mark.parametrize("replicated", [False, True]) +def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated): """With a real cross-process barrier every worker maps one shared inode, the name is gone while they run, and SIGKILL leaks nothing.""" num_workers = 2 @@ -927,7 +1006,7 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): procs = [ ctx.Process( target=_mp_barrier_construct_and_hold, - args=(iid, rank, num_workers, barrier, rank + 1, done_queue), + args=(iid, rank, num_workers, barrier, rank + 1, done_queue, replicated), ) for rank in range(num_workers) ] @@ -939,6 +1018,10 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): assert r["error"] is None, f"rank {r['rank']}: {r['error']}" assert len({r["inode"] for r in results}) == 1, "workers split onto two files" assert not any(r["path_exists"] for r in results) + assert all(r["contents_match"] for r in results) + assert sum(r["populate_calls"] for r in results) == ( + 1 if replicated else num_workers + ) assert not os.path.exists(path) finally: for p in procs: @@ -950,23 +1033,36 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid): assert restarted._creator is True, "restart must be able to create anew" -def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch): +@pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"]) +def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase): """A worker that dies before the rendezvous must still arrive at the barrier, or its peers block in the collective until it times out.""" - import vllm.v1.kv_offload.cpu.shared_offload_region as region - - monkeypatch.setattr( - region, - "check_shm_free_space", - MagicMock(side_effect=RuntimeError("Insufficient space")), + failure = MagicMock(side_effect=ValueError("Insufficient space")) + memory_check = failure if failure_phase == "memory" else None + if failure_phase == "shm": + monkeypatch.setattr(region_module, "check_shm_free_space", failure) + elif failure_phase == "population": + monkeypatch.setattr(region_module, "_get_populate_write_fn", failure) + path = f"/dev/shm/vllm_offload_{iid}.mmap" + seen_at_barrier = [] + barrier = MagicMock( + side_effect=lambda: seen_at_barrier.append(os.path.exists(path)) ) - barrier = MagicMock() - with pytest.raises(RuntimeError, match="Insufficient space"): - _make_region(iid, barrier=barrier) + with pytest.raises(ValueError, match="Insufficient space"): + SharedOffloadRegion( + engine_id=iid, + num_chunks=1, + rank=0, + kv_bytes_per_chunk=PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + barrier=barrier, + creator_memory_check=memory_check, + populate_only_on_creator=True, + ) barrier.assert_called_once_with() - assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap") + assert seen_at_barrier == [False] def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch): diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index dc75ba4ea1..2dadc8c980 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -12,6 +12,7 @@ from types import SimpleNamespace import pytest import torch +import vllm.v1.hisparse.binding as attn_utils_module from tests.v1.attention.utils import dense_kv_cache_views from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -239,6 +240,87 @@ def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyP assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set() +class _FakeSharedHostRegion: + def __init__(self) -> None: + self.cleanup_calls = 0 + self.base_tensor = torch.empty(1, dtype=torch.int8) + + def cleanup(self) -> None: + self.cleanup_calls += 1 + + +def test_profiling_cleanup_releases_tp_shared_region_once(monkeypatch): + """TP-shared profiling pools must use region-aware chunk cleanup.""" + region = _FakeSharedHostRegion() + runtime = SimpleNamespace( + _host_cache=object(), + registered_host_pool=region.base_tensor, + hot_backing=object(), + shared_host_region=region, + ) + forward_context = { + "layer": SimpleNamespace( + hisparse_cache=SimpleNamespace(runtime=runtime), + ) + } + released = [] + + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): + released.append((runtimes, pinned_host_pools, shared_host_region)) + + monkeypatch.setattr( + attn_utils_module, + "release_pinned_state", + release_pinned_state, + ) + + attn_utils_module.release_hisparse_profiling_cache(forward_context) + + assert released == [([runtime], [], region)] + + +@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"]) +def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): + """A failure after mmap allocation must not leak the shared registration.""" + region = _FakeSharedHostRegion() + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC + ), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), + ) + + def allocate(*args): + args[-1].shared_region = region + if failure_phase == "allocation": + raise RuntimeError("initialization failed") + return {} + + def bind(**kwargs): + if failure_phase == "binding": + raise RuntimeError("initialization failed") + return [] + + def buffers(*args, **kwargs): + raise RuntimeError("initialization failed") + + monkeypatch.setattr(attn_utils_module, "allocate_hisparse_kv_caches", allocate) + monkeypatch.setattr(attn_utils_module, "bind_hisparse_kv_caches", bind) + monkeypatch.setattr( + attn_utils_module, "initialize_hisparse_runtime_buffers", buffers + ) + with pytest.raises(RuntimeError, match="initialization failed"): + attn_utils_module.init_hisparse_kv_cache( + SimpleNamespace(), + torch.device("cpu"), + [], + vllm_config, + {}, + SimpleNamespace(), + ) + assert region.cleanup_calls == 1 + + def test_reshape_padded_kv_cache_strides_by_padded_page(): num_blocks = 3 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index bacd1cad84..7f4bea52c2 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -60,7 +60,11 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> No ) ] - monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool) + monkeypatch.setattr( + hisparse_binding, + "HiSparseHostPool", + lambda vllm_config, kv_cache_config: host_pool, + ) monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 8422a591c3..55cab287fd 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap from collections import deque from contextlib import nullcontext from types import SimpleNamespace @@ -188,6 +189,232 @@ def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary( assert handle.submit_layer_mirror.call_count == expected_calls +def _hisparse_parallel_config(**overrides): + values = { + "tensor_parallel_size": 2, + "pipeline_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "world_size": 2, + "distributed_executor_backend": "mp", + "nnodes_within_dp": 1, + "data_parallel_index": 3, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + config = SimpleNamespace(parallel_config=_hisparse_parallel_config()) + + assert hisparse_runtime_module.use_shared_hisparse_host_pool(config) + + unsupported = ( + {"tensor_parallel_size": 1, "world_size": 1}, + {"pipeline_parallel_size": 2, "world_size": 4}, + {"prefill_context_parallel_size": 2, "world_size": 4}, + {"decode_context_parallel_size": 2}, + {"world_size": 4}, + {"distributed_executor_backend": "ray"}, + {"nnodes_within_dp": 2}, + ) + for overrides in unsupported: + config.parallel_config = _hisparse_parallel_config(**overrides) + assert not hisparse_runtime_module.use_shared_hisparse_host_pool(config) + + +@pytest.mark.skip_global_cleanup +def test_hisparse_shared_host_pool_uses_one_replicated_mmap(monkeypatch): + page = mmap.PAGESIZE + tp_group = MagicMock() + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group) + + class FakeSharedOffloadRegion: + BLOCK_SIZE_ALIGNMENT = page + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.total_size_bytes = kwargs["num_chunks"] * kwargs["kv_bytes_per_chunk"] + self.base_tensor = torch.empty(self.total_size_bytes, dtype=torch.int8) + self.is_pinned = False + self.pinned_addresses = [] + self.view_sizes = [] + self.offset = 0 + + def create_next_canonical_view(self, size): + self.view_sizes.append(size) + view = torch.as_strided( + self.base_tensor, + size=(self.kwargs["num_chunks"], size), + stride=(self.kwargs["kv_bytes_per_chunk"], 1), + storage_offset=self.offset, + ) + self.offset += size + return view + + def cleanup(self): + pass + + monkeypatch.setattr( + hisparse_runtime_module, "SharedOffloadRegion", FakeSharedOffloadRegion + ) + registration_ranges = MagicMock(return_value=((0, page), (page, 4 * page))) + monkeypatch.setattr( + hisparse_runtime_module, + "_hisparse_registration_ranges", + registration_ranges, + ) + pinned: list[torch.Tensor] = [] + + def pin_tensor(tensor): + pinned.append(tensor) + + monkeypatch.setattr(hisparse_runtime_module, "pin_tensor", pin_tensor) + config = SimpleNamespace( + instance_id="instance", + parallel_config=_hisparse_parallel_config( + tensor_parallel_size=1, + world_size=1, + distributed_executor_backend="uni", + ), + ) + + pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools( + config, + [24, 40], + num_blocks=4, + host_block_stride=page, + use_shared_host_pool=True, + ) + + assert region is not None + assert private_pools == [] + assert region.kwargs == { + "engine_id": "hisparse_instance_dp3", + "num_chunks": 1, + "rank": 0, + "kv_bytes_per_chunk": 4 * page, + "cpu_page_size": 64, + "barrier": tp_group.barrier, + "creator_memory_check": hisparse_runtime_module.check_hisparse_host_memory, + "populate_only_on_creator": True, + } + assert region.view_sizes == [24, 40] + assert [pool.shape for pool in pools] == [(24,), (40,)] + registration_ranges.assert_called_once_with([24, 40], 4, page) + assert [ + (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes) + for tensor in pinned + ] == [(0, page), (page, 3 * page)] + assert region.is_pinned + + +@pytest.mark.parametrize("fail_registration", [None, 1, 2]) +def test_shared_host_pool_tracks_successful_registrations( + monkeypatch, fail_registration +): + """Only successful registrations may be unregistered on allocation failure.""" + page = mmap.PAGESIZE + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", MagicMock()) + backing = torch.empty(3 * page, dtype=torch.uint8) + region = MagicMock(base_tensor=backing, pinned_addresses=[], is_pinned=False) + monkeypatch.setattr( + hisparse_runtime_module, "SharedOffloadRegion", lambda **kw: region + ) + monkeypatch.setattr( + hisparse_runtime_module, + "_hisparse_registration_ranges", + lambda *args: ((0, page), (page, 3 * page)), + ) + cudart = MagicMock() + cudart.cudaHostRegister.side_effect = [ + SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2) + ] + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + config = SimpleNamespace( + instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0) + ) + context = ( + pytest.raises(RuntimeError, match="cudaHostRegister failed") + if fail_registration + else nullcontext() + ) + with context: + hisparse_runtime_module.allocate_hisparse_host_pools( + config, [3 * page], 3, page, use_shared_host_pool=True + ) + successful = 2 if fail_registration is None else fail_registration - 1 + assert region.pinned_addresses == [ + backing.data_ptr() + i * page for i in range(successful) + ] + assert region.is_pinned == bool(successful) + assert region.cleanup.call_count == int(fail_registration is not None) + attempts = 2 if fail_registration is None else fail_registration + assert ( + cudart.cudaHostRegister.call_args_list + == [ + call(backing.data_ptr(), page, 0), + call(backing.data_ptr() + page, 2 * page, 0), + ][:attempts] + ) + + +def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch): + """Aliased tensor configs must not count the shared backing once per view.""" + region = SimpleNamespace(base_tensor=torch.empty(64, dtype=torch.int8)) + allocate = MagicMock(return_value=([], [], region)) + monkeypatch.setattr( + hisparse_runtime_module, "allocate_hisparse_host_pools", allocate + ) + config = SimpleNamespace( + hisparse_shared_host_pool=True, + hisparse_host_num_blocks=4, + hisparse_host_block_stride=16, + kv_cache_tensors=[ + SimpleNamespace( + host_resident=True, + layers=["a", "b"], + offset=0, + layer_stride=12, + size=44, + ), + SimpleNamespace( + host_resident=True, layers=["c"], offset=24, layer_stride=20, size=44 + ), + ], + ) + vllm_config = SimpleNamespace() + pool = hisparse_runtime_module.HiSparseHostPool(vllm_config, config) + backing = pool.allocate(44) + allocate.assert_called_once_with( + vllm_config, [12, 12, 20], 4, 16, use_shared_host_pool=True + ) + assert backing.numel() == 44 + assert backing.data_ptr() == region.base_tensor.data_ptr() + assert pool.registered is region.base_tensor + + +def test_hisparse_registration_chunks_end_between_host_blocks(): + """Registration seams must not bisect a DMA-addressable host block.""" + page = mmap.PAGESIZE + ranges = hisparse_runtime_module._hisparse_registration_ranges( + tensor_sizes=[4 * 3 * page, 4 * 5 * page], + num_blocks=4, + host_block_stride=8 * page, + max_chunk_bytes=10 * page, + ) + + assert ranges == ( + (0, 9 * page), + (9 * page, 17 * page), + (17 * page, 27 * page), + (27 * page, 32 * page), + ) + + def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): backing = torch.full((10, 2, 3), -1, dtype=torch.float32) cache = backing[1:9] @@ -931,6 +1158,32 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc assert streams == [] +def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): + worker = object.__new__(HiSparseConnectorWorker) + worker._initialized = True + worker._slot_mapping_staging = None + worker.dma_stream = None + worker.cache_handles = [] + worker.pinned_host_pools = [] + worker.shared_host_region = object() + released = False + + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): + nonlocal released + assert runtimes == [] + assert pinned_host_pools == [] + assert shared_host_region is worker.shared_host_region + released = True + + monkeypatch.setattr( + hisparse_worker_module, "release_pinned_state", release_pinned_state + ) + + worker.shutdown() + + assert released + + class _TestReplaySSMMixer(MambaMixer2): def __init__(self): torch.nn.Module.__init__(self) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index 69721803be..b7ffe9ec97 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -147,6 +147,12 @@ def _select_written_row_mirrors( ) +def _is_hisparse_host_writer( + shared_host_region: SharedOffloadRegion | None, +) -> bool: + return shared_host_region is None or get_tensor_model_parallel_rank() == 0 + + def _create_hisparse_host_events( shared_host_region: SharedOffloadRegion | None, is_host_writer: bool, @@ -200,15 +206,26 @@ class HiSparseConnectorWorker: raise RuntimeError("HiSparse connector found no hot-cache handles.") hot_backings: dict[int, torch.Tensor] = {} registered_host_pools: dict[int, torch.Tensor] = {} + shared_host_regions: dict[int, SharedOffloadRegion] = {} for cache in cache_handles: hot_backing = cache.runtime.hot_backing hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing registered_pool = cache.runtime.registered_host_pool registered_host_pools[registered_pool.data_ptr()] = registered_pool + if (region := cache.runtime.shared_host_region) is not None: + shared_host_regions[id(region)] = region if len(hot_backings) != 1: raise RuntimeError("HiSparse hot tensors must share one GPU backing.") hot_backing = next(iter(hot_backings.values())) - pinned_host_pools = list(registered_host_pools.values()) + if len(shared_host_regions) > 1: + raise RuntimeError("HiSparse caches must share one host region.") + shared_host_region = next(iter(shared_host_regions.values()), None) + pinned_host_pools = ( + [] + if shared_host_region is not None + else list(registered_host_pools.values()) + ) + is_host_writer = _is_hisparse_host_writer(shared_host_region) resident = cache_handles[0].view assert resident is not None @@ -223,11 +240,14 @@ class HiSparseConnectorWorker: host_num_blocks, hot_backing.device, pinned_host_pools, + shared_host_region=shared_host_region, + is_host_writer=is_host_writer, ) except Exception: release_pinned_state( [cache.runtime for cache in cache_handles], pinned_host_pools, + shared_host_region, ) raise @@ -866,6 +886,8 @@ class HiSparseConnectorWorker: if self.dma_stream is not None: self.dma_stream.synchronize() release_pinned_state( - [cache.runtime for cache in self.cache_handles], self.pinned_host_pools + [cache.runtime for cache in self.cache_handles], + self.pinned_host_pools, + self.shared_host_region, ) self._initialized = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 42b9dd6aa3..c4e6d0f2a2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -17,6 +17,7 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, + UniformTypeKVCacheSpecs, iter_layer_specs, ) from vllm.v1.kv_offload.config import ( @@ -29,7 +30,7 @@ from vllm.v1.kv_offload.config import ( if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]: @@ -42,6 +43,19 @@ def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ... ) +def _group_kv_bytes_per_block(group: "KVCacheGroupSpec") -> int: + """Return the physical bytes occupied by one block of a cache group. + + Worker configs may retain ``UniformTypeKVCacheSpecs`` while scheduler + configs flatten that wrapper to one representative per-layer spec. Keep + the result invariant across those two representations. + """ + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.page_size_bytes + return spec.page_size_bytes * len(group.layer_names) + + def build_offloading_config( vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig", @@ -133,7 +147,7 @@ def build_offloading_config( worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks elif kv_cache_config.num_blocks > 0: worker_kv_bytes_per_block = sum( - group.kv_cache_spec.page_size_bytes for _, group in selected_groups + _group_kv_bytes_per_block(group) for _, group in selected_groups ) single_group_spec = ( diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 263e553d32..c23abc9cb4 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -120,27 +120,32 @@ def init_hisparse_kv_cache( block_tables: "BlockTables", ) -> dict[str, torch.Tensor]: """Allocate and bind HiSparse caches within the caller's allocation context.""" - host_pool = HiSparseHostPool() - kv_caches = allocate_hisparse_kv_caches( - kv_cache_config, - device, - vllm_config.cache_config.get_resolved_kv_cache_layout(), - kernel_block_sizes, - host_pool, - ) - cache_handles = bind_hisparse_kv_caches( - forward_context=forward_context, - kv_cache_config=kv_cache_config, - kv_caches=kv_caches, - block_tables=block_tables, - host_pool=host_pool, - ) - initialize_hisparse_runtime_buffers( - cache_handles, - max_num_reqs=vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, - ) - return kv_caches + host_pool = HiSparseHostPool(vllm_config, kv_cache_config) + try: + kv_caches = allocate_hisparse_kv_caches( + kv_cache_config, + device, + vllm_config.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + host_pool, + ) + cache_handles = bind_hisparse_kv_caches( + forward_context=forward_context, + kv_cache_config=kv_cache_config, + kv_caches=kv_caches, + block_tables=block_tables, + host_pool=host_pool, + ) + initialize_hisparse_runtime_buffers( + cache_handles, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, + ) + return kv_caches + except Exception: + if host_pool.shared_region is not None: + host_pool.shared_region.cleanup() + raise def _get_hisparse_cache( @@ -166,13 +171,24 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: if not runtimes: return - registered_pools = list( - { - runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool - for runtime in runtimes.values() - }.values() + shared_regions = { + id(runtime.shared_host_region): runtime.shared_host_region + for runtime in runtimes.values() + if runtime.shared_host_region is not None + } + assert len(shared_regions) <= 1 + shared_region = next(iter(shared_regions.values()), None) + registered_pools = ( + [] + if shared_region is not None + else list( + { + runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool + for runtime in runtimes.values() + }.values() + ) ) - release_pinned_state(list(runtimes.values()), registered_pools) + release_pinned_state(list(runtimes.values()), registered_pools, shared_region) for cache in cache_handles: cache.mirror_staging_cache = None cache.mirror_staging_slots = None @@ -257,6 +273,7 @@ def bind_hisparse_kv_caches( assert source_cache.untyped_storage().data_ptr() == ( host_pool.registered.untyped_storage().data_ptr() ) + cache_handle.runtime.shared_host_region = host_pool.shared_region cache_handle.runtime.bind_source_cache( source_cache, registered_host_pool=host_pool.registered, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index d74b0e8cfd..b3b45d3334 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -9,7 +9,11 @@ from vllm.config import VllmConfig from vllm.config.kv_transfer import hisparse_host_pool_gib from vllm.logger import init_logger from vllm.utils.math_utils import cdiv -from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig +from vllm.v1.hisparse.runtime import ( + ResolvedHiSparseConfig, + get_hisparse_host_block_stride, + use_shared_hisparse_host_pool, +) from vllm.v1.kv_cache_interface import ( HiSparseHotSpec, HiSparseResidentSpec, @@ -36,6 +40,8 @@ class HiSparseLayout: source_group: KVCacheGroupSpec device_groups: list[KVCacheGroupSpec] host_num_blocks: int + host_block_stride: int + shared_host_pool: bool def get_hisparse_kv_cache_groups( @@ -221,9 +227,12 @@ def create_hisparse_layout( regular_groups = groups[1:] gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups] - host_num_blocks = host_budget // sum( - spec.page_size_bytes for spec in source_specs.values() + shared_host_pool = use_shared_hisparse_host_pool(vllm_config) + host_block_stride = get_hisparse_host_block_stride( + sum(spec.page_size_bytes for spec in source_specs.values()), + use_shared_host_pool=shared_host_pool, ) + host_num_blocks = host_budget // host_block_stride if host_num_blocks <= 0: raise ValueError("HiSparse has no allocatable host blocks.") @@ -231,6 +240,8 @@ def create_hisparse_layout( source_group=source_group, device_groups=gpu_groups, host_num_blocks=host_num_blocks, + host_block_stride=host_block_stride, + shared_host_pool=shared_host_pool, ) @@ -339,6 +350,8 @@ def get_hisparse_kv_cache_config( kv_cache_tensors=kv_cache_tensors, kv_cache_groups=[*host_groups, *device_groups], hisparse_host_num_blocks=hisparse_layout.host_num_blocks, + hisparse_host_block_stride=hisparse_layout.host_block_stride, + hisparse_shared_host_pool=hisparse_layout.shared_host_pool, prefix_cache_retention_interval=( vllm_config.cache_config.prefix_cache_retention_interval ), diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 3dd8c017f5..314d3ba86b 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -4,6 +4,8 @@ from __future__ import annotations +import math +import mmap import time from collections.abc import Callable from dataclasses import dataclass @@ -13,17 +15,22 @@ import psutil import torch from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed import get_tp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import current_stream +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor logger = init_logger(__name__) +HOST_REGISTER_CHUNK_BYTES = 256 * 2**30 + if TYPE_CHECKING: from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup + from vllm.v1.kv_cache_interface import KVCacheConfig # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. FP8_DS_MLA_ROW_BYTES = 656 @@ -185,22 +192,195 @@ def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: class HiSparseHostPool: """Keep the single host backing and its exact CUDA registration range.""" - def __init__(self) -> None: + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: + self.vllm_config = vllm_config + self.kv_cache_config = kv_cache_config self.backing: torch.Tensor | None = None self.registered: torch.Tensor | None = None + self.shared_region: SharedOffloadRegion | None = None def allocate(self, size: int) -> torch.Tensor: assert self.backing is None, "HiSparse host pool is already allocated." - check_hisparse_host_memory(size) - self.backing, self.registered = allocate_pinned_host_pool(size) + config = self.kv_cache_config + if config.hisparse_shared_host_pool: + tensor_sizes = [ + size + for _, size in sorted( + (tensor.offset + i * tensor.layer_stride, tensor.layer_stride) + for tensor in config.kv_cache_tensors + if tensor.host_resident + for i in range(len(tensor.layers)) + ) + ] + assert config.hisparse_host_num_blocks is not None + assert config.hisparse_host_block_stride is not None + _, _, self.shared_region = allocate_hisparse_host_pools( + self.vllm_config, + tensor_sizes, + config.hisparse_host_num_blocks, + config.hisparse_host_block_stride, + use_shared_host_pool=True, + ) + assert self.shared_region is not None + self.registered = self.shared_region.base_tensor + self.backing = self.registered[:size] + else: + check_hisparse_host_memory(size) + self.backing, self.registered = allocate_pinned_host_pool(size) return self.backing +def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool: + """Whether replicated MLA host KV can share one local mmap.""" + parallel = vllm_config.parallel_config + return ( + current_platform.is_cuda_alike() + and parallel.tensor_parallel_size > 1 + and parallel.pipeline_parallel_size == 1 + and parallel.prefill_context_parallel_size == 1 + and parallel.decode_context_parallel_size == 1 + and parallel.world_size == parallel.tensor_parallel_size + and parallel.distributed_executor_backend == "mp" + and parallel.nnodes_within_dp == 1 + ) + + +def get_hisparse_host_block_stride( + page_size_bytes: int, *, use_shared_host_pool: bool +) -> int: + if use_shared_host_pool: + return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT) + return page_size_bytes + + +def _hisparse_registration_ranges( + tensor_sizes: list[int], + num_blocks: int, + host_block_stride: int, + max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES, +) -> tuple[tuple[int, int], ...]: + """Split a shared pool without bisecting any tensor's host block. + + CUDA batch copies reject descriptors spanning adjacent registrations. + """ + total_size = num_blocks * host_block_stride + if total_size % mmap.PAGESIZE: + raise ValueError("HiSparse host pool must be OS-page aligned.") + if max_chunk_bytes < mmap.PAGESIZE: + raise ValueError("HiSparse registration chunks must span at least one page.") + + tensors: list[tuple[int, int, int]] = [] + tensor_start = 0 + for tensor_size in tensor_sizes: + if tensor_size % num_blocks: + raise ValueError( + f"HiSparse host tensor size {tensor_size} is not divisible by " + f"{num_blocks} blocks." + ) + tensor_end = tensor_start + tensor_size + tensors.append((tensor_start, tensor_end, tensor_size // num_blocks)) + tensor_start = tensor_end + if tensor_start > total_size: + raise ValueError("HiSparse host tensors exceed the shared host pool.") + + def last_safe_boundary(lower: int, upper: int) -> int | None: + best = None + for start, end, block_bytes in tensors: + first_block = max(0, (lower - start) // block_bytes + 1) + last_block = min(num_blocks, (min(upper, end) - start) // block_bytes) + if first_block > last_block: + continue + divisor = math.gcd(block_bytes, mmap.PAGESIZE) + if start % divisor: + continue + period = mmap.PAGESIZE // divisor + if period == 1: + residue = 0 + else: + # Find block boundaries that are also OS-page boundaries. + residue = ( + (-start // divisor) * pow(block_bytes // divisor, -1, period) + ) % period + block = last_block - (last_block - residue) % period + if block >= first_block: + candidate = start + block * block_bytes + best = candidate if best is None else max(best, candidate) + + if upper >= tensor_start: + candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE + if candidate >= tensor_start and candidate > lower: + best = candidate if best is None else max(best, candidate) + return best + + ranges: list[tuple[int, int]] = [] + start = 0 + while total_size - start > max_chunk_bytes: + end = last_safe_boundary(start, start + max_chunk_bytes) + if end is None: + raise ValueError( + "A HiSparse host block is too large for the registration chunk limit." + ) + ranges.append((start, end)) + start = end + ranges.append((start, total_size)) + return tuple(ranges) + + +def allocate_hisparse_host_pools( + vllm_config: VllmConfig, + tensor_sizes: list[int], + num_blocks: int, + host_block_stride: int, + *, + use_shared_host_pool: bool, +) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]: + """Allocate private host tensors or one mmap shared by local TP ranks.""" + if not use_shared_host_pool: + check_hisparse_host_memory(sum(tensor_sizes)) + private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes] + return ( + [pool for pool, _ in private_pools], + [registered for _, registered in private_pools], + None, + ) + + region = SharedOffloadRegion( + engine_id=( + f"hisparse_{vllm_config.instance_id}_" + f"dp{vllm_config.parallel_config.data_parallel_index}" + ), + num_chunks=1, + rank=0, + kv_bytes_per_chunk=num_blocks * host_block_stride, + cpu_page_size=sum(tensor_sizes), + barrier=get_tp_group().barrier, + creator_memory_check=check_hisparse_host_memory, + populate_only_on_creator=True, + ) + try: + for start, end in _hisparse_registration_ranges( + tensor_sizes, num_blocks, host_block_stride + ): + tensor = region.base_tensor[start:end] + pin_tensor(tensor) + region.pinned_addresses.append(tensor.data_ptr()) + region.is_pinned = True + pools = [ + region.create_next_canonical_view(size).view(-1) for size in tensor_sizes + ] + except Exception: + region.cleanup() + raise + return pools, [], region + + def release_pinned_state( - runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor] + runtimes: list[HiSparseRuntime], + pinned_host_pools: list[torch.Tensor], + shared_host_region: SharedOffloadRegion | None = None, ) -> None: """Synchronize and release registered host KV pools.""" - if pinned_host_pools: + if pinned_host_pools or shared_host_region is not None: try: torch.accelerator.synchronize() except RuntimeError as e: @@ -208,7 +388,7 @@ def release_pinned_state( "HiSparse: CUDA context unusable at teardown (%s); leaving " "%d host-pool tensors pinned for kernel exit reclaim.", e, - len(pinned_host_pools), + len(pinned_host_pools) + int(shared_host_region is not None), ) return @@ -240,6 +420,8 @@ def release_pinned_state( del runtime._host_cache del runtime.registered_host_pool del runtime.hot_backing + if shared_host_region is not None: + shared_host_region.cleanup() def hisparse_prefill_staging_remap( @@ -536,6 +718,7 @@ class HiSparseRuntime: self.eager_host_mirror = config.eager_host_mirror self.resident_source_index = -1 self.request_state_indices: torch.Tensor | None = None + self.shared_host_region: SharedOffloadRegion | None = None @property def host_cache(self) -> torch.Tensor: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index fb839bcd68..e2869e047f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1435,6 +1435,12 @@ class KVCacheConfig: hisparse_host_num_blocks: int | None = None """Capacity of the dedicated HiSparse host-block manager, when enabled.""" + hisparse_host_block_stride: int | None = None + """Physical bytes between consecutive HiSparse host blocks.""" + + hisparse_shared_host_pool: bool = False + """Whether local TP ranks share one physical HiSparse host pool.""" + @cached_property def transfer_group_ids(self) -> tuple[int, ...]: """IDs of cache groups that participate in external KV transfer.""" diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 1a5dde21e6..46a1e7ae64 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import errno import mmap import os @@ -27,6 +28,10 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N while True: if os.fstat(fd).st_size >= expected_size: return + if os.fstat(fd).st_nlink == 0: + raise RuntimeError( + "Shared offload region creator failed during initialization." + ) if time.monotonic() > deadline: raise TimeoutError( f"Timed out waiting for mmap file to reach {expected_size} bytes" @@ -75,6 +80,9 @@ class SharedOffloadRegion: given, the path is unlinked once every worker has mapped the file, so the kernel reclaims the memory when the last worker exits, no matter how it exits; mappings taken before the unlink stay valid. + + Creator-only population pre-faults the entire region before the barrier + and requires that barrier to keep joiners from using unpopulated pages. """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE @@ -87,7 +95,12 @@ class SharedOffloadRegion: kv_bytes_per_chunk: int, cpu_page_size: int, barrier: Callable[[], None] | None = None, + *, + creator_memory_check: Callable[[int], None] | None = None, + populate_only_on_creator: bool = False, ) -> None: + if populate_only_on_creator and barrier is None: + raise ValueError("Creator-only population requires a barrier.") self.page_size = mmap.PAGESIZE assert kv_bytes_per_chunk % self.page_size == 0 @@ -112,25 +125,18 @@ class SharedOffloadRegion: # Joiner path — another worker won O_EXCL. Reopen and wait # for the file to reach expected size. self.fd = os.open(self.mmap_path, os.O_RDWR) - try: - _wait_for_file_size(self.fd, self.total_size_bytes) - except (TimeoutError, OSError): - os.close(self.fd) - raise + _wait_for_file_size(self.fd, self.total_size_bytes) logger.info("Opened existing mmap file %s", self.mmap_path) else: # Creator path. We won O_EXCL, so we own the file: any # failure here must clean up so concurrent joiners don't # land on a 0-byte stub and spin in _wait_for_file_size # for the full 30 s timeout. - try: - check_shm_free_space(self.total_size_bytes) - os.ftruncate(self.fd, self.total_size_bytes) - except (RuntimeError, OSError): - os.unlink(self.mmap_path) - os.close(self.fd) - raise self._creator = True + if creator_memory_check is not None: + creator_memory_check(self.total_size_bytes) + check_shm_free_space(self.total_size_bytes) + os.ftruncate(self.fd, self.total_size_bytes) logger.info( "Created mmap file %s (%.2f GB)", self.mmap_path, @@ -143,10 +149,21 @@ class SharedOffloadRegion: flags=mmap.MAP_SHARED, prot=mmap.PROT_READ | mmap.PROT_WRITE, ) + + if populate_only_on_creator and self._creator: + populate_write_fn = _get_populate_write_fn(self.mmap_obj) + populate_write_fn(self.mmap_obj, 0, self.total_size_bytes) except Exception: if self._creator: - os.unlink(self.mmap_path) + with contextlib.suppress(FileNotFoundError): + os.unlink(self.mmap_path) self._creator = False + if hasattr(self, "mmap_obj") and self.mmap_obj is not None: + self.mmap_obj.close() + self.mmap_obj = None + if hasattr(self, "fd") and self.fd is not None: + os.close(self.fd) + self.fd = None # Peers block inside the barrier until the collective times out if # we die before reaching it. Arrive anyway so every worker calls # barrier() exactly once and they fail on their own errors instead @@ -169,16 +186,28 @@ class SharedOffloadRegion: barrier() except Exception: if self._creator: - os.unlink(self.mmap_path) + with contextlib.suppress(FileNotFoundError): + os.unlink(self.mmap_path) self._creator = False self.mmap_obj.close() os.close(self.fd) + self.mmap_obj = None + self.fd = None raise if self._creator: os.unlink(self.mmap_path) self._creator = False logger.info("Unlinked mmap file %s", self.mmap_path) + self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) + self._views: list[torch.Tensor] = [] + self._canonical_offset = 0 + self.is_pinned: bool = False + self.pinned_addresses: list[int] = [] + + if populate_only_on_creator: + return + populate_write_fn = _get_populate_write_fn(self.mmap_obj) if rank is not None: @@ -205,10 +234,11 @@ class SharedOffloadRegion: "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0 ) - self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) - self._views: list[torch.Tensor] = [] - self._canonical_offset = 0 - self.is_pinned: bool = False + @property + def base_tensor(self) -> torch.Tensor: + if self._base is None: + raise RuntimeError("Shared offload region has been released.") + return self._base def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: """Allocate a strided int8 view for this worker, one canonical tensor. @@ -311,13 +341,18 @@ class SharedOffloadRegion: if self.is_pinned and self._base is not None: if current_platform.is_cuda_alike(): base_ptr = self._base.data_ptr() - result = torch.cuda.cudart().cudaHostUnregister(base_ptr) - if result.value != 0: - logger.warning( - "cudaHostUnregister failed for rank=%d (code=%d)", - self.rank, - result, - ) + 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() 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