diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 0d259c1c69..45fe7e6c4b 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -4,6 +4,7 @@ import unittest from unittest.mock import MagicMock +import numpy as np import torch from vllm.config import DeviceConfig, VllmConfig @@ -20,7 +21,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, ) -from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.outputs import LogprobsLists, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.structured_output import StructuredOutputManager @@ -84,6 +85,61 @@ def create_scheduler() -> Scheduler: class TestStreamingScheduler(unittest.TestCase): + def test_chunk_logprobs_do_not_depend_on_continuation_arrival(self): + """A queued continuation must not change the finishing chunk's logprobs.""" + for continuation_queued in (False, True): + with self.subTest(continuation_queued=continuation_queued): + scheduler = create_scheduler() + session = DummyRequest("session", prompt_token_ids=[1, 2, 3]) + session.sampling_params = SamplingParams( + stop_token_ids=[STOP_TOKEN], max_tokens=16, logprobs=0 + ) + scheduler.add_request(session) + scheduled = scheduler.schedule() + + continuation = DummyRequest("session", prompt_token_ids=[4, 5]) + assert continuation.sampling_params.logprobs is None + if continuation_queued: + scheduler.add_request(continuation) + + logprobs = LogprobsLists( + logprob_token_ids=np.array([[STOP_TOKEN]], dtype=np.int32), + logprobs=np.array([[-0.25]], dtype=np.float32), + sampled_token_ranks=np.array([1], dtype=np.int32), + ) + runner_output = ModelRunnerOutput( + req_ids=[session.request_id], + req_id_to_index={session.request_id: 0}, + sampled_token_ids=[[STOP_TOKEN]], + logprobs=logprobs, + prompt_logprobs_dict={}, + pooler_output=[], + ) + outputs = scheduler.update_from_output(scheduled, runner_output) + [output] = outputs[session.client_index].outputs + + if not continuation_queued: + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + scheduler.add_request(continuation) + assert session.status == RequestStatus.WAITING + assert session.sampling_params is continuation.sampling_params + assert output.new_token_ids == [STOP_TOKEN] + assert output.finish_reason == FinishReason.STOP + assert output.new_logprobs is not None, ( + "The finishing chunk requested logprobs; the continuation's " + "settings must not suppress them" + ) + np.testing.assert_array_equal( + output.new_logprobs.logprob_token_ids, logprobs.logprob_token_ids + ) + np.testing.assert_array_equal( + output.new_logprobs.logprobs, logprobs.logprobs + ) + np.testing.assert_array_equal( + output.new_logprobs.sampled_token_ranks, + logprobs.sampled_token_ranks, + ) + def test_add_request(self): scheduler = create_scheduler() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fa81626906..f05b12575c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2152,6 +2152,15 @@ class Scheduler(SchedulerInterface): self.kv_cache_manager.estimate_cached_tokens(request) ) + # Extract sample logprobs before stop handling can replace the + # sampling parameters for a streaming continuation. + if ( + request.sampling_params is not None + and request.sampling_params.num_logprobs is not None + and logprobs + ): + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) + finish_reason = None if stopped: # Capture finish_reason BEFORE _handle_stopped_request, which may @@ -2177,14 +2186,6 @@ class Scheduler(SchedulerInterface): else: stopped_preempted_reqs.add(request) - # Extract sample logprobs if needed. - if ( - request.sampling_params is not None - and request.sampling_params.num_logprobs is not None - and logprobs - ): - new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) - if self.return_sampling_mask: sampling_masks = model_runner_output.sampling_masks if new_token_ids and sampling_masks is not None: