[Feature][TTS] Port static chunk ramp to Moss-TTS streaming codec decode (#8124)

Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>

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]

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

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

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)