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