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 8, 15:15 UTC
[Bugfix][Core] Fix CPU memory leak from Request reference cycle in prefix caching (#34183)
Signed-off-by: Roger Wang <hey@rogerw.io>
Roger Wang committed
Feb 9, 2026 at 21:03 UTC
8a5e0e2b2bb925d162328927b7565514fa355da1
3 files changed
+14
-12
tests/v1/core/test_async_scheduler.py
+1
-1
@@ -236,7 +236,7 @@ def test_prefix_caching_for_multi_turn():
236
req._all_token_ids = req.prompt_token_ids.copy()
237
req.all_token_ids = ConstantList(req._all_token_ids)
238
req.block_hashes = []
239
- req.block_hashes = req.get_hash_new_full_blocks()
239
+ req.update_block_hashes()
240
241
# Schedule the next-turn requests.
242
for req in next_turn_requests:
vllm/v1/core/sched/scheduler.py
+2
-4
@@ -982,10 +982,8 @@ class Scheduler(SchedulerInterface):
982
983
session._all_token_ids.extend(update.prompt_token_ids or ())
984
session.prompt_token_ids.extend(update.prompt_token_ids or ())
985
- # Update block hashes for the new tokens
986
- # (mirrors Request.append_output_token_ids)
987
- if session.get_hash_new_full_blocks is not None:
988
- session.block_hashes.extend(session.get_hash_new_full_blocks())
985
+ # Update block hashes for the new tokens.
986
+ session.update_block_hashes()
987
session.num_prompt_tokens = len(session.prompt_token_ids)
988
session.arrival_time = update.arrival_time
989
session.sampling_params = update.sampling_params
vllm/v1/request.py
+11
-7
@@ -6,7 +6,6 @@ import time
6
from collections import deque
7
from collections.abc import Callable, Mapping
8
from dataclasses import dataclass
9
-from functools import partial
9
from typing import TYPE_CHECKING, Any
10
11
import torch
@@ -164,10 +163,11 @@ class Request:
163
self.num_external_computed_tokens = 0
164
165
self.block_hashes: list[BlockHash] = []
167
- self.get_hash_new_full_blocks: Callable[[], list[BlockHash]] | None = None
168
- if block_hasher is not None:
169
- self.get_hash_new_full_blocks = partial(block_hasher, self)
170
- self.block_hashes = self.get_hash_new_full_blocks()
166
+ # Store the block hasher without binding self to avoid creating a
167
+ # reference cycle (Request -> partial -> Request) that prevents
168
+ # immediate garbage collection via reference counting.
169
+ self._block_hasher: Callable[[Request], list[BlockHash]] | None = block_hasher
170
+ self.update_block_hashes()
171
172
self.skip_reading_prefix_cache = self.get_skip_reading_prefix_cache()
173
@@ -212,8 +212,12 @@ class Request:
212
self._output_token_ids.extend(token_ids)
213
self._all_token_ids.extend(token_ids)
214
215
- if self.get_hash_new_full_blocks is not None:
216
- self.block_hashes.extend(self.get_hash_new_full_blocks())
215
+ self.update_block_hashes()
216
+
217
+ def update_block_hashes(self) -> None:
218
+ """Compute block hashes for any new full blocks and append them."""
219
+ if self._block_hasher is not None:
220
+ self.block_hashes.extend(self._block_hasher(self))
221
222
@property
223
def use_structured_output(self) -> bool: