[Bugfix] MOSS-TTS: fix cross-request audio corruption under batching (#4415)

Signed-off-by: leo.yang <leo.yang.engineer@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/moss_tts/test_per_request_codes.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MOSS-TTS per-request code routing in the talker.
+
+Guards the cross-request audio-corruption class (RFC #4316 Issue 1, bug report
+#4355): under continuous batching every Stage-1 request must decode its *own*
+codes, not request 0's. The talker emits ``codes["audio"]`` as a batch-ordered
+**list** so the generic splitter (``to_payload_element``) routes
+``element[idx]`` to request ``idx``; the list must stay index-aligned with the
+batch (placeholder for skipped requests) or a shorter list silently falls back
+to entry 0 — the original corruption. The hidden-row mapping must also use the
+runner's real per-request token spans, not an equal split, or mixed
+prefill+decode steps sample each request's codes from the wrong rows.
+
+CPU-only and weight-free: ``make_omni_output`` is driven directly on a bare
+talker with ``_sample_audio_codes`` mocked to echo the hidden row it was given,
+so an assertion on the output value proves which row each request sampled from.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+N_VQ = 4
+
+
+@functools.lru_cache(maxsize=1)
+def _delay_talker_cls():
+    """Defer talker import (pulls vLLM model_executor) until first use."""
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSDelayTalkerForGeneration,
+    )
+
+    return MossTTSDelayTalkerForGeneration
+
+
+def _make_bare_delay_talker():
+    """A talker whose only behaviour is the routing/span logic under test.
+
+    ``_sample_audio_codes`` is replaced with a stub returning an (N_VQ,) tensor
+    filled with the value of the hidden row it received. With a ``hidden`` where
+    row r holds the scalar r, the emitted codes reveal which row each request
+    sampled.
+    """
+    cls = _delay_talker_cls()
+    talker = cls.__new__(cls)
+    talker._batch_state = None
+    talker.n_vq = N_VQ
+
+    def _stub_sample(last_h, state):
+        row_value = int(last_h.reshape(-1)[0].item())
+        return torch.full((N_VQ,), row_value, dtype=torch.long)
+
+    talker._sample_audio_codes = _stub_sample
+    return talker
+
+
+def _row_indexed_hidden(num_rows: int) -> torch.Tensor:
+    """(num_rows, 1) tensor where row r holds the scalar r."""
+    return torch.arange(num_rows, dtype=torch.float32).reshape(num_rows, 1)
+
+
+def _audio_payload(result):
+    return result.multimodal_outputs["codes"]["audio"]
+
+
+class TestPerRequestRouting:
+    def test_mixed_prefill_decode_routes_by_span(self):
+        """req 0 prefills 3 rows, req 1 decodes 1 row in the same step.
+
+        With real spans, req 0 samples its last row (2) and req 1 its last row
+        (3); the payload is a batch-aligned list, element[idx] per request.
+        """
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(4)
+        info_dicts = [{}, {}]
+        spans = [(0, 3), (3, 4)]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert isinstance(audio, list)
+        assert len(audio) == 2  # one entry per request, batch-aligned
+        assert audio[0].shape == (1, N_VQ)
+        assert torch.all(audio[0] == 2)  # req 0: last row of span [0,3)
+        assert torch.all(audio[1] == 3)  # req 1: last row of span [3,4)
+        # No cross-request leakage.
+        assert not torch.equal(audio[0], audio[1])
+
+    def test_no_request_samples_outside_its_span(self):
+        talker = _make_bare_delay_talker()
+        spans = [(0, 2), (2, 5), (5, 6)]
+        hidden = _row_indexed_hidden(6)
+        info_dicts = [{}, {}, {}]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert len(audio) == 3
+        sampled = []
+        for i, (start, end) in enumerate(spans):
+            value = int(audio[i].reshape(-1)[0].item())
+            assert start <= value < end, f"req {i} sampled row {value} outside [{start},{end})"
+            sampled.append(value)
+        assert len(set(sampled)) == len(sampled)  # distinct rows, no sharing
+
+    def test_skipped_request_keeps_list_aligned(self):
+        """An empty-span request gets an empty placeholder at its own index, so
+        later requests are not shifted onto request 0's codes."""
+        talker = _make_bare_delay_talker()
+        # req 1 has an empty span (start == end) and must not displace req 2.
+        spans = [(0, 1), (1, 1), (1, 2)]
+        hidden = _row_indexed_hidden(2)
+        info_dicts = [{}, {}, {}]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert len(audio) == 3
+        assert torch.all(audio[0] == 0)  # req 0: row 0
+        assert audio[1].numel() == 0  # req 1: empty placeholder, aligned
+        assert torch.all(audio[2] == 1)  # req 2: row 1, not shifted
+
+    def test_accumulates_into_per_request_history(self):
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(2)
+        prior = torch.full((1, N_VQ), 7, dtype=torch.long)
+        info_dicts = [{"audio_codes": {"accumulated": prior}}, {}]
+        spans = [(0, 1), (1, 2)]
+
+        result = talker.make_omni_output(
+            hidden,
+            runtime_additional_information=info_dicts,
+            request_token_spans=spans,
+        )
+        audio = _audio_payload(result)
+
+        assert audio[0].shape == (2, N_VQ)
+        assert torch.all(audio[0][0] == 7)  # prior history kept
+        assert torch.all(audio[0][1] == 0)  # this step's row 0 appended
+        assert torch.all(audio[1] == 1)
+
+    def test_runs_without_spans_backward_compat(self):
+        """Missing request_token_spans must not crash (equal-split fallback)
+        and still yields a batch-aligned list."""
+        talker = _make_bare_delay_talker()
+        hidden = _row_indexed_hidden(2)
+        info_dicts = [{}, {}]
+
+        result = talker.make_omni_output(hidden, runtime_additional_information=info_dicts)
+        audio = _audio_payload(result)
+        assert isinstance(audio, list)
+        assert len(audio) == 2
+
+    def test_no_codes_returns_empty_payload(self):
+        """Zero hidden rows → no codes → empty multimodal payload (no list)."""
+        talker = _make_bare_delay_talker()
+        hidden = torch.zeros((0, 1), dtype=torch.float32)
+        result = talker.make_omni_output(hidden, runtime_additional_information=[{}], request_token_spans=[(0, 0)])
+        assert result.multimodal_outputs == {}

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -553,7 +553,19 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         # rows align with hidden rows, which align with info_dicts in order.
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        audio_codes_list: list[torch.Tensor] = []
+        # Per-request (start, end) hidden-row spans from the runner. In mixed
+        # prefill+decode steps the per-request token counts differ, so an equal
+        # rows-per-request split would sample codes from the wrong rows; fall
+        # back to that split only when spans are unavailable.
+        spans = kwargs.get("request_token_spans")
+
+        # One accumulated-codes tensor per request, in batch order. The generic
+        # splitter (``to_payload_element``) routes ``per_req_codes[idx]`` to
+        # request ``idx``, so Stage 1 decodes each request's own codes instead
+        # of request 0's. Skipped requests get an empty placeholder so indices
+        # stay aligned (a shorter list would silently fall back to entry 0).
+        per_req_codes: list[torch.Tensor] = []
+        have_codes = False
 
         if hidden.numel() > 0:
             num_rows = hidden.shape[0]
@@ -561,10 +573,16 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
 
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
+                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
-                row_start = i * rows_per_req
-                row_end = min(row_start + rows_per_req, num_rows)
+                if spans is not None and i < len(spans):
+                    row_start, row_end = spans[i]
+                    row_end = min(int(row_end), num_rows)
+                else:
+                    row_start = i * rows_per_req
+                    row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
+                    per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
 
                 # Sample new audio codes from the last hidden state of this request.
@@ -582,21 +600,20 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
                     "current": new_codes,
                     "accumulated": updated_acc,
                 }
-                audio_codes_list.append(updated_acc)
+                per_req_codes.append(updated_acc)
+                have_codes = True
 
-        if not audio_codes_list:
+        if not have_codes:
             return OmniOutput(
                 text_hidden_states=hidden,
                 multimodal_outputs={},
             )
 
-        # Forward the accumulated audio codes (T, NQ) of the first request to
-        # Stage 1.  Multi-request batching for the codec is not yet supported
-        # by this adapter; the existing code already takes the first stage
-        # output anyway.
+        # Emit one accumulated-codes tensor per request as a list; the splitter
+        # picks ``element[idx]`` per request on the way to Stage 1.
         return OmniOutput(
             text_hidden_states=hidden,
-            multimodal_outputs={"codes": {"audio": audio_codes_list[0]}},
+            multimodal_outputs={"codes": {"audio": per_req_codes}},
         )
 
     # ------------------------------------------------------------------
@@ -933,15 +950,27 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
 
-        audio_codes_list: list[torch.Tensor] = []
+        # See delay talker: real per-request row spans from the runner, with the
+        # equal-split as fallback when unavailable.
+        spans = kwargs.get("request_token_spans")
+        # Per-request accumulated codes in batch order, pre-filled with empty
+        # placeholders so every skip path (stopped / eos / bos / pad) keeps
+        # indices aligned for to_payload_element's per-request routing. A
+        # shorter list would silently fall back to request 0's codes.
+        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
+        have_codes = False
         if hidden.numel() > 0 and info_dicts:
             num_rows = hidden.shape[0]
             rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     continue
