[Feature][TTS] Seed MOSS-TTS residual sampling from the deploy seed

tests/entrypoints/openai_api/test_serving_speech.py

@@ -5584,6 +5584,67 @@ def test_prepare_speech_generation_qwen3_default_seed_sets_tts_local_seed(
         assert stage0_params.extra_args["tts_local_seed"] == 42
         assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
 
+    @pytest.fixture
+    def moss_tts_server(self, mocker: MockerFixture):
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+            return_value=set(),
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+            return_value=None,
+        )
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.model_config = mocker.MagicMock(
+            model="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5",
+            hf_config=mocker.MagicMock(),
+        )
+        mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048, seed=42, extra_args=None)]
+        mock_engine_client.tts_batch_max_items = 32
+        mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+        mock_engine_client.tts_max_instructions_length = None
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                engine_args=SimpleNamespace(model_stage="moss_tts_local"),
+                tts_args={},
+            )
+        ]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        yield server
+        server.shutdown()
+
+    def test_prepare_speech_generation_moss_default_seed_sets_tts_local_seed(
+        self, moss_tts_server, mocker: MockerFixture
+    ):
+        """Deploy default seed should set tts_local_seed for MOSS-TTS.
+
+        When request.seed is None (e.g. vllm bench serve), the deploy config's
+        default_sampling_params.seed propagates to tts_local_seed via the MOSS
+        adapter's apply_sampling_overrides, matching the Qwen3-TTS pattern.
+        """
+        moss_tts_server._adapter.validate = mocker.MagicMock(return_value=None)
+        moss_tts_server._adapter.build = mocker.AsyncMock(
+            return_value=PreparedRequest(
+                prompt={"prompt": "hello"},
+                tts_params={},
+                model_type="moss_tts",
+            )
+        )
+        request = OpenAICreateSpeechRequest(input="hello")
+
+        asyncio.run(moss_tts_server._prepare_speech_generation(request))
+
+        stage0_params = moss_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+        assert stage0_params.extra_args["tts_local_seed"] == 42
+        assert moss_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
+
     def test_prepare_speech_generation_uses_adapter_model_type_label(
         self,
         voxtral_server,

tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py

@@ -37,8 +37,10 @@ class _Request:
     external_req_id: str
 
 
-def _manager(ramp: object = None, *, first: int = 1) -> _TransferManager:
-    extra = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp}
+def _manager(ramp: object = None, *, first: int = 1, adaptive: bool = False) -> _TransferManager:
+    extra: dict[str, object] = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp}
+    if adaptive:
+        extra["codec_chunk_adaptive"] = True
     return _TransferManager(_Connector({"extra": extra}))
 
 
