[BugFix] Finish async_chunk requests without pad-token injection (#3613)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -700,3 +700,232 @@ def test_wire_round_trip_struct_to_dict_contract():
     assert set(decoded.keys()) == set(expected.keys())
     assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
     assert set(decoded["codes"].keys()) == set(expected["codes"].keys())
+
+
+# ---------------------------------------------------------------
+# Deferred finish for upstream-completed requests
+# ---------------------------------------------------------------
+
+
+def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
+    """Build a mock scheduler with requests queued for deferred finish."""
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs}
+    cleanup_calls = []
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.ec_connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+    scheduler._pending_finish_reqs = list(pending_finish_reqs)
+
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = list(running)
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+    scheduler.requests = {r.request_id: r for r in running}
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=None,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={},
+    )
+    return scheduler, scheduler_output, model_runner_output, cleanup_calls
+
+
+def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
+    """A request whose upstream completed with no remaining tokens should
+    emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df1",
+        external_req_id="ext-df1",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=16,
+        num_prompt_tokens=16,
+        prompt_token_ids=list(range(1, 17)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[request],
+        pending_finish_reqs=[request],
+    )
+    scheduler._free_request.return_value = {"mock": "kv_params"}
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    scheduler._handle_stopped_request.assert_called_once_with(request)
+    scheduler._free_request.assert_called_once_with(request)
+    assert len(cleanup_calls) == 1
+    assert cleanup_calls[0][0] == ("req-df1", "ext-df1")
+
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].request_id == "req-df1"
+    assert eco.outputs[0].finish_reason == "stop"
+    assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"}
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_empty_prompt(mocker: MockerFixture):
+    """A request that never received any tokens (finished immediately upstream)
+    should still emit a FINISHED output and clean up."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df2",
+        external_req_id="ext-df2",
+        status=RequestStatus.WAITING,
+        is_finished=lambda: False,
+        num_computed_tokens=0,
+        num_prompt_tokens=0,
+        prompt_token_ids=[],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[],
+        pending_finish_reqs=[request],
+    )
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    scheduler._free_request.assert_called_once_with(request)
+    assert len(cleanup_calls) == 1
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].finish_reason == "stop"
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_skips_already_finished(mocker: MockerFixture):
+    """A request that was aborted between schedule() and update_from_output()
+    should be skipped without emitting output or freeing resources twice."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df3",
+        external_req_id="ext-df3",
+        status=RequestStatus.FINISHED_ABORTED,
+        is_finished=lambda: True,
+        num_computed_tokens=0,
+        num_prompt_tokens=0,
+        prompt_token_ids=[],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[],
+        pending_finish_reqs=[request],
+    )
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    scheduler._handle_stopped_request.assert_not_called()
+    scheduler._free_request.assert_not_called()
+    assert len(cleanup_calls) == 0
+    assert 0 not in result
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
+    """When _handle_stopped_request returns False (resumable request), the
+    output must still be emitted so the client stream doesn't hang."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df4",
+        external_req_id="ext-df4",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=16,
+        num_prompt_tokens=16,
+        prompt_token_ids=list(range(1, 17)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[request],
+        pending_finish_reqs=[request],
+    )
+    scheduler._handle_stopped_request.return_value = False
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    scheduler._handle_stopped_request.assert_called_once_with(request)
+    scheduler._free_request.assert_not_called()
+    assert len(cleanup_calls) == 0
+
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].request_id == "req-df4"
+    assert eco.outputs[0].finish_reason == "stop"
+    assert eco.outputs[0].kv_transfer_params is None
+    assert scheduler._pending_finish_reqs == []

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -44,6 +44,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self._pending_finish_reqs: list[Request] = []
         self.input_coordinator: OmniSchedulingCoordinator | None = None
         if uses_qwen3_omni_full_payload_input_coordinator(model_config):
             self.input_coordinator = OmniSchedulingCoordinator(
@@ -111,23 +112,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
-            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
                 if (
                     self.chunk_transfer_adapter is not None
                     and request.request_id in self.chunk_transfer_adapter.finished_requests
                 ):
-                    # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED.
-                    if len(request.prompt_token_ids) <= num_computed_tokens:
-                        request.prompt_token_ids.append(0)
-                        try:
-                            request._all_token_ids.append(0)  # type: ignore[attr-defined]
-                        except Exception:
-                            pass
-                    required_tokens = len(request.prompt_token_ids) - num_computed_tokens
-                else:
-                    req_index += 1
-                    continue
+                    self._pending_finish_reqs.append(request)
+                req_index += 1
+                continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -174,11 +166,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
                 if request.request_id in self.chunk_transfer_adapter.finished_requests:
-                    request.prompt_token_ids.append(0)
-                    try:
-                        request._all_token_ids.append(0)  # type: ignore[attr-defined]
-                    except Exception:
-                        pass
+                    self.waiting.pop_request()
+                    self._pending_finish_reqs.append(request)
+                    continue
                 else:
                     self.waiting.pop_request()
                     skipped_waiting_requests.prepend_request(request)
@@ -581,6 +571,36 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
+        # Finish async_chunk requests that schedule() collected because their
+        # upstream completed with no remaining codec tokens.
+        for request in self._pending_finish_reqs:
+            if request.is_finished():
+                continue
+            request.status = RequestStatus.FINISHED_STOPPED
+            finish_reason = request.get_finished_reason()
+            finished = self._handle_stopped_request(request)
+            kv_transfer_params = None
+            if finished:
+                kv_transfer_params = self._free_request(request)
+                if self.chunk_transfer_adapter is not None:
+                    self.chunk_transfer_adapter.cleanup(
+                        request.request_id,
+                        getattr(request, "external_req_id", None),
+                    )
+            outputs[request.client_index].append(
+                EngineCoreOutput(
+                    request_id=request.request_id,
+                    new_token_ids=[],
+                    finish_reason=finish_reason,
+                    stop_reason=request.stop_reason,
+                    events=request.take_events(),
+                    kv_transfer_params=kv_transfer_params,
+                    trace_headers=request.trace_headers,
+                )
+            )
+            stopped_running_reqs.add(request)
+        self._pending_finish_reqs.clear()
+
         # Remove the stopped requests from the running and waiting queues.
         if stopped_running_reqs:
             self.running = remove_all(self.running, stopped_running_reqs)