[bugfix] validate AR sampled-token logprob alignment (#5273)

Signed-off-by: hbhflw2000 <417911774@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>

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])

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]