From 5b3abe2f4d51560382146e955935f9e2e936e096 Mon Sep 17 00:00:00 2001 From: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Date: Tue, 29 Sep 2026 10:28:15 +0800 Subject: [PATCH] [Feature][TTS] Port static chunk ramp to Moss-TTS streaming codec decode (#8124) Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> --- .../test_moss_tts_chunk_ramp.py | 209 ++++++++++++++++++ vllm_omni/deploy/moss_tts_local.yaml | 4 + .../moss_tts/modeling_moss_tts_codec.py | 34 ++- .../stage_input_processors/moss_tts.py | 21 +- 4 files changed, 264 insertions(+), 4 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_moss_tts_chunk_ramp.py diff --git a/tests/model_executor/stage_input_processors/test_moss_tts_chunk_ramp.py b/tests/model_executor/stage_input_processors/test_moss_tts_chunk_ramp.py new file mode 100644 index 000000000..bf75b805f --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_tts_chunk_ramp.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""UT for Moss-TTS static chunk ramp in talker2codec_raw_async_chunk. + +Tests the Moss-TTS stage input processor's chunk-size mode selection: + 1. Static ramp (codec_chunk_ramp configured) + 2. Backward compat (not configured -> original IC/steady behavior) +and the codec-side warmup frame-size derivation, which must share the +processor's parse_chunk_ramp semantics (null / invalid / single-entry +values disable the ramp on both sides instead of raising at codec init). + +Framework utilities (parse_chunk_ramp, ramp_chunk_size) are already tested in +test_qwen3_tts_async_chunk.py -- this file only tests Moss-TTS adapter wiring. +""" + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import ( + _resolve_streaming_graph_frame_sizes, +) +from vllm_omni.model_executor.stage_input_processors.moss_tts import ( + talker2codec_raw_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_NQ = 12 # Moss-TTS Local v1.5 n_vq + + +def _frame(): + """One valid codec frame (n_vq codes, none equal to pad code 1024).""" + return torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=torch.long) + + +def _req(rid): + return SimpleNamespace(external_req_id=rid) + + +def _tm(*, chunk_frames=15, initial_chunk_frames=1, chunk_ramp=None): + extra = { + "codec_chunk_frames": chunk_frames, + "initial_codec_chunk_frames": initial_chunk_frames, + "codec_left_context_frames": 0, + } + if chunk_ramp is not None: + extra["codec_chunk_ramp"] = chunk_ramp + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + ramp_chunk_count=defaultdict(int), + request_payload={}, + connector=SimpleNamespace(config={"extra": extra}), + ) + + +def _feed(tm, rid, n_frames, *, finished=False): + """Call the processor with *n_frames* new frames appended.""" + if n_frames <= 0: + mm = None + else: + mm = {"codes": {"audio": _frame().reshape(1, -1).repeat(n_frames, 1)}} + return talker2codec_raw_async_chunk( + transfer_manager=tm, + multimodal_output=mm, + request=_req(rid), + is_finished=finished, + ) + + +class TestStaticRampEmission: + """Verify chunk sizes follow the ramp table [2, 4, 8, 15].""" + + def test_ramp_sequence(self): + tm = _tm(chunk_ramp=[2, 4, 8, 15]) + rid = "ramp-seq" + + # Chunk 0: need 2 frames + assert _feed(tm, rid, 1) is None # 1 < 2 -> hold + p0 = _feed(tm, rid, 1) # 2 >= 2 -> emit + assert p0 is not None + assert len(p0.codes.audio) == _NQ * 2 + tm.ramp_chunk_count[rid] = 1 + + # Chunk 1: need 4 frames + assert _feed(tm, rid, 3) is None + p1 = _feed(tm, rid, 1) + assert p1 is not None + assert len(p1.codes.audio) == _NQ * 4 + tm.ramp_chunk_count[rid] = 2 + + # Chunk 2: need 8 frames + assert _feed(tm, rid, 7) is None + p2 = _feed(tm, rid, 1) + assert p2 is not None + assert len(p2.codes.audio) == _NQ * 8 + tm.ramp_chunk_count[rid] = 3 + + # Chunk 3+: steady 15 + assert _feed(tm, rid, 14) is None + p3 = _feed(tm, rid, 1) + assert p3 is not None + assert len(p3.codes.audio) == _NQ * 15 + tm.ramp_chunk_count[rid] = 4 + + # Chunk 4: still 15 + assert _feed(tm, rid, 14) is None + p4 = _feed(tm, rid, 1) + assert p4 is not None + assert len(p4.codes.audio) == _NQ * 15 + + def test_ramp_finished_flush(self): + tm = _tm(chunk_ramp=[2, 4, 8, 15]) + rid = "ramp-flush" + + p0 = _feed(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + # 3 pending, finished -> flush all 3 + p_fin = _feed(tm, rid, 3, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + assert len(p_fin.codes.audio) == _NQ * 3 + + def test_ramp_steady_after_exhausted(self): + """Short ramp [2, 4] -> steady 15 after ramp exhausted.""" + tm = _tm(chunk_ramp=[2, 4], chunk_frames=15) + rid = "ramp-short" + + p0 = _feed(tm, rid, 2) + assert len(p0.codes.audio) == _NQ * 2 + tm.ramp_chunk_count[rid] = 1 + + p1 = _feed(tm, rid, 4) + assert len(p1.codes.audio) == _NQ * 4 + tm.ramp_chunk_count[rid] = 2 + + # Past ramp table -> steady + assert _feed(tm, rid, 14) is None + p2 = _feed(tm, rid, 1) + assert len(p2.codes.audio) == _NQ * 15 + + def test_ramp_string_config(self): + """YAML-style string value "2,4,8,15" parses the same as a list.""" + tm = _tm(chunk_ramp="2,4,8,15") + rid = "ramp-str" + + p0 = _feed(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _NQ * 2 + + +class TestBackwardCompat: + """Without ramp config, original IC/steady behavior is unchanged.""" + + def test_ic_then_steady(self): + tm = _tm(chunk_ramp=None, initial_chunk_frames=1, chunk_frames=15) + rid = "no-ramp" + + # IC=1: first chunk emits 1 frame + p0 = _feed(tm, rid, 1) + assert p0 is not None + assert len(p0.codes.audio) == _NQ * 1 + tm.put_req_chunk[rid] = 1 + + # Steady=15 + assert _feed(tm, rid, 14) is None + p1 = _feed(tm, rid, 1) + assert p1 is not None + assert len(p1.codes.audio) == _NQ * 15 + + def test_no_frames_returns_none(self): + tm = _tm(chunk_ramp=None, initial_chunk_frames=1) + rid = "empty" + assert _feed(tm, rid, 0) is None + + +class TestGraphFrameSizes: + """Codec warmup frame-size derivation shares parse_chunk_ramp semantics: + null / invalid / single-entry values disable the ramp (no raise at init).""" + + def _resolve(self, extra): + return _resolve_streaming_graph_frame_sizes(1, 15, extra) + + def test_absent_key(self): + assert self._resolve({}) == [1, 15] + assert self._resolve(None) == [1, 15] + + def test_explicit_null_disables(self): + # YAML `codec_chunk_ramp:` (present key, null value) must not raise + assert self._resolve({"codec_chunk_ramp": None}) == [1, 15] + + def test_invalid_values_disable(self): + assert self._resolve({"codec_chunk_ramp": "2,x,8"}) == [1, 15] + assert self._resolve({"codec_chunk_ramp": [4]}) == [1, 15] + assert self._resolve({"codec_chunk_ramp": [0, 4, 15]}) == [1, 15] + assert self._resolve({"codec_chunk_ramp": 5}) == [1, 15] + + def test_valid_ladder(self): + assert self._resolve({"codec_chunk_ramp": [2, 4, 8, 15]}) == [1, 2, 4, 8, 15] + + def test_string_form(self): + assert self._resolve({"codec_chunk_ramp": "2, 4, 8, 15"}) == [1, 2, 4, 8, 15] diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index ad16a6dea..41dca59f3 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -15,6 +15,10 @@ connectors: initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks codec_left_context_frames: 0 + # Chunk ramp (opt-in): gradual chunk-size ladder replacing the + # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last + # entry should equal codec_chunk_frames (steady state). + # codec_chunk_ramp: [2, 4, 8, 15] stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 6fd097d71..d87f01b4b 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -36,6 +36,7 @@ from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp logger = init_logger(__name__) @@ -225,6 +226,22 @@ class _MossCodecStreamSession: return out +def _resolve_streaming_graph_frame_sizes( + initial_frames: int, + steady_frames: int, + connector_extra_cfg: dict | None, +) -> list[int]: + """Streaming-graph frame sizes to pre-capture: ``{initial, steady}`` + ramp ladder. + + Ladder parsing is delegated to ``parse_chunk_ramp`` so the codec and the + stage input processor agree on one validated ``codec_chunk_ramp``: + null / malformed / single-entry values disable the ramp on both sides + (warn only) and never raise during codec init. + """ + ramp = parse_chunk_ramp(connector_extra_cfg or {}, steady=steady_frames) or [] + return sorted({s for s in (initial_frames, steady_frames, *ramp) if s > 0}) + + class MossTTSCodecDecoder(nn.Module): """Stage-1 decoder for all MOSS-TTS variants. @@ -283,8 +300,12 @@ class MossTTSCodecDecoder(nn.Module): self._stream_req_slots: dict[str, int] = {} self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config() - self._streaming_graph_frame_sizes = sorted( - {frames for frames in (self._initial_stream_chunk_frames, self._stream_chunk_frames) if frames > 0} + # codec_chunk_ramp ladder sizes join the streaming graph frame-size + # set so every ladder (B, T) combination is captured during warmup. + self._streaming_graph_frame_sizes = _resolve_streaming_graph_frame_sizes( + self._initial_stream_chunk_frames, + self._stream_chunk_frames, + self._connector_extra_cfg(), ) # ------------------------------------------------------------------ @@ -666,6 +687,15 @@ class MossTTSCodecDecoder(nn.Module): return int(extra_cfg[name]) return default + def _connector_extra_cfg(self) -> dict: + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) + else: + extra_cfg = getattr(connector_cfg, "extra", None) + return extra_cfg if isinstance(extra_cfg, dict) else {} + def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]: if getattr(self.vllm_config.model_config, "enforce_eager", True): return [] diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 3ba32271e..e07f552b1 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -14,6 +14,10 @@ from vllm.inputs import TokensPrompt as OmniTokensPrompt from vllm.logger import init_logger from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + parse_chunk_ramp, + ramp_chunk_size, +) logger = init_logger(__name__) @@ -284,8 +288,21 @@ def talker2codec_raw_async_chunk( initial_chunk_frames = chunk_frames pending = len(pending_frames) - emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 - threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + # codec_chunk_ramp (opt-in): static ladder of chunk sizes replacing the + # initial_codec_chunk_frames -> codec_chunk_frames cliff. The chunk index + # is maintained by the connector (ramp_chunk_count, incremented per + # emitted chunk); past the ladder, ramp_chunk_size falls back to steady. + if not hasattr(transfer_manager, "_moss_ramp_parsed"): + transfer_manager._moss_ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_frames) + _ramp = transfer_manager._moss_ramp_parsed + if _ramp is not None: + if not hasattr(transfer_manager, "ramp_chunk_count"): + transfer_manager.ramp_chunk_count = defaultdict(int) + _chunk_idx = int(transfer_manager.ramp_chunk_count.get(req_id, 0)) + threshold = ramp_chunk_size(_chunk_idx, _ramp, chunk_frames) + else: + emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 + threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames if pending <= 0: if is_finished: transfer_manager.code_prompt_token_ids.pop(req_id, None)