@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e19a3e172e

[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>

Lucas Wilkinson committed Sep 12, 2026 at 19:48 UTC e19a3e172ecc1d1ece8ecce9dfd91bf811a814a0

14 files changed +878 -85

docs/design/hisparse.md

+10

@@ -60,6 +60,16 @@ sees only device pools.

60 The source group has `block_pool_id=None`; device-pool consumers must narrow it

61 before indexing, so host ownership cannot masquerade as a numeric GPU pool.

63 +For single-node MP tensor parallelism, every TP worker maps the same pinned host

64 +pool and uses the same block and layer offsets. MLA source KV is replicated

65 +across TP ranks, so this stores one physical copy instead of one copy per rank.

66 +TP rank 0 writes the shared host pool; peers wait on its IPC events before

67 +reading it. Other executor and parallel layouts retain private per-rank pools.

68 +

69 +The shared layout backs the per-replica logical capacity with one physical pool;

70 +the private layout allocates one physical pool per rank. Physical pool size

71 +includes block-stride alignment.

72 +

73 ## Code boundary

75 ```text

tests/v1/core/test_kv_cache_utils.py

+24 -2

@@ -3,6 +3,7 @@

3 import copy

4 import hashlib

5 import importlib

6 +import mmap

7 import subprocess

8 import sys

9 from collections.abc import Callable

@@ -15,6 +16,7 @@ import pytest

16 import torch

18 import vllm.v1.core.kv_cache_utils as kv_cache_utils

19 +import vllm.v1.hisparse.runtime as hisparse_runtime_module

20 from vllm.config import (

21 CacheConfig,

22 KVTransferConfig,

@@ -90,7 +92,13 @@ pytestmark = pytest.mark.cpu_test

94 @pytest.mark.parametrize("gpu_block_size", [32, 64])

93 -def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size):

95 +@pytest.mark.parametrize("shared_host_pool", [False, True])

96 +def test_hisparse_hma_uses_resolved_gpu_block_size(

97 + monkeypatch, gpu_block_size, shared_host_pool

98 +):

99 + monkeypatch.setattr(

100 + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True

101 + )

102 specs = {

103 "model.layers.0.self_attn": MLAAttentionSpec(

104 block_size=gpu_block_size,

@@ -116,7 +124,15 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size):

124 hf_config=SimpleNamespace(index_topk=128),

125 max_model_len=gpu_block_size,

126 ),

119 - parallel_config=SimpleNamespace(decode_context_parallel_size=1),

127 + parallel_config=SimpleNamespace(

128 + tensor_parallel_size=2 if shared_host_pool else 1,

129 + pipeline_parallel_size=1,

130 + prefill_context_parallel_size=1,

131 + decode_context_parallel_size=1,

132 + world_size=2 if shared_host_pool else 1,

133 + distributed_executor_backend="mp",

134 + nnodes_within_dp=1,

135 + ),

136 cache_config=SimpleNamespace(

137 num_gpu_blocks_override=7,

138 prefix_cache_retention_interval=None,

@@ -137,6 +153,12 @@ def test_hisparse_hma_uses_resolved_gpu_block_size(monkeypatch, gpu_block_size):

153 assert cache_config.hisparse_host_num_blocks > 7

154

155 host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups

156 + assert cache_config.hisparse_shared_host_pool is shared_host_pool

157 + host_page = host_group.kv_cache_spec.page_size_bytes

158 + alignment = mmap.PAGESIZE if shared_host_pool else 1

159 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment

160 + assert cache_config.hisparse_host_block_stride == expected_host_stride

161 + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride

162 assert host_group.host_resident

163 assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups])

164 host_layers = set(host_group.layer_names)

tests/v1/kv_connector/unit/offloading_connector/test_config.py

+36

@@ -367,6 +367,42 @@ def test_hisparse_offloads_only_indexer_group():

367 assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1]

368

369

370 +def test_hisparse_partial_group_size_survives_scheduler_flattening():

371 + """Worker and scheduler representations must produce identical mmap rows."""

372 + layer_specs = {name: _full_attention_spec() for name in ("indexer.0", "indexer.1")}

373 + wrapped = UniformTypeKVCacheSpecs.from_specs(layer_specs)

374 + assert wrapped is not None

375 + source = KVCacheGroupSpec(

376 + ["source"],

377 + _full_attention_spec(),

378 + host_resident=True,

379 + role=KVCacheGroupRole.HISPARSE_SOURCE,

380 + )

381 + worker_indexer = KVCacheGroupSpec(

382 + list(layer_specs), wrapped, role=KVCacheGroupRole.HISPARSE_INDEXER

383 + )

384 + scheduler_indexer = KVCacheGroupSpec(

385 + list(layer_specs),

386 + next(iter(layer_specs.values())),

387 + role=KVCacheGroupRole.HISPARSE_INDEXER,

388 + )

389 +

390 + def make_config(indexer: KVCacheGroupSpec) -> KVCacheConfig:

391 + return KVCacheConfig(

392 + num_blocks=4,

393 + kv_cache_tensors=[],

394 + kv_cache_groups=[source, indexer],

395 + hisparse_host_num_blocks=4,

396 + )

397 +

398 + config = _make_vllm_config()

399 + worker = build_offloading_config(config, make_config(worker_indexer))

400 + scheduler = build_offloading_config(config, make_config(scheduler_indexer))

401 +

402 + assert worker.worker_kv_bytes_per_block == scheduler.worker_kv_bytes_per_block

403 + assert worker.worker_kv_bytes_per_block == wrapped.page_size_bytes

404 +

405 +

406 def test_zero_blocks_skips_tensor_layout_validation():

407 kv_cache_config = _make_sizing_kv_cache_config(packed=False)

408 kv_cache_config.num_blocks = 0

tests/v1/kv_offload/cpu/test_shared_offload_region.py

+113 -17

@@ -10,11 +10,12 @@ import os

10 import threading

11 import time

12 import uuid

13 -from unittest.mock import MagicMock

13 +from unittest.mock import MagicMock, call

15 import pytest

17 from vllm.utils.system_utils import get_mp_context

18 +from vllm.v1.kv_offload.cpu import shared_offload_region as region_module

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

20 SharedOffloadRegion,

21 _wait_for_file_size,

@@ -192,28 +193,43 @@ def _mp_barrier_construct_and_hold(

193 barrier,

194 fill_value: int,

195 done_queue,

196 + replicated: bool = False,

197 ) -> None:

198 """Construct with a real cross-process barrier, write, then hold the

199 mapping (no cleanup) until the parent SIGKILLs this process."""

200 try:

201 + populate_calls = 0

202 + get_populate_write_fn = region_module._get_populate_write_fn

203 +

204 + def track_population(mmap_obj):

205 + nonlocal populate_calls

206 + populate_calls += 1

207 + return get_populate_write_fn(mmap_obj)

208 +

209 + region_module._get_populate_write_fn = track_population

210 region = SharedOffloadRegion(

211 engine_id=engine_id,

212 num_chunks=2,

202 - rank=rank,

203 - kv_bytes_per_chunk=num_workers * PAGE_SIZE,

213 + rank=0 if replicated else rank,

214 + kv_bytes_per_chunk=PAGE_SIZE if replicated else num_workers * PAGE_SIZE,

215 cpu_page_size=PAGE_SIZE,

216 barrier=lambda: barrier.wait(30),

217 + populate_only_on_creator=replicated,

218 )

219 # The constructor's barrier precedes the creator's unlink.

220 barrier.wait(30)

221 t = region.create_next_worker_view(PAGE_SIZE)

210 - t[:, :] = fill_value

222 + if not replicated or rank == 0:

223 + t[:, :] = fill_value

224 + barrier.wait(30)

225 done_queue.put(

226 {

227 "rank": rank,

228 "error": None,

229 "inode": os.fstat(region.fd).st_ino,

230 "path_exists": os.path.exists(region.mmap_path),

231 + "populate_calls": populate_calls,

232 + "contents_match": bool((t == (1 if replicated else fill_value)).all()),

233 }

234 )

235 time.sleep(60) # hold the mapping until killed

@@ -752,6 +768,26 @@ def test_cleanup_idempotent(iid):

768 r.cleanup() # must be a no-op

769

770

771 +def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):

772 + """cleanup() must release every chunk from a split host registration."""

773 + r = _make_region(iid)

774 + cudart = MagicMock()

775 + cudart.cudaHostUnregister.return_value = MagicMock(value=0)

776 + monkeypatch.setattr(region_module, "current_platform", MagicMock())

777 + region_module.current_platform.is_cuda_alike.return_value = True

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

779 + r.pinned_addresses = [0x100000, 0x200000, 0x300000]

780 + r.is_pinned = True

781 +

782 + r.cleanup()

783 +

784 + assert cudart.cudaHostUnregister.call_args_list == [

785 + call(0x300000),

786 + call(0x200000),

787 + call(0x100000),

788 + ]

789 +

790 +

791 def test_cleanup_after_create_next_worker_view_releases_mmap(iid):

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

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

847 os.close(fd)

848

849

850 +@pytest.mark.skip_global_cleanup

851 +def test_wait_for_file_size_rejects_unlinked_file(tmp_path):

852 + path = tmp_path / "unlinked.mmap"

853 + fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600)

854 + try:

855 + os.unlink(path)

856 + with pytest.raises(RuntimeError, match="creator failed"):

857 + _wait_for_file_size(fd, PAGE_SIZE, timeout=1.0)

858 + finally:

859 + os.close(fd)

860 +

861 +

862 # ---------------------------------------------------------------------------

863 # Constructor — capacity validation

864 # ---------------------------------------------------------------------------

865

866

867 +@pytest.mark.skip_global_cleanup

868 +@pytest.mark.skipif(not os.path.isdir("/dev/shm"), reason="requires /dev/shm")

869 +def test_creator_memory_check_runs_only_for_creator(iid):

870 + checked_sizes: list[int] = []

871 + creator = SharedOffloadRegion(

872 + engine_id=iid,

873 + num_chunks=4,

874 + rank=0,

875 + kv_bytes_per_chunk=PAGE_SIZE,

876 + cpu_page_size=PAGE_SIZE,

877 + creator_memory_check=checked_sizes.append,

878 + )

879 + joiner: SharedOffloadRegion | None = None

880 + try:

881 + joiner = SharedOffloadRegion(

882 + engine_id=iid,

883 + num_chunks=4,

884 + rank=0,

885 + kv_bytes_per_chunk=PAGE_SIZE,

886 + cpu_page_size=PAGE_SIZE,

887 + creator_memory_check=checked_sizes.append,

888 + )

889 + assert checked_sizes == [4 * PAGE_SIZE]

890 + finally:

891 + if joiner is not None:

892 + joiner.cleanup()

893 + creator.cleanup()

894 + _cleanup_file(creator.mmap_path)

895 +

896 +

897 def test_insufficient_space_raises_clear_error(monkeypatch):

898 """A failed creator capacity check must clean up and give a clear error."""

899 import vllm.v1.kv_offload.cpu.shared_offload_region as region

@@ -916,7 +994,8 @@ def test_barrier_failure_unlinks_creator_and_raises(iid):

994 assert not os.path.exists(path)

995

996

919 -def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):

997 +@pytest.mark.parametrize("replicated", [False, True])

998 +def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated):

999 """With a real cross-process barrier every worker maps one shared inode,

1000 the name is gone while they run, and SIGKILL leaks nothing."""

1001 num_workers = 2

@@ -927,7 +1006,7 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):

1006 procs = [

1007 ctx.Process(

1008 target=_mp_barrier_construct_and_hold,

930 - args=(iid, rank, num_workers, barrier, rank + 1, done_queue),

1009 + args=(iid, rank, num_workers, barrier, rank + 1, done_queue, replicated),

1010 )

1011 for rank in range(num_workers)

1012 ]

@@ -939,6 +1018,10 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):

1018 assert r["error"] is None, f"rank {r['rank']}: {r['error']}"

1019 assert len({r["inode"] for r in results}) == 1, "workers split onto two files"

1020 assert not any(r["path_exists"] for r in results)

1021 + assert all(r["contents_match"] for r in results)

1022 + assert sum(r["populate_calls"] for r in results) == (

1023 + 1 if replicated else num_workers

1024 + )

1025 assert not os.path.exists(path)

1026 finally:

1027 for p in procs:

@@ -950,23 +1033,36 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):

1033 assert restarted._creator is True, "restart must be able to create anew"

1034

1035

953 -def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch):

1036 +@pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"])

1037 +def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase):

1038 """A worker that dies before the rendezvous must still arrive at the

1039 barrier, or its peers block in the collective until it times out."""

956 - import vllm.v1.kv_offload.cpu.shared_offload_region as region

957 -

958 - monkeypatch.setattr(

959 - region,

960 - "check_shm_free_space",

961 - MagicMock(side_effect=RuntimeError("Insufficient space")),

1040 + failure = MagicMock(side_effect=ValueError("Insufficient space"))

1041 + memory_check = failure if failure_phase == "memory" else None

1042 + if failure_phase == "shm":

1043 + monkeypatch.setattr(region_module, "check_shm_free_space", failure)

1044 + elif failure_phase == "population":

1045 + monkeypatch.setattr(region_module, "_get_populate_write_fn", failure)

1046 + path = f"/dev/shm/vllm_offload_{iid}.mmap"

1047 + seen_at_barrier = []

1048 + barrier = MagicMock(

1049 + side_effect=lambda: seen_at_barrier.append(os.path.exists(path))

1050 )

963 - barrier = MagicMock()

1051

965 - with pytest.raises(RuntimeError, match="Insufficient space"):

966 - _make_region(iid, barrier=barrier)

1052 + with pytest.raises(ValueError, match="Insufficient space"):

1053 + SharedOffloadRegion(

1054 + engine_id=iid,

1055 + num_chunks=1,

1056 + rank=0,

1057 + kv_bytes_per_chunk=PAGE_SIZE,

1058 + cpu_page_size=PAGE_SIZE,

1059 + barrier=barrier,

1060 + creator_memory_check=memory_check,

1061 + populate_only_on_creator=True,

1062 + )

1063

1064 barrier.assert_called_once_with()

969 - assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap")

1065 + assert seen_at_barrier == [False]

1066

1067

1068 def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch):

tests/v1/worker/test_attn_utils.py

+82

@@ -12,6 +12,7 @@ from types import SimpleNamespace

12 import pytest

13 import torch

15 +import vllm.v1.hisparse.binding as attn_utils_module

16 from tests.v1.attention.utils import dense_kv_cache_views

17 from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf

18 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

240 assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()

241

242

243 +class _FakeSharedHostRegion:

244 + def __init__(self) -> None:

245 + self.cleanup_calls = 0

246 + self.base_tensor = torch.empty(1, dtype=torch.int8)

247 +

248 + def cleanup(self) -> None:

249 + self.cleanup_calls += 1

250 +

251 +

252 +def test_profiling_cleanup_releases_tp_shared_region_once(monkeypatch):

253 + """TP-shared profiling pools must use region-aware chunk cleanup."""

254 + region = _FakeSharedHostRegion()

255 + runtime = SimpleNamespace(

256 + _host_cache=object(),

257 + registered_host_pool=region.base_tensor,

258 + hot_backing=object(),

259 + shared_host_region=region,

260 + )

261 + forward_context = {

262 + "layer": SimpleNamespace(

263 + hisparse_cache=SimpleNamespace(runtime=runtime),

264 + )

265 + }

266 + released = []

267 +

268 + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region):

269 + released.append((runtimes, pinned_host_pools, shared_host_region))

270 +

271 + monkeypatch.setattr(

272 + attn_utils_module,

273 + "release_pinned_state",

274 + release_pinned_state,

275 + )

276 +

277 + attn_utils_module.release_hisparse_profiling_cache(forward_context)

278 +

279 + assert released == [([runtime], [], region)]

280 +

281 +

282 +@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"])

283 +def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase):

284 + """A failure after mmap allocation must not leak the shared registration."""

285 + region = _FakeSharedHostRegion()

286 + vllm_config = SimpleNamespace(

287 + cache_config=SimpleNamespace(

288 + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC

289 + ),

290 + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),

291 + )

292 +

293 + def allocate(*args):

294 + args[-1].shared_region = region

295 + if failure_phase == "allocation":

296 + raise RuntimeError("initialization failed")

297 + return {}

298 +

299 + def bind(**kwargs):

300 + if failure_phase == "binding":

301 + raise RuntimeError("initialization failed")

302 + return []

303 +

304 + def buffers(*args, **kwargs):

305 + raise RuntimeError("initialization failed")

306 +

307 + monkeypatch.setattr(attn_utils_module, "allocate_hisparse_kv_caches", allocate)

308 + monkeypatch.setattr(attn_utils_module, "bind_hisparse_kv_caches", bind)

309 + monkeypatch.setattr(

310 + attn_utils_module, "initialize_hisparse_runtime_buffers", buffers

311 + )

312 + with pytest.raises(RuntimeError, match="initialization failed"):

313 + attn_utils_module.init_hisparse_kv_cache(

314 + SimpleNamespace(),

315 + torch.device("cpu"),

316 + [],

317 + vllm_config,

318 + {},

319 + SimpleNamespace(),

320 + )

321 + assert region.cleanup_calls == 1

322 +

323 +

324 def test_reshape_padded_kv_cache_strides_by_padded_page():

325 num_blocks = 3

326 spec = FullAttentionSpec(

tests/v1/worker/test_kv_cache_allocation_scope.py

+5 -1

@@ -60,7 +60,11 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> No

60 )

61 ]

63 - monkeypatch.setattr(hisparse_binding, "HiSparseHostPool", lambda: host_pool)

63 + monkeypatch.setattr(

64 + hisparse_binding,

65 + "HiSparseHostPool",

66 + lambda vllm_config, kv_cache_config: host_pool,

67 + )

68 monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate)

69 monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse)

tests/v1/worker/test_utils.py

+253

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import mmap

5 from collections import deque

6 from contextlib import nullcontext

7 from types import SimpleNamespace

@@ -188,6 +189,232 @@ def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary(

189 assert handle.submit_layer_mirror.call_count == expected_calls

190

191

192 +def _hisparse_parallel_config(**overrides):

193 + values = {

194 + "tensor_parallel_size": 2,

195 + "pipeline_parallel_size": 1,

196 + "prefill_context_parallel_size": 1,

197 + "decode_context_parallel_size": 1,

198 + "world_size": 2,

199 + "distributed_executor_backend": "mp",

200 + "nnodes_within_dp": 1,

201 + "data_parallel_index": 3,

202 + }

203 + values.update(overrides)

204 + return SimpleNamespace(**values)

205 +

206 +

207 +def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch):

208 + monkeypatch.setattr(

209 + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True

210 + )

211 + config = SimpleNamespace(parallel_config=_hisparse_parallel_config())

212 +

213 + assert hisparse_runtime_module.use_shared_hisparse_host_pool(config)

214 +

215 + unsupported = (

216 + {"tensor_parallel_size": 1, "world_size": 1},

217 + {"pipeline_parallel_size": 2, "world_size": 4},

218 + {"prefill_context_parallel_size": 2, "world_size": 4},

219 + {"decode_context_parallel_size": 2},

220 + {"world_size": 4},

221 + {"distributed_executor_backend": "ray"},

222 + {"nnodes_within_dp": 2},

223 + )

224 + for overrides in unsupported:

225 + config.parallel_config = _hisparse_parallel_config(**overrides)

226 + assert not hisparse_runtime_module.use_shared_hisparse_host_pool(config)

227 +

228 +

229 +@pytest.mark.skip_global_cleanup

230 +def test_hisparse_shared_host_pool_uses_one_replicated_mmap(monkeypatch):

231 + page = mmap.PAGESIZE

232 + tp_group = MagicMock()

233 + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group)

234 +

235 + class FakeSharedOffloadRegion:

236 + BLOCK_SIZE_ALIGNMENT = page

237 +

238 + def __init__(self, **kwargs):

239 + self.kwargs = kwargs

240 + self.total_size_bytes = kwargs["num_chunks"] * kwargs["kv_bytes_per_chunk"]

241 + self.base_tensor = torch.empty(self.total_size_bytes, dtype=torch.int8)

242 + self.is_pinned = False

243 + self.pinned_addresses = []

244 + self.view_sizes = []

245 + self.offset = 0

246 +

247 + def create_next_canonical_view(self, size):

248 + self.view_sizes.append(size)

249 + view = torch.as_strided(

250 + self.base_tensor,

251 + size=(self.kwargs["num_chunks"], size),

252 + stride=(self.kwargs["kv_bytes_per_chunk"], 1),

253 + storage_offset=self.offset,

254 + )

255 + self.offset += size

256 + return view

257 +

258 + def cleanup(self):

259 + pass

260 +

261 + monkeypatch.setattr(

262 + hisparse_runtime_module, "SharedOffloadRegion", FakeSharedOffloadRegion

263 + )

264 + registration_ranges = MagicMock(return_value=((0, page), (page, 4 * page)))

265 + monkeypatch.setattr(

266 + hisparse_runtime_module,

267 + "_hisparse_registration_ranges",

268 + registration_ranges,

269 + )

270 + pinned: list[torch.Tensor] = []

271 +

272 + def pin_tensor(tensor):

273 + pinned.append(tensor)

274 +

275 + monkeypatch.setattr(hisparse_runtime_module, "pin_tensor", pin_tensor)

276 + config = SimpleNamespace(

277 + instance_id="instance",

278 + parallel_config=_hisparse_parallel_config(

279 + tensor_parallel_size=1,

280 + world_size=1,

281 + distributed_executor_backend="uni",

282 + ),

283 + )

284 +

285 + pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools(

286 + config,

287 + [24, 40],

288 + num_blocks=4,

289 + host_block_stride=page,

290 + use_shared_host_pool=True,

291 + )

292 +

293 + assert region is not None

294 + assert private_pools == []

295 + assert region.kwargs == {

296 + "engine_id": "hisparse_instance_dp3",

297 + "num_chunks": 1,

298 + "rank": 0,

299 + "kv_bytes_per_chunk": 4 * page,

300 + "cpu_page_size": 64,

301 + "barrier": tp_group.barrier,

302 + "creator_memory_check": hisparse_runtime_module.check_hisparse_host_memory,

303 + "populate_only_on_creator": True,

304 + }

305 + assert region.view_sizes == [24, 40]

306 + assert [pool.shape for pool in pools] == [(24,), (40,)]

307 + registration_ranges.assert_called_once_with([24, 40], 4, page)

308 + assert [

309 + (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes)

310 + for tensor in pinned

311 + ] == [(0, page), (page, 3 * page)]

312 + assert region.is_pinned

313 +

314 +

315 +@pytest.mark.parametrize("fail_registration", [None, 1, 2])

316 +def test_shared_host_pool_tracks_successful_registrations(

317 + monkeypatch, fail_registration

318 +):

319 + """Only successful registrations may be unregistered on allocation failure."""

320 + page = mmap.PAGESIZE

321 + monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", MagicMock())

322 + backing = torch.empty(3 * page, dtype=torch.uint8)

323 + region = MagicMock(base_tensor=backing, pinned_addresses=[], is_pinned=False)

324 + monkeypatch.setattr(

325 + hisparse_runtime_module, "SharedOffloadRegion", lambda **kw: region

326 + )

327 + monkeypatch.setattr(

328 + hisparse_runtime_module,

329 + "_hisparse_registration_ranges",

330 + lambda *args: ((0, page), (page, 3 * page)),

331 + )

332 + cudart = MagicMock()

333 + cudart.cudaHostRegister.side_effect = [

334 + SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2)

335 + ]

336 + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart)

337 + config = SimpleNamespace(

338 + instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0)

339 + )

340 + context = (

341 + pytest.raises(RuntimeError, match="cudaHostRegister failed")

342 + if fail_registration

343 + else nullcontext()

344 + )

345 + with context:

346 + hisparse_runtime_module.allocate_hisparse_host_pools(

347 + config, [3 * page], 3, page, use_shared_host_pool=True

348 + )

349 + successful = 2 if fail_registration is None else fail_registration - 1

350 + assert region.pinned_addresses == [

351 + backing.data_ptr() + i * page for i in range(successful)

352 + ]

353 + assert region.is_pinned == bool(successful)

354 + assert region.cleanup.call_count == int(fail_registration is not None)

355 + attempts = 2 if fail_registration is None else fail_registration

356 + assert (

357 + cudart.cudaHostRegister.call_args_list

358 + == [

359 + call(backing.data_ptr(), page, 0),

360 + call(backing.data_ptr() + page, 2 * page, 0),

361 + ][:attempts]

362 + )

363 +

364 +

365 +def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch):

366 + """Aliased tensor configs must not count the shared backing once per view."""

367 + region = SimpleNamespace(base_tensor=torch.empty(64, dtype=torch.int8))

368 + allocate = MagicMock(return_value=([], [], region))

369 + monkeypatch.setattr(

370 + hisparse_runtime_module, "allocate_hisparse_host_pools", allocate

371 + )

372 + config = SimpleNamespace(

373 + hisparse_shared_host_pool=True,

374 + hisparse_host_num_blocks=4,

375 + hisparse_host_block_stride=16,

376 + kv_cache_tensors=[

377 + SimpleNamespace(

378 + host_resident=True,

379 + layers=["a", "b"],

380 + offset=0,

381 + layer_stride=12,

382 + size=44,

383 + ),

384 + SimpleNamespace(

385 + host_resident=True, layers=["c"], offset=24, layer_stride=20, size=44

386 + ),

387 + ],

388 + )

389 + vllm_config = SimpleNamespace()

390 + pool = hisparse_runtime_module.HiSparseHostPool(vllm_config, config)

391 + backing = pool.allocate(44)

392 + allocate.assert_called_once_with(

393 + vllm_config, [12, 12, 20], 4, 16, use_shared_host_pool=True

394 + )

395 + assert backing.numel() == 44

396 + assert backing.data_ptr() == region.base_tensor.data_ptr()

397 + assert pool.registered is region.base_tensor

398 +

399 +

400 +def test_hisparse_registration_chunks_end_between_host_blocks():

401 + """Registration seams must not bisect a DMA-addressable host block."""

402 + page = mmap.PAGESIZE

403 + ranges = hisparse_runtime_module._hisparse_registration_ranges(

404 + tensor_sizes=[4 * 3 * page, 4 * 5 * page],

405 + num_blocks=4,

406 + host_block_stride=8 * page,

407 + max_chunk_bytes=10 * page,

408 + )

409 +

410 + assert ranges == (

411 + (0, 9 * page),

412 + (9 * page, 17 * page),

413 + (17 * page, 27 * page),

414 + (27 * page, 32 * page),

415 + )

416 +

417 +

418 def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding():

419 backing = torch.full((10, 2, 3), -1, dtype=torch.float32)

420 cache = backing[1:9]

@@ -931,6 +1158,32 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc

1158 assert streams == []

1159

1160

1161 +def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch):

1162 + worker = object.__new__(HiSparseConnectorWorker)

1163 + worker._initialized = True

1164 + worker._slot_mapping_staging = None

1165 + worker.dma_stream = None

1166 + worker.cache_handles = []

1167 + worker.pinned_host_pools = []

1168 + worker.shared_host_region = object()

1169 + released = False

1170 +

1171 + def release_pinned_state(runtimes, pinned_host_pools, shared_host_region):

1172 + nonlocal released

1173 + assert runtimes == []

1174 + assert pinned_host_pools == []

1175 + assert shared_host_region is worker.shared_host_region

1176 + released = True

1177 +

1178 + monkeypatch.setattr(

1179 + hisparse_worker_module, "release_pinned_state", release_pinned_state

1180 + )

1181 +

1182 + worker.shutdown()

1183 +

1184 + assert released

1185 +

1186 +

1187 class _TestReplaySSMMixer(MambaMixer2):

1188 def __init__(self):

1189 torch.nn.Module.__init__(self)

vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py

+24 -2

@@ -147,6 +147,12 @@ def _select_written_row_mirrors(

147 )

148

149

150 +def _is_hisparse_host_writer(

151 + shared_host_region: SharedOffloadRegion | None,

152 +) -> bool:

153 + return shared_host_region is None or get_tensor_model_parallel_rank() == 0

154 +

155 +

156 def _create_hisparse_host_events(

157 shared_host_region: SharedOffloadRegion | None,

158 is_host_writer: bool,

@@ -200,15 +206,26 @@ class HiSparseConnectorWorker:

206 raise RuntimeError("HiSparse connector found no hot-cache handles.")

207 hot_backings: dict[int, torch.Tensor] = {}

208 registered_host_pools: dict[int, torch.Tensor] = {}

209 + shared_host_regions: dict[int, SharedOffloadRegion] = {}

210 for cache in cache_handles:

211 hot_backing = cache.runtime.hot_backing

212 hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing

213 registered_pool = cache.runtime.registered_host_pool

214 registered_host_pools[registered_pool.data_ptr()] = registered_pool

215 + if (region := cache.runtime.shared_host_region) is not None:

216 + shared_host_regions[id(region)] = region

217 if len(hot_backings) != 1:

218 raise RuntimeError("HiSparse hot tensors must share one GPU backing.")

219 hot_backing = next(iter(hot_backings.values()))

211 - pinned_host_pools = list(registered_host_pools.values())

220 + if len(shared_host_regions) > 1:

221 + raise RuntimeError("HiSparse caches must share one host region.")

222 + shared_host_region = next(iter(shared_host_regions.values()), None)

223 + pinned_host_pools = (

224 + []

225 + if shared_host_region is not None

226 + else list(registered_host_pools.values())

227 + )

228 + is_host_writer = _is_hisparse_host_writer(shared_host_region)

229

230 resident = cache_handles[0].view

231 assert resident is not None

@@ -223,11 +240,14 @@ class HiSparseConnectorWorker:

240 host_num_blocks,

241 hot_backing.device,

242 pinned_host_pools,

243 + shared_host_region=shared_host_region,

244 + is_host_writer=is_host_writer,

245 )

246 except Exception:

247 release_pinned_state(

248 [cache.runtime for cache in cache_handles],

249 pinned_host_pools,

250 + shared_host_region,

251 )

252 raise

253

@@ -866,6 +886,8 @@ class HiSparseConnectorWorker:

886 if self.dma_stream is not None:

887 self.dma_stream.synchronize()

888 release_pinned_state(

869 - [cache.runtime for cache in self.cache_handles], self.pinned_host_pools

889 + [cache.runtime for cache in self.cache_handles],

890 + self.pinned_host_pools,

891 + self.shared_host_region,

892 )

893 self._initialized = False

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

+16 -2

@@ -17,6 +17,7 @@ from vllm.v1.kv_cache_interface import (

17 MLAAttentionSpec,

18 SlidingWindowMLASpec,

19 SlidingWindowSpec,

20 + UniformTypeKVCacheSpecs,

21 iter_layer_specs,

22 )

23 from vllm.v1.kv_offload.config import (

@@ -29,7 +30,7 @@ from vllm.v1.kv_offload.config import (

31 if TYPE_CHECKING:

32 from vllm.config import VllmConfig

32 - from vllm.v1.kv_cache_interface import KVCacheConfig

33 + from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec

36 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, ...

43 )

46 +def _group_kv_bytes_per_block(group: "KVCacheGroupSpec") -> int:

47 + """Return the physical bytes occupied by one block of a cache group.

48 +

49 + Worker configs may retain ``UniformTypeKVCacheSpecs`` while scheduler

50 + configs flatten that wrapper to one representative per-layer spec. Keep

51 + the result invariant across those two representations.

52 + """

53 + spec = group.kv_cache_spec

54 + if isinstance(spec, UniformTypeKVCacheSpecs):

55 + return spec.page_size_bytes

56 + return spec.page_size_bytes * len(group.layer_names)

57 +

58 +

59 def build_offloading_config(

60 vllm_config: "VllmConfig",

61 kv_cache_config: "KVCacheConfig",

@@ -133,7 +147,7 @@ def build_offloading_config(

147 worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks

148 elif kv_cache_config.num_blocks > 0:

149 worker_kv_bytes_per_block = sum(

136 - group.kv_cache_spec.page_size_bytes for _, group in selected_groups

150 + _group_kv_bytes_per_block(group) for _, group in selected_groups

151 )

152

153 single_group_spec = (

vllm/v1/hisparse/binding.py

+44 -27

@@ -120,27 +120,32 @@ def init_hisparse_kv_cache(

120 block_tables: "BlockTables",

121 ) -> dict[str, torch.Tensor]:

122 """Allocate and bind HiSparse caches within the caller's allocation context."""

123 - host_pool = HiSparseHostPool()

124 - kv_caches = allocate_hisparse_kv_caches(

125 - kv_cache_config,

126 - device,

127 - vllm_config.cache_config.get_resolved_kv_cache_layout(),

128 - kernel_block_sizes,

129 - host_pool,

130 - )

131 - cache_handles = bind_hisparse_kv_caches(

132 - forward_context=forward_context,

133 - kv_cache_config=kv_cache_config,

134 - kv_caches=kv_caches,

135 - block_tables=block_tables,

136 - host_pool=host_pool,

137 - )

138 - initialize_hisparse_runtime_buffers(

139 - cache_handles,

140 - max_num_reqs=vllm_config.scheduler_config.max_num_seqs,

141 - max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,

142 - )

143 - return kv_caches

123 + host_pool = HiSparseHostPool(vllm_config, kv_cache_config)

124 + try:

125 + kv_caches = allocate_hisparse_kv_caches(

126 + kv_cache_config,

127 + device,

128 + vllm_config.cache_config.get_resolved_kv_cache_layout(),

129 + kernel_block_sizes,

130 + host_pool,

131 + )

132 + cache_handles = bind_hisparse_kv_caches(

133 + forward_context=forward_context,

134 + kv_cache_config=kv_cache_config,

135 + kv_caches=kv_caches,

136 + block_tables=block_tables,

137 + host_pool=host_pool,

138 + )

139 + initialize_hisparse_runtime_buffers(

140 + cache_handles,

141 + max_num_reqs=vllm_config.scheduler_config.max_num_seqs,

142 + max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,

143 + )

144 + return kv_caches

145 + except Exception:

146 + if host_pool.shared_region is not None:

147 + host_pool.shared_region.cleanup()

148 + raise

149

150

151 def _get_hisparse_cache(

@@ -166,13 +171,24 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None:

171 if not runtimes:

172 return

173

169 - registered_pools = list(

170 - {

171 - runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool

172 - for runtime in runtimes.values()

173 - }.values()

174 + shared_regions = {

175 + id(runtime.shared_host_region): runtime.shared_host_region

176 + for runtime in runtimes.values()

177 + if runtime.shared_host_region is not None

178 + }

179 + assert len(shared_regions) <= 1

180 + shared_region = next(iter(shared_regions.values()), None)

181 + registered_pools = (

182 + []

183 + if shared_region is not None

184 + else list(

185 + {

186 + runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool

187 + for runtime in runtimes.values()

188 + }.values()

189 + )

190 )

175 - release_pinned_state(list(runtimes.values()), registered_pools)

191 + release_pinned_state(list(runtimes.values()), registered_pools, shared_region)

192 for cache in cache_handles:

193 cache.mirror_staging_cache = None

194 cache.mirror_staging_slots = None

@@ -257,6 +273,7 @@ def bind_hisparse_kv_caches(

273 assert source_cache.untyped_storage().data_ptr() == (

274 host_pool.registered.untyped_storage().data_ptr()

275 )

276 + cache_handle.runtime.shared_host_region = host_pool.shared_region

277 cache_handle.runtime.bind_source_cache(

278 source_cache,

279 registered_host_pool=host_pool.registered,

vllm/v1/hisparse/layout.py

+16 -3

@@ -9,7 +9,11 @@ from vllm.config import VllmConfig

9 from vllm.config.kv_transfer import hisparse_host_pool_gib

10 from vllm.logger import init_logger

11 from vllm.utils.math_utils import cdiv

12 -from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig

12 +from vllm.v1.hisparse.runtime import (

13 + ResolvedHiSparseConfig,

14 + get_hisparse_host_block_stride,

15 + use_shared_hisparse_host_pool,

16 +)

17 from vllm.v1.kv_cache_interface import (

18 HiSparseHotSpec,

19 HiSparseResidentSpec,

@@ -36,6 +40,8 @@ class HiSparseLayout:

40 source_group: KVCacheGroupSpec

41 device_groups: list[KVCacheGroupSpec]

42 host_num_blocks: int

43 + host_block_stride: int

44 + shared_host_pool: bool

47 def get_hisparse_kv_cache_groups(

@@ -221,9 +227,12 @@ def create_hisparse_layout(

227 regular_groups = groups[1:]

228 gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups]

229

224 - host_num_blocks = host_budget // sum(

225 - spec.page_size_bytes for spec in source_specs.values()

230 + shared_host_pool = use_shared_hisparse_host_pool(vllm_config)

231 + host_block_stride = get_hisparse_host_block_stride(

232 + sum(spec.page_size_bytes for spec in source_specs.values()),

233 + use_shared_host_pool=shared_host_pool,

234 )

235 + host_num_blocks = host_budget // host_block_stride

236 if host_num_blocks <= 0:

237 raise ValueError("HiSparse has no allocatable host blocks.")

238

@@ -231,6 +240,8 @@ def create_hisparse_layout(

240 source_group=source_group,

241 device_groups=gpu_groups,

242 host_num_blocks=host_num_blocks,

243 + host_block_stride=host_block_stride,

244 + shared_host_pool=shared_host_pool,

245 )

246

247

@@ -339,6 +350,8 @@ def get_hisparse_kv_cache_config(

350 kv_cache_tensors=kv_cache_tensors,

351 kv_cache_groups=[*host_groups, *device_groups],

352 hisparse_host_num_blocks=hisparse_layout.host_num_blocks,

353 + hisparse_host_block_stride=hisparse_layout.host_block_stride,

354 + hisparse_shared_host_pool=hisparse_layout.shared_host_pool,

355 prefix_cache_retention_interval=(

356 vllm_config.cache_config.prefix_cache_retention_interval

357 ),

vllm/v1/hisparse/runtime.py

+189 -6

@@ -4,6 +4,8 @@

5 from __future__ import annotations

7 +import math

8 +import mmap

9 import time

10 from collections.abc import Callable

11 from dataclasses import dataclass

@@ -13,17 +15,22 @@ import psutil

15 import torch

17 from vllm.config import CUDAGraphMode, VllmConfig

18 +from vllm.distributed import get_tp_group

19 from vllm.forward_context import get_forward_context

20 from vllm.logger import init_logger

21 from vllm.platforms import current_platform

22 from vllm.utils.math_utils import round_up

23 from vllm.utils.torch_utils import current_stream

24 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

25 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor

27 logger = init_logger(__name__)

29 +HOST_REGISTER_CHUNK_BYTES = 256 * 2**30

30 +

31 if TYPE_CHECKING:

32 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup

33 + from vllm.v1.kv_cache_interface import KVCacheConfig

35 # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.

36 FP8_DS_MLA_ROW_BYTES = 656

@@ -185,22 +192,195 @@ def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]:

192 class HiSparseHostPool:

193 """Keep the single host backing and its exact CUDA registration range."""

194

188 - def __init__(self) -> None:

195 + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:

196 + self.vllm_config = vllm_config

197 + self.kv_cache_config = kv_cache_config

198 self.backing: torch.Tensor | None = None

199 self.registered: torch.Tensor | None = None

200 + self.shared_region: SharedOffloadRegion | None = None

201

202 def allocate(self, size: int) -> torch.Tensor:

203 assert self.backing is None, "HiSparse host pool is already allocated."

194 - check_hisparse_host_memory(size)

195 - self.backing, self.registered = allocate_pinned_host_pool(size)

204 + config = self.kv_cache_config

205 + if config.hisparse_shared_host_pool:

206 + tensor_sizes = [

207 + size

208 + for _, size in sorted(

209 + (tensor.offset + i * tensor.layer_stride, tensor.layer_stride)

210 + for tensor in config.kv_cache_tensors

211 + if tensor.host_resident

212 + for i in range(len(tensor.layers))

213 + )

214 + ]

215 + assert config.hisparse_host_num_blocks is not None

216 + assert config.hisparse_host_block_stride is not None

217 + _, _, self.shared_region = allocate_hisparse_host_pools(

218 + self.vllm_config,

219 + tensor_sizes,

220 + config.hisparse_host_num_blocks,

221 + config.hisparse_host_block_stride,

222 + use_shared_host_pool=True,

223 + )

224 + assert self.shared_region is not None

225 + self.registered = self.shared_region.base_tensor

226 + self.backing = self.registered[:size]

227 + else:

228 + check_hisparse_host_memory(size)

229 + self.backing, self.registered = allocate_pinned_host_pool(size)

230 return self.backing

231

232

233 +def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool:

234 + """Whether replicated MLA host KV can share one local mmap."""

235 + parallel = vllm_config.parallel_config

236 + return (

237 + current_platform.is_cuda_alike()

238 + and parallel.tensor_parallel_size > 1

239 + and parallel.pipeline_parallel_size == 1

240 + and parallel.prefill_context_parallel_size == 1

241 + and parallel.decode_context_parallel_size == 1

242 + and parallel.world_size == parallel.tensor_parallel_size

243 + and parallel.distributed_executor_backend == "mp"

244 + and parallel.nnodes_within_dp == 1

245 + )

246 +

247 +

248 +def get_hisparse_host_block_stride(

249 + page_size_bytes: int, *, use_shared_host_pool: bool

250 +) -> int:

251 + if use_shared_host_pool:

252 + return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT)

253 + return page_size_bytes

254 +

255 +

256 +def _hisparse_registration_ranges(

257 + tensor_sizes: list[int],

258 + num_blocks: int,

259 + host_block_stride: int,

260 + max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES,

261 +) -> tuple[tuple[int, int], ...]:

262 + """Split a shared pool without bisecting any tensor's host block.

263 +

264 + CUDA batch copies reject descriptors spanning adjacent registrations.

265 + """

266 + total_size = num_blocks * host_block_stride

267 + if total_size % mmap.PAGESIZE:

268 + raise ValueError("HiSparse host pool must be OS-page aligned.")

269 + if max_chunk_bytes < mmap.PAGESIZE:

270 + raise ValueError("HiSparse registration chunks must span at least one page.")

271 +

272 + tensors: list[tuple[int, int, int]] = []

273 + tensor_start = 0

274 + for tensor_size in tensor_sizes:

275 + if tensor_size % num_blocks:

276 + raise ValueError(

277 + f"HiSparse host tensor size {tensor_size} is not divisible by "

278 + f"{num_blocks} blocks."

279 + )

280 + tensor_end = tensor_start + tensor_size

281 + tensors.append((tensor_start, tensor_end, tensor_size // num_blocks))

282 + tensor_start = tensor_end

283 + if tensor_start > total_size:

284 + raise ValueError("HiSparse host tensors exceed the shared host pool.")

285 +

286 + def last_safe_boundary(lower: int, upper: int) -> int | None:

287 + best = None

288 + for start, end, block_bytes in tensors:

289 + first_block = max(0, (lower - start) // block_bytes + 1)

290 + last_block = min(num_blocks, (min(upper, end) - start) // block_bytes)

291 + if first_block > last_block:

292 + continue

293 + divisor = math.gcd(block_bytes, mmap.PAGESIZE)

294 + if start % divisor:

295 + continue

296 + period = mmap.PAGESIZE // divisor

297 + if period == 1:

298 + residue = 0

299 + else:

300 + # Find block boundaries that are also OS-page boundaries.

301 + residue = (

302 + (-start // divisor) * pow(block_bytes // divisor, -1, period)

303 + ) % period

304 + block = last_block - (last_block - residue) % period

305 + if block >= first_block:

306 + candidate = start + block * block_bytes

307 + best = candidate if best is None else max(best, candidate)

308 +

309 + if upper >= tensor_start:

310 + candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE

311 + if candidate >= tensor_start and candidate > lower:

312 + best = candidate if best is None else max(best, candidate)

313 + return best

314 +

315 + ranges: list[tuple[int, int]] = []

316 + start = 0

317 + while total_size - start > max_chunk_bytes:

318 + end = last_safe_boundary(start, start + max_chunk_bytes)

319 + if end is None:

320 + raise ValueError(

321 + "A HiSparse host block is too large for the registration chunk limit."

322 + )

323 + ranges.append((start, end))

324 + start = end

325 + ranges.append((start, total_size))

326 + return tuple(ranges)

327 +

328 +

329 +def allocate_hisparse_host_pools(

330 + vllm_config: VllmConfig,

331 + tensor_sizes: list[int],

332 + num_blocks: int,

333 + host_block_stride: int,

334 + *,

335 + use_shared_host_pool: bool,

336 +) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]:

337 + """Allocate private host tensors or one mmap shared by local TP ranks."""

338 + if not use_shared_host_pool:

339 + check_hisparse_host_memory(sum(tensor_sizes))

340 + private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes]

341 + return (

342 + [pool for pool, _ in private_pools],

343 + [registered for _, registered in private_pools],

344 + None,

345 + )

346 +

347 + region = SharedOffloadRegion(

348 + engine_id=(

349 + f"hisparse_{vllm_config.instance_id}_"

350 + f"dp{vllm_config.parallel_config.data_parallel_index}"

351 + ),

352 + num_chunks=1,

353 + rank=0,

354 + kv_bytes_per_chunk=num_blocks * host_block_stride,

355 + cpu_page_size=sum(tensor_sizes),

356 + barrier=get_tp_group().barrier,

357 + creator_memory_check=check_hisparse_host_memory,

358 + populate_only_on_creator=True,

359 + )

360 + try:

361 + for start, end in _hisparse_registration_ranges(

362 + tensor_sizes, num_blocks, host_block_stride

363 + ):

364 + tensor = region.base_tensor[start:end]

365 + pin_tensor(tensor)

366 + region.pinned_addresses.append(tensor.data_ptr())

367 + region.is_pinned = True

368 + pools = [

369 + region.create_next_canonical_view(size).view(-1) for size in tensor_sizes

370 + ]

371 + except Exception:

372 + region.cleanup()

373 + raise

374 + return pools, [], region

375 +

376 +

377 def release_pinned_state(

200 - runtimes: list[HiSparseRuntime], pinned_host_pools: list[torch.Tensor]

378 + runtimes: list[HiSparseRuntime],

379 + pinned_host_pools: list[torch.Tensor],

380 + shared_host_region: SharedOffloadRegion | None = None,

381 ) -> None:

382 """Synchronize and release registered host KV pools."""

203 - if pinned_host_pools:

383 + if pinned_host_pools or shared_host_region is not None:

384 try:

385 torch.accelerator.synchronize()

386 except RuntimeError as e:

@@ -208,7 +388,7 @@ def release_pinned_state(

388 "HiSparse: CUDA context unusable at teardown (%s); leaving "

389 "%d host-pool tensors pinned for kernel exit reclaim.",

390 e,

211 - len(pinned_host_pools),

391 + len(pinned_host_pools) + int(shared_host_region is not None),

392 )

393 return

394

@@ -240,6 +420,8 @@ def release_pinned_state(

420 del runtime._host_cache

421 del runtime.registered_host_pool

422 del runtime.hot_backing

423 + if shared_host_region is not None:

424 + shared_host_region.cleanup()

425

426

427 def hisparse_prefill_staging_remap(

@@ -536,6 +718,7 @@ class HiSparseRuntime:

718 self.eager_host_mirror = config.eager_host_mirror

719 self.resident_source_index = -1

720 self.request_state_indices: torch.Tensor | None = None

721 + self.shared_host_region: SharedOffloadRegion | None = None

722

723 @property

724 def host_cache(self) -> torch.Tensor:

vllm/v1/kv_cache_interface.py

@@ -1435,6 +1435,12 @@ class KVCacheConfig:

1435 hisparse_host_num_blocks: int | None = None

1436 """Capacity of the dedicated HiSparse host-block manager, when enabled."""

1437

1438 + hisparse_host_block_stride: int | None = None

1439 + """Physical bytes between consecutive HiSparse host blocks."""

1440 +

1441 + hisparse_shared_host_pool: bool = False

1442 + """Whether local TP ranks share one physical HiSparse host pool."""

1443 +

1444 @cached_property

1445 def transfer_group_ids(self) -> tuple[int, ...]:

1446 """IDs of cache groups that participate in external KV transfer."""

vllm/v1/kv_offload/cpu/shared_offload_region.py

+60 -25

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import contextlib

4 import errno

5 import mmap

6 import os

@@ -27,6 +28,10 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N

28 while True:

29 if os.fstat(fd).st_size >= expected_size:

30 return

31 + if os.fstat(fd).st_nlink == 0:

32 + raise RuntimeError(

33 + "Shared offload region creator failed during initialization."

34 + )

35 if time.monotonic() > deadline:

36 raise TimeoutError(

37 f"Timed out waiting for mmap file to reach {expected_size} bytes"

@@ -75,6 +80,9 @@ class SharedOffloadRegion:

80 given, the path is unlinked once every worker has mapped the file, so

81 the kernel reclaims the memory when the last worker exits, no matter

82 how it exits; mappings taken before the unlink stay valid.

83 +

84 + Creator-only population pre-faults the entire region before the barrier

85 + and requires that barrier to keep joiners from using unpopulated pages.

86 """

88 BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE

@@ -87,7 +95,12 @@ class SharedOffloadRegion:

95 kv_bytes_per_chunk: int,

96 cpu_page_size: int,

97 barrier: Callable[[], None] | None = None,

98 + *,

99 + creator_memory_check: Callable[[int], None] | None = None,

100 + populate_only_on_creator: bool = False,

101 ) -> None:

102 + if populate_only_on_creator and barrier is None:

103 + raise ValueError("Creator-only population requires a barrier.")

104 self.page_size = mmap.PAGESIZE

105 assert kv_bytes_per_chunk % self.page_size == 0

106

@@ -112,25 +125,18 @@ class SharedOffloadRegion:

125 # Joiner path — another worker won O_EXCL. Reopen and wait

126 # for the file to reach expected size.

127 self.fd = os.open(self.mmap_path, os.O_RDWR)

115 - try:

116 - _wait_for_file_size(self.fd, self.total_size_bytes)

117 - except (TimeoutError, OSError):

118 - os.close(self.fd)

119 - raise

128 + _wait_for_file_size(self.fd, self.total_size_bytes)

129 logger.info("Opened existing mmap file %s", self.mmap_path)

130 else:

131 # Creator path. We won O_EXCL, so we own the file: any

132 # failure here must clean up so concurrent joiners don't

133 # land on a 0-byte stub and spin in _wait_for_file_size

134 # for the full 30 s timeout.

126 - try:

127 - check_shm_free_space(self.total_size_bytes)

128 - os.ftruncate(self.fd, self.total_size_bytes)

129 - except (RuntimeError, OSError):

130 - os.unlink(self.mmap_path)

131 - os.close(self.fd)

132 - raise

135 self._creator = True

136 + if creator_memory_check is not None:

137 + creator_memory_check(self.total_size_bytes)

138 + check_shm_free_space(self.total_size_bytes)

139 + os.ftruncate(self.fd, self.total_size_bytes)

140 logger.info(

141 "Created mmap file %s (%.2f GB)",

142 self.mmap_path,

@@ -143,10 +149,21 @@ class SharedOffloadRegion:

149 flags=mmap.MAP_SHARED,

150 prot=mmap.PROT_READ | mmap.PROT_WRITE,

151 )

152 +

153 + if populate_only_on_creator and self._creator:

154 + populate_write_fn = _get_populate_write_fn(self.mmap_obj)

155 + populate_write_fn(self.mmap_obj, 0, self.total_size_bytes)

156 except Exception:

157 if self._creator:

148 - os.unlink(self.mmap_path)

158 + with contextlib.suppress(FileNotFoundError):

159 + os.unlink(self.mmap_path)

160 self._creator = False

161 + if hasattr(self, "mmap_obj") and self.mmap_obj is not None:

162 + self.mmap_obj.close()

163 + self.mmap_obj = None

164 + if hasattr(self, "fd") and self.fd is not None:

165 + os.close(self.fd)

166 + self.fd = None

167 # Peers block inside the barrier until the collective times out if

168 # we die before reaching it. Arrive anyway so every worker calls

169 # barrier() exactly once and they fail on their own errors instead

@@ -169,16 +186,28 @@ class SharedOffloadRegion:

186 barrier()

187 except Exception:

188 if self._creator:

172 - os.unlink(self.mmap_path)

189 + with contextlib.suppress(FileNotFoundError):

190 + os.unlink(self.mmap_path)

191 self._creator = False

192 self.mmap_obj.close()

193 os.close(self.fd)

194 + self.mmap_obj = None

195 + self.fd = None

196 raise

197 if self._creator:

198 os.unlink(self.mmap_path)

199 self._creator = False

200 logger.info("Unlinked mmap file %s", self.mmap_path)

201

202 + self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)

203 + self._views: list[torch.Tensor] = []

204 + self._canonical_offset = 0

205 + self.is_pinned: bool = False

206 + self.pinned_addresses: list[int] = []

207 +

208 + if populate_only_on_creator:

209 + return

210 +

211 populate_write_fn = _get_populate_write_fn(self.mmap_obj)

212

213 if rank is not None:

@@ -205,10 +234,11 @@ class SharedOffloadRegion:

234 "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0

235 )

236

208 - self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)

209 - self._views: list[torch.Tensor] = []

210 - self._canonical_offset = 0

211 - self.is_pinned: bool = False

237 + @property

238 + def base_tensor(self) -> torch.Tensor:

239 + if self._base is None:

240 + raise RuntimeError("Shared offload region has been released.")

241 + return self._base

242

243 def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor:

244 """Allocate a strided int8 view for this worker, one canonical tensor.

@@ -311,13 +341,18 @@ class SharedOffloadRegion:

341 if self.is_pinned and self._base is not None:

342 if current_platform.is_cuda_alike():

343 base_ptr = self._base.data_ptr()

314 - result = torch.cuda.cudart().cudaHostUnregister(base_ptr)

315 - if result.value != 0:

316 - logger.warning(

317 - "cudaHostUnregister failed for rank=%d (code=%d)",

318 - self.rank,

319 - result,

320 - )

344 + addresses = self.pinned_addresses or [base_ptr]

345 + for address in reversed(addresses):

346 + result = torch.cuda.cudart().cudaHostUnregister(address)

347 + if result.value != 0:

348 + logger.warning(

349 + "cudaHostUnregister failed for rank=%d, "

350 + "address=%#x (code=%d)",

351 + self.rank,

352 + address,

353 + result.value,

354 + )

355 + self.pinned_addresses.clear()

356 self.is_pinned = False

357 # Release views before _base: each view holds a _base reference and a

358 # direct StorageImpl reference. Freeing views first lets both refcounts