From b4cfbc24d33ca17bc764a75ffe749654654521c1 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Tue, 7 Jul 2026 14:32:55 +0700 Subject: [PATCH] [Bugfix][Core] Fix host memory leak from undrained new_block_ids (#44490) Signed-off-by: Ting Sun Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: chaunceyjiang --- vllm/v1/core/kv_cache_coordinator.py | 1 + vllm/v1/core/sched/scheduler.py | 7 +++--- vllm/v1/core/single_type_kv_cache_manager.py | 25 ++++++++++---------- 3 files changed, 18 insertions(+), 15 deletions(-) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index a759d7a80a..5829e553f9 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -115,6 +115,7 @@ class KVCacheCoordinator(ABC): dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, scheduler_block_size=self.scheduler_block_size, + needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing, ) for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f75bfae2b5..b315269b86 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1078,10 +1078,11 @@ class Scheduler(SchedulerInterface): self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. + new_attn_block_ids = self.kv_cache_manager.take_new_block_ids() new_block_ids_to_zero = ( - (self.kv_cache_manager.take_new_block_ids() or None) - if self.needs_kv_cache_zeroing - else None + (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None ) # Dynamic speculative decoding: compute optimal K diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 642fe3e6a0..5d554b5925 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -45,6 +45,7 @@ class SingleTypeKVCacheManager(ABC): scheduler_block_size: int, dcp_world_size: int = 1, pcp_world_size: int = 1, + needs_kv_cache_zeroing: bool = False, max_admission_blocks_per_request: int | None = None, ) -> None: """ @@ -55,6 +56,8 @@ class SingleTypeKVCacheManager(ABC): kv_cache_group_id: The id of the kv cache group of this manager. scheduler_block_size: The scheduling granularity (LCM of all group block sizes); a multiple of this manager's ``block_size``. + needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs + newly allocated block IDs from this manager. max_admission_blocks_per_request: Recycling-aware per-request block cap used by `get_num_blocks_to_allocate`. Only set for spec types that recycle blocks across chunks (SWA, @@ -73,6 +76,14 @@ class SingleTypeKVCacheManager(ABC): self.block_pool = block_pool self.enable_caching = enable_caching self._max_admission_blocks_per_request = max_admission_blocks_per_request + # Record newly allocated block ids only when worker-side zeroing will + # consume them and this manager holds a spec type that gets zeroed. + self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + HiddenStateCacheSpec, + ) self.new_block_ids: list[int] = [] # Mapping from request ID to blocks to track the blocks allocated @@ -268,12 +279,7 @@ class SingleTypeKVCacheManager(ABC): cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) ) req_blocks.extend(allocated_blocks) - if type(self.kv_cache_spec) in ( - FullAttentionSpec, - TQFullAttentionSpec, - MLAAttentionSpec, - HiddenStateCacheSpec, - ): + if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -301,12 +307,7 @@ class SingleTypeKVCacheManager(ABC): else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) - if type(self.kv_cache_spec) in ( - FullAttentionSpec, - TQFullAttentionSpec, - MLAAttentionSpec, - HiddenStateCacheSpec, - ): + if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in new_blocks) return new_blocks