[5/N] Share HiSparse host cache across TP ranks (reopens #52760) (#56629)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Chao Lei <leichao139636@163.com>
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Chao Lei <leichao139636@163.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>

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

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)

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

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):

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(

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)
 

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)

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

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 = (

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,

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
         ),

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:

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."""

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