From 8def699a7a8292cb0bb580a1e6db24ea7c0b80a5 Mon Sep 17 00:00:00 2001
From: S1ro1 <matej.sirovatka@gmail.com>
Date: Fri, 18 Sep 2026 00:19:27 +0000
Subject: [PATCH] [Bugfix][HiSparse] Account for GPU copies during prefix
 admission

Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
Co-authored-by: OpenAI Codex
---
 tests/v1/core/test_prefix_caching.py         | 45 ++++++++++++++++++++
 vllm/v1/core/single_type_kv_cache_manager.py |  5 +++
 vllm/v1/hisparse/coordinator.py              | 11 +++++
 3 files changed, 61 insertions(+)

diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py
index 229f9c17f9ee..f6ce6648da60 100644
--- a/tests/v1/core/test_prefix_caching.py
+++ b/tests/v1/core/test_prefix_caching.py
@@ -1101,6 +1101,51 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():
     )
 
 
+@pytest.mark.parametrize("free_blocks", [7, 8, 9, 10])
+def test_hisparse_prefix_copy_admission_under_memory_pressure(free_blocks):
+    """Prefix-copy pins must leave enough GPU blocks for the hot region."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    resident_ids = [
+        block.block_id for block in manager.get_blocks("original").blocks[2]
+    ]
+    manager.free(original)
+    pool = manager.block_pool
+    pressure = pool.get_new_blocks(pool.get_num_free_blocks() - free_blocks)
+    evicted_ids = {block.block_id for block in pressure}
+    expected_resident_ids = [
+        pool.null_block.block_id if block_id in evicted_ids else block_id
+        for block_id in resident_ids[:3]
+    ]
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    computed, num_computed, _ = manager.get_computed_blocks(resumed)
+    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
+    allocated = manager.allocate_slots(
+        resumed,
+        num_new_tokens=len(tokens) - num_computed,
+        num_new_computed_tokens=num_computed,
+        new_computed_blocks=computed,
+    )
+    if free_blocks < 10:
+        assert allocated is None
+        assert pool.get_num_free_blocks() == free_blocks
+        pool.free_blocks(pressure)
+        allocated = manager.allocate_slots(
+            resumed,
+            num_new_tokens=len(tokens) - num_computed,
+            num_new_computed_tokens=num_computed,
+            new_computed_blocks=computed,
+        )
+    assert allocated is not None
+    assert [
+        block.block_id for block in manager.get_blocks("resumed").blocks[2][:3]
+    ] == expected_resident_ids
+
+
 def test_hisparse_host_backed_request_accepts_local_prefix_hit():
     """Local group-completion hits need no external host import allocation."""
     manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index 835a434c1e7c..4f1ae1665c0b 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -2285,6 +2285,11 @@ def get_num_blocks_to_allocate(
                 assert self.coordinator is not None
                 assert self.coordinator.gpu_pool is not None
                 return self.coordinator.gpu_pool.num_gpu_blocks + 1
+        if self.coordinator is not None and request_id not in self.num_cached_block:
+            num_host_pages = cdiv(num_local_computed_tokens, self.block_size)
+            return self.coordinator.get_num_copy_blocks_to_pin(
+                new_computed_blocks[:num_host_pages]
+            )
         return 0
 
     def allocate_new_blocks(
diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py
index 056e19479a33..7fd068d5c819 100644
--- a/vllm/v1/hisparse/coordinator.py
+++ b/vllm/v1/hisparse/coordinator.py
@@ -235,6 +235,17 @@ def get_host_block_pool(self) -> BlockPool | None:
     # GPU copies of published host pages
     # ------------------------------------------------------------------
 
+    def get_num_copy_blocks_to_pin(self, host_blocks: Sequence[KVCacheBlock]) -> int:
+        """Count free GPU copies that adopting a host prefix will pin."""
+        blocks_to_pin: set[int] = set()
+        for host_block in host_blocks:
+            if host_block.is_null or host_block.block_hash is None:
+                continue
+            for block in self.copies.get(host_block.block_hash, ()):
+                if block.ref_cnt == 0 and not block.is_null:
+                    blocks_to_pin.add(block.block_id)
+        return len(blocks_to_pin)
+
     def _adopt_copies(
         self,
         request_id: str,