[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/engine/test_output_processor_mrv2_text.py

@@ -7,8 +7,9 @@ from unittest.mock import MagicMock
 
 import pytest
 import torch
-from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.engine import EngineCoreRequest, FinishReason
+from vllm.v1.engine.logprobs import LogprobsProcessor
 
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.output_modality import OutputModality
@@ -16,6 +17,7 @@ from vllm_omni.engine.output_processor import (
     MultimodalOutputProcessor,
     OmniRequestState,
 )
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -109,6 +111,37 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
     assert not state.mm_accumulated.is_empty
 
 
+def test_fixed_token_scores_survive_output_processing_and_stage_wrapping():
+    processor, state = _make_processor(OutputModality.TEXT, _Detokenizer())
+    request = EngineCoreRequest(
+        request_id="r",
+        prompt_token_ids=[1, 2, 3],
+        mm_features=None,
+        sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+    state.logprobs_processor = LogprobsProcessor.from_new_request(None, request)
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    processed = processor.process_outputs(
+        [
+            OmniEngineCoreOutput(
+                request_id="r",
+                new_token_ids=[42],
+                prompt_token_id_logprobs=scores,
+                finish_reason=FinishReason.STOP,
+            )
+        ]
+    )
+
+    (raw_output,) = processed.request_outputs
+    final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0)
+    torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores)
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_generation_stage_mm_delivered_only_with_terminal(streaming):
     # Non-terminal chunks accumulate; the terminal output carries the audio.