Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 45 additions & 0 deletions tests/v1/core/test_prefix_caching.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions vllm/v1/core/single_type_kv_cache_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
11 changes: 11 additions & 0 deletions vllm/v1/hisparse/coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading