[Bugfix][EC Connector] Don't stop an encoder-instance request before its images are encoded (#50275)
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Tianyu Guo committed
Aug 5, 2026 at 13:53 UTC
50c51682a18cf2417c229ee4253fb2a089d5cb78
4 files changed
+121
-5
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+2
-2
@@ -108,8 +108,8 @@ async def fanout_encoder_primer(
108
"messages": [
109
{"role": "user", "content": [item]},
110
],
111
- # Only need 1 token so the server actually runs the encoder path
112
- "max_tokens": 1,
111
+ # No max_tokens cap: the encoder instance never samples, it finishes
112
+ # once the prompt is encoded and its embeddings are published.
113
"stream": False,
114
}
115
tasks.append(
tests/v1/core/test_scheduler.py
+103
-1
@@ -38,7 +38,12 @@ from vllm.v1.kv_cache_interface import (
38
KVCacheGroupSpec,
39
MambaSpec,
40
)
41
-from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
41
+from vllm.v1.outputs import (
42
+ DraftTokenIds,
43
+ KVConnectorOutput,
44
+ ModelRunnerOutput,
45
+ make_empty_encoder_model_runner_output,
46
+)
47
from vllm.v1.request import Request, RequestStatus
48
from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
49
@@ -5730,3 +5735,100 @@ def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
5735
num_scheduled = output.num_scheduled_tokens[replay.request_id]
5736
# Must resume at the convergent boundary (block 0), not the deep FA hit.
5737
assert replay.num_tokens - num_scheduled == block_size
5738
+
5739
+
5740
+def _make_encoder_instance_request(scheduler, text_prefix=8, image_tokens=16):
5741
+ """One request whose image sits after some text, as an EPD proxy sends it.
5742
+
5743
+ `create_scheduler` builds a text-only model, so `compute_mm_encoder_budget`
5744
+ leaves the encoder cache at zero capacity; give it room so the admitted path
5745
+ is reachable.
5746
+ """
5747
+ ecm = scheduler.encoder_cache_manager
5748
+ ecm.cache_size = ecm.num_free_slots = ecm.num_freeable_slots = 8 * image_tokens
5749
+ scheduler.max_num_encoder_input_tokens = 8 * image_tokens
5750
+
5751
+ (request,) = create_requests(
5752
+ num_requests=1,
5753
+ num_tokens=text_prefix + image_tokens + 4,
5754
+ # Completion must hinge on the prompt being encoded, not on max_tokens.
5755
+ max_tokens=16,
5756
+ mm_hashes_list=[["img-hash-0"]],
5757
+ mm_positions=[[PlaceholderRange(offset=text_prefix, length=image_tokens)]],
5758
+ )
5759
+ scheduler.add_request(request)
5760
+ return request
5761
+
5762
+
5763
+def test_encoder_instance_defers_stop_until_prompt_is_consumed():
5764
+ """An encoder instance must not finish a request that has not encoded yet.
5765
+
5766
+ When the encoder cache cannot admit the multi-modal item, the scheduler
5767
+ schedules only the tokens before it. Finishing that step would hand the
5768
+ client a successful completion for a request whose image was never
5769
+ encoded, so no embedding is ever published to the EC connector.
5770
+ """
5771
+ scheduler = create_scheduler(
5772
+ max_num_seqs=8,
5773
+ max_num_batched_tokens=1024,
5774
+ use_ec_connector=True,
5775
+ ec_role="ec_producer",
5776
+ )
5777
+ request = _make_encoder_instance_request(scheduler)
5778
+ req_id = request.request_id
5779
+
5780
+ # Encoder budget exhausted: the item cannot be admitted this step.
5781
+ scheduler.encoder_cache_manager.can_allocate = lambda *a, **k: False
5782
+
5783
+ output = scheduler.schedule()
5784
+ assert output.num_scheduled_tokens[req_id] > 0
5785
+ assert not output.scheduled_encoder_inputs.get(req_id)
5786
+ assert output.num_scheduled_tokens[req_id] < request.num_prompt_tokens
5787
+
5788
+ scheduler.update_from_output(
5789
+ output,
5790
+ make_empty_encoder_model_runner_output(output),
5791
+ )
5792
+
5793
+ assert not request.is_finished(), (
5794
+ "request finished before its image was encoded, so no embedding was "
5795
+ "published to the EC connector"
5796
+ )
5797
+ assert request.num_output_tokens == 0
5798
+
5799
+ # With the budget restored the item is encoded and only then does the
5800
+ # request finish.
5801
+ del scheduler.encoder_cache_manager.can_allocate # unshadow the real method
5802
+
5803
+ output = scheduler.schedule()
5804
+ assert output.scheduled_encoder_inputs.get(req_id) == [0]
5805
+
5806
+ scheduler.update_from_output(
5807
+ output,
5808
+ make_empty_encoder_model_runner_output(output),
5809
+ )
5810
+ assert request.is_finished()
5811
+
5812
+
5813
+def test_encoder_instance_finishes_request_once_prompt_is_consumed():
5814
+ """The unblocked path completes in a single step, without sampling."""
5815
+ scheduler = create_scheduler(
5816
+ max_num_seqs=8,
5817
+ max_num_batched_tokens=1024,
5818
+ use_ec_connector=True,
5819
+ ec_role="ec_producer",
5820
+ )
5821
+ request = _make_encoder_instance_request(scheduler)
5822
+ req_id = request.request_id
5823
+
5824
+ output = scheduler.schedule()
5825
+ assert output.scheduled_encoder_inputs.get(req_id) == [0]
5826
+ assert output.num_scheduled_tokens[req_id] == request.num_prompt_tokens
5827
+
5828
+ scheduler.update_from_output(
5829
+ output,
5830
+ make_empty_encoder_model_runner_output(output),
5831
+ )
5832
+ assert request.status == RequestStatus.FINISHED_STOPPED
5833
+ # The encoder instance publishes an embedding, not tokens.
5834
+ assert request.num_output_tokens == 0
vllm/v1/core/sched/scheduler.py
+12
@@ -94,6 +94,7 @@ class Scheduler(SchedulerInterface):
94
)
95
self.structured_output_manager = structured_output_manager
96
self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder
97
+ self.is_encoder_only = vllm_config.is_encoder_only
98
99
# include_finished_set controls whether a separate set of finished
100
# request ids should be included in the EngineCoreOutputs returned
@@ -1813,6 +1814,17 @@ class Scheduler(SchedulerInterface):
1814
# Pooling stops as soon as there is output.
1815
request.status = RequestStatus.FINISHED_STOPPED
1816
stopped = True
1817
+ elif (
1818
+ self.is_encoder_only
1819
+ and request.num_computed_tokens >= request.num_prompt_tokens
1820
+ ):
1821
+ # An encoder instance runs the encoder and publishes the
1822
+ # embeddings instead of sampling, so it stops as soon as the
1823
+ # whole prompt is consumed. Encoder inputs are never scheduled
1824
+ # past a multi-modal item the encoder cache could not admit, so
1825
+ # a consumed prompt also means every item in it was encoded.
1826
+ request.status = RequestStatus.FINISHED_STOPPED
1827
+ stopped = True
1828
1829
if new_token_ids and self.structured_output_manager.should_advance(
1830
request, new_token_ids=new_token_ids
vllm/v1/outputs.py
+4
-2
@@ -358,8 +358,10 @@ def make_empty_encoder_model_runner_output(
358
# Give every request its own contiguous index
359
req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)}
360
361
- # No tokens generated yet ⇒ one empty list per request
362
- sampled_token_ids: list[list[int]] = [[0] for _ in req_ids]
361
+ # An encoder instance never samples, so it emits no tokens at all. The
362
+ # scheduler finishes these requests once their prompt is fully encoded
363
+ # (see `Scheduler.update_from_output`).
364
+ sampled_token_ids: list[list[int]] = [[] for _ in req_ids]
365
366
# Pooler outputs are not available yet ⇒ use None placeholders
367
pooler_output: list[torch.Tensor | None] = [None for _ in req_ids]