Skip to content
Merged
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
18 changes: 18 additions & 0 deletions tests/v1/core/test_kv_cache_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
hash_block_tokens,
init_none_hash,
make_block_hash_with_group_id,
null_head_len,
tensor_data,
to_event_extra_keys,
)
Expand Down Expand Up @@ -2156,3 +2157,20 @@ def test_wrapped_mamba_group_requires_block_zeroing():
)

assert config.needs_kv_cache_zeroing


def test_null_head_len():
null = KVCacheBlock(0, is_null=True)
real = [KVCacheBlock(i) for i in range(1, 4)]
tables = [
[],
real,
[null] * 5,
[null] * 5 + real,
[null] * 2 + real[:1] + [null] * 3 + real[1:], # not one contiguous head
]
for table in tables:
expected = next((i for i, b in enumerate(table) if not b.is_null), len(table))
assert null_head_len(table) == expected
blocks = KVCacheBlocks(([null] * 4 + real, [null] + real[:1] + [null]))
assert blocks.get_unhashed_block_ids_all_groups() == [[1, 2, 3], [1]]
32 changes: 32 additions & 0 deletions tests/v1/core/test_single_type_kv_cache_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -1012,3 +1012,35 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
f"but allocator pulled {len(new_blocks)}"
)
total_computed = num_tokens


def test_pop_blocks_for_free_trims_null_head():
block_size = 2
spec = SlidingWindowSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
sliding_window=4,
)
block_pool = BlockPool(
num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size
)
manager = get_sliding_window_manager(spec, block_pool)
null = block_pool.null_block
real = block_pool.get_new_blocks(2)
manager.req_to_blocks["swa"] = [null] * 3 + real
manager.req_to_blocks["dense"] = dense = block_pool.get_new_blocks(2)

assert manager.pop_blocks_for_free("swa") == real
assert manager.pop_blocks_for_free("dense") is dense # no null head, no copy
assert "swa" not in manager.req_to_blocks
assert all(b.ref_cnt == 1 for b in real + dense)

# The null block is never pooled, so touch / free_blocks leave it alone.
null_ref_cnt = null.ref_cnt
num_free_blocks = block_pool.get_num_free_blocks()
block_pool.touch([null])
block_pool.free_blocks([null])
assert null.ref_cnt == null_ref_cnt
assert block_pool.get_num_free_blocks() == num_free_blocks
8 changes: 6 additions & 2 deletions vllm/v1/core/block_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -759,9 +759,11 @@ def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
"""
for block in blocks:
if block.is_null:
continue
# ref_cnt=0 means this block is in the free list (i.e. eviction
# candidate), so remove it.
if block.ref_cnt == 0 and not block.is_null:
if block.ref_cnt == 0:
self.free_block_queue.remove(block)
block.ref_cnt += 1
if self.metrics_collector:
Expand All @@ -785,11 +787,13 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
blocks_to_evict_first = []
other_pools: dict[BlockPool, list[KVCacheBlock]] = {}
for block in ordered_blocks:
if block.is_null: # never pooled; its ref_cnt is not maintained
continue
if block.pool is not None and block.pool is not self:
other_pools.setdefault(block.pool, []).append(block)
continue
block.ref_cnt -= 1
if block.ref_cnt == 0 and not block.is_null:
if block.ref_cnt == 0:
if block.block_hash is None or not self.enable_caching:
# LIFO reuse of non-cached blocks for better GPU locality.
blocks_to_evict_first.append(block)
Expand Down
4 changes: 2 additions & 2 deletions vllm/v1/core/kv_cache_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
get_kv_cache_coordinator,
)
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy
from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy, trim_null_head
from vllm.v1.core.single_type_kv_cache_manager import MambaManager
from vllm.v1.kv_cache_interface import (
AttentionSpec,
Expand Down Expand Up @@ -117,7 +117,7 @@ def get_unhashed_block_ids_all_groups(self) -> list[list[int]]:
return [
[
block.block_id
for block in group
for block in trim_null_head(group)
if block.block_hash is None and not block.is_null
]
for group in self.blocks
Expand Down
36 changes: 35 additions & 1 deletion vllm/v1/core/kv_cache_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,15 @@
import os
from collections.abc import Callable, Iterable, Iterator, Sequence
from dataclasses import dataclass, field, replace
from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload
from typing import (
TYPE_CHECKING,
Any,
NamedTuple,
NewType,
TypeAlias,
TypeVar,
overload,
)

from vllm import envs
from vllm.config import VllmConfig
Expand Down Expand Up @@ -216,6 +224,32 @@ def __repr__(self) -> str:
)


def null_head_len(blocks: Sequence[KVCacheBlock]) -> int:
"""Number of leading null blocks, e.g. a sliding-window table's skipped prefix."""
if not blocks or not blocks[0].is_null:
return 0
lo, hi = 1, len(blocks)
while lo < hi:
mid = (lo + hi) // 2
if blocks[mid].is_null:
lo = mid + 1
else:
hi = mid
# All null entries are one shared object, so count() checks the run at C speed.
if blocks[:lo].count(blocks[0]) == lo:
return lo
return next(i for i, b in enumerate(blocks) if not b.is_null)


_BlockSeq = TypeVar("_BlockSeq", list[KVCacheBlock], Sequence[KVCacheBlock])


def trim_null_head(blocks: _BlockSeq) -> _BlockSeq:
"""``blocks`` without its leading null blocks; ``blocks`` itself if none."""
nhl = null_head_len(blocks)
return blocks[nhl:] if nhl else blocks


class KVCacheBlockCopy(NamedTuple):
src_block_id: int
dst_block_id: int
Expand Down
8 changes: 5 additions & 3 deletions vllm/v1/core/single_type_kv_cache_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
BlockHashWithGroupId,
KVCacheBlock,
resolve_block_hashes,
trim_null_head,
)
from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool
from vllm.v1.kv_cache_interface import (
Expand Down Expand Up @@ -569,14 +570,15 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
request_id: The request ID.
Returns:
The request's blocks in allocation order.
The request's blocks in allocation order, without the leading
null blocks (e.g. a sliding-window table's skipped prefix).
"""
# Default to [] in case a request is freed (aborted) before alloc.
req_blocks = self.req_to_blocks.pop(request_id, [])
self.num_cached_block.pop(request_id, None)
self._partial_hit_reqs.pop(request_id, None)
return req_blocks
return trim_null_head(req_blocks)

def free(self, request_id: str) -> None:
"""Free the blocks for the request.
Expand Down Expand Up @@ -2538,7 +2540,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
blocks = self.get_residency_row(request_id)
self.residency_changes.pop(request_id, None)
super().pop_blocks_for_free(request_id)
return blocks
return trim_null_head(blocks)

def adopt_resident_page(
self, request_id: str, block_idx: int, block: KVCacheBlock
Expand Down
Loading