@@ -135,3 +137,91 @@ def test_graph_shapes_and_fast_path_follow_sender(
     assert codec._stream_max_step_frames == max_step
     assert codec._first_chunk_fast
     assert codec._initial_stream_chunk_frames == effective_first
+
+
+# ---------------------------------------------------------------------------
+# Adaptive controller
+# ---------------------------------------------------------------------------
+
+
+def test_adaptive_chunk_0_uses_ic_threshold() -> None:
+    """Chunk 0 under adaptive uses the same IC/steady threshold as default."""
+    manager = _manager(adaptive=True, first=1)
+    packet = _emit(manager, "a", 0)
+    assert packet is not None
+    assert packet.meta.codec_chunk_frames == 1
+    assert not bool(packet.meta.finished)
+
+
+def test_adaptive_finish_flushes_remaining() -> None:
+    """Adaptive flushes all remaining frames on finish."""
+    manager = _manager(adaptive=True, first=1)
+    # Chunk 0: emit 1 frame (IC threshold = 1).
+    first = _emit(manager, "a", 0)
+    assert first is not None
+    assert first.meta.codec_chunk_frames == 1
+    # Feed 1 more frame (pending=1, adaptive target>=2, hold — no emit).
+    held = _emit(manager, "a", 1)
+    assert held is None
+    # Finish with 1 pending → flush.
+    final = _emit(manager, "a", None, finished=True)
+    assert final is not None
+    assert bool(final.meta.finished)
+    assert final.meta.codec_chunk_frames == 1
+    assert "a" not in manager.code_prompt_token_ids
+
+
+def test_adaptive_empty_finish_returns_sentinel() -> None:
+    """Adaptive returns an empty-finished sentinel when no frames remain."""
+    manager = _manager(adaptive=True, first=1)
+    # Emit chunk 0, then finish with no pending frames.
+    first = _emit(manager, "a", 0)
+    assert first is not None
+    # Flush remaining with finish + no new frames.
+    final = _emit(manager, "a", None, finished=True)
+    assert final is not None
+    assert bool(final.meta.finished)
+    assert final.meta.codec_chunk_frames == 0
+    assert final.codes.audio.numel() == 0
+
+
+def test_adaptive_takes_precedence_over_ramp() -> None:
+    """When both codec_chunk_ramp and codec_chunk_adaptive are set, adaptive wins."""
+    manager = _manager(ramp=[4, 8, 15], adaptive=True, first=1)
+    # Chunk 0: adaptive uses IC=1, not ramp's first entry (4).
+    packet = _emit(manager, "a", 0)
+    assert packet is not None
+    assert packet.meta.codec_chunk_frames == 1
+
+
+def test_adaptive_requests_progress_independently() -> None:
+    """Two requests under adaptive do not interfere."""
+    manager = _manager(adaptive=True, first=1)
+    # Start both requests — chunk 0 emits 1 frame each.
+    a0 = _emit(manager, "a", 0)
+    b0 = _emit(manager, "b", 0)
+    assert a0 is not None and b0 is not None
+    assert a0.meta.codec_chunk_frames == 1
+    assert b0.meta.codec_chunk_frames == 1
+    # Feed 1 more frame to 'a' (held by adaptive target >= 2).
+    assert _emit(manager, "a", 1) is None
+    # Finish 'a' with 1 pending → flush 1.
+    final_a = _emit(manager, "a", None, finished=True)
+    assert final_a is not None
+    assert bool(final_a.meta.finished)
+    assert final_a.meta.codec_chunk_frames == 1
+    # 'b' can still be finished independently.
+    final_b = _emit(manager, "b", None, finished=True)
+    assert final_b is not None
+    assert bool(final_b.meta.finished)
+
+
+def test_adaptive_states_present_during_request() -> None:
+    """The adaptive controller is created on chunk 0 and persists until finish."""
+    manager = _manager(adaptive=True, first=1)
+    _emit(manager, "a", 0)
+    assert hasattr(manager, "_adaptive_states")
+    assert "a" in manager._adaptive_states
+    # Finish cleans up processor-side pending frames.
+    _emit(manager, "a", None, finished=True)
+    assert "a" not in manager.code_prompt_token_ids

vllm_omni/deploy/moss_tts_local_v1.yaml

@@ -26,6 +26,10 @@ connectors:
       # 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]
+      # Adaptive chunk ramp (opt-in): buffer-feedback controller that
+      # auto-tunes chunk sizes across hardware. Takes precedence over
+      # codec_chunk_ramp when both are set.
+      # codec_chunk_adaptive: true
 
 stages:
   - stage_id: 0

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -647,7 +647,24 @@ def apply_sampling_overrides(
         prompt: dict | None = None,
         request_id: str | None = None,
     ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)
+        sampling_params_list = apply_max_new_tokens(sampling_params_list, request)
+
+        # Shallow-copy stage 0 so the shared defaults stay immutable.
+        import copy
+
+        sampling_params_list = [copy.copy(params) for params in sampling_params_list]
+
+        # Propagate the deploy/YAML stage seed to tts_local_seed. The generic
+        # serving layer applies an explicit request.seed afterwards, so the
+        # request value still has higher precedence.
+        stage0_params = sampling_params_list[0]
+        default_seed = getattr(stage0_params, "seed", None)
+        if default_seed is not None:
+            if stage0_params.extra_args is None:
+                stage0_params.extra_args = {}
+            stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed))
+
+        return sampling_params_list
 
 
 @register_tts_adapter

vllm_omni/model_executor/stage_input_processors/moss_tts.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import time
 from collections import defaultdict
 from collections.abc import Mapping
 from typing import Any
@@ -18,7 +19,13 @@
 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
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
+    AdaptiveChunkController,
+    compute_adaptive_emit,
+    parse_adaptive_config,
+    parse_chunk_ramp,
+    ramp_chunk_size,
+)
 
 logger = init_logger(__name__)
 
