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
fix(v1/kv_cache): resolve async KV transfer bug in cascade attention (#23485)
Signed-off-by: Ayush Satyam <ayushsatyam146@gmail.com>
Ayush Satyam committed
Oct 8, 2025 at 10:16 UTC
cd9890544b9773882a274f944f82c0414df9c991
4 files changed
+41
-72
vllm/v1/core/kv_cache_coordinator.py
+9
-16
@@ -148,27 +148,22 @@ class KVCacheCoordinator(ABC):
148
for manager in self.single_type_managers:
149
manager.free(request_id)
150
151
- def get_num_common_prefix_blocks(
152
- self, request_id: str, num_running_requests: int
153
- ) -> list[int]:
151
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
152
"""
155
- Get the number of common prefix blocks for all requests in the RUNNING
156
- state for each kv cache group.
153
+ Get the number of common prefix blocks for all requests with allocated
154
+ KV cache for each kv cache group.
155
156
Args:
159
- request_id: The request ID.
160
- num_running_requests: The total number of requests in the RUNNING
161
- state.
157
+ running_request_id: The request ID of any running request, used to
158
+ identify the common prefix blocks.
159
160
Returns:
164
- list[int]: The number of common prefix blocks for all requests in
165
- the RUNNING state for each kv cache group.
161
+ list[int]: The number of common prefix blocks for each kv cache group.
162
"""
167
- num_blocks_per_group = [
168
- manager.get_num_common_prefix_blocks(request_id, num_running_requests)
163
+ return [
164
+ manager.get_num_common_prefix_blocks(running_request_id)
165
for manager in self.single_type_managers
166
]
171
- return num_blocks_per_group
167
168
def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
169
"""
@@ -226,9 +221,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
221
)
222
self.num_single_type_manager = len(self.single_type_managers)
223
229
- def get_num_common_prefix_blocks(
230
- self, request_id: str, num_running_requests: int
231
- ) -> list[int]:
224
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
225
return [0] * self.num_single_type_manager
226
227
def find_longest_cache_hit(
vllm/v1/core/kv_cache_manager.py
+18
-28
@@ -10,7 +10,7 @@ from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
10
from vllm.v1.core.kv_cache_utils import KVCacheBlock
11
from vllm.v1.kv_cache_interface import KVCacheConfig
12
from vllm.v1.metrics.stats import PrefixCacheStats
13
-from vllm.v1.request import Request, RequestStatus
13
+from vllm.v1.request import Request
14
15
logger = init_logger(__name__)
16
@@ -344,49 +344,39 @@ class KVCacheManager:
344
self.prefix_cache_stats.reset = True
345
return True
346
347
- def get_num_common_prefix_blocks(
348
- self,
349
- request: Request,
350
- num_running_requests: int,
351
- ) -> list[int]:
352
- """Calculate the number of common prefix blocks shared by all requests
353
- in the RUNNING state for each kv cache group.
347
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
348
+ """Calculate the number of common prefix blocks for each kv cache group.
349
355
- The function determines this by selecting any request and iterating
356
- through its blocks. A block is considered a common prefix block if its
357
- `ref_cnt` equals the total number of requests in the RUNNING state.
350
+ The function selects a running request and iterates through its blocks.
351
+ A block is considered a common prefix block if ALL requests with
352
+ allocated KV cache share it (i.e., ref_cnt equals the number of entries
353
+ in req_to_blocks).
354
359
- NOTE(woosuk): The number of requests in the RUNNING state is **greater
355
+ NOTE(woosuk): The number of requests with allocated KV cache is **greater
356
than or equal to** the number of requests scheduled in the current step.
361
- This is because the RUNNING state only indicates that:
357
+ This is because having allocated KV cache only indicates that:
358
1. The request has not yet finished, and
359
2. The request holds its blocks unfreed.
360
365
- While all scheduled requests must be in the RUNNING state, the inverse
366
- is not necessarily true. There may be RUNNING requests that are not
367
- scheduled in the current step.
361
+ While all scheduled requests must have allocated KV cache, the inverse
362
+ is not necessarily true. There may be requests with allocated KV cache
363
+ that are not scheduled in the current step.
364
365
This can result in an edge case where the number of common prefix blocks
366
is 0, even though all scheduled requests share a common prefix. This
371
- occurs because there may be unscheduled RUNNING requests that do not
372
- share the common prefix. Currently, this case cannot be easily detected,
373
- so the function returns 0 in such cases.
367
+ occurs because there may be unscheduled requests that do not share the
368
+ common prefix. Currently, this case cannot be easily detected, so the
369
+ function returns 0 in such cases.
370
371
Args:
376
- request: Any request in the RUNNING state, used to identify the
377
- common prefix blocks.
378
- num_running_requests: The total number of requests in the RUNNING
379
- state. This can be different from the number of scheduled
380
- requests in the current step.
372
+ running_request_id: The request ID of any running request, used to
373
+ identify the common prefix blocks.
374
375
Returns:
376
list[int]: The number of common prefix blocks for each kv cache
377
group.
378
"""
386
- assert request.status == RequestStatus.RUNNING
387
- return self.coordinator.get_num_common_prefix_blocks(
388
- request.request_id, num_running_requests
389
- )
379
+ return self.coordinator.get_num_common_prefix_blocks(running_request_id)
380
381
def take_events(self) -> list[KVCacheEvent]:
382
"""Take the KV cache events from the block pool.
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -597,7 +597,7 @@ class Scheduler(SchedulerInterface):
597
any_request = self.running[0]
598
num_common_prefix_blocks = (
599
self.kv_cache_manager.get_num_common_prefix_blocks(
600
- any_request, len(self.running)
600
+ any_request.request_id
601
)
602
)
603
vllm/v1/core/single_type_kv_cache_manager.py
+13
-27
@@ -182,21 +182,17 @@ class SingleTypeKVCacheManager(ABC):
182
self.num_cached_block.pop(request_id, None)
183
184
@abstractmethod
185
- def get_num_common_prefix_blocks(
186
- self, request_id: str, num_running_requests: int
187
- ) -> int:
185
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
186
"""
189
- Get the number of common prefix blocks for all requests in the RUNNING
190
- state.
187
+ Get the number of common prefix blocks for all requests with allocated
188
+ KV cache.
189
190
Args:
193
- request_id: The request ID.
194
- num_running_requests: The total number of requests in the RUNNING
195
- state.
191
+ running_request_id: The request ID.
192
193
Returns:
198
- The number of common prefix blocks for all requests in the RUNNING
199
- state.
194
+ The number of common prefix blocks for all requests with allocated
195
+ KV cache.
196
"""
197
198
raise NotImplementedError
@@ -302,13 +298,11 @@ class FullAttentionManager(SingleTypeKVCacheManager):
298
# No need to remove blocks for full attention.
299
pass
300
305
- def get_num_common_prefix_blocks(
306
- self, request_id: str, num_running_requests: int
307
- ) -> int:
308
- blocks = self.req_to_blocks[request_id]
301
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
302
+ blocks = self.req_to_blocks[running_request_id]
303
num_common_blocks = 0
304
for block in blocks:
311
- if block.ref_cnt == num_running_requests:
305
+ if block.ref_cnt == len(self.req_to_blocks):
306
num_common_blocks += 1
307
else:
308
break
@@ -408,9 +402,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
402
blocks[i] = self._null_block
403
self.block_pool.free_blocks(removed_blocks)
404
411
- def get_num_common_prefix_blocks(
412
- self, request_id: str, num_running_requests: int
413
- ) -> int:
405
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
406
"""
407
NOTE(Chen): The prefix blocks are null blocks for sliding window layers.
408
So it's not correct to count ref_cnt like FullAttentionManager. Return
@@ -544,9 +536,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
536
blocks[i] = self._null_block
537
self.block_pool.free_blocks(removed_blocks)
538
547
- def get_num_common_prefix_blocks(
548
- self, request_id: str, num_running_requests: int
549
- ) -> int:
539
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
540
"""
541
cascade attention is not supported by chunked local attention.
542
"""
@@ -596,9 +586,7 @@ class MambaManager(SingleTypeKVCacheManager):
586
# (for which find_longest_cache_hit returns block_pool.null_block)
587
pass
588
599
- def get_num_common_prefix_blocks(
600
- self, request_id: str, num_running_requests: int
601
- ) -> int:
589
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
590
"""
591
cascade attention is not supported by mamba
592
"""
@@ -648,9 +636,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
636
# requests, so this method is not relevant.
637
raise ValueError("Should not be called as prefix caching is disabled.")
638
651
- def get_num_common_prefix_blocks(
652
- self, request_id: str, num_running_requests: int
653
- ) -> int:
639
+ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
640
# Cross-attention blocks contain request-specific encoder states
641
# and are not shared between different requests
642
return 0