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 10, 22:45 UTC
[Core][Hybrid allocator + connector] Support hybrid allocator + kv cache connector (#30166)
Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu> Co-authored-by: KuntaiDu <kuntai@uchicago.edu>
Yifan Qiao committed
Dec 26, 2025 at 18:25 UTC
52bf0665168c539d2d061a664ad62b18a12e80bb
8 files changed
+360
-109
tests/v1/core/test_single_type_kv_cache_manager.py
+62
-8
@@ -21,13 +21,23 @@ from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowS
21
pytestmark = pytest.mark.cpu_test
22
23
24
-def get_sliding_window_manager(sliding_window_spec, block_pool):
25
- return SlidingWindowManager(sliding_window_spec, block_pool, kv_cache_group_id=0)
24
+def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
25
+ return SlidingWindowManager(
26
+ sliding_window_spec,
27
+ block_pool,
28
+ enable_caching=enable_caching,
29
+ kv_cache_group_id=0,
30
+ )
31
32
28
-def get_chunked_local_attention_manager(chunked_local_attention_spec, block_pool):
33
+def get_chunked_local_attention_manager(
34
+ chunked_local_attention_spec, block_pool, enable_caching=True
35
+):
36
return ChunkedLocalAttentionManager(
30
- chunked_local_attention_spec, block_pool, kv_cache_group_id=0
37
+ chunked_local_attention_spec,
38
+ block_pool,
39
+ enable_caching=enable_caching,
40
+ kv_cache_group_id=0,
41
)
42
43
@@ -332,11 +342,53 @@ def test_get_num_blocks_to_allocate():
342
]
343
344
assert (
335
- manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1) == 20
345
+ manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1, 0)
346
+ == 20
347
)
348
assert (
338
- manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2) == 15
349
+ manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2, 0)
350
+ == 15
351
+ )
352
+
353
+
354
+def test_evictable_cached_blocks_not_double_allocated():
355
+ block_size = 2
356
+ sliding_window_length = 2 * block_size
357
+ sliding_window_spec = SlidingWindowSpec(
358
+ block_size=block_size,
359
+ num_kv_heads=1,
360
+ head_size=1,
361
+ dtype=torch.float32,
362
+ sliding_window=sliding_window_length,
363
+ )
364
+
365
+ block_pool = BlockPool(
366
+ num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
367
+ )
368
+ manager = get_sliding_window_manager(sliding_window_spec, block_pool)
369
+
370
+ request_id = "req"
371
+ evictable_block = block_pool.blocks[1] # ref_cnt == 0, eviction candidate
372
+
373
+ num_blocks_to_allocate = manager.get_num_blocks_to_allocate(
374
+ request_id=request_id,
375
+ num_tokens=2 * block_size,
376
+ new_computed_blocks=[evictable_block],
377
+ total_computed_tokens=block_size,
378
+ )
379
+ # Free capacity check should count evictable cached blocks, but allocation
380
+ # should only allocate the truly new block.
381
+ assert num_blocks_to_allocate == 2
382
+
383
+ manager.allocate_new_computed_blocks(
384
+ request_id,
385
+ [evictable_block],
386
+ num_local_computed_tokens=block_size,
387
+ num_external_computed_tokens=0,
388
)
389
+ new_blocks = manager.allocate_new_blocks(request_id, num_tokens=4)
390
+ assert len(new_blocks) == 1
391
+ assert len(manager.req_to_blocks[request_id]) == 2
392
393
394
def test_chunked_local_attention_get_num_blocks_to_allocate():
@@ -359,8 +411,10 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
411
]
412
413
assert (
362
- manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1) == 20
414
+ manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1, 0)
415
+ == 20
416
)
417
assert (
365
- manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2) == 15
418
+ manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2, 0)
419
+ == 15
420
)
vllm/v1/core/block_pool.py
+13
-10
@@ -254,6 +254,10 @@ class BlockPool:
254
[] if self.enable_kv_cache_events else None
255
)
256
for i, blk in enumerate(new_full_blocks):
257
+ # Some blocks may be null blocks when enabling sparse attention like
258
+ # sliding window attention. We skip null blocks here.
259
+ if blk.is_null:
260
+ continue
261
assert blk.block_hash is None
262
block_hash = new_block_hashes[i]
263
@@ -361,7 +365,7 @@ class BlockPool:
365
)
366
return True
367
364
- def touch(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None:
368
+ def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
369
"""Touch a block increases its reference count by 1, and may remove
370
the block from the free queue. This is used when a block is hit by
371
another request with the same prefix.
@@ -369,15 +373,14 @@ class BlockPool:
373
Args:
374
blocks: A list of blocks to touch.
375
"""
372
- for blocks_per_group in blocks:
373
- for block in blocks_per_group:
374
- # ref_cnt=0 means this block is in the free list (i.e. eviction
375
- # candidate), so remove it.
376
- if block.ref_cnt == 0 and not block.is_null:
377
- self.free_block_queue.remove(block)
378
- block.ref_cnt += 1
379
- if self.metrics_collector:
380
- self.metrics_collector.on_block_accessed(block)
376
+ for block in blocks:
377
+ # ref_cnt=0 means this block is in the free list (i.e. eviction
378
+ # candidate), so remove it.
379
+ if block.ref_cnt == 0 and not block.is_null:
380
+ self.free_block_queue.remove(block)
381
+ block.ref_cnt += 1
382
+ if self.metrics_collector:
383
+ self.metrics_collector.on_block_accessed(block)
384
385
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
386
"""Free a list of blocks. The blocks should be ordered by their
vllm/v1/core/kv_cache_coordinator.py
+35
-11
@@ -60,6 +60,7 @@ class KVCacheCoordinator(ABC):
60
get_manager_for_kv_cache_spec(
61
kv_cache_spec=kv_cache_group.kv_cache_spec,
62
block_pool=self.block_pool,
63
+ enable_caching=enable_caching,
64
kv_cache_group_id=i,
65
dcp_world_size=dcp_world_size,
66
pcp_world_size=pcp_world_size,
@@ -73,6 +74,7 @@ class KVCacheCoordinator(ABC):
74
num_tokens: int,
75
new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
76
num_encoder_tokens: int,
77
+ total_computed_tokens: int,
78
) -> int:
79
"""
80
Get the number of blocks needed to be allocated for the request.
@@ -85,9 +87,10 @@ class KVCacheCoordinator(ABC):
87
prefix caching.
88
num_encoder_tokens: The number of encoder tokens for allocating
89
blocks for cross-attention.
90
+ total_computed_tokens: Include both local and external tokens.
91
92
Returns:
90
- The number of blocks.
93
+ The number of blocks to allocate.
94
"""
95
num_blocks_to_allocate = 0
96
for i, manager in enumerate(self.single_type_managers):
@@ -95,30 +98,48 @@ class KVCacheCoordinator(ABC):
98
# For cross-attention, we issue a single static allocation
99
# of blocks based on the number of encoder input tokens.
100
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
98
- request_id, num_encoder_tokens, []
101
+ request_id, num_encoder_tokens, [], 0
102
)
103
else:
104
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
102
- request_id, num_tokens, new_computed_blocks[i]
105
+ request_id,
106
+ num_tokens,
107
+ new_computed_blocks[i],
108
+ total_computed_tokens,
109
)
110
return num_blocks_to_allocate
111
106
- def save_new_computed_blocks(
107
- self, request_id: str, new_computed_blocks: tuple[Sequence[KVCacheBlock], ...]
112
+ def allocate_new_computed_blocks(
113
+ self,
114
+ request_id: str,
115
+ new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
116
+ num_local_computed_tokens: int,
117
+ num_external_computed_tokens: int,
118
) -> None:
119
"""
110
- Add the new computed blocks to the request.
120
+ Add the new computed blocks to the request. Optionally allocate new
121
+ blocks for external computed tokens (if any).
122
123
Args:
124
request_id: The request ID.
125
new_computed_blocks: The new computed blocks just hitting the
126
prefix cache.
127
+ num_local_computed_tokens: The number of local computed tokens.
128
+ num_external_computed_tokens: The number of external computed tokens.
129
"""
130
for i, manager in enumerate(self.single_type_managers):
118
- manager.save_new_computed_blocks(request_id, new_computed_blocks[i])
131
+ manager.allocate_new_computed_blocks(
132
+ request_id,
133
+ new_computed_blocks[i],
134
+ num_local_computed_tokens,
135
+ num_external_computed_tokens,
136
+ )
137
138
def allocate_new_blocks(
121
- self, request_id: str, num_tokens: int, num_encoder_tokens: int = 0
139
+ self,
140
+ request_id: str,
141
+ num_tokens: int,
142
+ num_encoder_tokens: int = 0,
143
) -> tuple[list[KVCacheBlock], ...]:
144
"""
145
Allocate new blocks for the request to give it at least `num_tokens`
@@ -184,17 +205,20 @@ class KVCacheCoordinator(ABC):
205
for manager in self.single_type_managers
206
]
207
187
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
208
+ def remove_skipped_blocks(
209
+ self, request_id: str, total_computed_tokens: int
210
+ ) -> None:
211
"""
212
Remove the blocks that are no longer needed from `blocks` and replace
213
the removed blocks with null_block.
214
215
Args:
216
request_id: The request ID.
194
- num_computed_tokens: The number of tokens that have been computed.
217
+ total_computed_tokens: The total number of computed tokens, including
218
+ local computed tokens and external computed tokens.
219
"""
220
for manager in self.single_type_managers:
197
- manager.remove_skipped_blocks(request_id, num_computed_tokens)
221
+ manager.remove_skipped_blocks(request_id, total_computed_tokens)
222
223
def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
224
"""
vllm/v1/core/kv_cache_manager.py
+108
-42
@@ -210,6 +210,7 @@ class KVCacheManager:
210
num_new_computed_tokens: int = 0,
211
new_computed_blocks: KVCacheBlocks | None = None,
212
num_lookahead_tokens: int = 0,
213
+ num_external_computed_tokens: int = 0,
214
delay_cache_blocks: bool = False,
215
num_encoder_tokens: int = 0,
216
) -> KVCacheBlocks | None:
@@ -217,16 +218,16 @@ class KVCacheManager:
218
219
Args:
220
request: The request to allocate slots.
220
- num_new_tokens: The number of tokens to allocate, including external
221
- tokens. Note that this does not include tokens that have
222
- already been computed locally (i.e. new_computed_blocks).
221
+ num_new_tokens: The number of new tokens to be allocated and computed.
222
num_new_computed_tokens: The number of new computed tokens just
223
hitting the prefix caching, excluding external tokens.
224
new_computed_blocks: The cached blocks for the above new computed
226
- tokens.
225
+ tokens, grouped as a tuple by kv cache groups.
226
num_lookahead_tokens: The number of speculative tokens to allocate.
227
This is used by spec decode proposers with kv-cache such
228
as eagle.
229
+ num_external_computed_tokens: The number of tokens that their
230
+ KV caches are not cached by vLLM but cached by the connector.
231
delay_cache_blocks: Whether to skip caching the blocks. This is
232
used by P/D when allocating blocks used in a KV transfer
233
which will complete in a future step.
@@ -236,29 +237,81 @@ class KVCacheManager:
237
238
Blocks layout:
239
```
239
- -----------------------------------------------------------------------
240
- | < computed > | < new computed > | < new > | < pre-allocated > |
241
- -----------------------------------------------------------------------
242
- | < required > |
243
- --------------------------------------------------
244
- | < full > |
245
- ------------------------------------------------
246
- | <new full> |
247
- --------------
240
+ ----------------------------------------------------------------------
241
+ | < comp > | < new_comp > | < ext_comp > | < new > | < lookahead > |
242
+ ----------------------------------------------------------------------
243
+ | < to be computed > |
244
+ ----------------------------------------------------------------------
245
+ | < to be allocated > |
246
+ ----------------------------------------------------------------------
247
+ | < to be cached (roughly, |
248
+ | details below)> |
249
+ ----------------------------------------------------------------------
250
+ | Prefix-cached tokens from either vLLM |
251
+ | or connector. Can be safely removed if |
252
+ | they are outside sliding window. |
253
+ ----------------------------------------------------------------------
254
+ | < cached by vLLM > | not cached by |
255
+ | vLLM, but |
256
+ | ref_cnt | ref_cnt not | cached by |
257
+ | increased| increased yet| connector |
258
+ ----------------------------------------------------------------------
259
```
249
- The following *_blocks are illustrated in this layout.
260
+
261
+ Abbrivations:
262
+
263
+ ```
264
+ comp = request.num_computed_tokens
265
+ new_comp = num_new_computed_tokens
266
+ = len(new_computed_blocks) * block_size
267
+ ext_comp = num_external_computed_tokens, cached by the connector
268
+ new = num_new_tokens, including unverified draft tokens
269
+ lookahead = num_lookahead_tokens
270
+ ```
271
+
272
+ NOTE: for new tokens which include both verified and unverified draft
273
+ tokens, we only cache the verified tokens (by capping the number at
274
+ `request.num_tokens`).
275
+
276
+ The allocation has three stages:
277
+ - Free unnecessary blocks in `comp` and check
278
+ if we have sufficient free blocks (return None if not).
279
+ - Handle prefix tokens (`comp + new_comp + ext_comp`):
280
+ - Free unnecessary blocks (e.g. outside sliding window)
281
+ - Allocate new blocks for `ext_comp` tokens inside
282
+ sliding window
283
+ - Allocate new blocks for tokens to be computed (`new + lookahead`)
284
285
Returns:
286
A list of new allocated blocks.
287
"""
254
- if num_new_tokens == 0:
255
- raise ValueError("num_new_tokens must be greater than 0")
288
+ # When loading KV data asynchronously, we may have zero new tokens to
289
+ # compute while still allocating slots for externally computed tokens.
290
+ if num_new_tokens == 0 and num_external_computed_tokens == 0:
291
+ raise ValueError(
292
+ "num_new_tokens must be greater than 0 when there are no "
293
+ "external computed tokens"
294
+ )
295
296
if new_computed_blocks is not None:
297
new_computed_block_list = new_computed_blocks.blocks
298
else:
299
new_computed_block_list = self.empty_kv_cache_blocks.blocks
300
301
+ # The number of computed tokens is the number of computed tokens plus
302
+ # the new prefix caching hits
303
+ num_local_computed_tokens = (
304
+ request.num_computed_tokens + num_new_computed_tokens
305
+ )
306
+ total_computed_tokens = min(
307
+ num_local_computed_tokens + num_external_computed_tokens,
308
+ self.max_model_len,
309
+ )
310
+ num_tokens_need_slot = min(
311
+ total_computed_tokens + num_new_tokens + num_lookahead_tokens,
312
+ self.max_model_len,
313
+ )
314
+
315
# Free the blocks that are skipped during the attention computation
316
# (e.g., tokens outside the sliding window).
317
# We can do this even if we cannot schedule this request due to
@@ -266,15 +319,7 @@ class KVCacheManager:
319
# Should call this function before allocating new blocks to reduce
320
# the number of evicted blocks.
321
self.coordinator.remove_skipped_blocks(
269
- request.request_id, request.num_computed_tokens
270
- )
271
-
272
- # The number of computed tokens is the number of computed tokens plus
273
- # the new prefix caching hits
274
- num_computed_tokens = request.num_computed_tokens + num_new_computed_tokens
275
- num_tokens_need_slot = min(
276
- num_computed_tokens + num_new_tokens + num_lookahead_tokens,
277
- self.max_model_len,
322
+ request.request_id, total_computed_tokens
323
)
324
325
num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
@@ -282,25 +327,25 @@ class KVCacheManager:
327
num_tokens=num_tokens_need_slot,
328
new_computed_blocks=new_computed_block_list,
329
num_encoder_tokens=num_encoder_tokens,
330
+ total_computed_tokens=num_local_computed_tokens
331
+ + num_external_computed_tokens,
332
)
333
334
if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
335
# Cannot allocate new blocks
336
return None
337
291
- # Touch the computed blocks to make sure they won't be evicted.
292
- if self.enable_caching:
293
- self.block_pool.touch(new_computed_block_list)
294
- else:
295
- assert not any(new_computed_block_list), (
296
- "Computed blocks should be empty when prefix caching is disabled"
297
- )
298
-
299
- if new_computed_block_list is not self.empty_kv_cache_blocks.blocks:
338
+ if (
339
+ new_computed_block_list is not self.empty_kv_cache_blocks.blocks
340
+ or num_external_computed_tokens > 0
341
+ ):
342
# Append the new computed blocks to the request blocks until now to
343
# avoid the case where the new blocks cannot be allocated.
302
- self.coordinator.save_new_computed_blocks(
303
- request.request_id, new_computed_block_list
344
+ self.coordinator.allocate_new_computed_blocks(
345
+ request_id=request.request_id,
346
+ new_computed_blocks=new_computed_block_list,
347
+ num_local_computed_tokens=num_local_computed_tokens,
348
+ num_external_computed_tokens=num_external_computed_tokens,
349
)
350
351
new_blocks = self.coordinator.allocate_new_blocks(
@@ -312,12 +357,14 @@ class KVCacheManager:
357
if not self.enable_caching or delay_cache_blocks:
358
return self.create_kv_cache_blocks(new_blocks)
359
315
- # NOTE(woosuk): We want to commit (cache) up to num_computed_tokens +
316
- # num_new_tokens, but must exclude "non-committable" tokens (e.g.,
317
- # draft tokens that could be rejected). Therefore, we cap the number
318
- # at `request.num_tokens`, ensuring only "finalized" tokens are cached.
360
+ # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
361
+ # + num_external_computed_tokens + num_new_tokens, but must exclude
362
+ # "non-committable" tokens (e.g., draft tokens that could be rejected).
363
+ # Therefore, we cap the number at `request.num_tokens`, ensuring only
364
+ # "finalized" tokens are cached.
365
num_tokens_to_cache = min(
320
- num_computed_tokens + num_new_tokens, request.num_tokens
366
+ total_computed_tokens + num_new_tokens,
367
+ request.num_tokens,
368
)
369
self.coordinator.cache_blocks(request, num_tokens_to_cache)
370
@@ -333,6 +380,19 @@ class KVCacheManager:
380
"""
381
self.coordinator.free(request.request_id)
382
383
+ def remove_skipped_blocks(
384
+ self, request_id: str, total_computed_tokens: int
385
+ ) -> None:
386
+ """Remove the blocks that are no longer needed from `blocks` and replace
387
+ the removed blocks with null_block.
388
+
389
+ Args:
390
+ request_id: The request ID.
391
+ total_computed_tokens: The total number of computed tokens, including
392
+ local computed tokens and external computed tokens.
393
+ """
394
+ self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens)
395
+
396
def evict_blocks(self, block_ids: set[int]) -> None:
397
"""evict blocks from the prefix cache by their block IDs.
398
@@ -408,7 +468,13 @@ class KVCacheManager:
468
return self.get_blocks(request_id).get_block_ids()
469
470
def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
411
- """Cache the blocks for the request, if enabled."""
471
+ """Cache the blocks for the request, if enabled.
472
+
473
+ Args:
474
+ request: The request to cache the blocks.
475
+ num_computed_tokens: The number of computed tokens, including tokens
476
+ that are already cached and tokens to be cached.
477
+ """
478
if self.enable_caching:
479
self.coordinator.cache_blocks(request, num_computed_tokens)
480
vllm/v1/core/sched/scheduler.py
+12
-4
@@ -587,10 +587,11 @@ class Scheduler(SchedulerInterface):
587
588
new_blocks = self.kv_cache_manager.allocate_slots(
589
request,
590
- num_new_tokens + num_external_computed_tokens,
591
- num_new_local_computed_tokens,
592
- new_computed_blocks,
590
+ num_new_tokens,
591
+ num_new_computed_tokens=num_new_local_computed_tokens,
592
+ new_computed_blocks=new_computed_blocks,
593
num_lookahead_tokens=effective_lookahead_tokens,
594
+ num_external_computed_tokens=num_external_computed_tokens,
595
delay_cache_blocks=load_kv_async,
596
num_encoder_tokens=num_encoder_tokens,
597
)
@@ -606,7 +607,7 @@ class Scheduler(SchedulerInterface):
607
if self.connector is not None:
608
self.connector.update_state_after_alloc(
609
request,
609
- new_computed_blocks + new_blocks,
610
+ self.kv_cache_manager.get_blocks(request.request_id),
611
num_external_computed_tokens,
612
)
613
@@ -1580,6 +1581,13 @@ class Scheduler(SchedulerInterface):
1581
if self.connector is None:
1582
return False, None
1583
1584
+ # Free any out-of-window prefix blocks before we hand the block table to
1585
+ # the connector.
1586
+ self.kv_cache_manager.remove_skipped_blocks(
1587
+ request_id=request.request_id,
1588
+ total_computed_tokens=request.num_tokens,
1589
+ )
1590
+
1591
block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
1592
1593
if not isinstance(self.connector, SupportsHMA):
vllm/v1/core/single_type_kv_cache_manager.py
+122
-31
@@ -30,6 +30,7 @@ class SingleTypeKVCacheManager(ABC):
30
self,
31
kv_cache_spec: KVCacheSpec,
32
block_pool: BlockPool,
33
+ enable_caching: bool,
34
kv_cache_group_id: int,
35
dcp_world_size: int = 1,
36
pcp_world_size: int = 1,
@@ -48,6 +49,7 @@ class SingleTypeKVCacheManager(ABC):
49
self.block_size *= dcp_world_size * pcp_world_size
50
self.kv_cache_spec = kv_cache_spec
51
self.block_pool = block_pool
52
+ self.enable_caching = enable_caching
53
54
# Mapping from request ID to blocks to track the blocks allocated
55
# for each request, so that we can free the blocks when the request
@@ -68,6 +70,7 @@ class SingleTypeKVCacheManager(ABC):
70
request_id: str,
71
num_tokens: int,
72
new_computed_blocks: Sequence[KVCacheBlock],
73
+ total_computed_tokens: int,
74
) -> int:
75
"""
76
Get the number of blocks needed to be allocated for the request.
@@ -78,46 +81,121 @@ class SingleTypeKVCacheManager(ABC):
81
tokens that are already allocated).
82
new_computed_blocks: The new computed blocks just hitting the
83
prefix caching.
84
+ total_computed_tokens: Include both local and external computed
85
+ tokens.
86
87
Returns:
83
- The number of blocks.
88
+ The number of blocks to allocate.
89
"""
90
91
num_required_blocks = cdiv(num_tokens, self.block_size)
87
- num_new_blocks = (
88
- num_required_blocks
89
- - len(new_computed_blocks)
90
- - len(self.req_to_blocks[request_id])
92
+ num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
93
+
94
+ if request_id in self.num_cached_block:
95
+ # Fast-path: a running request won't have any new prefix-cache hits.
96
+ assert len(new_computed_blocks) == 0
97
+ # NOTE: With speculative decoding, request's blocks may be allocated
98
+ # for draft tokens which are later rejected. In this case,
99
+ # num_required_blocks may be smaller than num_req_blocks.
100
+ return max(num_required_blocks - num_req_blocks, 0)
101
+
102
+ num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
103
+ num_local_computed_blocks = len(new_computed_blocks) + num_req_blocks
104
+ # Number of whole blocks that are skipped by the attention window.
105
+ # If nothing is skipped, this is 0.
106
+ num_skipped_blocks = num_skipped_tokens // self.block_size
107
+ # We need blocks for the non-skipped suffix. If there are still
108
+ # local-computed blocks inside the window, they contribute to the
109
+ # required capacity; otherwise, skipped blocks dominate.
110
+ num_new_blocks = max(
111
+ num_required_blocks - max(num_skipped_blocks, num_local_computed_blocks),
112
+ 0,
113
)
92
- # If a computed block of a request is an eviction candidate (in the
93
- # free queue and ref_cnt == 0), it will be changed from a free block
94
- # to a computed block when the request is allocated, so we also count
95
- # it as needed to be allocated.
96
- num_evictable_computed_blocks = sum(
97
- blk.ref_cnt == 0 and not blk.is_null for blk in new_computed_blocks
114
+
115
+ # Among the `new_computed_blocks`, the first `num_skipped_blocks` worth
116
+ # of blocks are skipped; `num_req_blocks` of those may already be in
117
+ # `req_to_blocks`, so only skip the remainder from `new_computed_blocks`.
118
+ num_skipped_new_computed_blocks = max(0, num_skipped_blocks - num_req_blocks)
119
+
120
+ # If a computed block is an eviction candidate (in the free queue and
121
+ # ref_cnt == 0), it will be removed from the free queue when touched by
122
+ # the allocated request, so we must count it in the free-capacity check.
123
+ num_evictable_blocks = sum(
124
+ blk.ref_cnt == 0 and not blk.is_null
125
+ for blk in new_computed_blocks[num_skipped_new_computed_blocks:]
126
)
99
- return num_new_blocks + num_evictable_computed_blocks
127
+ return num_new_blocks + num_evictable_blocks
128
101
- def save_new_computed_blocks(
102
- self, request_id: str, new_computed_blocks: Sequence[KVCacheBlock]
129
+ def allocate_new_computed_blocks(
130
+ self,
131
+ request_id: str,
132
+ new_computed_blocks: Sequence[KVCacheBlock],
133
+ num_local_computed_tokens: int,
134
+ num_external_computed_tokens: int,
135
) -> None:
136
"""
105
- Add the new computed blocks to the request.
137
+ Add the new computed blocks to the request. This involves three steps:
138
+ 1. Touch the computed blocks to make sure they won't be evicted.
139
+ 1.5. (Optional) For sliding window, skip blocks are padded with null blocks.
140
+ 2. Add the remaining computed blocks.
141
+ 3. (Optional) For KV connectors, allocate new blocks for external computed
142
+ tokens (if any).
143
144
Args:
145
request_id: The request ID.
146
new_computed_blocks: The new computed blocks just hitting the
147
prefix cache.
148
+ num_local_computed_tokens: The number of local computed tokens.
149
+ num_external_computed_tokens: The number of external computed tokens.
150
"""
112
- if request_id not in self.num_cached_block:
113
- # A new request.
114
- req_blocks = self.req_to_blocks[request_id]
115
- assert len(req_blocks) == 0
116
- req_blocks.extend(new_computed_blocks)
117
- self.num_cached_block[request_id] = len(new_computed_blocks)
118
- else:
119
- # A running request. Should not have new computed blocks.
151
+
152
+ if request_id in self.num_cached_block:
153
+ # Fast-path: a running request won't have any new prefix-cache hits.
154
+ # It should not have any new computed blocks.
155
assert len(new_computed_blocks) == 0
156
+ return
157
+
158
+ # A new request.
159
+ req_blocks = self.req_to_blocks[request_id]
160
+ assert len(req_blocks) == 0
161
+ num_total_computed_tokens = (
162
+ num_local_computed_tokens + num_external_computed_tokens
163
+ )
164
+ num_skipped_tokens = self.get_num_skipped_tokens(num_total_computed_tokens)
165
+ num_skipped_blocks = num_skipped_tokens // self.block_size
166
+ if num_skipped_blocks > 0:
167
+ # It is possible that all new computed blocks are skipped when
168
+ # num_skipped_blocks > len(new_computed_blocks).
169
+ new_computed_blocks = new_computed_blocks[num_skipped_blocks:]
170
+ # Some external computed tokens may be skipped too.
171
+ num_external_computed_tokens = min(
172
+ num_total_computed_tokens - num_skipped_tokens,
173
+ num_external_computed_tokens,
174
+ )
175
+
176
+ # Touch the computed blocks to make sure they won't be evicted.
177
+ if self.enable_caching:
178
+ self.block_pool.touch(new_computed_blocks)
179
+ else:
180
+ assert not any(new_computed_blocks), (
181
+ "Computed blocks should be empty when prefix caching is disabled"
182
+ )
183
+
184
+ # Skip blocks are padded with null blocks.
185
+ req_blocks.extend([self._null_block] * num_skipped_blocks)
186
+ # Add the remaining computed blocks.
187
+ req_blocks.extend(new_computed_blocks)
188
+ # All cached hits (including skipped nulls) are already cached; mark
189
+ # them so cache_blocks() will not try to re-cache blocks that already
190
+ # have a block_hash set.
191
+ self.num_cached_block[request_id] = len(req_blocks)
192
+
193
+ if num_external_computed_tokens > 0:
194
+ # Allocate new blocks for external computed tokens.
195
+ allocated_blocks = self.block_pool.get_new_blocks(
196
+ cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
197
+ )
198
+ req_blocks.extend(allocated_blocks)
199
200
def allocate_new_blocks(
201
self, request_id: str, num_tokens: int
@@ -252,7 +330,9 @@ class SingleTypeKVCacheManager(ABC):
330
331
raise NotImplementedError
332
255
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
333
+ def remove_skipped_blocks(
334
+ self, request_id: str, total_computed_tokens: int
335
+ ) -> None:
336
"""
337
Remove and free the blocks that are no longer needed for attention computation.
338
The removed blocks should be replaced by null_block.
@@ -262,18 +342,24 @@ class SingleTypeKVCacheManager(ABC):
342
343
Args:
344
request_id: The request ID.
265
- num_computed_tokens: The number of tokens that have been computed.
345
+ total_computed_tokens: The total number of computed tokens, including
346
+ local computed tokens and external computed tokens.
347
"""
348
# Remove the blocks that will be skipped during attention computation.
268
- num_skipped_tokens = self.get_num_skipped_tokens(num_computed_tokens)
349
+ num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
350
if num_skipped_tokens <= 0:
351
# This indicates that ALL tokens are inside attention window.
352
# Thus we do not need to free any blocks outside attention window.
353
# A typical case is full attention that we never free any token
354
# before the request is finished.
355
return
275
- num_skipped_blocks = num_skipped_tokens // self.block_size
356
blocks = self.req_to_blocks[request_id]
357
+ num_skipped_blocks = num_skipped_tokens // self.block_size
358
+ # `num_skipped_tokens` may include tokens that haven't been allocated yet
359
+ # (e.g., when the attention window moves into the external computed tokens
360
+ # range), so we must cap to the number of blocks that currently exist for
361
+ # this request.
362
+ num_skipped_blocks = min(num_skipped_blocks, len(blocks))
363
removed_blocks: list[KVCacheBlock] = []
364
# Because the block starts from index 0, the num_skipped_block-th block
365
# corresponds to index num_skipped_blocks - 1.
@@ -486,7 +572,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
572
Returns:
573
The number of tokens that will be skipped for attention computation.
574
"""
489
- return num_computed_tokens - self.sliding_window + 1
575
+ return max(0, num_computed_tokens - self.sliding_window + 1)
576
577
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
578
"""
@@ -711,6 +797,7 @@ class MambaManager(SingleTypeKVCacheManager):
797
request_id: str,
798
num_tokens: int,
799
new_computed_blocks: Sequence[KVCacheBlock],
800
+ total_computed_tokens: int,
801
) -> int:
802
# Allocate extra `num_speculative_blocks` blocks for
803
# speculative decoding (MTP/EAGLE) with linear attention.
@@ -721,7 +808,7 @@ class MambaManager(SingleTypeKVCacheManager):
808
* self.kv_cache_spec.num_speculative_blocks
809
)
810
return super().get_num_blocks_to_allocate(
724
- request_id, num_tokens, new_computed_blocks
811
+ request_id, num_tokens, new_computed_blocks, total_computed_tokens
812
)
813
814
def allocate_new_blocks(
@@ -749,8 +836,12 @@ class MambaManager(SingleTypeKVCacheManager):
836
class CrossAttentionManager(SingleTypeKVCacheManager):
837
"""Manager for cross-attention KV cache in encoder-decoder models."""
838
752
- def save_new_computed_blocks(
753
- self, request_id: str, new_computed_blocks: Sequence[KVCacheBlock]
839
+ def allocate_new_computed_blocks(
840
+ self,
841
+ request_id: str,
842
+ new_computed_blocks: Sequence[KVCacheBlock],
843
+ num_local_computed_tokens: int,
844
+ num_external_computed_tokens: int,
845
) -> None:
846
# We do not cache blocks for cross-attention to be shared between
847
# requests, so `new_computed_blocks` should always be empty.
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -624,7 +624,7 @@ class Worker(WorkerBase):
624
output = self.model_runner.execute_model(
625
scheduler_output, intermediate_tensors
626
)
627
- if isinstance(output, (ModelRunnerOutput, NoneType)):
627
+ if isinstance(output, ModelRunnerOutput | NoneType):
628
return output
629
630
assert isinstance(output, IntermediateTensors)
vllm/v1/worker/tpu_worker.py
+7
-2
@@ -304,6 +304,13 @@ class TPUWorker:
304
305
def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
306
"""Allocate GPU KV cache with the specified kv_cache_config."""
307
+ # Init kv cache connector here, because it requires
308
+ # `kv_cache_config`.
309
+ # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
310
+ # because `initialize_kv_cache` will inject kv cache groups not
311
+ # related to kv cache connector (e.g. kv cache sharing layers).
312
+ ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
313
+
314
self.model_runner.initialize_kv_cache(kv_cache_config)
315
316
def check_health(self) -> None:
@@ -336,8 +343,6 @@ class TPUWorker:
343
parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size
344
)
345
339
- ensure_kv_transfer_initialized(vllm_config)
340
-
346
def shutdown(self) -> None:
347
self.model_runner.ensure_kv_transfer_shutdown()
348