iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][HiSparse] Reuse pinned host pool after graph profiling (#61668)

Browse Source 
Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Itay Etelis
2026-10-15 12:33:03 -07:00
committed by GitHub
co-authored by Matthew Bonanni Claude Opus 5.5
parent 105605b8e9
commit 49f944be3e
6 changed files with 87 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_attn_utils.py
+47 -25
View File
@@ -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():
tests/v1/worker/test_kv_cache_allocation_scope.py
+1 -1
View File
@@ -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)
tests/v1/worker/test_utils.py
+5
View File
@@ -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):
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
+3
View File
@@ -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
vllm/v1/hisparse/binding.py
+11 -29
View File
@@ -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(
vllm/v1/hisparse/runtime.py
+20 -1
View File
@@ -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