[Bugfix][Scheduler] Preserve logprobs across streaming continuations (#57447)
Signed-off-by: 0xsensei <prblmslvr.aditya@gmail.com>
0xsensei committed
Sep 29, 2026 at 01:20 UTC
2840ca7a332525fe73bbc55ef508722372875823
2 files changed
+66
-9
tests/v1/streaming_input/test_scheduler_streaming.py
+57
-1
@@ -4,6 +4,7 @@
4
import unittest
5
from unittest.mock import MagicMock
6
7
+import numpy as np
8
import torch
9
10
from vllm.config import DeviceConfig, VllmConfig
@@ -20,7 +21,7 @@ from vllm.v1.kv_cache_interface import (
21
KVCacheConfig,
22
KVCacheGroupSpec,
23
)
23
-from vllm.v1.outputs import ModelRunnerOutput
24
+from vllm.v1.outputs import LogprobsLists, ModelRunnerOutput
25
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
26
from vllm.v1.structured_output import StructuredOutputManager
27
@@ -84,6 +85,61 @@ def create_scheduler() -> Scheduler:
85
86
87
class TestStreamingScheduler(unittest.TestCase):
88
+ def test_chunk_logprobs_do_not_depend_on_continuation_arrival(self):
89
+ """A queued continuation must not change the finishing chunk's logprobs."""
90
+ for continuation_queued in (False, True):
91
+ with self.subTest(continuation_queued=continuation_queued):
92
+ scheduler = create_scheduler()
93
+ session = DummyRequest("session", prompt_token_ids=[1, 2, 3])
94
+ session.sampling_params = SamplingParams(
95
+ stop_token_ids=[STOP_TOKEN], max_tokens=16, logprobs=0
96
+ )
97
+ scheduler.add_request(session)
98
+ scheduled = scheduler.schedule()
99
+
100
+ continuation = DummyRequest("session", prompt_token_ids=[4, 5])
101
+ assert continuation.sampling_params.logprobs is None
102
+ if continuation_queued:
103
+ scheduler.add_request(continuation)
104
+
105
+ logprobs = LogprobsLists(
106
+ logprob_token_ids=np.array([[STOP_TOKEN]], dtype=np.int32),
107
+ logprobs=np.array([[-0.25]], dtype=np.float32),
108
+ sampled_token_ranks=np.array([1], dtype=np.int32),
109
+ )
110
+ runner_output = ModelRunnerOutput(
111
+ req_ids=[session.request_id],
112
+ req_id_to_index={session.request_id: 0},
113
+ sampled_token_ids=[[STOP_TOKEN]],
114
+ logprobs=logprobs,
115
+ prompt_logprobs_dict={},
116
+ pooler_output=[],
117
+ )
118
+ outputs = scheduler.update_from_output(scheduled, runner_output)
119
+ [output] = outputs[session.client_index].outputs
120
+
121
+ if not continuation_queued:
122
+ assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
123
+ scheduler.add_request(continuation)
124
+ assert session.status == RequestStatus.WAITING
125
+ assert session.sampling_params is continuation.sampling_params
126
+ assert output.new_token_ids == [STOP_TOKEN]
127
+ assert output.finish_reason == FinishReason.STOP
128
+ assert output.new_logprobs is not None, (
129
+ "The finishing chunk requested logprobs; the continuation's "
130
+ "settings must not suppress them"
131
+ )
132
+ np.testing.assert_array_equal(
133
+ output.new_logprobs.logprob_token_ids, logprobs.logprob_token_ids
134
+ )
135
+ np.testing.assert_array_equal(
136
+ output.new_logprobs.logprobs, logprobs.logprobs
137
+ )
138
+ np.testing.assert_array_equal(
139
+ output.new_logprobs.sampled_token_ranks,
140
+ logprobs.sampled_token_ranks,
141
+ )
142
+
143
def test_add_request(self):
144
scheduler = create_scheduler()
145
vllm/v1/core/sched/scheduler.py
+9
-8
@@ -2152,6 +2152,15 @@ class Scheduler(SchedulerInterface):
2152
self.kv_cache_manager.estimate_cached_tokens(request)
2153
)
2154
2155
+ # Extract sample logprobs before stop handling can replace the
2156
+ # sampling parameters for a streaming continuation.
2157
+ if (
2158
+ request.sampling_params is not None
2159
+ and request.sampling_params.num_logprobs is not None
2160
+ and logprobs
2161
+ ):
2162
+ new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
2163
+
2164
finish_reason = None
2165
if stopped:
2166
# Capture finish_reason BEFORE _handle_stopped_request, which may
@@ -2177,14 +2186,6 @@ class Scheduler(SchedulerInterface):
2186
else:
2187
stopped_preempted_reqs.add(request)
2188
2180
- # Extract sample logprobs if needed.
2181
- if (
2182
- request.sampling_params is not None
2183
- and request.sampling_params.num_logprobs is not None
2184
- and logprobs
2185
- ):
2186
- new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
2187
-
2189
if self.return_sampling_mask:
2190
sampling_masks = model_runner_output.sampling_masks
2191
if new_token_ids and sampling_masks is not None: