[Core] Add token_ids output type and enforce canonical names (#7843)

Signed-off-by: Jungle430 <junglece430@gmail.com>
Co-authored-by: Codex <noreply@openai.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -184,7 +184,7 @@ def _make_llm_metadata(
         default_sampling_params=types.SimpleNamespace(name=f"sp-{stage_id}-{replica_id}"),
         custom_process_input_func=None,
         engine_input_source=[] if stage_id == 0 else [stage_id - 1],
-        engine_output_type="token_ids",
+        engine_output_type="text",
         replica_id=replica_id,
         is_comprehension=is_comprehension,
     )
@@ -1330,6 +1330,32 @@ def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_e
     ]
 
 
+@pytest.mark.parametrize("engine_output_type", ["", "tokens", "text+speech", "AUDIO"])
+def test_stage_runtime_rejects_invalid_output_type_during_initialization(mocker, engine_output_type):
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True))
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg)
+    stage_plan.replicas[0].metadata.engine_output_type = engine_output_type
+    stage_client = mocker.Mock()
+    mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan])
+    initialize_replicas = mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [stage_client]})
+
+    # Keep output-processor construction real: no request is needed to parse
+    # the type, even though the engine replica is already initialized here.
+    with pytest.raises(ValueError, match="Unknown modality"):
+        runtime.initialize()
+
+    initialize_replicas.assert_called_once()
+    stage_client.shutdown.assert_called_once_with()
+    assert runtime.stage_pools == []
+
+
 def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/engine/test_output_modality.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for Phase 1 foundation types (RFC #1601).
 
 Note: Uses importlib to load modules directly, bypassing the vllm_omni
@@ -7,10 +9,14 @@ package __init__ which requires the vllm base package.
 import importlib.util
 import sys
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 import torch
 
+if TYPE_CHECKING:
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -19,6 +25,7 @@ _OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs"
 
 def _load_module(name: str, filepath: Path):
     spec = importlib.util.spec_from_file_location(name, filepath)
+    assert spec is not None and spec.loader is not None
     mod = importlib.util.module_from_spec(spec)
     sys.modules[name] = mod
     spec.loader.exec_module(mod)
@@ -37,45 +44,155 @@ _mm_mod = _load_module(
     "vllm_omni.outputs.mm_outputs",
     _OUTPUTS_DIR / "mm_outputs.py",
 )
+_accumulation_mod = _load_module(
+    "vllm_omni.outputs.multimodal_accumulation",
+    _OUTPUTS_DIR / "multimodal_accumulation.py",
+)
 
 OutputModality = _om_mod.OutputModality
 TensorAccumulationStrategy = _om_mod.TensorAccumulationStrategy
 get_accumulation_strategy = _om_mod.get_accumulation_strategy
 register_key_accumulation_strategy = _om_mod.register_key_accumulation_strategy
-MultimodalPayload = _mm_mod.MultimodalPayload
+if not TYPE_CHECKING:
+    MultimodalPayload = _mm_mod.MultimodalPayload
 MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput
 
 
-def test_output_modality_parsing_and_flags():
-    """Test OutputModality enum: from_string, aliases, compounds, properties, and accumulation strategy."""
-    # Defaults
-    assert OutputModality.from_string(None) == OutputModality.TEXT
-    assert OutputModality.from_string("") == OutputModality.TEXT
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (None, OutputModality.TEXT),
+        ("text", OutputModality.TEXT),
+        ("image", OutputModality.IMAGE),
+        ("audio", OutputModality.AUDIO),
+        ("latent", OutputModality.LATENT),
+        ("token_ids", OutputModality.TOKEN_IDS),
+        ("text+token_ids", OutputModality.TEXT | OutputModality.TOKEN_IDS),
+        ("token_ids,latent", OutputModality.TOKEN_IDS | OutputModality.LATENT),
+        ("text+image", OutputModality.TEXT | OutputModality.IMAGE),
+        ("text,audio", OutputModality.TEXT | OutputModality.AUDIO),
+        ("latent+image,audio", OutputModality.LATENT | OutputModality.IMAGE | OutputModality.AUDIO),
+        ("image+text", OutputModality.TEXT | OutputModality.IMAGE),
+        ("audio+audio", OutputModality.AUDIO),
+    ],
+)
+def test_output_modality_parses_canonical_names(value, expected):
+    assert OutputModality.from_string(value) == expected
 
-    # Direct names and case insensitivity
-    assert OutputModality.from_string("image") == OutputModality.IMAGE
-    assert OutputModality.from_string("Audio") == OutputModality.AUDIO
 
