[Bugfix][KV Cache] Prevent negative external block allocation (#52707)
Signed-off-by: haic0 <149741444+haic0@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
haic0 committed
Aug 28, 2026 at 11:10 UTC
53bf99050213da20fe29bb6dbb49eed5fad813dc
2 files changed
+46
-2
tests/v1/core/test_single_type_kv_cache_manager.py
+43
@@ -14,11 +14,13 @@ from vllm.v1.core.kv_cache_utils import (
14
)
15
from vllm.v1.core.single_type_kv_cache_manager import (
16
ChunkedLocalAttentionManager,
17
+ FullAttentionManager,
18
RSWAManager,
19
SlidingWindowManager,
20
)
21
from vllm.v1.kv_cache_interface import (
22
ChunkedLocalAttentionSpec,
23
+ FullAttentionSpec,
24
RSWASpec,
25
SlidingWindowSpec,
26
)
@@ -26,6 +28,47 @@ from vllm.v1.kv_cache_interface import (
28
pytestmark = pytest.mark.cpu_test
29
30
31
+def test_external_computed_blocks_do_not_corrupt_free_pool():
32
+ block_size = 4
33
+ spec = FullAttentionSpec(
34
+ block_size=block_size,
35
+ num_kv_heads=1,
36
+ head_size=1,
37
+ dtype=torch.float32,
38
+ )
39
+ block_pool = BlockPool(
40
+ num_gpu_blocks=10,
41
+ enable_caching=False,
42
+ hash_block_size=block_size,
43
+ )
44
+ manager = FullAttentionManager(
45
+ spec,
46
+ block_pool=block_pool,
47
+ enable_caching=False,
48
+ kv_cache_group_id=0,
49
+ scheduler_block_size=block_size,
50
+ )
51
+ request_id = "request"
52
+ manager.allocate_new_blocks(
53
+ request_id,
54
+ num_tokens=3 * block_size,
55
+ num_tokens_main_model=3 * block_size,
56
+ )
57
+ num_free_blocks = block_pool.get_num_free_blocks()
58
+
59
+ # Speculative allocations can exceed the blocks implied by the eventual
60
+ # external computed-token count. This must not request a negative number
61
+ # of blocks, which inflates the free-queue counter.
62
+ manager.allocate_external_computed_blocks(
63
+ request_id,
64
+ num_local_computed_tokens=0,
65
+ num_external_computed_tokens=block_size,
66
+ )
67
+
68
+ assert block_pool.get_num_free_blocks() == num_free_blocks
69
+ assert len(manager.req_to_blocks[request_id]) == 3
70
+
71
+
72
def get_sliding_window_manager(
73
sliding_window_spec,
74
block_pool,
vllm/v1/core/single_type_kv_cache_manager.py
+3
-2
@@ -317,9 +317,10 @@ class SingleTypeKVCacheManager(ABC):
317
return
318
319
req_blocks = self.req_to_blocks[request_id]
320
- allocated_blocks = self.block_pool.get_new_blocks(
321
- cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
320
+ num_new_blocks = max(
321
+ 0, cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
322
)
323
+ allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks)
324
req_blocks.extend(allocated_blocks)
325
if self._record_new_block_ids:
326
self.new_block_ids.extend(b.block_id for b in allocated_blocks)