diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 394063843..92bb30a27 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/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") diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index b043f0c0b..ecb4c4278 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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