[Perf][Ming-TTS] Add initial latent chunk for lower streaming TTFP (#5011)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>

tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.ming_tts.constants import LATENT_DIM, PATCH_SIZE
+from vllm_omni.model_executor.stage_input_processors.ming_tts import (
+    MING_EMIT_PATCH_COUNT_KEY,
+    llm2audio_vae_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_transfer_manager(*, chunk_size: int = 5, initial_chunk_size: int = 2):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        request_payload={},
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "latent_chunk_size": chunk_size,
+                    "initial_latent_chunk_size": initial_chunk_size,
+                    "latent_left_context": 0,
+                }
+            }
+        ),
+    )
+
+
+def _make_request(req_id: str = "req", *, finished: bool = False):
+    return SimpleNamespace(
+        external_req_id=req_id,
+        is_finished=lambda: finished,
+    )
+
+
+def _make_output(value: float):
+    return {
+        "ming_has_patch": torch.tensor([1], dtype=torch.bool),
+        "ming_latent_patch": torch.full((1, PATCH_SIZE, LATENT_DIM), value),
+    }
+
+
+def _append_patch(tm, req_id: str, idx: int, *, finished: bool = False):
+    return llm2audio_vae_async_chunk(
+        tm,
+        _make_output(float(idx)),
+        _make_request(req_id, finished=finished),
+        is_finished=finished,
+    )
+
+
+def test_ming_async_chunk_emits_small_initial_chunk_then_steady_chunks():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    assert _append_patch(tm, "req", 0) is None
+
+    first = _append_patch(tm, "req", 1)
+    assert first is not None
+    assert first.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
+    assert first.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
+    assert first.meta.finished.item() is False
+
+    for idx in range(2, 6):
+        assert _append_patch(tm, "req", idx) is None
+
+    second = _append_patch(tm, "req", 6)
+    assert second is not None
+    assert second.latent.shape == (5, PATCH_SIZE, LATENT_DIM)
+    assert second.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 5
+
+
+def test_ming_async_chunk_flushes_short_final_chunk():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    payload = _append_patch(tm, "req", 0, finished=True)
+
+    assert payload is not None
+    assert payload.latent.shape == (1, PATCH_SIZE, LATENT_DIM)
+    assert payload.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 1
+    assert payload.meta.finished.item() is True
+
+
+def test_ming_async_chunk_flushes_leftover_after_initial_chunk():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    assert _append_patch(tm, "req", 0) is None
+    assert _append_patch(tm, "req", 1) is not None
+    assert _append_patch(tm, "req", 2) is None
+
+    final = _append_patch(tm, "req", 3, finished=True)
+
+    assert final is not None
+    assert final.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
+    assert final.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
+    assert final.meta.finished.item() is True

vllm_omni/deploy/ming_tts.yaml

@@ -15,6 +15,7 @@ connectors:
     name: SharedMemoryConnector
     extra:
       latent_chunk_size: 25
+      initial_latent_chunk_size: 4
       latent_left_context: 0
 
 stages:

vllm_omni/deploy/ming_tts_moe.yaml

@@ -23,6 +23,7 @@ connectors:
     name: SharedMemoryConnector
     extra:
       latent_chunk_size: 25
+      initial_latent_chunk_size: 4
       latent_left_context: 0
 
 stages:

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -22,6 +22,7 @@ from .constants import (
     DEFAULT_SIGMA,
     DEFAULT_TEMPERATURE,
     HISTORY_PATCH_SIZE,
+    INITIAL_LATENT_CHUNK_SIZE,
     KEY_CFG,
     KEY_CHUNK_ID,
     KEY_DECODE_STEP,
@@ -239,6 +240,7 @@ class MingTTSConfig:
     max_decode_steps: int = MAX_DECODE_STEPS
 
     latent_chunk_size: int = LATENT_CHUNK_SIZE
+    initial_latent_chunk_size: int = INITIAL_LATENT_CHUNK_SIZE
     latent_left_context: int = LATENT_LEFT_CONTEXT
 
     text_eos_token_id: int = TEXT_EOS_TOKEN_ID

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -58,6 +58,7 @@ DEFAULT_TEMPERATURE = 0.0
 
 # Connector / Stage-2 streaming defaults (runtime tuning)
 LATENT_CHUNK_SIZE = 25
+INITIAL_LATENT_CHUNK_SIZE = 4
 LATENT_LEFT_CONTEXT = 0
 MAX_DECODE_STEPS = 200
 

vllm_omni/model_executor/models/ming_tts/validation.py

@@ -172,6 +172,9 @@ def validate_ming_tts_config(cfg: Any) -> None:
 
     if cfg.latent_chunk_size <= 0:
         raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
+    initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0)
+    if initial_latent_chunk_size < 0:
+        raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.")
     if cfg.latent_left_context < 0:
         raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
     if cfg.max_decode_steps <= 0:

