diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 56447264a..f9b08ca15 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock +import numpy as np import pytest import torch from vllm.outputs import PoolingRequestOutput @@ -250,6 +251,25 @@ def test_no_detokenizer_make_request_output(): assert AUDIO in result.outputs[0].multimodal_output +def test_no_detokenizer_make_request_output_with_routed_experts(): + """make_request_output accepts the routed_experts arg that the multimodal + output channel (_process_mm_only_outputs) passes positionally, and attaches + it to the completion output. + + Regression: the call site passes 6 positional args + (..., kv_transfer_params, routed_experts) but the override previously took + only 5, raising TypeError on every generation-stage output. + """ + s = _make_no_detok_state(RequestOutputKind.CUMULATIVE) + s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO) + routed_experts = np.zeros((2, 3), dtype=np.int32) + # Mirror the exact call shape of _process_mm_only_outputs. + result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts) + assert result is not None + assert not isinstance(result, PoolingRequestOutput) + assert result.outputs[0].routed_experts is routed_experts + + def test_no_detokenizer_stream_interval_skipped(): """stream_interval > 1 does not crash when detokenizer is None.""" kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5} diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index b84a9eb3a..c9044e29e 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -250,6 +250,7 @@ class OmniRequestState(RequestState): finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, + routed_experts: Any = None, ) -> OmniRequestOutput | PoolingRequestOutput | None: """Create a request output from generation results. @@ -263,6 +264,8 @@ class OmniRequestState(RequestState): finish_reason: Optional finish reason indicating why generation stopped stop_reason: Optional stop reason (token ID or stop string) kv_transfer_params: Optional KV cache transfer parameters + routed_experts: Optional MoE routed-expert ids for this step, + attached to the completion output for generation stages Returns: OmniRequestOutput or PoolingRequestOutput if output should be @@ -310,7 +313,7 @@ class OmniRequestState(RequestState): external_req_id = self.external_req_id - output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) if self.parent_req is None: outputs = [output] @@ -332,6 +335,7 @@ class OmniRequestState(RequestState): token_ids: list[int], finish_reason: FinishReason | None, stop_reason: int | str | None, + routed_experts: Any = None, ) -> MultimodalCompletionOutput | CompletionOutput: """Create a completion output with multimodal data attached. @@ -352,6 +356,7 @@ class OmniRequestState(RequestState): cumulative_logprob=None, finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, + routed_experts=routed_experts, ) else: base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)