@@ -240,6 +247,9 @@ def talker2codec_raw_async_chunk(
     a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern
     de-delay is applied on this path. An optional ``codec_chunk_ramp`` selects
     successive chunk sizes and takes precedence over the initial chunk size.
+    An optional ``codec_chunk_adaptive`` enables a buffer-feedback controller
+    that auto-tunes chunk sizes across hardware; it takes precedence over the
+    static ramp when both are configured.
     """
     external_req_id = getattr(request, "external_req_id", None)
     req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request)))
@@ -301,11 +311,131 @@ def talker2codec_raw_async_chunk(
     if not hasattr(transfer_manager, "_moss_chunk_ramp"):
         transfer_manager._moss_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_frames)
     ramp = transfer_manager._moss_chunk_ramp
+
+    # --- Adaptive controller (takes precedence over static ramp) ---
+    # Parse once per transfer_manager; config is static for its lifetime.
+    if not hasattr(transfer_manager, "_adaptive_parsed"):
+        transfer_manager._adaptive_parsed = parse_adaptive_config(cfg)
+    (
+        adaptive_enabled,
+        adaptive_min,
+        adaptive_margin,
+        adaptive_divisor,
+        adaptive_delta_min,
+    ) = transfer_manager._adaptive_parsed
+
+    if adaptive_enabled:
+        _adaptive_states = getattr(transfer_manager, "_adaptive_states", None)
+        if _adaptive_states is None:
+            _adaptive_states = {}
+            transfer_manager._adaptive_states = _adaptive_states
+
+        ctrl = _adaptive_states.get(req_id)
+        chunk_index = int(transfer_manager.ramp_chunk_count.get(req_id, 0))
+
+        if ctrl is None or chunk_index == 0:
+            # Chunk 0: use IC/steady threshold (same as original logic).
+            if pending <= 0:
+                if is_finished:
+                    transfer_manager.code_prompt_token_ids.pop(req_id, None)
+                    transfer_manager.request_payload.pop(req_id, None)
+                    return OmniPayloadStruct(
+                        codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                        meta=MetaStruct(
+                            req_id=[req_id],
+                            left_context_size=0,
+                            codec_chunk_frames=0,
+                            codec_left_context_frames=0,
+                            stream_finished=torch.tensor(True, dtype=torch.bool),
+                            finished=torch.tensor(True, dtype=torch.bool),
+                        ),
+                        request_id=req_id,
+                    )
+                return None
+            if not is_finished and pending < threshold:
+                return None
+            emit_frames = pending if is_finished else threshold
+            target_size = emit_frames
+
+            now = time.monotonic()
+            ctrl = AdaptiveChunkController(
+                first_emit_time=now,
+                last_emit_time=now,
+                ramp_divisor=adaptive_divisor,
+                ramp_delta_min=adaptive_delta_min,
+            )
+            _adaptive_states[req_id] = ctrl
+        else:
+            now = time.monotonic()
+            target_size = ctrl.compute_next_chunk_size(
+                now,
+                adaptive_min,
+                chunk_frames,
+                adaptive_margin,
+            )
+            emit, emit_frames = compute_adaptive_emit(
+                pending,
+                0,
+                target_size,
+                is_finished,
+            )
+            if not emit:
+                return None
+            if emit_frames == 0:
+                ctrl.log_summary(req_id)
+                transfer_manager.code_prompt_token_ids.pop(req_id, None)
+                transfer_manager.request_payload.pop(req_id, None)
+                return OmniPayloadStruct(
+                    codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                    meta=MetaStruct(
+                        req_id=[req_id],
+                        left_context_size=0,
+                        codec_chunk_frames=0,
+                        codec_left_context_frames=0,
+                        stream_finished=torch.tensor(True, dtype=torch.bool),
+                        finished=torch.tensor(True, dtype=torch.bool),
+                    ),
+                    request_id=req_id,
+                )
+
+        # Common emit path for adaptive.
+        actual_emit = min(emit_frames, pending)
+        chunk_rows = pending_frames[:actual_emit]
+        del pending_frames[:actual_emit]
+        chunk_np = np.stack([np.asarray(row, dtype=np.int64) for row in chunk_rows])
+        finished = bool(is_finished and len(pending_frames) == 0)
+
+        ctrl.record_emit(time.monotonic(), actual_emit, 0, target_size)
+        if finished:
+            ctrl.log_summary(req_id)
+
+        codec_flat = torch.from_numpy(np.ascontiguousarray(chunk_np.T).reshape(-1))
+
+        if finished:
+            transfer_manager.code_prompt_token_ids.pop(req_id, None)
+            transfer_manager.request_payload.pop(req_id, None)
+
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=codec_flat),
+            meta=MetaStruct(
+                req_id=[req_id],
+                left_context_size=0,
+                codec_chunk_frames=int(chunk_np.shape[0]),
+                codec_left_context_frames=0,
+                code_flat_numel=int(codec_flat.numel()),
+                stream_finished=torch.tensor(finished, dtype=torch.bool),
+                finished=torch.tensor(finished, dtype=torch.bool),
+            ),
+            request_id=req_id,
+        )
+
+    # Static ramp override (only when adaptive is not enabled).
     if ramp is not None:
         # The ladder is indexed by the connector's segment-local counter, which
         # restarts at each segment boundary; put_req_chunk is request-global.
         ramp_index = int(transfer_manager.ramp_chunk_count.get(req_id, 0))
         threshold = ramp_chunk_size(ramp_index, ramp, chunk_frames)
+
     if pending <= 0:
         if is_finished:
             transfer_manager.code_prompt_token_ids.pop(req_id, None)