A high-throughput and memory-efficient inference and serving engine for LLMs
[BugFix] Avoid prematurely freeing cached mm encoder outputs (#45347)
Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 12, 2026 at 15:39 UTC
1a369783e9a09cfd9ebed9799a7b8bbffdc9896f
2 files changed
+182
-3
tests/v1/core/test_scheduler.py
+174
@@ -4435,6 +4435,180 @@ def test_eagle3_mm_encoder_cache_with_shift():
4435
)
4436
4437
4438
+def test_free_encoder_inputs_respects_unconfirmed_placeholders():
4439
+ """Regression test for issue #38551 (rollback path): under async
4440
+ scheduling with speculative decoding, num_computed_tokens is advanced
4441
+ optimistically and can be rolled back when in-flight draft tokens are
4442
+ rejected. Freeing an encoder input as soon as num_computed_tokens passes
4443
+ the end of its placeholder range allows a later rollback to rewind back
4444
+ into the range, after which the worker's MM-embedding gather reads an
4445
+ evicted entry and crashes the engine with "Encoder cache miss". The
4446
+ scheduler must retain the input until the *confirmed* progress
4447
+ (num_computed_tokens - num_output_placeholders) passes the range end, so
4448
+ that no pending rejection can rewind into the range."""
4449
+ scheduler = create_scheduler(
4450
+ model="llava-hf/llava-1.5-7b-hf",
4451
+ num_speculative_tokens=3,
4452
+ )
4453
+ mm_start_pos = 50
4454
+ mm_length = 100
4455
+ mm_positions = [
4456
+ [PlaceholderRange(offset=mm_start_pos, length=mm_length)],
4457
+ ]
4458
+ request = create_requests(
4459
+ num_requests=1,
4460
+ num_tokens=mm_start_pos + mm_length + 100,
4461
+ mm_positions=mm_positions,
4462
+ )[0]
4463
+ manager = scheduler.encoder_cache_manager
4464
+ manager.allocate(request, 0)
4465
+ mm_end = mm_start_pos + mm_length
4466
+
4467
+ # One optimistically-scheduled in-flight step advanced num_computed_tokens
4468
+ # by 1 sampled + 3 draft tokens; none are confirmed yet, so all 4 are
4469
+ # still output placeholders that a rejection could rewind.
4470
+ request.num_output_placeholders = 4
4471
+
4472
+ # Optimistic progress reaches the end of the MM range, but the confirmed
4473
+ # position (mm_end + 1 - 4) is still inside it: a rejection could rewind
4474
+ # back into the range, so the entry must be retained.
4475
+ request.num_computed_tokens = mm_end + 1
4476
+ scheduler._free_encoder_inputs(request)
4477
+ assert manager.get_cached_input_ids(request) == {0}
4478
+
4479
+ # Confirmed position still inside the range.
4480
+ request.num_computed_tokens = mm_end + 3
4481
+ scheduler._free_encoder_inputs(request)
4482
+ assert manager.get_cached_input_ids(request) == {0}
4483
+
4484
+ # Confirmed position (mm_end + 4 - 4) now reaches the range end: even if
4485
+ # every unconfirmed token is rejected, progress cannot rewind into the
4486
+ # range, so the entry is freed.
4487
+ request.num_computed_tokens = mm_end + 4
4488
+ scheduler._free_encoder_inputs(request)
4489
+ assert manager.get_cached_input_ids(request) == set()
4490
+
4491
+
4492
+def test_free_encoder_inputs_unchanged_without_spec_decode():
4493
+ """Without speculative decoding, encoder inputs are freed as soon as
4494
+ num_computed_tokens passes the placeholder range, as before."""
4495
+ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
4496
+ mm_positions = [[PlaceholderRange(offset=50, length=100)]]
4497
+ request = create_requests(
4498
+ num_requests=1,
4499
+ num_tokens=250,
4500
+ mm_positions=mm_positions,
4501
+ )[0]
4502
+ manager = scheduler.encoder_cache_manager
4503
+ manager.allocate(request, 0)
4504
+
4505
+ request.num_computed_tokens = 149
4506
+ scheduler._free_encoder_inputs(request)
4507
+ assert manager.get_cached_input_ids(request) == {0}
4508
+
4509
+ request.num_computed_tokens = 150
4510
+ scheduler._free_encoder_inputs(request)
4511
+ assert manager.get_cached_input_ids(request) == set()
4512
+
4513
+
4514
+def test_encoder_cache_retained_across_preemption_and_resume():
4515
+ """Regression guard for issue #38551 (preemption path).
4516
+
4517
+ A request preempted under KV pressure resets num_computed_tokens to 0
4518
+ and drops its encoder references (scheduler._preempt_request calls
4519
+ encoder_cache_manager.free). Because that only moves the entry into
4520
+ `freeable` (it is not evicted), the worker still holds it: the scheduler
4521
+ must NOT report the mm_hash as freed. On resume, re-requesting the
4522
+ encoder input must pull the still-cached entry back out of `freeable`
4523
+ without scheduling a recompute, keeping the scheduler and worker
4524
+ consistent. The spec-rollback retention margin does not gate this path,
4525
+ so it is covered separately here."""
4526
+ scheduler = create_scheduler(
4527
+ model="llava-hf/llava-1.5-7b-hf",
4528
+ num_speculative_tokens=3,
4529
+ )
4530
+ mm_positions = [[PlaceholderRange(offset=50, length=100)]]
4531
+ request = create_requests(
4532
+ num_requests=1,
4533
+ num_tokens=250,
4534
+ mm_hashes_list=[["img_a"]],
4535
+ mm_positions=mm_positions,
4536
+ )[0]
4537
+ manager = scheduler.encoder_cache_manager
4538
+ mm_hash = request.mm_features[0].identifier
4539
+
4540
+ # Prefill scheduled and computed the encoder input; it is pinned.
4541
+ manager.allocate(request, 0)
4542
+ assert manager.get_cached_input_ids(request) == {0}
4543
+
4544
+ # Preemption drops the request's encoder references (scheduler.py:
4545
+ # _preempt_request -> encoder_cache_manager.free) and resets progress.
4546
+ manager.free(request)
4547
+ request.num_computed_tokens = 0
4548
+ # The entry is now ref-free but only `freeable` (not evicted): the
4549
+ # worker still holds it, so nothing must be reported as freed.
4550
+ assert mm_hash in manager.cached
4551
+ assert mm_hash in manager.freeable
4552
+ assert manager.get_freed_mm_hashes() == []
4553
+
4554
+ # Resume re-requests the encoder output. The still-cached entry is pulled
4555
+ # back out of `freeable` with no recompute and no worker-side free.
4556
+ assert manager.check_and_update_cache(request, 0) is True
4557
+ assert mm_hash not in manager.freeable
4558
+ assert manager.get_cached_input_ids(request) == {0}
4559
+ assert manager.get_freed_mm_hashes() == []
4560
+
4561
+
4562
+def test_encoder_cache_recomputed_when_evicted_during_preemption():
4563
+ """Companion to the retention case (issue #38551, preemption path).
4564
+
4565
+ If a preempted request's retained encoder entry IS evicted under memory
4566
+ pressure before it resumes, the scheduler reports the mm_hash as freed
4567
+ (so the worker drops it) and a resume must schedule a recompute rather
4568
+ than assume the worker still holds it. check_and_update_cache must
4569
+ return False so the encoder input is re-scheduled."""
4570
+ scheduler = create_scheduler(
4571
+ model="llava-hf/llava-1.5-7b-hf",
4572
+ num_speculative_tokens=3,
4573
+ )
4574
+ mm_positions = [[PlaceholderRange(offset=50, length=100)]]
4575
+ request = create_requests(
4576
+ num_requests=1,
4577
+ num_tokens=250,
4578
+ mm_hashes_list=[["img_a"]],
4579
+ mm_positions=mm_positions,
4580
+ )[0]
4581
+ manager = scheduler.encoder_cache_manager
4582
+ mm_hash = request.mm_features[0].identifier
4583
+
4584
+ manager.allocate(request, 0)
4585
+ # Preemption drops references; the entry becomes freeable.
4586
+ manager.free(request)
4587
+ request.num_computed_tokens = 0
4588
+ assert mm_hash in manager.freeable
4589
+
4590
+ # A new request with a different image hits memory pressure and evicts
4591
+ # the freeable entry to make room.
4592
+ other = create_requests(
4593
+ num_requests=1,
4594
+ num_tokens=250,
4595
+ mm_hashes_list=[["img_b"]],
4596
+ mm_positions=mm_positions,
4597
+ req_ids=["1"],
4598
+ )[0]
4599
+ manager.num_free_slots = 50 # force eviction of the freeable entry
4600
+ assert manager.can_allocate(
4601
+ other, 0, encoder_compute_budget=10_000, num_embeds_to_schedule=0
4602
+ )
4603
+
4604
+ # The evicted entry is reported to the worker, which drops it.
4605
+ assert mm_hash not in manager.cached
4606
+ assert manager.get_freed_mm_hashes() == [mm_hash]
4607
+
4608
+ # On resume the original request must recompute (cache miss is correct).
4609
+ assert manager.check_and_update_cache(request, 0) is False
4610
+
4611
+
4612
@pytest.mark.parametrize("use_kv_connector", [False, True])
4613
def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector):
4614
"""Test that ensure_cache_available() returning False defers the request.
vllm/v1/core/sched/scheduler.py
+8
-3
@@ -1808,9 +1808,14 @@ class Scheduler(SchedulerInterface):
1808
# we know we're done with the encoder input. Cross Attention
1809
# KVs have been calculated and cached already.
1810
self.encoder_cache_manager.free_encoder_input(request, input_id)
1811
- elif start_pos + num_tokens <= request.num_computed_tokens:
1812
- # The encoder output is already processed and stored
1813
- # in the decoder's KV cache.
1811
+ elif (
1812
+ start_pos + num_tokens
1813
+ <= request.num_computed_tokens - request.num_output_placeholders
1814
+ ):
1815
+ # The encoder output is already processed and stored in the
1816
+ # decoder's KV cache, and progress is far enough past the
1817
+ # placeholder range that no pending draft-token rejection can
1818
+ # roll num_computed_tokens back into it.
1819
self.encoder_cache_manager.free_encoder_input(request, input_id)
1820
1821
def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: