[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>

docs/configuration/stage_configs.md

@@ -35,7 +35,7 @@ Common `StagePipelineConfig` fields include:
 | `final_output` / `final_output_type` | Whether the stage produces a user-visible output and its modality. |
 | `owns_tokenizer` | Whether this stage owns the pipeline tokenizer. |
 | `model_arch`, `hf_config_name` | Stage-specific model architecture and nested HF config selector. |
-| `engine_output_type` | Runtime output representation such as `text`, `latent`, or `audio`. |
+| `engine_output_type` | Canonical runtime output type: `text`, `image`, `audio`, `latent`, or `token_ids`. |
 | `custom_process_input_func` | Processor applied to this stage's incoming payload. |
 | `custom_process_next_stage_input_func` | Processor used for full-payload handoff to the next stage. |
 | `async_chunk_process_next_stage_input_func` | Processor used for async chunk handoff. |
@@ -45,6 +45,50 @@ To add or change topology, define and register a new pipeline variant. Use
 deploy YAML only for runtime placement, resource sizing, connectors, and other
 deployment overrides.
 
+### Engine output types
+
+`engine_output_type` is parsed when the LLM stage's output processor is built
+during service initialization. Invalid values fail startup before requests
+are accepted; engine replicas may already have been initialized at that point
+and are cleaned up on failure.
+
+Use the exact lowercase names listed above. Combinations such as `text+audio`
+or `text,audio` remain supported, with every component canonical. Omitted or
+`None` values retain the default behavior (`TEXT` in the output processor).
+Empty strings, whitespace, case variants, and aliases are rejected.
+
+For existing custom pipelines or direct callers, replace old aliases explicitly:
+
+| Old value | Canonical value |
+| ------- | ------- |
+| `speech`, `wav`, `waveform` | `audio` |
+| `images`, `pixel_values`, `pixels` | `image` |
+| `latents` | `latent` |
+| `tokens` | `text` for text output, or `token_ids` for discrete IDs |
+
+`token_ids` is a separate type, not a `text` alias. It describes discrete
+vocabulary or codebook indices, distinct from continuous latent representations.
+Existing models using `latent` retain their historical payload contracts.
+GLM-Image's AR stage keeps `engine_output_type="token_ids"`; generated IDs still
+travel through `cumulative_token_ids`, and source-image IDs through
+`ids.prior_image`.
+
+Bare primary tensors and generic producer keys (`model_outputs` or `hidden`)
+use the declared type as their default output key, including `token_ids`.
+The normalizer does not inspect or convert their dtype. Distinct auxiliary
+outputs must use explicit field names: `model_outputs` and `hidden` both map
+to the same default key and must not represent two independent outputs.
+GLM-Image disables auxiliary hidden output for its default GPU execution.
+Some runner paths still attach it; that field is then named `token_ids` even
+though it contains hidden states. The GLM-Image bridge ignores this auxiliary
+field and reads generated IDs from `cumulative_token_ids` instead.
+
+Token-ID tensors use the existing dimension-0 accumulation default, subject to
+registered per-key overrides, and remain available after DELTA emissions.
+This is separate from the existing accumulation of `cumulative_token_ids`.
+`final_output_type` and Diffusion's `output_type` are separate settings and are
+unchanged by these rules.
+
 ## Deploy configuration schema
 
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:

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(

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for GLM-Image stage input processor."""
 
+from collections.abc import Mapping
 from types import SimpleNamespace
 
 import pytest
 import torch
 
+from vllm_omni.data_entry_keys import unflatten_payload
+from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
 from vllm_omni.model_executor.stage_input_processors.glm_image import (
     _first_source_image,
     _has_source_image,
@@ -15,6 +18,7 @@ from vllm_omni.model_executor.stage_input_processors.glm_image import (
     ar2diffusion,
     compute_max_tokens,
 )
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -24,7 +28,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 # =============================================================================
 
 
-def _source_output(token_ids: list[int], mm_output: dict | None = None):
+def _source_output(token_ids: list[int], mm_output: Mapping | None = None):
     """Create a minimal AR output mock."""
     return SimpleNamespace(
         outputs=[SimpleNamespace(token_ids=token_ids, cumulative_token_ids=token_ids)],
@@ -276,6 +280,84 @@ class TestParseGeneratedTokens:
 
 
 class TestAr2Diffusion:
+    @pytest.mark.parametrize("include_hidden", [False, True], ids=["tokens_only", "with_generic_hidden"])
+    def test_t2i_token_id_output_preserves_cumulative_tokens(self, include_hidden):
+        """DiT consumes the cumulative token sequence, not the generic tensor payload."""
+        output_type = GLM_IMAGE_PIPELINE.stages[0].engine_output_type
+        assert output_type == "token_ids"
+
+        # Default GPU execution omits hidden; some runner paths still attach it.
+        # The shared normalizer renames it without changing the tensor's values.
+        mm_output = None
+        if include_hidden:
+            hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+            mm_output = MultimodalPayload.from_raw({"hidden": hidden}, output_type)
+            assert mm_output is not None
+            assert set(mm_output) == {"token_ids"}
+            torch.testing.assert_close(mm_output["token_ids"], hidden)
+
+        # At 64x64 the AR sequence contains a 256-token preview, four target
+        # image tokens, and EOS. The final emitted delta contains only EOS.
+        token_ids = [9] * 256 + [1, 2, 3, 4] + [16385]
+        source_output = _source_output(token_ids, mm_output)
+        source_output.outputs[0].token_ids = [16385]
+        prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 64, "target_w": 64}}
+
+        result = ar2diffusion([source_output], prompt=prompt)
+
+        assert result is not None
+        assert (result["prompt"], result["height"], result["width"]) == ("a cat", 64, 64)
+        torch.testing.assert_close(
+            result["extra"]["prior_token_ids"],
+            torch.tensor([1, 1, 2, 2, 1, 1, 2, 2, 3, 3, 4, 4, 3, 3, 4, 4]),
+        )
+        assert result["extra"]["prior_token_image_ids"] is None
+
+    @pytest.mark.parametrize("include_hidden", [False, True], ids=["explicit_ids_only", "with_generic_hidden"])
+    def test_i2i_token_id_output_preserves_source_image_tokens(self, include_hidden):
+        """Explicit image IDs survive normalization with or without hidden output."""
+        from PIL import Image
+
+        output_type = GLM_IMAGE_PIPELINE.stages[0].engine_output_type
+        assert output_type == "token_ids"
+
+        prior_image_ids = torch.tensor([5, 6, 7, 8])
+        raw_output = {"ids.prior_image": prior_image_ids}
+        hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        if include_hidden:
+            raw_output["hidden"] = hidden
+        mm_output = MultimodalPayload.from_raw(raw_output, output_type)
+        assert mm_output is not None
+        # The output processor restores dotted runner keys before emitting
+        # the payload consumed by the stage input processor.
+        mm_output = MultimodalPayload.from_dict(unflatten_payload(mm_output.to_dict()))
+        assert mm_output is not None
+        assert set(mm_output) == ({"ids", "token_ids"} if include_hidden else {"ids"})
+        torch.testing.assert_close(mm_output["ids"]["prior_image"], prior_image_ids)
+        if include_hidden:
+            torch.testing.assert_close(mm_output["token_ids"], hidden)
+
+        source_output = _source_output([1, 2, 3, 4, 16385], mm_output)
+        source_output.outputs[0].token_ids = [16385]
+        image = Image.new("RGB", (64, 64))
+        prompt = {
+            "prompt": "edit this",
+            "mm_processor_kwargs": {"target_h": 64, "target_w": 64},
+            "multi_modal_data": {"image": image},
+        }
+
+        result = ar2diffusion([source_output], prompt=prompt, requires_multimodal_data=True)
+
+        assert result is not None
+        assert (result["prompt"], result["height"], result["width"]) == ("edit this", 64, 64)
+        assert result["pil_image"] is image
+        torch.testing.assert_close(
+            result["extra"]["prior_token_ids"],
+            torch.tensor([1, 1, 2, 2, 1, 1, 2, 2, 3, 3, 4, 4, 3, 3, 4, 4]),
+        )
+        assert len(result["extra"]["prior_token_image_ids"]) == 1
+        torch.testing.assert_close(result["extra"]["prior_token_image_ids"][0], prior_image_ids)
+
     def test_basic_t2i(self):
         """Test basic text-to-image pipeline: AR -> Diffusion."""
         # 1024x1024 t2i: small(256) + large(1024) + EOS

vllm_omni/config/model.py

@@ -105,7 +105,7 @@ class OmniModelConfig(ModelConfig):
          worker_type: Model Type, e.g., "ar" or "generation"
          engine_output_type: Optional output type specification for the engine.
              Used to route outputs to appropriate processors (e.g., "image",
-             "audio", "latents"). If None, output type is inferred.
+             "audio", "latent", "token_ids"). If None, output type is inferred.
          stage_connector_config: Stage connector configuration dictionary.
              Contains "name" (connector name), "extra" (extra connector config).
          task_type: Model-defined startup task type. Each model validates its

vllm_omni/config/stage_config.py

@@ -1050,7 +1050,7 @@ def _build_engine_args(
     engine_args["retains_state_across_chunks"] = ps.retains_state_across_chunks
     if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch:
         engine_args.setdefault("model_class_name", ps.model_arch)
-    if ps.engine_output_type:
+    if ps.engine_output_type is not None:
         engine_args["engine_output_type"] = ps.engine_output_type
     if next_stage_proc:
         engine_args["custom_process_next_stage_input_func"] = next_stage_proc

vllm_omni/engine/arg_utils.py

@@ -156,7 +156,7 @@ class OmniEngineArgs(EngineArgs):
             (default: "Qwen2_5OmniForConditionalGeneration")
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
-            "audio", "latents"). If None, output type is inferred.
+            "audio", "latent", "token_ids"). If None, output type is inferred.
         hf_config_name: Optional key for HF config subkey to be extracted
             for this stage, e.g., talker_config; If None, the default
             HF config will be used.

vllm_omni/outputs/output_modality.py

@@ -10,21 +10,8 @@ for type-safe multimodal output routing and tensor merging.
 
 from __future__ import annotations
 
-import re
 from enum import Enum, Flag, auto
 
-_MODALITY_ALIASES: dict[str, str] = {
-    "speech": "audio",
-    "images": "image",
-    "latents": "latent",
-    "wav": "audio",
-    "waveform": "audio",
-    "pixel_values": "image",
-    "pixels": "image",
-    "token_ids": "text",
-    "tokens": "text",
-}
-
 
 class OutputModalityNames(str, Enum):
     """Keys for output modalities.
@@ -41,13 +28,17 @@ class OutputModalityNames(str, Enum):
     IMAGE = "image"
     AUDIO = "audio"
     LATENT = "latent"
+    TOKEN_IDS = "token_ids"
 
 
 # Specify which output modalities may be drained when handling delta messages.
-# For some types, e.g., latents, we need to be careful to ensure the full context
-# is passed as the stream yields due to assumptions in the I/O processing and model
-# when async chunk isn't enabled.
-NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT}
+# Retain intermediate outputs such as latents and token IDs across emissions
+# so downstream stages can receive full context when async chunk isn't enabled.
+NON_DRAINABLE_MODALITIES = {
+    OutputModalityNames.TEXT,
+    OutputModalityNames.LATENT,
+    OutputModalityNames.TOKEN_IDS,
+}
 DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES}
 
 
