diff --git a/tests/v1/streaming_input/test_async_llm_streaming.py b/tests/v1/streaming_input/test_async_llm_streaming.py index 9dac5ee652..2d4ee7f115 100644 --- a/tests/v1/streaming_input/test_async_llm_streaming.py +++ b/tests/v1/streaming_input/test_async_llm_streaming.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from vllm.engine.protocol import StreamingInput +from vllm.exceptions import VLLMValidationError from vllm.outputs import STREAM_FINISHED, RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.async_llm import AsyncLLM @@ -107,6 +108,25 @@ def make_output(request_id: str, finished: bool) -> RequestOutput: ) +@pytest.mark.parametrize( + ("is_mm_encoder_only", "is_diffusion"), + [(True, False), (False, True)], + ids=["encoder-only", "diffusion"], +) +def test_resumable_requests_rejected_for_non_decoder_models( + is_mm_encoder_only: bool, is_diffusion: bool +): + llm = MagicMock() + llm.vllm_config.is_mm_encoder_only = is_mm_encoder_only + llm.model_config.is_diffusion = is_diffusion + + with pytest.raises( + VLLMValidationError, + match="Resumable requests are not supported", + ): + AsyncLLM._validate_resumable_request(llm) + + @pytest.mark.asyncio async def test_generate_with_async_generator(): """Test generate with an async input generator. @@ -180,6 +200,8 @@ async def test_empty_input_stream_finishes_without_engine_request(): llm.get_supported_tasks = AsyncMock(return_value=("generate",)) llm._run_output_handler = MagicMock() llm._add_request = AsyncMock() + llm.vllm_config = MagicMock(is_mm_encoder_only=False) + llm.model_config = MagicMock(is_diffusion=False) final_req = MagicMock() final_req.request_id = "request-int" @@ -204,6 +226,7 @@ async def test_empty_input_stream_finishes_without_engine_request(): def _make_streaming_llm(num_chunks: int, has_request: bool = True): llm = AsyncLLM.__new__(AsyncLLM) llm._validate_streaming_input_sampling_params = MagicMock() + llm._validate_resumable_request = MagicMock() llm.get_supported_tasks = AsyncMock(return_value=("generate",)) llm._run_output_handler = MagicMock() llm._add_request = AsyncMock() diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index c6df585a58..471772d055 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -9,6 +9,7 @@ import pytest import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils +from tests.v1.core.utils import create_scheduler as create_async_scheduler from tests.v1.core.utils import create_scheduler as create_spec_scheduler from vllm.config import DeviceConfig, VllmConfig from vllm.multimodal.inputs import ( @@ -179,6 +180,7 @@ class TestStreamingScheduler(unittest.TestCase): request_id="session", prompt_token_ids=[1, 2, 3], ) + session.append_output_token_ids(10) session.num_computed_tokens = len(session.prompt_token_ids) session.max_tokens = 2 @@ -195,6 +197,7 @@ class TestStreamingScheduler(unittest.TestCase): assert session.sampling_params.max_tokens == 10 assert session.max_tokens == 10 + session.append_output_token_ids(11) session.num_computed_tokens = len(session.prompt_token_ids) new_request2 = DummyRequest( @@ -216,6 +219,7 @@ class TestStreamingScheduler(unittest.TestCase): request_id="session", prompt_token_ids=[1, 2, 3], ) + session.append_output_token_ids(10) session.num_computed_tokens = len(session.prompt_token_ids) new_request = DummyRequest( @@ -246,6 +250,7 @@ class TestStreamingScheduler(unittest.TestCase): prompt_token_ids=[1, 2, 3], mm_features=[mm_feature], ) + session.append_output_token_ids(10) session.num_computed_tokens = len(session.prompt_token_ids) mm_feature = MultiModalFeatureSpec( @@ -316,6 +321,7 @@ class TestStreamingScheduler(unittest.TestCase): ) scheduler.add_request(session) session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + session.append_output_token_ids(10) session.num_computed_tokens = len(session.prompt_token_ids) next_request = DummyRequest( @@ -716,3 +722,54 @@ def test_next_chunk_fully_computed_after_stop_on_accepted_draft(async_scheduling assert list(session.prompt_token_ids) == [1, 2, 3, 10, 7, 8] assert output.num_scheduled_tokens["session"] == 2 + + +@pytest.mark.parametrize("chunk_queued_before_stop", [False, True]) +def test_session_resumes_after_stop_while_preempted(chunk_queued_before_stop): + """A session preempted with a step in flight can stop when that step's + output is delivered. It must stay schedulable and keep its tokens.""" + scheduler = create_async_scheduler(async_scheduling=True) + session = DummyRequest("session", prompt_token_ids=list(range(10))) + scheduler.add_request(session) + scheduler.update_from_output(scheduler.schedule(), _model_output("session", [10])) + in_flight = scheduler.schedule() + next_chunk = DummyRequest("session", prompt_token_ids=[20, 21]) + if chunk_queued_before_stop: + scheduler.add_request(next_chunk) + + scheduler.running.remove(session) + scheduler._preempt_request(session, timestamp=0.0) + scheduler.update_from_output(in_flight, _model_output("session", [STOP_TOKEN])) + # The parked variant holds a streaming-input slot until the next chunk. + assert scheduler.num_waiting_for_streaming_input == ( + 0 if chunk_queued_before_stop else 1 + ) + if not chunk_queued_before_stop: + scheduler.add_request(next_chunk) + output = scheduler.schedule() + + assert scheduler.num_waiting_for_streaming_input == 0 + expected = list(range(10)) + [10, 20, 21] + assert list(session.all_token_ids) == expected + assert list(session.prompt_token_ids) == expected + assert output.num_scheduled_tokens["session"] == len(expected) + + +def test_non_resumable_request_stays_dequeued_after_stop_while_preempted(): + """Only resumable requests continue after a stop; a request that finishes + while preempted must not be put back in a queue.""" + scheduler = create_async_scheduler(async_scheduling=True) + request = DummyRequest("req", resumable=False, prompt_token_ids=list(range(10))) + scheduler.add_request(request) + scheduler.update_from_output(scheduler.schedule(), _model_output("req", [10])) + in_flight = scheduler.schedule() + + scheduler.running.remove(request) + scheduler._preempt_request(request, timestamp=0.0) + scheduler.update_from_output(in_flight, _model_output("req", [STOP_TOKEN])) + + assert request.status == RequestStatus.FINISHED_STOPPED + assert request not in scheduler.waiting + assert request not in scheduler.kv_holding_waiting + assert request not in scheduler.deferred_waiting + assert "req" not in scheduler.schedule().num_scheduled_tokens diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c7619a2f66..ee9175304d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1732,14 +1732,14 @@ class Scheduler(SchedulerInterface): Discards the last sampled output token from the prior input chunk. """ - # Current streaming input behaviour: Keep only computed output tokens - # (discard final sampled output token). - num_computed_tokens = session.num_computed_tokens + # Keep all output tokens but the final sampled one. + keep_end = session.num_tokens - 1 + assert keep_end >= session.num_computed_tokens kept_output_tokens = session._all_token_ids[ - session.num_prompt_tokens : num_computed_tokens + session.num_prompt_tokens : keep_end ] - del session._all_token_ids[num_computed_tokens:] - del session.block_hashes[num_computed_tokens // self.hash_block_size :] + del session._all_token_ids[keep_end:] + del session.block_hashes[keep_end // self.hash_block_size :] session._output_token_ids.clear() assert session.prompt_token_ids is not None # Extend prompt with kept output tokens. @@ -2155,6 +2155,8 @@ class Scheduler(SchedulerInterface): # to avoid expensive operations inside the loop. stopped_running_reqs: set[Request] = set() stopped_preempted_reqs: set[Request] = set() + # Stopped resumable requests to re-enqueue, in output order. + continuing_reqs: list[Request] = [] for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 request = self.requests.get(req_id) @@ -2320,6 +2322,8 @@ class Scheduler(SchedulerInterface): kv_transfer_params["remote_prefill_cached_tokens"] = ( prefill_stats.num_cached_tokens ) + else: + continuing_reqs.append(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) @@ -2371,14 +2375,16 @@ class Scheduler(SchedulerInterface): assert not prompt_logprobs_tensors assert prompt_token_id_logprobs is None - # Remove the stopped requests from the running and waiting queues. + # Remove stopped requests from their queues, then enqueue the ones + # that continue, so each ends up in exactly one queue. if stopped_running_reqs: self.running = remove_all(self.running, stopped_running_reqs) if stopped_preempted_reqs: - # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) self.kv_holding_waiting.remove_requests(stopped_preempted_reqs) self.deferred_waiting.difference_update(stopped_preempted_reqs) + for request in continuing_reqs: + self._enqueue_waiting_request(request) error_req_ids = set(self.grammar_compile_error_reqs) self.grammar_compile_error_reqs.clear() @@ -2568,7 +2574,6 @@ class Scheduler(SchedulerInterface): request.status = RequestStatus.WAITING_FOR_STREAMING_REQ self.num_waiting_for_streaming_input += 1 - self._enqueue_waiting_request(request) return False def _update_request_with_output( diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 8b7f913b43..3cd9486898 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -599,6 +599,7 @@ class AsyncLLM(EngineClient): session_id: str | None = None, kv_hints: KvHintsEnvelope | None = None, ) -> RequestOutputCollector: + self._validate_resumable_request() self._validate_streaming_input_sampling_params(sampling_params) inputs = dict( @@ -693,6 +694,13 @@ class AsyncLLM(EngineClient): queue._input_stream_task = asyncio.create_task(handle_inputs()) return queue + def _validate_resumable_request(self) -> None: + if self.vllm_config.is_mm_encoder_only or self.model_config.is_diffusion: + raise VLLMValidationError( + "Resumable requests are not supported for encoder-only or " + "diffusion models." + ) + @staticmethod def _validate_streaming_input_sampling_params( params: SamplingParams | PoolingParams,