From 49f944be3e9fa12cbff5f30ea5ed8e1efa534e20 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:33:03 +0300 Subject: [PATCH] [Perf][HiSparse] Reuse pinned host pool after graph profiling (#61668) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Signed-off-by: Matthew Bonanni Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- tests/v1/worker/test_attn_utils.py | 72 ++++++++++++------- .../worker/test_kv_cache_allocation_scope.py | 2 +- tests/v1/worker/test_utils.py | 5 ++ .../kv_connector/v1/hisparse/worker.py | 3 + vllm/v1/hisparse/binding.py | 40 +++-------- vllm/v1/hisparse/runtime.py | 21 +++++- 6 files changed, 87 insertions(+), 56 deletions(-) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 8d30403f30..0f1537403e 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -9,12 +9,14 @@ never addressed by the logical view. from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import numpy as np import pytest import torch import vllm.v1.hisparse.binding as attn_utils_module +import vllm.v1.hisparse.runtime as hisparse_runtime_module from tests.v1.attention.utils import ( BatchSpec, create_common_attn_metadata, @@ -403,49 +405,63 @@ class _FakeSharedHostRegion: 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() +def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch): + """The real KV cache must reuse the profiling host pool, not pin it again.""" + pinned = torch.empty(8, dtype=torch.int8) + allocate_pinned = MagicMock(return_value=(pinned, pinned)) + monkeypatch.setattr( + hisparse_runtime_module, "allocate_pinned_host_pool", allocate_pinned + ) + monkeypatch.setattr( + hisparse_runtime_module, "check_hisparse_host_memory", lambda size: None + ) + vllm_config = SimpleNamespace(instance_id="profiling-reuse") + kv_cache_config = SimpleNamespace(hisparse_shared_host_pool=False) + pool = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) + backing = pool.allocate(8) runtime = SimpleNamespace( - _host_cache=object(), - registered_host_pool=region.base_tensor, - hot_backing=object(), - shared_host_region=region, + _host_cache=backing, registered_host_pool=pinned, hot_backing=object() ) forward_context = { - "layer": SimpleNamespace( - hisparse_cache=SimpleNamespace(runtime=runtime), - ) + "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)] + assert not hasattr(runtime, "_host_cache") + reused = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) + assert reused is pool + assert reused.allocate(8) is backing + allocate_pinned.assert_called_once_with(8) + hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) @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.""" +@pytest.mark.parametrize("shared", [False, True]) +def test_init_hisparse_rolls_back_host_pool(monkeypatch, failure_phase, shared): + """Failed initialization must release shared or private registrations.""" region = _FakeSharedHostRegion() + registered = torch.empty(4096, dtype=torch.int8) + cudart = MagicMock() + cudart.cudaHostUnregister.return_value.value = 0 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) vllm_config = SimpleNamespace( cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC ), scheduler_config=SimpleNamespace(max_num_seqs=1), + instance_id="rollback", ) + pools = [] def allocate(*args): - args[-1].shared_region = region + pools.append(args[-1]) + if shared: + args[-1].shared_region = region + else: + args[-1].registered = registered + args[-1].backing = registered[:8] if failure_phase == "allocation": raise RuntimeError("initialization failed") return {} @@ -472,7 +488,13 @@ def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): {}, SimpleNamespace(), ) - assert region.cleanup_calls == 1 + assert region.cleanup_calls == int(shared) + if shared: + cudart.cudaHostUnregister.assert_not_called() + else: + cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr()) + assert attn_utils_module.get_hisparse_host_pool(vllm_config, None) is not pools[0] + hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) def test_reshape_padded_kv_cache_strides_by_padded_page(): diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 52cf2a3741..b5b288d9b5 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -62,7 +62,7 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> No monkeypatch.setattr( hisparse_binding, - "HiSparseHostPool", + "get_hisparse_host_pool", lambda vllm_config, kv_cache_config: host_pool, ) monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a8f49c1b88..b2e9e9054f 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1257,6 +1257,8 @@ def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): worker.cache_handles = [] worker.pinned_host_pools = [] worker.shared_host_region = object() + worker.vllm_config = SimpleNamespace(instance_id="shutdown") + pool = hisparse_runtime_module.get_hisparse_host_pool(worker.vllm_config, None) released = False def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): @@ -1273,6 +1275,9 @@ def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): worker.shutdown() assert released + new_pool = hisparse_runtime_module.get_hisparse_host_pool(worker.vllm_config, None) + assert new_pool is not pool + hisparse_runtime_module.drop_hisparse_host_pool(worker.vllm_config) class _TestReplaySSMMixer(MambaMixer2): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index c53be27079..6e3d3415c3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -28,6 +28,7 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, + drop_hisparse_host_pool, release_pinned_state, update_hisparse_residency, ) @@ -254,6 +255,7 @@ class HiSparseConnectorWorker: pinned_host_pools, shared_host_region, ) + drop_hisparse_host_pool(self.vllm_config) raise def initialize( @@ -972,4 +974,5 @@ class HiSparseConnectorWorker: self.pinned_host_pools, self.shared_host_region, ) + drop_hisparse_host_pool(self.vllm_config) self._initialized = False diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 1420a45396..6a02dd571c 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -17,6 +17,8 @@ from vllm.v1.hisparse.layout import ( from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, HiSparseHostPool, + drop_hisparse_host_pool, + get_hisparse_host_pool, initialize_hisparse_runtime_buffers, release_pinned_state, ) @@ -126,7 +128,7 @@ 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(vllm_config, kv_cache_config) + host_pool = get_hisparse_host_pool(vllm_config, kv_cache_config) try: kv_caches = allocate_hisparse_kv_caches( kv_cache_config, @@ -150,6 +152,9 @@ def init_hisparse_kv_cache( except Exception: if host_pool.shared_region is not None: host_pool.shared_region.cleanup() + elif host_pool.registered is not None: + release_pinned_state([], [host_pool.registered]) + drop_hisparse_host_pool(vllm_config) raise @@ -163,37 +168,14 @@ def _get_hisparse_cache( def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: - cache_handles = [ - cache - for layer in forward_context.values() - if (cache := getattr(layer, "hisparse_cache", None)) is not None - ] + """Detach the profiling caches but keep the host pool for the real KV cache.""" runtimes = { id(cache.runtime): cache.runtime - for cache in cache_handles - if hasattr(cache.runtime, "_host_cache") + for layer in forward_context.values() + if (cache := getattr(layer, "hisparse_cache", None)) is not None + and hasattr(cache.runtime, "_host_cache") } - if not runtimes: - return - - 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, shared_region) + release_pinned_state(list(runtimes.values()), []) def bind_hisparse_kv_caches( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 3cc45c680e..87916ba4ad 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -223,7 +223,9 @@ class HiSparseHostPool: self.shared_region: SharedOffloadRegion | None = None def allocate(self, size: int) -> torch.Tensor: - assert self.backing is None, "HiSparse host pool is already allocated." + if self.backing is not None: + assert self.backing.numel() == size, "HiSparse host pool size changed." + return self.backing config = self.kv_cache_config if config.hisparse_shared_host_pool: tensor_sizes = [ @@ -253,6 +255,23 @@ class HiSparseHostPool: return self.backing +_host_pools: dict[str, HiSparseHostPool] = {} + + +def get_hisparse_host_pool( + vllm_config: VllmConfig, kv_cache_config: KVCacheConfig +) -> HiSparseHostPool: + """Return the engine's host pool, which outlives the profiling KV cache.""" + instance_id = vllm_config.instance_id + if instance_id not in _host_pools: + _host_pools[instance_id] = HiSparseHostPool(vllm_config, kv_cache_config) + return _host_pools[instance_id] + + +def drop_hisparse_host_pool(vllm_config: VllmConfig) -> None: + _host_pools.pop(vllm_config.instance_id, None) + + 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