[Bugfix][Core] Retain encoder cache references for repeated multimodal inputs (#59942)

Signed-off-by: bzsuni <bingzhe.sun@daocloud.io>
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>
Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com>
This commit is contained in:
bzsuni
2026-10-05 15:52:46 +00:00
committed by GitHub
co-authored by cjackal
parent edde9d2ce4
commit 528772a4bf
2 changed files with 65 additions and 25 deletions
+27
View File
@@ -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(
+38 -25
View File
@@ -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(