vllm_omni/model_executor/stage_input_processors/ming_tts.py

@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    INITIAL_LATENT_CHUNK_SIZE,
     KEY_CHUNK_ID,
     KEY_REQUEST_ID,
     LATENT_CHUNK_SIZE,
@@ -116,7 +117,7 @@ def _decode_stop_reason(value: Any) -> str | None:
     return MING_STOP_REASON_BY_CODE.get(int(value))
 
 
-def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
+def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]:
     connector = getattr(transfer_manager, "connector", None)
     raw_cfg = getattr(connector, "config", {}) or {}
     cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, Mapping) else {}
@@ -131,9 +132,19 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
             LATENT_LEFT_CONTEXT,
         )
     chunk_size = int(cfg.get("latent_chunk_size", LATENT_CHUNK_SIZE))
+    initial_chunk_size = int(cfg.get("initial_latent_chunk_size", INITIAL_LATENT_CHUNK_SIZE))
     left_context = int(cfg.get("latent_left_context", LATENT_LEFT_CONTEXT))
     if chunk_size <= 0:
         raise ValueError(f"Invalid Ming latent_chunk_size={chunk_size}")
+    if initial_chunk_size < 0:
+        raise ValueError(f"Invalid Ming initial_latent_chunk_size={initial_chunk_size}")
+    if initial_chunk_size > chunk_size:
+        logger.warning(
+            "Ming initial_latent_chunk_size=%d > latent_chunk_size=%d, clamping to latent_chunk_size.",
+            initial_chunk_size,
+            chunk_size,
+        )
+        initial_chunk_size = chunk_size
     # Stage-2 VAE caches past_key_values and stream_state by request_id.
     # Replaying left-context latents would double-feed cached decoder state.
     if left_context != 0:
@@ -142,7 +153,7 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
             "Ming boundary continuity is handled by per-request decoder state cache, not "
             f"latent replay. Got latent_left_context={left_context}."
         )
-    return chunk_size, left_context
+    return chunk_size, initial_chunk_size, left_context
 
 
 def _build_chunk_observability(
@@ -196,7 +207,7 @@ def llm2audio_vae_async_chunk(
     if patch is not None:
         transfer_manager.code_prompt_token_ids[request_id].append(patch)
 
-    chunk_size, _ = _get_async_chunk_config(transfer_manager)
+    chunk_size, initial_chunk_size, _ = _get_async_chunk_config(transfer_manager)
 
     patches = transfer_manager.code_prompt_token_ids[request_id]
     seen_patch_len = int(state.get("seen_patch_len", 0))
@@ -225,11 +236,17 @@ def llm2audio_vae_async_chunk(
             )
         return None
 
-    chunk_length = length % chunk_size
-    if chunk_length != 0 and not finished:
-        return None
-
-    emit_count = chunk_length if chunk_length != 0 else chunk_size
+    use_first_chunk = 0 < initial_chunk_size < chunk_size and seen_patch_len == 0
+    if finished:
+        emit_count = length
+    elif use_first_chunk:
+        if length < initial_chunk_size:
+            return None
+        emit_count = initial_chunk_size
+    else:
+        if length < chunk_size:
+            return None
+        emit_count = chunk_size
     emit_patches = list(new_patches[:emit_count])
     state["seen_patch_len"] = seen_patch_len + len(emit_patches)
     latent_patches = torch.stack(emit_patches, dim=0)