-    # Aliases
-    assert OutputModality.from_string("speech") == OutputModality.AUDIO
-    assert OutputModality.from_string("latents") == OutputModality.LATENT
-    assert OutputModality.from_string("pixel_values") == OutputModality.IMAGE
+@pytest.mark.parametrize(
+    "value",
+    [
+        "speech",
+        "images",
+        "latents",
+        "wav",
+        "waveform",
+        "pixel_values",
+        "pixels",
+        "tokens",
+        "text+speech",
+        "token_ids+tokens",
+        "TOKEN_IDS",
+        "token_ids ",
+        "Audio",
+        "TEXT",
+        "text+Image",
+        " audio",
+        "audio ",
+        "text +audio",
+        "text, audio",
+        "text\taudio",
+        "audio\n",
+        "",
+        " ",
+        "+audio",
+        "audio,",
+        "text++audio",
+        "text,,audio",
+        "text+,audio",
+        "video",
+        "text+unknown",
+    ],
+)
+def test_output_modality_rejects_noncanonical_names(value):
+    with pytest.raises(ValueError, match="Unknown modality"):
+        OutputModality.from_string(value)
 
-    # Compound
+
+def test_output_modality_flags_and_accumulation_strategy():
     compound = OutputModality.from_string("text+image")
     assert compound.has_text and compound.has_multimodal
 
     # Flag properties
     assert OutputModality.TEXT.has_text and not OutputModality.TEXT.has_multimodal
     assert OutputModality.IMAGE.has_multimodal and not OutputModality.IMAGE.has_text
+    assert OutputModality.TOKEN_IDS.has_multimodal and not OutputModality.TOKEN_IDS.has_text
+    assert OutputModality.TOKEN_IDS != OutputModality.LATENT
 
     # Accumulation strategy
     assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST
     assert get_accumulation_strategy(OutputModality.IMAGE) == TensorAccumulationStrategy.CONCAT_DIM0
+    assert get_accumulation_strategy(OutputModality.TOKEN_IDS) == TensorAccumulationStrategy.CONCAT_DIM0
 
