diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py new file mode 100644 index 000000000..eabb4fdc0 --- /dev/null +++ b/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 == {} diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index e7e100a18..a03b4c1d1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/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)}, }, ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 349e1da04..cdf2286bf 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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)