[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:
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user