[Bugfix][Core] Fix host memory leak from undrained new_block_ids (#44490)
Signed-off-by: Ting Sun <suntcrick@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: chaunceyjiang <chaunceyjiang@gmail.com>
Ting SUN committed
Jul 7, 2026 at 14:32 UTC
b4cfbc24d33ca17bc764a75ffe749654654521c1
3 files changed
+18
-15
vllm/v1/core/kv_cache_coordinator.py
+1
@@ -115,6 +115,7 @@ class KVCacheCoordinator(ABC):
115
dcp_world_size=dcp_world_size,
116
pcp_world_size=pcp_world_size,
117
scheduler_block_size=self.scheduler_block_size,
118
+ needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing,
119
)
120
for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
121
)
vllm/v1/core/sched/scheduler.py
+4
-3
@@ -1078,10 +1078,11 @@ class Scheduler(SchedulerInterface):
1078
self.prev_step_scheduled_req_ids.clear()
1079
self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
1080
1081
+ # Drain new attention block ids every step so the manager-side list
1082
+ # does not grow unbounded; only kv-cache zeroing consumes them.
1083
+ new_attn_block_ids = self.kv_cache_manager.take_new_block_ids()
1084
new_block_ids_to_zero = (
1082
- (self.kv_cache_manager.take_new_block_ids() or None)
1083
- if self.needs_kv_cache_zeroing
1084
- else None
1085
+ (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None
1086
)
1087
1088
# Dynamic speculative decoding: compute optimal K
vllm/v1/core/single_type_kv_cache_manager.py
+13
-12
@@ -45,6 +45,7 @@ class SingleTypeKVCacheManager(ABC):
45
scheduler_block_size: int,
46
dcp_world_size: int = 1,
47
pcp_world_size: int = 1,
48
+ needs_kv_cache_zeroing: bool = False,
49
max_admission_blocks_per_request: int | None = None,
50
) -> None:
51
"""
@@ -55,6 +56,8 @@ class SingleTypeKVCacheManager(ABC):
56
kv_cache_group_id: The id of the kv cache group of this manager.
57
scheduler_block_size: The scheduling granularity (LCM of all group
58
block sizes); a multiple of this manager's ``block_size``.
59
+ needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs
60
+ newly allocated block IDs from this manager.
61
max_admission_blocks_per_request: Recycling-aware per-request
62
block cap used by `get_num_blocks_to_allocate`. Only set for
63
spec types that recycle blocks across chunks (SWA,
@@ -73,6 +76,14 @@ class SingleTypeKVCacheManager(ABC):
76
self.block_pool = block_pool
77
self.enable_caching = enable_caching
78
self._max_admission_blocks_per_request = max_admission_blocks_per_request
79
+ # Record newly allocated block ids only when worker-side zeroing will
80
+ # consume them and this manager holds a spec type that gets zeroed.
81
+ self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in (
82
+ FullAttentionSpec,
83
+ TQFullAttentionSpec,
84
+ MLAAttentionSpec,
85
+ HiddenStateCacheSpec,
86
+ )
87
self.new_block_ids: list[int] = []
88
89
# Mapping from request ID to blocks to track the blocks allocated
@@ -268,12 +279,7 @@ class SingleTypeKVCacheManager(ABC):
279
cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
280
)
281
req_blocks.extend(allocated_blocks)
271
- if type(self.kv_cache_spec) in (
272
- FullAttentionSpec,
273
- TQFullAttentionSpec,
274
- MLAAttentionSpec,
275
- HiddenStateCacheSpec,
276
- ):
282
+ if self._record_new_block_ids:
283
self.new_block_ids.extend(b.block_id for b in allocated_blocks)
284
285
def allocate_new_blocks(
@@ -301,12 +307,7 @@ class SingleTypeKVCacheManager(ABC):
307
else:
308
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
309
req_blocks.extend(new_blocks)
304
- if type(self.kv_cache_spec) in (
305
- FullAttentionSpec,
306
- TQFullAttentionSpec,
307
- MLAAttentionSpec,
308
- HiddenStateCacheSpec,
309
- ):
310
+ if self._record_new_block_ids:
311
self.new_block_ids.extend(b.block_id for b in new_blocks)
312
return new_blocks
313