A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Core] Reuse empty block lists whenever possible in KVCacheBlocks to mitigate GC costs (#24964)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Oct 14, 2025 at 12:58 UTC
acaa2c0a4a53dbb57f85f1042b1a6f1e3f24cef5
5 files changed
+53
-26
vllm/v1/core/block_pool.py
+2
-2
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable
3
+from collections.abc import Iterable, Sequence
4
from typing import Any
5
6
from vllm.distributed.kv_events import (
@@ -328,7 +328,7 @@ class BlockPool:
328
)
329
return True
330
331
- def touch(self, blocks: tuple[list[KVCacheBlock], ...]) -> None:
331
+ def touch(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None:
332
"""Touch a block increases its reference count by 1, and may remove
333
the block from the free queue. This is used when a block is hit by
334
another request with the same prefix.
vllm/v1/core/kv_cache_coordinator.py
+3
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC, abstractmethod
4
+from collections.abc import Sequence
5
6
from vllm.v1.core.block_pool import BlockPool
7
from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
@@ -51,7 +52,7 @@ class KVCacheCoordinator(ABC):
52
self,
53
request_id: str,
54
num_tokens: int,
54
- new_computed_blocks: tuple[list[KVCacheBlock], ...],
55
+ new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
56
num_encoder_tokens: int,
57
) -> int:
58
"""
@@ -84,7 +85,7 @@ class KVCacheCoordinator(ABC):
85
return num_blocks_to_allocate
86
87
def save_new_computed_blocks(
87
- self, request_id: str, new_computed_blocks: tuple[list[KVCacheBlock], ...]
88
+ self, request_id: str, new_computed_blocks: tuple[Sequence[KVCacheBlock], ...]
89
) -> None:
90
"""
91
Add the new computed blocks to the request.
vllm/v1/core/kv_cache_manager.py
+36
-15
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import itertools
5
+from collections.abc import Sequence
6
from dataclasses import dataclass
7
from typing import Literal, overload
8
@@ -23,7 +25,7 @@ class KVCacheBlocks:
25
structure from the Scheduler.
26
"""
27
26
- blocks: tuple[list[KVCacheBlock], ...]
28
+ blocks: tuple[Sequence[KVCacheBlock], ...]
29
"""
30
`blocks[i][j]` refers to the i-th kv_cache_group
31
and the j-th block of tokens.We don't use block of
@@ -31,12 +33,20 @@ class KVCacheBlocks:
33
kv_cache_groups have the same number of blocks, which is true for now but
34
will be broken if we want to give different block_size to different
35
kv_cache_groups in the future.
36
+
37
+ Each single type KVCacheBlocks could be represented as:
38
+ - list[KVCacheBlock] for more than one KVCacheBlock
39
+ - an empty tuple for requests without KVCacheBlock
40
+ (a precomputed KVCacheBlocks is in KVCacheManager to avoid GC overhead)
41
"""
42
43
def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
44
"""Adds two KVCacheBlocks instances."""
45
return KVCacheBlocks(
39
- tuple(blk1 + blk2 for blk1, blk2 in zip(self.blocks, other.blocks))
46
+ tuple(
47
+ list(itertools.chain(blk1, blk2))
48
+ for blk1, blk2 in zip(self.blocks, other.blocks)
49
+ )
50
)
51
52
@overload
@@ -74,8 +84,10 @@ class KVCacheBlocks:
84
return [block.block_id for block in self.blocks[0] if block.block_hash is None]
85
86
def new_empty(self) -> "KVCacheBlocks":
77
- """Creates a new KVCacheBlocks instance with no blocks."""
78
- return KVCacheBlocks(tuple([] for _ in range(len(self.blocks))))
87
+ """
88
+ Creates a new KVCacheBlocks instance with no blocks.
89
+ """
90
+ return KVCacheBlocks(tuple(() for _ in range(len(self.blocks))))
91
92
93
class KVCacheManager:
@@ -131,6 +143,15 @@ class KVCacheManager:
143
self.block_pool = self.coordinator.block_pool
144
self.kv_cache_config = kv_cache_config
145
146
+ # Pre-constructed KVCacheBlocks with no blocks, callers should use this
147
+ # via create_kv_cache_blocks instead of creating new ones to avoid GC
148
+ # overhead.
149
+ #
150
+ # We use nested tuples to ensure the empty KVCacheBlocks is immutable.
151
+ self.empty_kv_cache_blocks = KVCacheBlocks(
152
+ tuple(() for _ in range(self.num_kv_cache_groups))
153
+ )
154
+
155
@property
156
def usage(self) -> float:
157
"""Get the KV cache usage.
@@ -170,7 +191,7 @@ class KVCacheManager:
191
request.sampling_params is not None
192
and request.sampling_params.prompt_logprobs is not None
193
):
173
- return self.create_empty_block_list(), 0
194
+ return self.empty_kv_cache_blocks, 0
195
196
# NOTE: When all tokens hit the cache, we must recompute the last token
197
# to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
@@ -198,7 +219,7 @@ class KVCacheManager:
219
self.prefix_cache_stats.queries += request.num_tokens
220
self.prefix_cache_stats.hits += num_new_computed_tokens
221
201
- return KVCacheBlocks(computed_blocks), num_new_computed_tokens
222
+ return (self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens)
223
224
def allocate_slots(
225
self,
@@ -251,9 +272,7 @@ class KVCacheManager:
272
if new_computed_blocks is not None:
273
new_computed_block_list = new_computed_blocks.blocks
274
else:
254
- new_computed_block_list = tuple(
255
- [] for _ in range(len(self.kv_cache_config.kv_cache_groups))
256
- )
275
+ new_computed_block_list = self.empty_kv_cache_blocks.blocks
276
277
# Free the blocks that are skipped during the attention computation
278
# (e.g., tokens outside the sliding window).
@@ -305,7 +324,7 @@ class KVCacheManager:
324
# P/D: delay caching blocks if we have to recv from
325
# remote. Update state for locally cached blocks.
326
if not self.enable_caching or delay_cache_blocks:
308
- return KVCacheBlocks(new_blocks)
327
+ return self.create_kv_cache_blocks(new_blocks)
328
329
# NOTE(woosuk): We want to commit (cache) up to num_computed_tokens +
330
# num_new_tokens, but must exclude "non-committable" tokens (e.g.,
@@ -316,7 +335,7 @@ class KVCacheManager:
335
)
336
self.coordinator.cache_blocks(request, num_tokens_to_cache)
337
319
- return KVCacheBlocks(new_blocks)
338
+ return self.create_kv_cache_blocks(new_blocks)
339
340
def free(self, request: Request) -> None:
341
"""Free the blocks allocated for the request.
@@ -388,7 +407,7 @@ class KVCacheManager:
407
408
def get_blocks(self, request_id: str) -> KVCacheBlocks:
409
"""Get the blocks of a request."""
391
- return KVCacheBlocks(self.coordinator.get_blocks(request_id))
410
+ return self.create_kv_cache_blocks(self.coordinator.get_blocks(request_id))
411
412
def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:
413
"""Get the block ids of a request."""
@@ -399,6 +418,8 @@ class KVCacheManager:
418
if self.enable_caching:
419
self.coordinator.cache_blocks(request, num_computed_tokens)
420
402
- def create_empty_block_list(self) -> KVCacheBlocks:
403
- """Creates a new KVCacheBlocks instance with no blocks."""
404
- return KVCacheBlocks(tuple([] for _ in range(self.num_kv_cache_groups)))
421
+ def create_kv_cache_blocks(
422
+ self, blocks: tuple[list[KVCacheBlock], ...]
423
+ ) -> KVCacheBlocks:
424
+ # Only create new KVCacheBlocks for non-empty blocks
425
+ return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
vllm/v1/core/sched/scheduler.py
+1
-3
@@ -421,9 +421,7 @@ class Scheduler(SchedulerInterface):
421
# KVTransfer: WAITING reqs have num_computed_tokens > 0
422
# after async KV recvs are completed.
423
else:
424
- new_computed_blocks = (
425
- self.kv_cache_manager.create_empty_block_list()
426
- )
424
+ new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks
425
num_new_local_computed_tokens = 0
426
num_computed_tokens = request.num_computed_tokens
427
vllm/v1/core/single_type_kv_cache_manager.py
+11
-4
@@ -3,6 +3,7 @@
3
import itertools
4
from abc import ABC, abstractmethod
5
from collections import defaultdict
6
+from collections.abc import Sequence
7
8
from vllm.utils import cdiv
9
from vllm.v1.core.block_pool import BlockPool
@@ -61,7 +62,10 @@ class SingleTypeKVCacheManager(ABC):
62
self._null_block = block_pool.null_block
63
64
def get_num_blocks_to_allocate(
64
- self, request_id: str, num_tokens: int, new_computed_blocks: list[KVCacheBlock]
65
+ self,
66
+ request_id: str,
67
+ num_tokens: int,
68
+ new_computed_blocks: Sequence[KVCacheBlock],
69
) -> int:
70
"""
71
Get the number of blocks needed to be allocated for the request.
@@ -93,7 +97,7 @@ class SingleTypeKVCacheManager(ABC):
97
return num_new_blocks + num_evictable_computed_blocks
98
99
def save_new_computed_blocks(
96
- self, request_id: str, new_computed_blocks: list[KVCacheBlock]
100
+ self, request_id: str, new_computed_blocks: Sequence[KVCacheBlock]
101
) -> None:
102
"""
103
Add the new computed blocks to the request.
@@ -593,7 +597,10 @@ class MambaManager(SingleTypeKVCacheManager):
597
return 0
598
599
def get_num_blocks_to_allocate(
596
- self, request_id: str, num_tokens: int, new_computed_blocks: list[KVCacheBlock]
600
+ self,
601
+ request_id: str,
602
+ num_tokens: int,
603
+ new_computed_blocks: Sequence[KVCacheBlock],
604
) -> int:
605
# Allocate extra `num_speculative_blocks` blocks for
606
# speculative decoding (MTP/EAGLE) with linear attention.
@@ -625,7 +632,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
632
"""Manager for cross-attention KV cache in encoder-decoder models."""
633
634
def save_new_computed_blocks(
628
- self, request_id: str, new_computed_blocks: list[KVCacheBlock]
635
+ self, request_id: str, new_computed_blocks: Sequence[KVCacheBlock]
636
) -> None:
637
# We do not cache blocks for cross-attention to be shared between
638
# requests, so `new_computed_blocks` should always be empty.