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

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

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)