[Bugfix] Trim AR sampled logprob rows with new_token_ids on mid-step stop (#5564)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -105,6 +105,87 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
+def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
+    """Spec decode can sample tokens past a stop; the trailing tokens are
+    trimmed by _update_request_with_output, and the emitted logprob rows must
+    be trimmed with them (upstream slices logprobs after the trim)."""
+    request = _Request("req")
+    request.num_computed_tokens = 8
+    scheduler = SimpleNamespace(
+        perf_metrics=None,
+        connector=None,
+        chunk_transfer_adapter=None,
+        requests={request.request_id: request},
+        running=[request],
+        waiting=_RequestQueue(),
+        skipped_waiting=_RequestQueue(),
+        structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
+        transfer_triggered_requests=set(),
+        active_kv_transfers=set(),
+        pending_stop_after_extraction=set(),
+        waiting_for_transfer_free=set(),
+        _new_prompt_len_snapshot={},
+        finished_req_ids=set(),
+        finished_req_ids_dict=defaultdict(set),
+        kv_cache_manager=SimpleNamespace(take_events=lambda: None),
+        kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
+    )
+
+    def update_request_trimming(req, token_ids):
+        # Mirror vllm Scheduler._update_request_with_output: the first sampled
+        # token hits EOS, so the spec-accepted tokens behind it are trimmed
+        # in place and the request stops.
+        del token_ids[1:]
+        req.status = RequestStatus.FINISHED_STOPPED
+        return token_ids, True
+
+    def free_request(req):
+        scheduler.requests.pop(req.request_id, None)
+        scheduler.finished_req_ids.add(req.request_id)
+        scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
+        return None, None
+
+    scheduler._update_request_with_output = update_request_trimming
+    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
+    scheduler._handle_stopped_request = lambda _request: True
+    scheduler._free_request = free_request
+    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
+    scheduler.make_stats = lambda *args, **kwargs: None
+    scheduler._capture_omni_connector_output = lambda _output: None
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 3},
+        scheduled_spec_decode_tokens={"req": [8, 9]},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[7, 8, 9]],
+        logprobs=_logprobs(
+            [[7, 1], [8, 2], [9, 3]],
+            [[-0.7, -1.0], [-0.8, -1.1], [-0.9, -1.2]],
+            cumulative_rows=[0],
+        ),
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req": 0},
+        routed_experts=None,
+    )
+
+    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+    (output,) = outputs[0].outputs
+
+    assert output.new_token_ids == [7]
+    assert output.finish_reason is FinishReason.STOP
+    token_rows = np.asarray(output.new_logprobs.logprob_token_ids)
+    value_rows = np.asarray(output.new_logprobs.logprobs)
+    assert token_rows.shape[0] == len(output.new_token_ids)
+    assert value_rows.shape[0] == len(output.new_token_ids)
+    np.testing.assert_array_equal(token_rows[:, 0], [7])
+
+
 def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     """A bad runner response must not bring down unrelated batch requests."""
     bad = _Request("bad")

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -471,7 +471,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Check for stop and update request status.
             if new_token_ids:
+                num_sampled_tokens = len(new_token_ids)
                 new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
+                if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens:
+                    # A mid-step stop (e.g. spec-decode tokens sampled past
+                    # EOS) trims new_token_ids after the validation slice
+                    # above; re-slice so the emitted logprob rows stay 1:1
+                    # with the emitted tokens, as upstream vLLM does by
+                    # slicing after the trim.
+                    new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
             elif request.pooling_params and pooler_output is not None:
                 # Pooling stops as soon as there is output.
                 request.status = RequestStatus.FINISHED_STOPPED