diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1489315ec2..6a0b6393a2 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import ( hash_block_tokens, init_none_hash, make_block_hash_with_group_id, + null_head_len, tensor_data, to_event_extra_keys, ) @@ -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]] diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index cd4581dcef..5f9f1abf4c 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -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 diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 3b0102fe84..41839bd1f7 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -759,9 +759,11 @@ class BlockPool: """ 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: @@ -785,11 +787,13 @@ class BlockPool: 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) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 4f472170a8..0ac1c38818 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -14,7 +14,7 @@ from vllm.v1.core.kv_cache_coordinator import ( 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, @@ -117,7 +117,7 @@ class KVCacheBlocks: 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 diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 5c51a8d131..aadc9cda57 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -8,7 +8,15 @@ import math 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 @@ -216,6 +224,32 @@ class KVCacheBlock: ) +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 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 28a7a1106f..a583abb61e 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -16,6 +16,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHashWithGroupId, KVCacheBlock, resolve_block_hashes, + trim_null_head, ) from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool from vllm.v1.kv_cache_interface import ( @@ -569,14 +570,15 @@ class SingleTypeKVCacheManager(ABC): 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. @@ -2538,7 +2540,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): 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