diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 57e4721b0..c2c0eb733 100644 --- a/docs/configuration/stage_configs.md +++ b/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: diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 1a92c7b6e..befa22857 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index e5c5d36c1..64020b633 100644 --- a/tests/engine/test_output_modality.py +++ b/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 ===") diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 9fbecc1c1..cc54f07b6 100644 --- a/tests/engine/test_output_processor.py +++ b/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) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 9b8b70586..bdf0150b0 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/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, diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py index 0b21f7573..697410569 100644 --- a/tests/engine/test_stage_metadata.py +++ b/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( diff --git a/tests/model_executor/stage_input_processors/test_glm_image.py b/tests/model_executor/stage_input_processors/test_glm_image.py index 8dc333fb6..2fef42cc0 100644 --- a/tests/model_executor/stage_input_processors/test_glm_image.py +++ b/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 diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 13115f00e..3c43b53a5 100644 --- a/vllm_omni/config/model.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 167664b37..b243ea882 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a46fc49e2..4bea0e05b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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. diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 1e8489e99..0a9cdc09c 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/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 diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index cac8d9990..ca09e93f4 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/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