From 528772a4bf589fd52eec40bc8ca79c84d99206b5 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:52:46 +0800 Subject: [PATCH] [Bugfix][Core] Retain encoder cache references for repeated multimodal inputs (#59942) Signed-off-by: bzsuni Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- tests/v1/core/test_scheduler.py | 27 ++++++++++++++ vllm/v1/core/sched/scheduler.py | 63 ++++++++++++++++++++------------- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 2451ea5ec0..66368a299d 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -688,6 +688,33 @@ def test_throttle_capacity_bound_guard_admits(): assert "b" in output.num_scheduled_tokens +def test_same_step_duplicate_encoder_input_stays_cached(): + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + max_num_batched_tokens=1024, + max_model_len=2048, + ) + request = create_requests( + 1, + num_tokens=2000, + req_ids=["repeated"], + mm_hashes_list=[["image", "image", "image"]], + mm_positions=[ + [PlaceholderRange(offset=offset, length=576) for offset in (0, 600, 1300)] + ], + )[0] + scheduler.add_request(request) + + output = scheduler.schedule() + assert output.scheduled_encoder_inputs == {request.request_id: [0]} + _model_output(scheduler, output, [[]]) + + # The second occurrence is partially consumed; the third is not scheduled. + cache = scheduler.encoder_cache_manager + assert cache.get_cached_input_ids(request) == {1} + assert "image" not in cache.freeable + + def test_no_mm_input_chunking(): # Disable multimodal input chunking. scheduler = create_scheduler( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b390f1aadb..e25b5adda4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -709,6 +709,7 @@ class Scheduler(SchedulerInterface): # Schedule encoder inputs. encoder_inputs_to_schedule = None external_load_encoder_input: list[int] = [] + duplicate_encoder_inputs: list[int] = [] new_encoder_compute_budget = encoder_compute_budget if request.has_encoder_inputs: ( @@ -716,6 +717,7 @@ class Scheduler(SchedulerInterface): num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) = self._try_schedule_encoder_inputs( request, request.num_computed_tokens, @@ -856,17 +858,13 @@ class Scheduler(SchedulerInterface): # Encoder-related. if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule - # Allocate the encoder cache. - for i in encoder_inputs_to_schedule: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget - if external_load_encoder_input: - for i in external_load_encoder_input: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) + if encoder_inputs_to_schedule or external_load_encoder_input: + self._allocate_encoder_inputs( + request, + (encoder_inputs_to_schedule or []) + external_load_encoder_input, + duplicate_encoder_inputs, + ) # Record the LoRAs in scheduled_running_reqs scheduled_loras: set[int] = set() @@ -1065,6 +1063,7 @@ class Scheduler(SchedulerInterface): encoder_inputs_to_schedule = None external_load_encoder_input = [] + duplicate_encoder_inputs = [] new_encoder_compute_budget = encoder_compute_budget pad_spec_decode = False @@ -1166,6 +1165,7 @@ class Scheduler(SchedulerInterface): num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) = self._try_schedule_encoder_inputs( request, num_computed_tokens, @@ -1335,18 +1335,14 @@ class Scheduler(SchedulerInterface): # Encoder-related. if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule - # Allocate the encoder cache. - for i in encoder_inputs_to_schedule: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget - # Allocate for external load encoder cache - if external_load_encoder_input: - for i in external_load_encoder_input: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) + if encoder_inputs_to_schedule or external_load_encoder_input: + self._allocate_encoder_inputs( + request, + (encoder_inputs_to_schedule or []) + + external_load_encoder_input, + duplicate_encoder_inputs, + ) # re-queue requests skipped in this pass ahead of older skipped items. if step_skipped_kv_holding: @@ -1742,6 +1738,20 @@ class Scheduler(SchedulerInterface): self.encoder_cache_mismatch_reqs.add(request.request_id) return True + def _allocate_encoder_inputs( + self, + request: Request, + input_ids: list[int], + duplicate_input_ids: list[int], + ) -> None: + for input_id in input_ids: + self.encoder_cache_manager.allocate(request, input_id) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, input_id) + + for input_id in duplicate_input_ids: + self.encoder_cache_manager.check_and_update_cache(request, input_id) + def _try_schedule_encoder_inputs( self, request: Request, @@ -1749,7 +1759,7 @@ class Scheduler(SchedulerInterface): num_new_tokens: int, encoder_compute_budget: int, shift_computed_tokens: int = 0, - ) -> tuple[list[int], int, int, list[int]]: + ) -> tuple[list[int], int, int, list[int], list[int]]: """Determine which encoder inputs need to be scheduled in the current step, and update `num_new_tokens` and encoder token budget accordingly. @@ -1770,12 +1780,13 @@ class Scheduler(SchedulerInterface): blocks and externally cached blocks (via KVConnector). """ if num_new_tokens == 0 or not request.has_encoder_inputs: - return [], num_new_tokens, encoder_compute_budget, [] + return [], num_new_tokens, encoder_compute_budget, [], [] encoder_inputs_to_schedule: list[int] = [] mm_features = request.mm_features assert mm_features is not None assert len(mm_features) > 0 external_load_encoder_input = [] + duplicate_encoder_inputs = [] # NOTE: since scheduler operates on the request level (possibly with # multiple encoder inputs per request), we need to create temporary @@ -1832,9 +1843,10 @@ class Scheduler(SchedulerInterface): scheduled_num_encoder_embeds, num_encoder_embeds, ): - return [], 0, encoder_compute_budget, [] + return [], 0, encoder_compute_budget, [], [] # The same encoder input has already been scheduled in the # current step. + duplicate_encoder_inputs.append(i) continue cached_num_encoder_embeds = ( @@ -1849,7 +1861,7 @@ class Scheduler(SchedulerInterface): num_encoder_embeds, ) ): - return [], 0, encoder_compute_budget, [] + return [], 0, encoder_compute_budget, [], [] if self.encoder_cache_manager.check_and_update_cache(request, i): # The encoder input is already computed and cached from a @@ -1925,6 +1937,7 @@ class Scheduler(SchedulerInterface): num_new_tokens, encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) def _make_scheduled_encoder_input_stats(