@@ -59,6 +50,10 @@ class OutputModality(Flag):
     Single:   ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ...
     Compound: ``OutputModality.TEXT | OutputModality.IMAGE``  (text+image)
 
+    TODO: Describe modality, representation, and update/retention rules per
+    named output field instead of sharing one stage-level type. Audit legacy
+    LATENT uses when introducing that contract.
+
     Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's
     native path (pooling_output → PoolingRequestOutput), handled entirely
     by the base OutputProcessor. vLLM-Omni's layer does not participate.
@@ -68,29 +63,43 @@ class OutputModality(Flag):
     IMAGE = auto()
     AUDIO = auto()
     LATENT = auto()
+    """Continuous latent tensors used for further computation or decoding.
+
+    Existing LATENT producers may carry other representations; their current
+    payload contracts remain unchanged.
+    """
+
+    TOKEN_IDS = auto()
+    """Discrete vocabulary/codebook indices, e.g. GLM-Image's AR image tokens.
+
+    Unlike continuous latents, these values identify discrete symbols.
+    The distinction is semantic, not inferred from the tensor dtype.
+    """
 
     @classmethod
     def from_string(cls, s: str | None) -> OutputModality:
-        """Parse a free-text modality string into an OutputModality flag.
+        """Parse canonical modality names into an OutputModality flag.
 
