[Bugfix] Per-key accumulation strategy for Qwen3-TTS codec frame outputs (#7608)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

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)

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",

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:

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:

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")