-                row_start = i * rows_per_req
-                row_end = min(row_start + rows_per_req, num_rows)
+                if spans is not None and i < len(spans):
+                    row_start, row_end = spans[i]
+                    row_end = min(int(row_end), num_rows)
+                else:
+                    row_start = i * rows_per_req
+                    row_end = min(row_start + rows_per_req, num_rows)
                 if row_start >= row_end:
                     continue
 
@@ -1008,23 +1037,25 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
                 state["step"] = int(state.get("step", 0)) + 1
-                audio_codes_list.append(updated_acc)
+                per_req_codes[i] = updated_acc
+                have_codes = True
 
-        if not audio_codes_list:
+        if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Forward the first request's accumulated codes (multi-request
-        # batching for the codec is not yet wired up — same as delay).
-        # The realtime variant emits raw codes (no delay pattern), so we
-        # signal that to the chunk processor via a 1-D bool tensor in
-        # ``meta.finished`` (already a tensor field on the schema). The
-        # processor reads its truth value as "skip de-delay" only for
-        # realtime; the delay path doesn't populate ``meta`` at all.
-        device = audio_codes_list[0].device
+        # Emit one accumulated-codes tensor per request as a list; the splitter
+        # routes ``element[idx]`` to request idx, so each Stage-1 request decodes
+        # its own codes instead of request 0's.
+        # The realtime variant emits raw codes (no delay pattern), so we signal
+        # that to the chunk processor via a 1-D bool tensor in ``meta.finished``
+        # (already a tensor field on the schema). The processor reads its truth
+        # value as "skip de-delay" only for realtime; the delay path doesn't
+        # populate ``meta`` at all.
+        device = hidden.device
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
-                "codes": {"audio": audio_codes_list[0]},
+                "codes": {"audio": per_req_codes},
                 "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
             },
         )

vllm_omni/worker/gpu_model_runner.py

@@ -1313,6 +1313,20 @@ class OmniGPUModelRunner(GPUModelRunner):
 
             traceback.print_exc()
 
+        # Per-request (start, end) hidden-row spans so make_omni_output can map
+        # flat hidden rows to the right request in mixed prefill+decode steps,
+        # instead of assuming an equal rows-per-request split (which samples the
+        # wrong rows whenever per-request token counts differ).
+        nstp = self._omni_num_scheduled_tokens_np
+        if nstp is not None and len(nstp) == len(self.input_batch.req_ids):
+            try:
+                model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp)
+            except Exception as e:
+                # Visible on purpose: the fallback is the equal rows-per-request
+                # split, which can re-introduce the cross-request corruption this
+                # plumbing fixes — a silent failure here must not pass unnoticed.
+                logger.warning("[OMNI] Failed to compute request_token_spans: %s", e)
+
         if self._omni_query_start_loc_model_kwarg:
             try:
                 num_reqs = len(self.input_batch.req_ids)