diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index ff15e5824..e5c5d36c1 100644 --- a/tests/engine/test_output_modality.py +++ b/tests/engine/test_output_modality.py @@ -41,6 +41,7 @@ _mm_mod = _load_module( 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 MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput @@ -156,3 +157,90 @@ def test_output_modality_printed_examples(capsys): captured = capsys.readouterr() assert "OutputModality Parsing" in captured.out assert "MultimodalPayload" in captured.out + + +def test_get_accumulation_strategy_key_override_precedence(): + """A registered per-key override wins over the modality-wide default; + any other key under the same modality keeps using that default.""" + assert get_accumulation_strategy(OutputModality.AUDIO, "unregistered.key") == TensorAccumulationStrategy.CONCAT_LAST + + register_key_accumulation_strategy("__test_override__.frames", TensorAccumulationStrategy.CONCAT_DIM0) + assert ( + get_accumulation_strategy(OutputModality.AUDIO, "__test_override__.frames") + == TensorAccumulationStrategy.CONCAT_DIM0 + ) + # A plain waveform key under the same AUDIO modality is unaffected. + assert get_accumulation_strategy(OutputModality.AUDIO, "waveform") == TensorAccumulationStrategy.CONCAT_LAST + assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST + + +def _talker_codec_frame_payload(audio_key: str, ref_key: str) -> MultimodalPayload: + """Synthetic stand-in for a TTS talker's accumulated per-step output. + + ``audio_key`` grows one ``[1, num_codebooks]`` row per decode step on top + of a ``[5, num_codebooks]`` prefill block (Qwen3-TTS's ``codes.audio``); + ``ref_key`` re-emits the same constant ``[100, num_codebooks]`` + reference-context matrix unchanged at every step (Qwen3-TTS's + ``codes.ref``). Neither is a waveform chunk. + """ + payload = MultimodalPayload() + payload.tensors[audio_key] = [ + torch.zeros(5, 16), + torch.full((1, 16), 1.0), + torch.full((1, 16), 2.0), + torch.full((1, 16), 3.0), + ] + ref_block = torch.arange(100 * 16, dtype=torch.float32).reshape(100, 16) + payload.tensors[ref_key] = [ref_block.clone(), ref_block.clone(), ref_block.clone()] + return payload + + +def test_consolidate_tensors_default_audio_strategy_corrupts_codec_frame_keys(): + """Reproduces both failure modes of forcing codec-frame keys through the + AUDIO modality's waveform-tuned default (CONCAT_LAST) with no per-key + override registered: the audio-frames key's dim-0 growth makes + CONCAT_LAST's dim=-1 concat raise (mismatched dim 0), which the + flatten-chunks fallback turns into a 1-D blob instead of the intended + ``[8, 16]``; the ref-frames key's matching last dim lets CONCAT_LAST + "succeed" by silently concatenating three identical ``[100, 16]`` blocks + into ``[100, 48]`` instead of keeping a single copy. + """ + payload = _talker_codec_frame_payload("unregistered.codes.audio", "unregistered.codes.ref") + + payload.consolidate_tensors(OutputModality.AUDIO) + + assert tuple(payload.tensors["unregistered.codes.audio"].shape) != (8, 16) + assert tuple(payload.tensors["unregistered.codes.ref"].shape) == (100, 48) + + +def test_consolidate_tensors_with_qwen3_tts_key_overrides_is_correct(): + """With the per-key overrides Qwen3-TTS's pipeline module registers + (``codes.audio`` -> CONCAT_DIM0, ``codes.ref`` -> REPLACE), both keys + consolidate correctly under the AUDIO modality despite not being + waveform chunks.""" + register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0) + register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE) + + payload = _talker_codec_frame_payload("codes.audio", "codes.ref") + expected_ref = payload.tensors["codes.ref"][0] + + payload.consolidate_tensors(OutputModality.AUDIO) + + assert tuple(payload.tensors["codes.audio"].shape) == (8, 16) + assert tuple(payload.tensors["codes.ref"].shape) == (100, 16) + assert torch.equal(payload.tensors["codes.ref"], expected_ref) + + +def test_consolidate_tensors_raises_with_key_name_instead_of_silently_keeping_last(): + """A concat failure under a non-CONCAT_LAST strategy must be surfaced + with the offending key name rather than resolved by silently keeping + only the last chunk. Previously the except-branch only special-cased + the literal key ``"audio"``, so any other mismatched key (e.g. a + codec-frame key routed to CONCAT_DIM0) was corrupted instead of + failing loudly. + """ + payload = MultimodalPayload() + payload.tensors["latent.frames"] = [torch.zeros(2, 8), torch.zeros(3, 9)] + + with pytest.raises(RuntimeError, match="latent.frames"): + payload.consolidate_tensors(OutputModality.LATENT) diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 6dfe8736d..dd33aeca8 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -10,9 +10,23 @@ from vllm_omni.config.stage_config import ( StageExecutionType, StagePipelineConfig, ) +from vllm_omni.outputs.output_modality import ( + TensorAccumulationStrategy, + register_key_accumulation_strategy, +) _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" +# The talker's codec-frame keys are not waveform chunks, so they must not go +# through the AUDIO modality's waveform-tuned default (CONCAT_LAST) whenever +# a stage carrying them consolidates its output under that modality (e.g. an +# AUDIO-typed final stage that forwards them directly instead of routing +# through code2wav): ``codes.audio`` grows one row per decode step along +# dim 0, and ``codes.ref`` re-emits the constant reference-context matrix +# unchanged at every step. +register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0) +register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE) + QWEN3_TTS_PIPELINE = PipelineConfig( model_type="qwen3_tts", default_deploy_config_name="qwen3_tts.yaml", diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py index b1393ba19..45d1fc6fc 100644 --- a/vllm_omni/outputs/mm_outputs.py +++ b/vllm_omni/outputs/mm_outputs.py @@ -16,7 +16,11 @@ import torch from vllm.logger import init_logger from vllm.outputs import CompletionOutput -from vllm_omni.outputs.output_modality import TensorAccumulationStrategy +from vllm_omni.outputs.output_modality import ( + OutputModality, + TensorAccumulationStrategy, + get_accumulation_strategy, +) from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu logger = init_logger(__name__) @@ -46,21 +50,24 @@ def _consolidate_tensor_list( tensor_list: list[torch.Tensor], strategy: TensorAccumulationStrategy, ) -> torch.Tensor: - """Concatenate a deferred tensor list, with fallbacks on shape mismatch.""" + """Concatenate a deferred tensor list, with a fallback for CONCAT_LAST only. + + A shape mismatch under CONCAT_LAST is expected for genuine audio waveform + chunks (their non-last dims can vary between async-chunk steps), so that + case falls back to flattening each chunk before concatenating. Any other + strategy failing to concatenate means the tensors are not shaped the way + that strategy assumes (e.g. a key was assigned the wrong strategy), so it + is surfaced instead of silently keeping only the last chunk. + """ try: return _cat_tensors(tensor_list, strategy) - except RuntimeError: - # [TODO: this part is for async chunk, not for audio only. can we come up with a design - # that identify this by whether using async chunk instead of modality?] - if key != "audio": - logger.warning("Error concatenating tensor for key %s; keeping last tensor", key) - return tensor_list[-1] - # Audio chunks may have mismatched shapes; retry along the last dim, - # then fall back to flattening each chunk. - try: - return torch.cat(tensor_list, dim=-1) - except RuntimeError: - return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) + except RuntimeError as exc: + if strategy is not TensorAccumulationStrategy.CONCAT_LAST: + raise RuntimeError( + f"Error concatenating tensor list for key {key!r} under strategy {strategy}: {exc}" + ) from exc + logger.warning("CONCAT_LAST failed for key %s; flattening each chunk before concatenating", key) + return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None: @@ -167,12 +174,17 @@ class MultimodalPayload(Mapping): self.metadata[key] = value return self - def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None: + def consolidate_tensors(self, modality: OutputModality) -> None: """Concatenate deferred tensor lists into single tensors. - Tensors are generated content accumulated as chunks, so lists are - concatenated according to *strategy* (e.g. audio chunks along the - time dimension, latent frames along the batch dimension). + Tensors are generated content accumulated as chunks, so each key's + list is concatenated according to the strategy + ``get_accumulation_strategy(modality, key)`` resolves for it (e.g. + audio waveform chunks along the time dimension, latent frames along + the batch dimension). Most keys share *modality*'s default, but a + key can be registered for a different strategy -- e.g. a codec-frame + matrix that grows along dim 0 rather than the waveform-tuned default + for its modality -- via ``register_key_accumulation_strategy``. """ # Relocate scalar metadata tensors (e.g. sample rate) that from_dict # routed into .tensors, so they take the REPLACE path via .metadata @@ -185,6 +197,7 @@ class MultimodalPayload(Mapping): for key, value in list(self.tensors.items()): if _is_tensor_list(value): + strategy = get_accumulation_strategy(modality, key) self.tensors[key] = _consolidate_tensor_list(key, value, strategy) def consolidate_metadata(self) -> None: diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 7cdbc1f19..412e61619 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -115,8 +115,43 @@ class TensorAccumulationStrategy(Enum): """Replace previous tensor entirely with the latest one.""" -def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: - """Determine tensor merge strategy from the multimodal flags.""" +# Per-tensor-key overrides for the modality-wide default below. A modality +# is a coarse hint (e.g. AUDIO tensors are usually waveform chunks that grow +# along the last dimension), but some models emit tensors under that same +# modality whose accumulation semantics differ -- e.g. discrete codec-frame +# matrices that grow along dim 0, or a constant reference/context tensor +# re-emitted unchanged at every step. Pipelines that produce such keys +# register the correct strategy here via ``register_key_accumulation_strategy`` +# instead of forcing every tensor of the modality through one default. +_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = {} + + +def register_key_accumulation_strategy(key: str, strategy: TensorAccumulationStrategy) -> None: + """Register a per-key override for ``get_accumulation_strategy``. + + Args: + key: The flattened multimodal output key this override applies to, + e.g. ``"codes.audio"`` as produced by + ``vllm_omni.data_entry_keys.flatten_payload`` from a nested + ``{"codes": {"audio": tensor}}`` payload. + strategy: The accumulation strategy to use for this key, regardless + of the modality its request is otherwise associated with. + """ + _KEY_ACCUMULATION_OVERRIDES[key] = strategy + + +def get_accumulation_strategy( + modality: OutputModality, + key: str | None = None, +) -> TensorAccumulationStrategy: + """Determine the tensor merge strategy for one output key. + + A registered per-key override (see ``register_key_accumulation_strategy``) + always wins; otherwise the strategy falls back to the modality-wide + default. + """ + if key is not None and key in _KEY_ACCUMULATION_OVERRIDES: + return _KEY_ACCUMULATION_OVERRIDES[key] if OutputModality.AUDIO in modality: return TensorAccumulationStrategy.CONCAT_LAST if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index bd60e4793..5d2ede43d 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -28,7 +28,7 @@ from vllm_omni.outputs.multimodal_accumulation import ( is_non_final_delta_audio_chunk, replace_snapshot_keys, ) -from vllm_omni.outputs.output_modality import OutputModality, get_accumulation_strategy +from vllm_omni.outputs.output_modality import OutputModality logger = init_logger(__name__) @@ -140,8 +140,8 @@ class OmniRequestState(RequestState): """Consolidate accumulated tensor lists into single tensors. Uses TensorAccumulationStrategy derived from the output modality - to determine concatenation behavior. Metadata values always use - REPLACE (keep latest). + (with any per-key override applied) to determine concatenation + behavior. Metadata values always use REPLACE (keep latest). """ if self.mm_accumulated.is_empty: return @@ -150,10 +150,9 @@ class OmniRequestState(RequestState): modality = OutputModality.from_string(self.mm_type) except (ValueError, KeyError): modality = OutputModality.TEXT - strategy = get_accumulation_strategy(modality) try: - self.mm_accumulated.consolidate_tensors(strategy) + self.mm_accumulated.consolidate_tensors(modality) self.mm_accumulated.consolidate_metadata() except (RuntimeError, TypeError, KeyError): logger.exception("Error consolidating multimodal tensors")