diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py new file mode 100644 index 000000000..155bd894f --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -0,0 +1,181 @@ +"""Tests for the AR sampled-token logprob contract.""" + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import numpy as np +import pytest +from vllm.v1.engine import FinishReason +from vllm.v1.outputs import LogprobsLists +from vllm.v1.request import RequestStatus + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _logprobs( + token_ids: list[list[int]], + values: list[list[float]], + *, + cumulative_rows: list[int] | None = None, +) -> LogprobsLists: + return LogprobsLists( + logprob_token_ids=np.asarray(token_ids, dtype=np.int32), + logprobs=np.asarray(values, dtype=np.float32), + sampled_token_ranks=np.zeros(len(token_ids), dtype=np.int32), + cu_num_generated_tokens=cumulative_rows, + ) + + +def test_slices_requested_rows_and_preserves_sampled_token_values() -> None: + all_logprobs = _logprobs( + [[3, 8], [4, 9], [11, 2], [12, 1]], + [[-0.3, -1.0], [-0.4, -1.1], [-0.11, -2.0], [-0.12, -2.1]], + cumulative_rows=[0, 2], + ) + + sliced = _slice_sampled_logprobs(all_logprobs, req_index=1, sampled_token_ids=[11, 12]) + + np.testing.assert_array_equal(sliced.logprob_token_ids[:, 0], [11, 12]) + np.testing.assert_allclose(sliced.logprobs[:, 0], [-0.11, -0.12]) + + +def test_rejects_missing_logprobs() -> None: + with pytest.raises(RuntimeError, match="model runner returned none"): + _slice_sampled_logprobs(None, req_index=0, sampled_token_ids=[7]) + + +def test_rejects_row_count_mismatch() -> None: + logprobs = _logprobs([[7, 1]], [[-0.7, -1.0]]) + + with pytest.raises(RuntimeError, match="row count does not match"): + _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7, 8]) + + +def test_rejects_sampled_token_mismatch() -> None: + logprobs = _logprobs([[9, 7]], [[-0.9, -1.0]]) + + with pytest.raises(RuntimeError, match="generated_token=7 logprob_token=9"): + _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7]) + + +def test_rejects_non_finite_sampled_logprob() -> None: + logprobs = _logprobs([[7, 1]], [[float("nan"), -1.0]]) + + with pytest.raises(RuntimeError, match="non-finite values at rows \\[0\\]"): + _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7]) + + +class _Request: + def __init__(self, request_id: str) -> None: + self.request_id = request_id + self.client_index = 0 + self.status = RequestStatus.RUNNING + self.sampling_params = SimpleNamespace(num_logprobs=1) + self.num_computed_tokens = 0 + self.num_output_placeholders = 0 + self.has_encoder_inputs = False + self.pooling_params = None + self.resumable = True + self.stop_reason = None + self.trace_headers = None + self.num_nans_in_logits = 0 + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + def get_finished_reason(self): + return RequestStatus.get_finished_reason(self.status) + + def take_events(self): + return None + + def take_prefill_stats(self): + return None + + +class _RequestQueue(list): + def remove_requests(self, requests) -> None: + for request in requests: + if request in self: + self.remove(request) + + +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") + good = _Request("good") + scheduler = SimpleNamespace( + perf_metrics=None, + connector=None, + chunk_transfer_adapter=None, + requests={bad.request_id: bad, good.request_id: good}, + running=[bad, good], + 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), + ) + update_calls: list[str] = [] + + def update_request(request, token_ids): + update_calls.append(request.request_id) + return token_ids, False + + def free_request(request): + scheduler.requests.pop(request.request_id, None) + scheduler.finished_req_ids.add(request.request_id) + scheduler.finished_req_ids_dict[request.client_index].add(request.request_id) + return None + + scheduler._update_request_with_output = update_request + 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={"bad": 1, "good": 1}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + model_runner_output = SimpleNamespace( + sampled_token_ids=[[7], [8]], + logprobs=_logprobs( + [[9, 1], [8, 2]], + [[-0.9, -1.0], [-0.8, -1.0]], + cumulative_rows=[0, 1], + ), + prompt_logprobs_dict={}, + pooler_output=None, + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={"bad": 0, "good": 1}, + routed_experts=None, + ) + + outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) + output_by_id = {output.request_id: output for output in outputs[0].outputs} + + assert update_calls == ["good"] + assert bad.status is RequestStatus.FINISHED_ERROR + assert bad.stop_reason is not None and "misaligned" in bad.stop_reason + assert "bad" not in scheduler.requests + assert output_by_id["bad"].finish_reason is FinishReason.ERROR + assert output_by_id["bad"].new_token_ids == [] + assert output_by_id["good"].new_token_ids == [8] + np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8]) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 22917197d..d87fc0c8a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -5,6 +5,7 @@ from collections.abc import Iterable from time import time from typing import Any +import numpy as np from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats @@ -36,6 +37,51 @@ from vllm_omni.outputs import OmniConnectorOutput logger = init_logger(__name__) +class SampledLogprobContractError(RuntimeError): + """The model runner returned unusable sampled-token logprobs.""" + + +def _slice_sampled_logprobs(logprobs: Any, req_index: int, sampled_token_ids: list[int]) -> Any: + """Slice and validate the sampled-token logprobs for one AR request.""" + if logprobs is None: + raise SampledLogprobContractError("AR logprobs were requested, but the model runner returned none") + + sliced = logprobs.slice_request(req_index, len(sampled_token_ids)) + token_rows = np.asarray(sliced.logprob_token_ids) + value_rows = np.asarray(sliced.logprobs) + expected_rows = len(sampled_token_ids) + + if token_rows.ndim != 2 or value_rows.ndim != 2: + raise SampledLogprobContractError( + "AR sampled-token logprobs must be rank-2 arrays, " + f"got token_ids={token_rows.shape} logprobs={value_rows.shape}" + ) + if token_rows.shape[0] != expected_rows or value_rows.shape[0] != expected_rows: + raise SampledLogprobContractError( + "AR sampled-token logprob row count does not match generated tokens: " + f"tokens={expected_rows} token_id_rows={token_rows.shape[0]} " + f"logprob_rows={value_rows.shape[0]}" + ) + if expected_rows == 0: + return sliced + if token_rows.shape[1] == 0 or value_rows.shape[1] == 0: + raise SampledLogprobContractError("AR sampled-token logprob rows are empty") + + sampled = np.asarray(sampled_token_ids) + if not np.array_equal(token_rows[:, 0], sampled): + mismatch = np.flatnonzero(token_rows[:, 0] != sampled) + first = int(mismatch[0]) + raise SampledLogprobContractError( + "AR sampled-token logprobs are misaligned: " + f"row={first} generated_token={int(sampled[first])} " + f"logprob_token={int(token_rows[first, 0])}" + ) + if not np.isfinite(value_rows[:, 0]).all(): + bad_rows = np.flatnonzero(~np.isfinite(value_rows[:, 0])).tolist() + raise SampledLogprobContractError(f"AR sampled-token logprobs contain non-finite values at rows {bad_rows}") + return sliced + + class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): """Synchronous AutoRegressive scheduler for vLLM-Omni. This class is also used as a base class for the OmniARAsyncScheduler and holds most of the @@ -362,6 +408,27 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): req_index = model_runner_output.req_id_to_index[req_id] generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] + status_before_stop = request.status + new_logprobs = None + logprob_validation_failed = False + + # Validate before mutating request token state. A bad runner output + # is request-local: terminate only this request and keep processing + # the rest of the batch. + if ( + generated_token_ids + and request.sampling_params is not None + and request.sampling_params.num_logprobs is not None + ): + try: + new_logprobs = _slice_sampled_logprobs(logprobs, req_index, generated_token_ids) + except SampledLogprobContractError as exc: + logger.error("Invalid AR sampled-token logprobs for request %s: %s", req_id, exc) + request.status = RequestStatus.FINISHED_ERROR + request.stop_reason = str(exc) + request.resumable = False + generated_token_ids = [] + logprob_validation_failed = True scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id) if scheduled_spec_token_ids and generated_token_ids: @@ -391,16 +458,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if request.has_encoder_inputs: self._free_encoder_inputs(request) - stopped = False + stopped = logprob_validation_failed is_segment_finished = False finished = False - new_logprobs = None new_token_ids = generated_token_ids pooler_output = pooler_outputs[req_index] if pooler_outputs else None mm_output = mm_outputs[req_index] if mm_outputs else None inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None kv_transfer_params = None - status_before_stop = request.status finish_reason = None routed_experts = None @@ -471,10 +536,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): 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 num_nans_in_logits is not None and req_id in num_nans_in_logits: request.num_nans_in_logits = num_nans_in_logits[req_id]