diff --git a/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py new file mode 100644 index 000000000..bff3e31e3 --- /dev/null +++ b/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 diff --git a/vllm_omni/deploy/ming_tts.yaml b/vllm_omni/deploy/ming_tts.yaml index 1a2dc99ec..684a47f98 100644 --- a/vllm_omni/deploy/ming_tts.yaml +++ b/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: diff --git a/vllm_omni/deploy/ming_tts_moe.yaml b/vllm_omni/deploy/ming_tts_moe.yaml index fe7806b1d..48c5e9aa0 100644 --- a/vllm_omni/deploy/ming_tts_moe.yaml +++ b/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: diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py index a83989a03..fdac2a30e 100644 --- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py index a189e4285..e0508bda7 100644 --- a/vllm_omni/model_executor/models/ming_tts/constants.py +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_tts/validation.py b/vllm_omni/model_executor/models/ming_tts/validation.py index c4ccf9dcd..3f3da6cf0 100644 --- a/vllm_omni/model_executor/models/ming_tts/validation.py +++ b/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: diff --git a/vllm_omni/model_executor/stage_input_processors/ming_tts.py b/vllm_omni/model_executor/stage_input_processors/ming_tts.py index 0f2effdcb..058cda788 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_tts.py +++ b/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)