-    # Unknown raises
-    with pytest.raises(ValueError, match="Unknown modality"):
-        OutputModality.from_string("video")
+
+@pytest.mark.parametrize("producer_key", [None, "model_outputs", "hidden", "token_ids"])
+def test_token_id_payload_preserves_discrete_primary_output(producer_key):
+    token_ids = torch.tensor([[12, 85], [206, 7]], dtype=torch.long)
+    raw = token_ids if producer_key is None else {producer_key: token_ids}
+
+    payload = MultimodalPayload.from_raw(raw, "token_ids")
+
+    assert payload is not None
+    assert set(payload) == {"token_ids"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+
+
+def test_token_id_payload_keeps_explicit_latents_and_image_ids_separate():
+    token_ids = torch.tensor([12, 85, 206], dtype=torch.long)
+    latent = torch.tensor([[0.1, 0.2], [0.3, 0.4]])
+    prior_image = torch.tensor([4, 5], dtype=torch.long)
+
+    payload = MultimodalPayload.from_raw(
+        {"model_outputs": token_ids, "latent": latent, "ids.prior_image": prior_image},
+        "token_ids",
+    )
+
+    assert payload is not None
+    assert set(payload) == {"token_ids", "latent", "ids.prior_image"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+    torch.testing.assert_close(payload["latent"], latent)
+    torch.testing.assert_close(payload["ids.prior_image"], prior_image)
+
+
+def test_token_id_payload_consolidates_dim0_and_respects_key_overrides():
+    reference_key = "__test_token_ids__.reference"
+    register_key_accumulation_strategy(reference_key, TensorAccumulationStrategy.REPLACE)
+    first_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
+    last_ids = torch.tensor([[5, 6]], dtype=torch.long)
+    first_ref = torch.tensor([[7, 8]], dtype=torch.long)
+    last_ref = torch.tensor([[9, 10]], dtype=torch.long)
+    payload = MultimodalPayload()
+    payload.tensors["token_ids"] = [first_ids, last_ids]
+    payload.tensors[reference_key] = [first_ref, last_ref]
+
+    payload.consolidate_tensors(OutputModality.TOKEN_IDS)
+
+    torch.testing.assert_close(payload["token_ids"], torch.cat([first_ids, last_ids], dim=0))
+    torch.testing.assert_close(payload[reference_key], last_ref)
+
+
+def test_delta_drain_retains_token_ids_and_latents():
+    token_ids = torch.tensor([12, 85], dtype=torch.long)
+    latent = torch.tensor([[0.1, 0.2]])
+    payload = MultimodalPayload.from_dict({"token_ids": token_ids, "latent": latent, "audio": torch.ones(4)})
+    assert payload is not None
+
+    _accumulation_mod.drain_delta_payload(payload)
+
+    assert set(payload) == {"token_ids", "latent"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+    torch.testing.assert_close(payload["latent"], latent)
 
 
 def test_multimodal_payload_and_completion_output():
@@ -107,7 +224,7 @@ def test_multimodal_payload_and_completion_output():
 def test_output_modality_printed_examples(capsys):
     """Printed examples for output modality types."""
     print("\n=== OutputModality Parsing ===")
-    for s in [None, "", "image", "Audio", "speech", "latents", "pixel_values", "text+image"]:
+    for s in [None, "text", "image", "audio", "latent", "text+image", "text,audio"]:
         print(f"  from_string({s!r:20s}) -> {OutputModality.from_string(s)}")
 
     print("\n=== Flag Properties ===")

tests/engine/test_output_processor.py

@@ -3,6 +3,7 @@
 """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
 
 from dataclasses import dataclass
+from enum import Flag
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -81,6 +82,26 @@ def test_output_modality_name_uses_string_value():
     assert str(OutputModalityNames.AUDIO) == "audio"
 
 
+@pytest.mark.parametrize(
+    ("name", "expected"),
+    [
+        ("TOKEN_IDS", "token_ids"),
+        ("TEXT|TOKEN_IDS", "token_ids"),
+        ("LATENT|TOKEN_IDS", "latent"),
+        ("IMAGE|TOKEN_IDS", "image"),
+        ("AUDIO|TOKEN_IDS", "audio"),
+    ],
+)
+def test_modality_type_string_handles_reloaded_token_id_enum(name, expected):
+    # A separate Flag class reproduces enum identity changes after module reload.
+    ReloadedOutputModality = Flag("ReloadedOutputModality", ["TEXT", "IMAGE", "AUDIO", "LATENT", "TOKEN_IDS"])
+    value = ReloadedOutputModality(0)
+    for member in name.split("|"):
+        value |= ReloadedOutputModality[member]
+
+    assert output_processor._modality_to_type_string(value) == expected
+
+
 def test_init_empty_dict():
     """Ensure mm_accumulated is initially empty."""
     assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
@@ -574,6 +595,92 @@ def _make_no_detok_state(output_kind: RequestOutputKind):
     return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
 
 
+def test_token_id_delta_retains_ids_and_latents_until_final_consolidation(mocker):
+    state = _make_state(RequestOutputKind.DELTA)
+    state.detokenizer = mocker.MagicMock(
+        output_token_ids=[10],
+        get_next_output_text=mocker.MagicMock(return_value=""),
+        num_output_tokens=mocker.MagicMock(return_value=1),
+    )
+    first_ids = torch.tensor([[1, 2]], dtype=torch.long)
+    last_ids = torch.tensor([[3, 4]], dtype=torch.long)
+    first_latent = torch.tensor([[0.1, 0.2, 0.3]])
+    last_latent = torch.tensor([[0.4, 0.5, 0.6]])
+    state.add_multimodal_tensor({"model_outputs": first_ids, "latent": first_latent}, mm_type="token_ids")
+
+    first = state.make_request_output([10], None, None, None)
+
+    assert first is not None and not isinstance(first, PoolingRequestOutput)
+    first_completion = first.outputs[0]
+    assert set(state.mm_accumulated) == {"token_ids", "latent"}
+    torch.testing.assert_close(first_completion.multimodal_output["token_ids"], first_ids)
+    torch.testing.assert_close(first_completion.multimodal_output["latent"], first_latent)
+    assert first_completion.cumulative_token_ids == [10]
+
+    state.detokenizer.output_token_ids.append(11)
+    state.add_multimodal_tensor({"model_outputs": last_ids, "latent": last_latent}, mm_type="token_ids")
+    final = state.make_request_output([11], None, FinishReason.STOP, None)
+
+    assert final is not None and not isinstance(final, PoolingRequestOutput)
+    completion = final.outputs[0]
+    assert set(state.mm_accumulated) == {"token_ids", "latent"}
+    torch.testing.assert_close(completion.multimodal_output["token_ids"], torch.cat([first_ids, last_ids], dim=0))
+    torch.testing.assert_close(completion.multimodal_output["latent"], torch.cat([first_latent, last_latent], dim=0))
+    # Completion tokens remain on their own cumulative path, independent of
+    # tensor payload IDs; the earlier completion keeps its original snapshot.
+    assert completion.cumulative_token_ids == [10, 11]
+    assert first_completion.cumulative_token_ids == [10]
+
+
+@pytest.mark.parametrize(
+    ("engine_output_type", "expected_key", "expected_values"),
+    [
+        (None, "text", [[0, 1, 2], [3, 4, 5]]),
+        ("text", "text", [[0, 1, 2], [3, 4, 5]]),
+        ("image", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("latent", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("token_ids", "token_ids", [[0, 1, 2], [3, 4, 5]]),
+        ("audio", "audio", [[0, 1, 2, 3, 4, 5]]),
+        ("text+image", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("text,latent", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("text+audio", "audio", [[0, 1, 2, 3, 4, 5]]),
+        ("text+token_ids", "token_ids", [[0, 1, 2], [3, 4, 5]]),
+        ("latent,token_ids", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("image+token_ids", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("audio+token_ids", "audio", [[0, 1, 2, 3, 4, 5]]),
+    ],
+)
+def test_engine_output_type_controls_payload_key_and_concatenation(engine_output_type, expected_key, expected_values):
+    """Canonical config values retain output naming and tensor merge behavior."""
+    processor = MultimodalOutputProcessor(
+        tokenizer=None,
+        log_stats=False,
+        engine_core_output_type=engine_output_type,
+    )
+    state = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id].append(state.request_id)
+
+    first = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[10],
+        multimodal_output={"model_outputs": torch.tensor([[0, 1, 2]])},
+    )
+    assert not processor.process_outputs([first]).request_outputs
+
+    last = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[11],
+        finish_reason=FinishReason.STOP,
+        multimodal_output={"model_outputs": torch.tensor([[3, 4, 5]])},
+    )
+    result = processor.process_outputs([last])
+    assert len(result.request_outputs) == 1
+    payload = result.request_outputs[0].outputs[0].multimodal_output
+    assert set(payload) == {expected_key}
+    torch.testing.assert_close(payload[expected_key], torch.tensor(expected_values))
+
+
 def test_no_detokenizer_completion_output():
     """_new_completion_output works when detokenizer is None (generation stages)."""
     s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)

tests/engine/test_stage_engine_args.py

@@ -226,7 +226,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
                 execution_type=StageExecutionType.LLM_AR,
                 requires_multimodal_data=True,
                 hf_config_name="thinker_config",
-                engine_output_type="hidden_states",
+                engine_output_type="latent",
                 model_arch="Qwen3OmniMoeForConditionalGeneration",
                 model_subdir="ar-model",
                 tokenizer_subdir="ar-tokenizer",
@@ -238,7 +238,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
                 model_stage="talker",
                 execution_type=StageExecutionType.LLM_GENERATION,
                 input_sources=(0,),
-                engine_output_type="audio_tokens",
+                engine_output_type="latent",
             ),
             StagePipelineConfig(
                 stage_id=2,

tests/engine/test_stage_metadata.py

@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import copy
 import operator
+from dataclasses import replace
 
 import pytest
 from vllm.sampling_params import SamplingParams
@@ -33,7 +37,7 @@ def _metadata_inputs() -> tuple[PipelineConfig, DeployConfig]:
                 execution_type=StageExecutionType.LLM_AR,
                 owns_tokenizer=True,
                 requires_multimodal_data=True,
-                engine_output_type="token_ids",
+                engine_output_type="text",
                 custom_process_input_func="operator.add",
                 prompt_expand_func="operator.mul",
             ),
@@ -42,7 +46,7 @@ def _metadata_inputs() -> tuple[PipelineConfig, DeployConfig]:
                 model_stage="talker",
                 execution_type=StageExecutionType.LLM_GENERATION,
                 input_sources=(0,),
-                engine_output_type="audio_tokens",
+                engine_output_type="latent",
                 custom_process_input_func="operator.sub",
                 sync_process_input_func="operator.floordiv",
             ),
@@ -160,6 +164,26 @@ def test_extract_stage_metadata_preserves_legacy_one_argument_api():
     assert metadata.custom_process_input_func is operator.add
 
 
+@pytest.mark.parametrize("engine_output_type", [None, "", "latent", "token_ids", "text+token_ids", "text+audio"])
+def test_extract_stage_metadata_preserves_declared_engine_output_type(engine_output_type):
+    pipeline, deploy = _metadata_inputs()
+    pipeline = replace(
+        pipeline,
+        stages=(replace(pipeline.stages[0], engine_output_type=engine_output_type), *pipeline.stages[1:]),
+    )
+    omni_config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=copy.deepcopy(deploy))
+    legacy_config = merge_pipeline_deploy(pipeline, deploy)[0].to_omegaconf()
+
+    # An explicit empty value must reach the output processor, where it is
+    # rejected during startup, instead of silently becoming the None default.
+    assert legacy_config.engine_args.get("engine_output_type") == engine_output_type
+    assert extract_legacy_stage_metadata(legacy_config).engine_output_type == engine_output_type
+    assert (
+        extract_stage_metadata_from_omni_stage_config(omni_config.stage_by_id(0)).engine_output_type
+        == engine_output_type
+    )
+
+
 def test_extract_stage_metadata_defaults_missing_engine_input_source():
     pipeline, deploy = _metadata_inputs()
     legacy_config = merge_pipeline_deploy(