diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 32284eeac..e2ad98f56 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Unit tests for Omni AR streaming-session async placeholder handling.""" from __future__ import annotations from types import SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -11,6 +15,8 @@ import pytest # isort: off import vllm_omni # noqa: F401 - import for side effects (patch vLLM) from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler @@ -51,6 +57,65 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate: ) +def _run_resumable_segment_stop(session: Request) -> None: + sched = MagicMock() + sched.requests = {session.request_id: session} + sched.perf_metrics = None + sched.structured_output_manager.should_advance.return_value = False + sched._update_request_with_output.return_value = ([42], True) + sched._handle_stopped_request.return_value = False + sched.chunk_transfer_adapter = None + sched.running = [session] + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {session.request_id: 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {session.request_id: 0} + model_runner_output.routed_experts = None + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + +@pytest.mark.parametrize("outstanding_async_tokens", [0, 1, 2]) +def test_resumable_segment_stop_reconciles_async_placeholders( + outstanding_async_tokens: int, +) -> None: + """A segment stop discards and rolls back only in-flight async tokens.""" + session = _make_request() + session.status = RequestStatus.RUNNING + session.resumable = True + session.append_output_token_ids([7, 8]) + session.num_computed_tokens = session.num_tokens + outstanding_async_tokens + session.num_output_placeholders = outstanding_async_tokens + session.spec_token_ids = [-1] * outstanding_async_tokens + + _run_resumable_segment_stop(session) + + assert session.async_tokens_to_discard == outstanding_async_tokens + assert session.num_computed_tokens == session.num_tokens + assert session.num_output_placeholders == 0 + assert session.spec_token_ids == [] + assert session._output_token_ids == [] + + def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None: sched = _make_scheduler(stage_id=0) session = _make_request() diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 93654f27f..a9702eeb3 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -438,8 +438,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Downstream async-chunk stages receive real payloads from the # connector. This update only resumes polling for the next segment. self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id) - request.async_tokens_to_discard = 1 - request.num_output_placeholders = 0 + outstanding_async_tokens = request.num_output_placeholders + if outstanding_async_tokens > 0: + # Discard only outputs that are already in flight and + # roll back their optimistic computed-token accounting. + request.async_tokens_to_discard = outstanding_async_tokens + request.num_computed_tokens -= outstanding_async_tokens + request.num_output_placeholders = 0 request.spec_token_ids = [] request._output_token_ids.clear() if finished: