[Model] Batch MiniCPM-o Talker codec sampling (#5792)

Signed-off-by: Zhou248 <62542847+Zhou248@users.noreply.github.com>

tests/model_executor/models/minicpmo_4_5/test_talker_batching.py

@@ -12,7 +12,9 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
     MiniCPMO45OmniForConditionalGeneration,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
+    _REPETITION_PENALTY_CHUNK_SIZE,
     MiniCPMO45OmniTTSForConditionalGeneration,
+    _apply_batched_repetition_penalty,
     _max_audio_tokens,
     _restore_weight_norm_weight,
 )
@@ -73,6 +75,50 @@ def _routed(output, index: int):
     )
 
 
+def _reference_repetition_penalty(
+    logits: torch.Tensor,
+    history: torch.Tensor,
+    *,
+    penalty: float,
+    window_size: int,
+) -> torch.Tensor:
+    if penalty == 1.0 or history.numel() == 0:
+        return logits
+    recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long)
+    frequencies = torch.bincount(recent, minlength=logits.shape[-1]).to(dtype=logits.dtype)
+    alpha = torch.pow(torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype), frequencies)
+    return torch.where(logits < 0, logits * alpha, logits / alpha)
+
+
+def _make_sampling_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
+    talker = _make_talker()
+    talker._codec_temperature = 0.8
+    talker._codec_top_k = 5
+    talker._codec_top_p = 0.85
+    talker._codec_repetition_penalty = 1.05
+    talker._request_audio_states = {
+        "req-a": {"min_tokens": 2},
+        "req-b": {"min_tokens": 0},
+    }
+    talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)])
+    with torch.no_grad():
+        talker.head_code[0].weight.copy_(
+            torch.tensor(
+                [
+                    [0.1, -0.3],
+                    [0.2, 0.4],
+                    [-0.5, 0.2],
+                    [0.7, -0.1],
+                    [-0.2, 0.8],
+                    [0.6, 0.3],
+                    [-0.4, -0.7],
+                    [0.3, 0.5],
+                ]
+            )
+        )
+    return talker
+
+
 @pytest.mark.parametrize(
     ("condition_tokens", "expected"),
     [(3, 64), (100, 1000), (1000, 2048)],
@@ -101,16 +147,153 @@ def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() ->
     assert torch.equal(restored, linear.weight)
 
 
+def test_batched_repetition_penalty_matches_request_local_rows() -> None:
+    logits = torch.tensor(
+        [
+            [-2.0, -1.0, 1.0, 2.0, 3.0, 4.0],
+            [4.0, 3.0, 2.0, 1.0, -1.0, -2.0],
+            [1.0, -1.0, 2.0, -2.0, 3.0, -3.0],
+        ]
+    )
+    histories = [
+        torch.tensor([1, 1, 2, 5]),
+        torch.tensor([0, 4, 4]),
+        torch.empty(0, dtype=torch.long),
+    ]
+
+    actual = _apply_batched_repetition_penalty(
+        logits,
+        histories,
+        penalty=1.2,
+        window_size=3,
+    )
+    expected = torch.cat(
+        [
+            _reference_repetition_penalty(
+                logits[row : row + 1],
+                history,
+                penalty=1.2,
+                window_size=3,
+            )
+            for row, history in enumerate(histories)
+        ],
+        dim=0,
+    )
+
+    assert torch.equal(actual, expected)
+
+
+def test_batched_repetition_penalty_matches_rows_across_chunks(mocker) -> None:
+    batch_size = 2 * _REPETITION_PENALTY_CHUNK_SIZE + 1
+    vocab_size = 11
+    logits = torch.arange(batch_size * vocab_size, dtype=torch.float32).reshape(batch_size, vocab_size) - 100
+    histories = [
+        torch.tensor([(row + offset) % vocab_size for offset in range(row % 7)], dtype=torch.long)
+        for row in range(batch_size)
+    ]
+
+    expected = torch.cat(
+        [
+            _reference_repetition_penalty(
+                logits[row : row + 1],
+                history,
+                penalty=1.2,
+                window_size=5,
+            )
+            for row, history in enumerate(histories)
+        ],
+        dim=0,
+    )
+    bincount = mocker.spy(torch, "bincount")
+    actual = _apply_batched_repetition_penalty(
+        logits,
+        histories,
+        penalty=1.2,
+        window_size=5,
+    )
+
+    assert torch.equal(actual, expected)
+    assert [call.kwargs["minlength"] for call in bincount.call_args_list] == [
+        _REPETITION_PENALTY_CHUNK_SIZE * vocab_size,
+        _REPETITION_PENALTY_CHUNK_SIZE * vocab_size,
+        vocab_size,
+    ]
+
+
+def test_batched_codec_sampling_is_request_order_independent() -> None:
+    talker = _make_sampling_talker()
+    histories = [torch.tensor([1, 1, 2]), torch.tensor([3, 4, 4])]
+    hidden = torch.tensor([[0.25, -0.75], [-0.5, 0.5]])
+
+    forward_order = talker._sample_audio_codes(
+        hidden,
+        histories,
+        ["req-a", "req-b"],
+        [0, 0],
+    )
+    talker._request_generators = {}
+    reverse_order = talker._sample_audio_codes(
+        hidden.flip(0),
+        list(reversed(histories)),
+        ["req-b", "req-a"],
+        [0, 0],
+    )
+
+    assert torch.equal(forward_order[0], reverse_order[1])
+    assert torch.equal(forward_order[1], reverse_order[0])
+
+
+def test_batched_codec_sampling_preserves_request_rng_after_compaction() -> None:
+    talker = _make_sampling_talker()
+    history_a = torch.tensor([1, 1, 2])
+    history_b = torch.tensor([3, 4, 4])
+    hidden_a = torch.tensor([[0.25, -0.75]])
+    hidden_b = torch.tensor([[-0.5, 0.5]])
+
+    first_batch = talker._sample_audio_codes(
+        torch.cat([hidden_a, hidden_b]),
+        [history_a, history_b],
+        ["req-a", "req-b"],
+        [0, 0],
+    )
+    compacted_b = talker._sample_audio_codes(
+        hidden_b,
+        [history_b],
+        ["req-b"],
+        [1],
+    )
+
+    talker._request_generators = {}
+    standalone_b_first = talker._sample_audio_codes(
+        hidden_b,
+        [history_b],
+        ["req-b"],
+        [0],
+    )[0]
+    standalone_b_second = talker._sample_audio_codes(
+        hidden_b,
+        [history_b],
+        ["req-b"],
+        [1],
+    )[0]
+
+    assert torch.equal(first_batch[1], standalone_b_first)
+    assert torch.equal(compacted_b[0], standalone_b_second)
+
+
 def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None:
     talker = _make_talker()
     seen: list[tuple[str, list[float], list[int]]] = []
 
-    def sample(hidden, history, request_id, step):
-        assert step == 0
-        seen.append((request_id, hidden.reshape(-1).tolist(), history.tolist()))
-        return torch.tensor(2 if request_id == "req-a" else 3)
+    def sample(hidden, histories, request_ids, steps):
+        assert steps == [0, 0]
+        seen.extend(
+            (request_id, hidden[row].reshape(-1).tolist(), histories[row].tolist())
+            for row, request_id in enumerate(request_ids)
+        )
+        return torch.tensor([2, 3])
 
-    mocker.patch.object(talker, "_sample_audio_code", side_effect=sample)
+    batched_sample = mocker.patch.object(talker, "_sample_audio_codes", side_effect=sample)
     infos = [
         {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}},
         {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}},
@@ -126,6 +309,7 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N
         ("req-a", [2.0, 0.0], [1]),
         ("req-b", [3.0, 0.0], []),
     ]
+    batched_sample.assert_called_once()
     assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2]
     assert infos[1]["audio_codes"]["accumulated"].tolist() == [3]
     assert set(output.multimodal_outputs) == {"codes", "meta"}
@@ -140,7 +324,7 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N
 
 def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
     talker = _make_talker()
-    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
+    mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2, 2]))
     infos = [
         {
             "request_id": "req-a",
@@ -186,7 +370,7 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
 
 def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
     talker = _make_talker()
-    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
+    mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2]))
     info = {
         "request_id": "req-missing-fence",
         "native_duplex": True,
@@ -203,7 +387,7 @@ def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
 
 def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None:
     talker = _make_talker()
-    sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
+    sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2]))
     infos = [
         {
             "request_id": "req-prefill",
@@ -225,7 +409,7 @@ def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) ->
     )
 
     sample.assert_called_once()
-    assert sample.call_args.args[2] == "req-decode"
+    assert sample.call_args.args[2] == ["req-decode"]
     assert infos[0]["audio_state"]["step"] == 0
     assert infos[0]["audio_codes"]["accumulated"].numel() == 0
     assert infos[1]["audio_state"]["step"] == 5
@@ -235,7 +419,7 @@ def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) ->
 
 def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None:
     talker = _make_talker()
-    sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(7))
+    sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([7]))
     info = {
         "request_id": "req-stop",
         "audio_state": {"step": 3},
@@ -265,7 +449,7 @@ def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None:
 
 def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None:
     talker = _make_talker()
-    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3))
+    mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3]))
     info = {
         "request_id": "req-limit",
         "audio_state": {"step": 1, "max_tokens": 2},
@@ -288,7 +472,7 @@ def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None:
 
 def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None:
     talker = _make_talker()
-    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3))
+    mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3]))
     first_info = {
         "request_id": "req-local-state",
         "audio_state": {"step": 1, "max_tokens": 3},

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py

@@ -11,7 +11,8 @@ Pipeline:
   4. Continuously generate request-aligned discrete audio-code deltas
 """
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -32,6 +33,7 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 _REPETITION_WINDOW = 16
+_REPETITION_PENALTY_CHUNK_SIZE = 16
 _MIN_AUDIO_TOKENS = 64
 _MAX_AUDIO_TOKENS = 2048
 _AUDIO_TOKENS_PER_TEXT_TOKEN = 10
@@ -46,6 +48,17 @@ _CODEC_MIN_TOKENS = 50
 _DUPLEX_CODEC_TOKENS_PER_CHUNK = 26
 
 
+@dataclass(slots=True)
+class _PendingCodecSample:
+    output_index: int
+    hidden_row: torch.Tensor
+    codes: torch.Tensor
+    request_id: str
+    step: int
+    state: dict[str, Any]
+    info: dict[str, Any]
+
+
 def _max_audio_tokens(condition_tokens: int) -> int:
     """Bound codec generation with a conservative text-length estimate.
 
@@ -65,20 +78,48 @@ def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor)
     return torch._weight_norm(weight_v, weight_g, dim=0)
 
 
-def _apply_repetition_penalty(
+def _apply_batched_repetition_penalty(
     logits: torch.Tensor,
-    history: torch.Tensor,
+    histories: Sequence[torch.Tensor],
     *,
     penalty: float,
     window_size: int,
 ) -> torch.Tensor:
-    """Match MiniCPMTTS' frequency-aware repetition penalty."""
-    if penalty == 1.0 or history.numel() == 0:
+    """Apply request-local frequency penalties to a batch of codec logits."""
+    if logits.ndim != 2:
+        raise ValueError(f"batched codec logits must be 2D, got shape {tuple(logits.shape)}")
+    batch_size, vocab_size = logits.shape
+    if len(histories) != batch_size:
+        raise ValueError(f"expected {batch_size} codec histories, got {len(histories)}")
+    if penalty == 1.0:
         return logits
-    recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long)
-    frequencies = torch.bincount(recent, minlength=logits.shape[-1]).to(dtype=logits.dtype)
-    alpha = torch.pow(torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype), frequencies)
-    return torch.where(logits < 0, logits * alpha, logits / alpha)
+
+    if batch_size == 0:
+        return logits
+
+    penalized = logits.clone()
+    penalty_tensor = torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype)
+    for start in range(0, batch_size, _REPETITION_PENALTY_CHUNK_SIZE):
+        end = min(start + _REPETITION_PENALTY_CHUNK_SIZE, batch_size)
+        chunk_logits = logits[start:end]
+        encoded_rows: list[torch.Tensor] = []
+        for local_row, history in enumerate(histories[start:end]):
+            recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long)
+            if recent.numel() > 0:
+                encoded_rows.append(recent + local_row * vocab_size)
+        if not encoded_rows:
+            continue
+
+        # Bound the int64 bincount workspace independently of request concurrency.
+        encoded = encoded_rows[0] if len(encoded_rows) == 1 else torch.cat(encoded_rows)
+        frequencies = torch.bincount(
+            encoded,
+            minlength=(end - start) * vocab_size,
+        ).reshape(end - start, vocab_size)
+        alpha = torch.pow(penalty_tensor, frequencies.to(dtype=logits.dtype))
+        penalized[start:end] = torch.where(chunk_logits < 0, chunk_logits * alpha, chunk_logits / alpha)
+
+    return penalized
 
 
 def _apply_top_k_top_p(
@@ -370,28 +411,47 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             self._request_generators[request_id] = generator
         return generator
 
-    def _sample_audio_code(
+    def _sample_audio_codes(
         self,
-        hidden_state: torch.Tensor,
-        history: torch.Tensor,
-        request_id: str,
-        step: int,
+        hidden_states: torch.Tensor,
+        histories: Sequence[torch.Tensor],
+        request_ids: Sequence[str],
+        steps: Sequence[int],
     ) -> torch.Tensor:
-        logits = self.head_code[0](hidden_state).float() / self._codec_temperature
+        batch_size = int(hidden_states.shape[0])
+        if not (len(histories) == len(request_ids) == len(steps) == batch_size):
+            raise ValueError(
+                "MiniCPM-o batched codec sampling requires one history, request id, "
+                f"and step per hidden row, got batch={batch_size}, histories={len(histories)}, "
+                f"request_ids={len(request_ids)}, steps={len(steps)}"
+            )
+        if batch_size == 0:
+            return torch.empty(0, dtype=torch.long, device=hidden_states.device)
+
+        logits = self.head_code[0](hidden_states).float() / self._codec_temperature
         eos_id = self._num_audio_tokens - 1
-        logits = _apply_repetition_penalty(
+        logits = _apply_batched_repetition_penalty(
             logits,
-            history,
+            histories,
             penalty=self._codec_repetition_penalty,
             window_size=_REPETITION_WINDOW,
         )
         request_states = getattr(self, "_request_audio_states", {})
-        state = request_states.get(request_id)
-        min_tokens = (
-            int(state.get("min_tokens", self._codec_min_tokens)) if isinstance(state, dict) else self._codec_min_tokens
+        mask_eos_values: list[bool] = []
+        for request_id, step in zip(request_ids, steps, strict=True):
+            state = request_states.get(request_id)
+            min_tokens = (
+                int(state.get("min_tokens", self._codec_min_tokens))
+                if isinstance(state, dict)
+                else self._codec_min_tokens
+            )
+            mask_eos_values.append(step < min_tokens)
+        mask_eos = torch.tensor(
+            mask_eos_values,
+            dtype=torch.bool,
+            device=logits.device,
         )
-        if step < min_tokens:
-            logits[..., eos_id] = float("-inf")
+        logits[:, eos_id].masked_fill_(mask_eos, float("-inf"))
         logits = _apply_top_k_top_p(
             logits,
             top_k=self._codec_top_k,
@@ -399,11 +459,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             min_tokens_to_keep=3,
         )
         probabilities = torch.softmax(logits, dim=-1)
-        return torch.multinomial(
-            probabilities,
-            num_samples=1,
-            generator=self._request_generator(request_id, probabilities.device),
-        ).reshape(())
+        # torch.multinomial accepts one generator for the entire call. Keep it
+        # per row so compaction, request ordering, or another request finishing
+        # cannot perturb request-local RNG streams; the expensive projection
+        # and logit transforms above remain fully batched.
+        return torch.cat(
+            [
+                torch.multinomial(
+                    probabilities[row : row + 1],
+                    num_samples=1,
+                    generator=self._request_generator(request_id, probabilities.device),
+                )
+                for row, request_id in enumerate(request_ids)
+            ],
+            dim=0,
+        ).reshape(-1)
 
     def make_omni_output(
         self,
@@ -426,15 +496,18 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             )
         emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos)
 
-        stop_rows: list[torch.Tensor] = []
-        codec_deltas: list[torch.Tensor] = []
-        terminal_flags: list[torch.Tensor] = []
+        # Rows default to continue. Only previously finished or newly terminal
+        # requests stop; prefill/ineligible rows stay aligned as False.
+        stop_flags = [False] * len(infos)
         native_duplex_flags: list[torch.Tensor] = []
         duplex_epochs: list[torch.Tensor] = []
         duplex_turn_ids: list[torch.Tensor] = []
         segment_texts_utf8: list[torch.Tensor] = []
         turn_end_flags: list[torch.Tensor] = []
         empty_delta = hidden.new_empty((0, 1), dtype=torch.long)
+        codec_deltas = [empty_delta for _ in infos]
+        terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos]
+        pending_samples: list[_PendingCodecSample] = []
         for index, info in enumerate(infos):
             info_dict = info if isinstance(info, dict) else {}
             native_duplex = info_dict.get("native_duplex") is True
@@ -480,16 +553,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool))
 
             if not isinstance(info, dict):
-                stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
-                codec_deltas.append(empty_delta)
-                terminal_flags.append(torch.tensor(False, dtype=torch.bool))
                 continue
             start, end = spans[index]
             end = min(int(end), int(hidden.shape[0]))
             if int(start) >= end:
-                stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
-                codec_deltas.append(empty_delta)
-                terminal_flags.append(torch.tensor(False, dtype=torch.bool))
                 continue
             request_id = str(info.get("request_id", index))
             request_states = getattr(self, "_request_audio_states", None)
@@ -501,17 +568,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 state = dict(info.get("audio_state", {}) or {})
                 request_states[request_id] = state
             if state.get("finished"):
-                stop_rows.append(hidden.new_tensor([float("-inf"), 0.0]))
-                codec_deltas.append(empty_delta)
-                terminal_flags.append(torch.tensor(False, dtype=torch.bool))
+                stop_flags[index] = True
                 continue
             if not sample_eligible[index]:
                 # vLLM computes a logit row for incomplete chunked prefills but
                 # discards its sampled token. Advancing codec/RNG state here
                 # would make output depend on prefill chunking and compaction.
-                stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
-                codec_deltas.append(empty_delta)
-                terminal_flags.append(torch.tensor(False, dtype=torch.bool))
                 continue
             codes = state.get("codes")
             if not isinstance(codes, torch.Tensor):
@@ -521,8 +583,42 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
             else:
                 codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1)
             step = int(state.get("step", 0))
-            sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step)
-            sampled_id = int(sampled.item())
+            pending_samples.append(
+                _PendingCodecSample(
+                    output_index=index,
+                    hidden_row=hidden[end - 1 : end],
+                    codes=codes,
+                    request_id=request_id,
+                    step=step,
+                    state=state,
+                    info=info,
+                )
+            )
+
+        if pending_samples:
+            active_hidden_rows = [pending.hidden_row for pending in pending_samples]
+            active_hidden = (
+                active_hidden_rows[0] if len(active_hidden_rows) == 1 else torch.cat(active_hidden_rows, dim=0)
+            )
+            sampled_batch = self._sample_audio_codes(
+                active_hidden,
+                [pending.codes for pending in pending_samples],
+                [pending.request_id for pending in pending_samples],
+                [pending.step for pending in pending_samples],
+            )
+            # One batched device-to-host synchronization replaces one .item()
+            # synchronization per request.
+            sampled_ids = sampled_batch.detach().to(device="cpu").tolist()
+        else:
+            sampled_batch = hidden.new_empty((0,), dtype=torch.long)
+            sampled_ids = []
+
+        for row, pending in enumerate(pending_samples):
+            sampled = sampled_batch[row].reshape(())
+            sampled_id = int(sampled_ids[row])
+            codes = pending.codes
+            state = pending.state
+            info = pending.info
             is_eos = sampled_id == self._num_audio_tokens - 1
             state["step"] = int(state.get("step", 0)) + 1
             reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048))
@@ -541,11 +637,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 "current": sampled.reshape(1),
                 "accumulated": codes,
             }
-            codec_deltas.append(delta)
-            terminal_flags.append(torch.tensor(finished, dtype=torch.bool))
-            stop_rows.append(hidden.new_tensor([float("-inf"), 0.0] if finished else [0.0, float("-inf")]))
+            codec_deltas[pending.output_index] = delta
+            terminal_flags[pending.output_index] = torch.tensor(finished, dtype=torch.bool)
+            stop_flags[pending.output_index] = finished
 
-        self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2))
+        self._batch_stop_logits = (
+            hidden.new_tensor([[float("-inf"), 0.0] if finished else [0.0, float("-inf")] for finished in stop_flags])
+            if stop_flags
+            else hidden.new_empty((0, 2))
+        )
         # Lists are deliberate: the runner routes element i to request i,
         # preserving compaction alignment while emitting only this step's code.
         meta_outputs = {"finished": terminal_flags}