-        Handles common aliases and compound strings separated by + or ,.
+        Names must be lowercase and contain no whitespace. Compound names
+        may be separated by + or ,. None preserves the default TEXT modality.
 
         Examples::
 
             OutputModality.from_string("text+image")
             # → OutputModality.TEXT | OutputModality.IMAGE
         """
-        if not s or not s.strip():
+        if s is None:
             return cls.TEXT
 
-        parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())]
         result = cls(0)
-        for p in parts:
-            p = _MODALITY_ALIASES.get(p, p)
+        for part in s.replace(",", "+").split("+"):
             try:
-                result |= cls[p.upper()]
-            except KeyError:
-                raise ValueError(f"Unknown modality: {p!r}. Supported: {[(m.name or '').lower() for m in cls]}")
+                name = OutputModalityNames(part)
+            except ValueError:
+                raise ValueError(
+                    f"Unknown modality: {part!r}. Supported: {[m.value for m in OutputModalityNames]}"
+                ) from None
+            result |= cls[name.name]
         return result
 
     @property
@@ -106,7 +115,7 @@ class TensorAccumulationStrategy(Enum):
     """Strategy for merging incremental multimodal tensors."""
 
     CONCAT_DIM0 = "concat_dim0"
-    """Concatenate along dimension 0. Used for image/latent tensors."""
+    """Concatenate along dimension 0. Used for image/latent/token-ID tensors."""
 
     CONCAT_LAST = "concat_last"
     """Concatenate along the last dimension. Used for audio waveforms."""
@@ -166,4 +175,4 @@ def get_accumulation_strategy(
         return TensorAccumulationStrategy.CONCAT_LAST
     if OutputModality.IMAGE in modality or OutputModality.LATENT in modality:
         return TensorAccumulationStrategy.CONCAT_DIM0
-    return TensorAccumulationStrategy.CONCAT_DIM0  # default
+    return TensorAccumulationStrategy.CONCAT_DIM0  # TEXT / TOKEN_IDS default

vllm_omni/outputs/output_processor.py

@@ -46,6 +46,8 @@ def _modality_to_type_string(modality: OutputModality) -> str:
             return "image"
         if OutputModality.LATENT in modality:
             return "latent"
+        if OutputModality.TOKEN_IDS in modality:
+            return "token_ids"
     except TypeError:
         # Flag identity mismatch (e.g. after module reload in tests).
         name = getattr(modality, "name", "") or ""
@@ -56,6 +58,8 @@ def _modality_to_type_string(modality: OutputModality) -> str:
             return "image"
         if "latent" in lowered:
             return "latent"
+        if "token_ids" in lowered:
+            return "token_ids"
     return "text"
 
 
@@ -410,7 +414,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             log_stats: Whether to log statistics
             stream_interval: Stream interval for output generation
             engine_core_output_type: Optional output type string (e.g.,
-                "image", "audio", "latent"). Converted to OutputModality
+                "image", "audio", "latent", "token_ids"). Converted to OutputModality
                 internally. Kept for backward compatibility with
                 stage_init_utils.
             output_modality: Type-safe output modality flag. Used to tag