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 13, 17:00 UTC
[Core][Hybrid allocator + connector 2/n] Unify `remove_skipped_blocks` by `get_last_useful_token` (#25431)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
Nov 5, 2025 at 16:12 UTC
efe73e9b57574934d5d30b889d66ef54b1f16711
1 file changed
+112
-76
vllm/v1/core/single_type_kv_cache_manager.py
+112
-76
@@ -243,18 +243,53 @@ class SingleTypeKVCacheManager(ABC):
243
244
raise NotImplementedError
245
246
- @abstractmethod
246
def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
247
"""
249
- Remove the blocks that are no longer needed from `blocks` and free the
250
- blocks. The removed blocks should be replaced by null_block.
251
- Need to be customized for each attention type.
248
+ Remove and free the blocks that are no longer needed for attention computation.
249
+ The removed blocks should be replaced by null_block.
250
+
251
+ This function depends on `get_num_skipped_tokens`, which need to be implemented
252
+ differently for each attention type.
253
254
Args:
255
request_id: The request ID.
256
num_computed_tokens: The number of tokens that have been computed.
257
"""
257
- raise NotImplementedError
258
+ # Remove the blocks that will be skipped during attention computation.
259
+ num_skipped_tokens = self.get_num_skipped_tokens(num_computed_tokens)
260
+ if num_skipped_tokens <= 0:
261
+ # This indicates that ALL tokens are inside attention window.
262
+ # Thus we do not need to free any blocks outside attention window.
263
+ # A typical case is full attention that we never free any token
264
+ # before the request is finished.
265
+ return
266
+ num_skipped_blocks = num_skipped_tokens // self.block_size
267
+ blocks = self.req_to_blocks[request_id]
268
+ removed_blocks: list[KVCacheBlock] = []
269
+ # Because the block starts from index 0, the num_skipped_block-th block
270
+ # corresponds to index num_skipped_blocks - 1.
271
+ for i in range(num_skipped_blocks - 1, -1, -1):
272
+ if blocks[i] == self._null_block:
273
+ # If the block is already a null block, the blocks before it
274
+ # should also have been set to null blocks by the previous calls
275
+ # to this function.
276
+ break
277
+ removed_blocks.append(blocks[i])
278
+ blocks[i] = self._null_block
279
+ self.block_pool.free_blocks(removed_blocks)
280
+
281
+ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
282
+ """
283
+ Get the number of tokens that will be skipped for attention computation.
284
+
285
+ Args:
286
+ num_computed_tokens: The number of tokens that have been computed.
287
+
288
+ Returns:
289
+ The number of tokens that will be skipped for attention computation.
290
+ """
291
+ # The default behavior is to not skip any tokens.
292
+ return 0
293
294
295
class FullAttentionManager(SingleTypeKVCacheManager):
@@ -298,10 +333,6 @@ class FullAttentionManager(SingleTypeKVCacheManager):
333
computed.pop()
334
return computed_blocks
335
301
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
302
- # No need to remove blocks for full attention.
303
- pass
304
-
336
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
337
blocks = self.req_to_blocks[running_request_id]
338
num_common_blocks = 0
@@ -389,28 +420,33 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
420
computed.pop()
421
return computed_blocks
422
392
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
393
- # Remove the blocks that are no longer be in the sliding window and
394
- # skipped during the attention computation.
395
- last_useful_token = num_computed_tokens - self.sliding_window + 1
396
- last_useful_block = last_useful_token // self.block_size
397
- if last_useful_block <= 0:
398
- # Early return if tokens are not enough to fill the sliding window
399
- return
400
- blocks = self.req_to_blocks[request_id]
401
- if blocks[last_useful_block - 1] == self._null_block:
402
- # Early return if there are no blocks to remove
403
- return
404
- removed_blocks: list[KVCacheBlock] = []
405
- for i in range(last_useful_block - 1, -1, -1):
406
- if blocks[i] == self._null_block:
407
- # If the block is already a null block, the blocks before it
408
- # should also have been set to null blocks by the previous calls
409
- # to this function.
410
- break
411
- removed_blocks.append(blocks[i])
412
- blocks[i] = self._null_block
413
- self.block_pool.free_blocks(removed_blocks)
423
+ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
424
+ """
425
+ Get the number of tokens that will be skipped for attention computation.
426
+
427
+ For sliding window, this corresponds to the tokens that are prior to
428
+ the current sliding window.
429
+
430
+ Example:
431
+ sliding_window=4, num_computed_tokens=7
432
+
433
+ Tokens: [ 0 1 2 3 4 5 6 7 ]
434
+ | ---- computed -----|
435
+ ^ next token to be computed
436
+ |-----------| sliding window for next token
437
+ |--skipped---|
438
+
439
+ The current window contains tokens 4~7. Tokens 0~3 will be skipped for
440
+ attention computation since they are outside the sliding window.
441
+ Thus, get_num_skipped_tokens(7) == 4.
442
+
443
+ Args:
444
+ num_computed_tokens: The number of tokens that have been computed.
445
+
446
+ Returns:
447
+ The number of tokens that will be skipped for attention computation.
448
+ """
449
+ return num_computed_tokens - self.sliding_window + 1
450
451
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
452
"""
@@ -511,40 +547,51 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
547
break
548
return computed_blocks
549
514
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
515
- # Remove the blocks that are no longer be in the chunked attention
516
- # window and skipped during the attention computation.
517
-
518
- # [chunk 0][chunk 1]local_attention_start_idx ... current
519
- # we computed previous number of chunks to get the idx of
520
- # current chunk window starting offset,
521
- # e.g. for computed 1024 tokens, the 1024th token (0 indexed)
522
- # is in the second chunk, there are 1 prev chunk, the start idx
523
- # is 1024. for 1023, it will be 0.
524
- num_cached_block = self.num_cached_block.get(request_id, 0)
525
- local_attention_start_idx = (
526
- (num_computed_tokens)
527
- // self.attention_chunk_size
528
- * self.attention_chunk_size
529
- )
530
- first_useful_block_idx = local_attention_start_idx // self.block_size
531
- if num_cached_block > 0:
532
- # Make sure we don't delete the last cached block
533
- first_useful_block_idx = min(first_useful_block_idx, num_cached_block - 1)
534
- # if block size = 128, 0 -> block 0, 1024 (= 128 * 8) ->
535
- # block 8, 372 (= 128 * 2 + 116) -> block 2
536
- blocks = self.req_to_blocks[request_id]
537
- removed_blocks: list[KVCacheBlock] = []
538
- # we need to keep the last block to get the previous hash key
539
- for i in range(first_useful_block_idx - 1, -1, -1):
540
- if blocks[i] == self._null_block:
541
- # If the block is already a null block, the blocks before it
542
- # should also have been set to null blocks by the previous calls
543
- # to this function.
544
- break
545
- removed_blocks.append(blocks[i])
546
- blocks[i] = self._null_block
547
- self.block_pool.free_blocks(removed_blocks)
550
+ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
551
+ """
552
+ Get the number of tokens that will be skipped for attention computation.
553
+
554
+ For chunked local attention, this corresponds to the tokens that are on
555
+ the left side of the current chunk.
556
+
557
+ Example 1:
558
+ chunk size = 8, num_computed_tokens = 13
559
+ Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ...
560
+ | ----- computed ---------------|
561
+ ^^ next token to be computed
562
+ |----------------| <-- attention window for
563
+ next token
564
+ |--- skipped -----|
565
+ Output: get_num_skipped_tokens(13) == 8
566
+
567
+ Example 2:
568
+ chunk size = 8, num_computed_tokens = 8
569
+ Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ...
570
+ | --- computed ---|
571
+ ^ next token to be computed
572
+ |--| <-- attention window for next token
573
+ | --- skipped ----|
574
+ Output: get_num_skipped_tokens(8) == 8
575
+
576
+ Example 3:
577
+ chunk size = 8, num_computed_tokens = 7
578
+ Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ...
579
+ |---computed---|
580
+ ^ next token to be computed
581
+ |-----------------| <-- attention window for next token
582
+ no token should be skipped.
583
+ Output: get_num_skipped_tokens(7) == 0
584
+
585
+ Args:
586
+ num_computed_tokens: The number of tokens that have been computed.
587
+
588
+ Returns:
589
+ The number of tokens that will be skipped for attention computation.
590
+ """
591
+ num_skipped_tokens = (
592
+ num_computed_tokens // self.attention_chunk_size
593
+ ) * self.attention_chunk_size
594
+ return num_skipped_tokens
595
596
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
597
"""
@@ -590,12 +637,6 @@ class MambaManager(SingleTypeKVCacheManager):
637
638
return computed_blocks
639
593
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
594
- # Here unused blocks may be freed up for running requests.
595
- # TODO(@s3woz) Free up all blocks that aren't needed by Mamba2
596
- # (for which find_longest_cache_hit returns block_pool.null_block)
597
- pass
598
-
640
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
641
"""
642
cascade attention is not supported by mamba
@@ -676,11 +717,6 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
717
# Return empty blocks to indicate no cache hits
718
raise NotImplementedError("CrossAttentionManager does not support caching")
719
679
- def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
680
- # Cross-attention blocks represent encoder states which are needed
681
- # for the entire decoding process, so no blocks should be skipped
682
- pass
683
-
720
721
spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
722
FullAttentionSpec: FullAttentionManager,