diff --git a/docs/design/feature/async_chunk_design.md b/docs/design/feature/async_chunk_design.md index 40b8f215a..5ae5f5967 100644 --- a/docs/design/feature/async_chunk_design.md +++ b/docs/design/feature/async_chunk_design.md @@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros For qwen3-omni: - **Thinker → Talker**: Per decode step (typically chunk_size=1) -- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA +- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override. - **Code2Wav**: Streaming decode with code2wav chunk_size With `async_chunk`: diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index ca5f7a2ac..2f0695298 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -96,7 +96,7 @@ Content-Type: application/json | `language` | string | "Auto" | Language (see supported languages below) | | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | integer | 2048 | Maximum tokens to generate | -| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) | +| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. | **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md index afadd39f4..352c076ad 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_tts.md +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -99,7 +99,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str ``` Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames` -and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation +in the stage config). The initial chunk size is dynamically selected based on server load for reduced +TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation completes. This demonstrates that audio data is available progressively rather than only at the end. > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged. diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md index c198f2c6e..ce2dd588e 100644 --- a/examples/offline_inference/qwen3_tts/README.md +++ b/examples/offline_inference/qwen3_tts/README.md @@ -96,7 +96,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str ``` Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames` -and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation +in the stage config). The initial chunk size is dynamically selected based on server load for reduced +TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation completes. This demonstrates that audio data is available progressively rather than only at the end. > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged. diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index 5088a4ca6..6984ebd42 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -289,7 +289,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w | `language` | string | "Auto" | Language (see supported languages below) | | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | int | 2048 | Maximum tokens to generate | -| `initial_codec_chunk_frames` | int | null | Initial chunk size for reduced TTFA (overrides stage config) | +| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. | | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) | **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 1b6fe2e4b..8d6d21430 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -7,16 +7,20 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + compute_dynamic_initial_chunk_size, + max_ic_for_chunk_size, +) from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( talker2code2wav, talker2code2wav_async_chunk, ) -_FRAME = [1, 2, 3, 4] # 4-codebook frame -_Q = len(_FRAME) # num quantizers +_FRAME = [1, 2, 3, 4] +_Q = len(_FRAME) -def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = None): +def _req(rid, *, finished, initial_codec_chunk_frames=None): ai = None if initial_codec_chunk_frames is not None: entry = SimpleNamespace(list_data=[initial_codec_chunk_frames]) @@ -28,9 +32,10 @@ def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = N ) -def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0): +def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1): return SimpleNamespace( code_prompt_token_ids=defaultdict(list), + scheduler_max_num_seqs=max_num_seqs, put_req_chunk=defaultdict(int), request_payload={}, connector=SimpleNamespace( @@ -38,16 +43,14 @@ def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0): "extra": { "codec_chunk_frames": chunk_frames, "codec_left_context_frames": left_context, - "initial_codec_chunk_frames": initial_chunk, } } ), ) -def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None): +def _call(tm, rid, *, n_frames, finished=False, req_ic=None): tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)] - tm.put_req_chunk[rid] = put_req return talker2code2wav_async_chunk( transfer_manager=tm, pooling_output={"audio_codes": torch.zeros((0,))}, @@ -56,84 +59,68 @@ def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None): ) -def test_does_not_emit_empty_chunk_when_not_finished(): +def test_empty_returns_none(): tm = _tm() - request = _req("rid-empty", finished=False) - - payload = talker2code2wav_async_chunk( + p = talker2code2wav_async_chunk( transfer_manager=tm, pooling_output={"audio_codes": torch.zeros((0,))}, - request=request, + request=_req("r", finished=False), ) - - assert payload is None + assert p is None -def test_flushes_tail_when_finished_without_pooler_output(): +def test_eof_marker_when_finished_empty(): tm = _tm() - rid = "rid-tail" - tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(24)] - request = _req(rid, finished=True) - - payload = talker2code2wav_async_chunk( + p = talker2code2wav_async_chunk( transfer_manager=tm, pooling_output=None, - request=request, + request=_req("r", finished=True), + is_finished=True, ) - - assert payload is not None - assert payload["finished"].item() is True - assert len(payload["code_predictor_codes"]) == _Q * 24 + assert p == {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)} -def test_emits_eof_marker_when_finished_with_no_frames(): +def test_flush_on_finish(): tm = _tm() - request = _req("rid-eof", finished=True) - - payload = talker2code2wav_async_chunk( + tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(24)] + p = talker2code2wav_async_chunk( transfer_manager=tm, pooling_output=None, - request=request, + request=_req("r", finished=True), + is_finished=True, ) - - assert payload == { - "code_predictor_codes": [], - "finished": torch.tensor(True, dtype=torch.bool), - } + assert p is not None + assert p["finished"].item() is True + assert len(p["code_predictor_codes"]) == _Q * 24 _CASES = [ - # Normal path (initial=0): emit at chunk_size boundaries - ((25, 25, 0), (24, 0, False), None), - ((25, 25, 0), (25, 0, False), (0, 25)), - # Initial-chunk phase: hold, first emit, second emit - ((25, 25, 10), (9, 0, False), None), - ((25, 25, 10), (10, 0, False), (0, 10)), - ((25, 25, 10), (20, 1, False), (10, 20)), - # Non-divisible: holds at chunk boundary - ((25, 25, 12), (25, 2, False), None), - # Normal phase: offset by initial_coverage (chunk//initial * initial) - ((25, 25, 10), (45, 2, False), (20, 45)), - # Second normal emit (offset must stay stable) - ((25, 25, 10), (70, 3, False), (25, 50)), - # initial >= chunk clamps to chunk_size (behaves as normal) - ((25, 25, 30), (25, 0, False), (0, 25)), - # finished=True flushes IC tail - ((25, 25, 10), (5, 0, True), (0, 5)), - # finished=True flushes non-divisible IC residual - ((25, 25, 12), (25, 2, True), (24, 25)), - # finished=True flushes normal phase tail - ((25, 25, 10), (30, 2, True), (20, 30)), + # Dynamic IC=16, cs=25, initial_coverage=16 + ((25, 25, 0), 24, False, None), # IC phase: 24%16!=0 -> hold + ((25, 25, 0), 25, False, None), # transition: adjusted=9, hold (no replay) + ((25, 25, 0), 41, False, (16, 41)), # first normal emit, lc=16 + # Per-request IC=10, cs=25, initial_coverage=20 + ((25, 25, 10), 9, False, None), # IC: hold + ((25, 25, 10), 10, False, (0, 10)), # IC: emit at boundary + ((25, 25, 10), 25, False, None), # transition: hold (no replay) + ((25, 25, 10), 45, False, (20, 45)), # first normal emit, lc=20 + ((25, 25, 10), 5, True, (0, 5)), # finished flushes IC tail + ((25, 25, 10), 33, True, (20, 33)), # finished flushes normal tail + # IC=8, cs=16: IC divides chunk_size evenly (edge case) + ((16, 25, 8), 8, False, (0, 8)), # IC: emit + ((16, 25, 8), 16, False, None), # transition: hold (no replay) + ((16, 25, 8), 24, False, (8, 24)), # first normal emit, lc=8 + # Per-request override: IC=15 at n_frames=10 -> 10%15!=0 -> hold + ((25, 25, 15), 10, False, None), ] -@pytest.mark.parametrize("config, state, expected", _CASES) -def test_streaming_decoding_with_variable_initial(config, state, expected): - chunk_frames, left_context, initial_chunk = config - n_frames, put_req, finished = state - - tm = _tm(chunk_frames=chunk_frames, left_context=left_context, initial_chunk=initial_chunk) - payload = _call(tm, "r", n_frames=n_frames, put_req=put_req, finished=finished) +@pytest.mark.parametrize("config, n_frames, finished, expected", _CASES) +def test_streaming_phases(config, n_frames, finished, expected): + chunk_frames, left_context, req_ic_val = config + tm = _tm(chunk_frames=chunk_frames, left_context=left_context) + req_ic = req_ic_val if req_ic_val > 0 else None + payload = _call(tm, "r", n_frames=n_frames, finished=finished, req_ic=req_ic) if expected is None: assert payload is None @@ -144,22 +131,86 @@ def test_streaming_decoding_with_variable_initial(config, state, expected): assert len(payload["code_predictor_codes"]) == _Q * exp_window -def test_per_request_override_activates_initial_phase(): - tm = _tm(initial_chunk=0) - payload = _call(tm, "r-override", n_frames=10, req_ic=10) - assert payload is not None - assert payload["left_context_size"] == 0 - assert len(payload["code_predictor_codes"]) == _Q * 10 +def test_dynamic_ic_adapts_to_load(): + # chunk_size=25 -> max_ic=16, steps=[2,4,8,16] + tm = _tm(max_num_seqs=8) + + # Low load (1/8) -> IC=2 -> emit at 2 + p1 = _call(tm, "r", n_frames=2) + assert p1 is not None + assert len(p1["code_predictor_codes"]) == _Q * 2 + + # High load: add 4 others -> active=5/8 -> IC=8 -> emit at 8 + for i in range(4): + tm.code_prompt_token_ids[f"other-{i}"] = [[0]] + p2 = _call(tm, "r", n_frames=8) + assert p2 is not None + assert len(p2["code_predictor_codes"]) == _Q * 8 + + # Requests past initial phase still count in load factor + tm2 = _tm(max_num_seqs=4) + for i in range(3): + tm2.code_prompt_token_ids[f"long-{i}"] = [[0]] * 50 # well past cs=25 + # active=4/4=1.0 -> IC=16 + p3 = _call(tm2, "new", n_frames=16) + assert p3 is not None + assert len(p3["code_predictor_codes"]) == _Q * 16 -def test_per_request_override_wins_over_stage_config(): - tm = _tm(initial_chunk=5) - payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15) - assert payload is None +def test_ic_load_change_mid_request(): + """IC stateless: load spike mid-request shifts initial_coverage.""" + tm = _tm(chunk_frames=25, left_context=25, max_num_seqs=8) + + # Low load -> IC=2 -> emit at frame 2 + p1 = _call(tm, "r", n_frames=2) + assert p1 is not None + + # Spike load: 6 others -> IC=16 -> initial_coverage=16 + for i in range(6): + tm.code_prompt_token_ids[f"other-{i}"] = [[0]] * 10 + + # adjusted=25-16=9, 9%25!=0 -> hold + assert _call(tm, "r", n_frames=25) is None + + # First normal emit at 16+25=41 + p3 = _call(tm, "r", n_frames=41) + assert p3 is not None + assert p3["left_context_size"] == 16 + + +@pytest.mark.parametrize( + "active,max_bs,max_ic,expected", + [ + (0, 4, 32, 2), # zero load -> min step + (2, 4, 32, 8), # mid load + (4, 4, 32, 32), # full load + (10, 4, 16, 16), # over capacity, capped + (0, 4, 1, 1), # max_ic below min step + (0, 0, 16, 2), # zero capacity edge case + ], +) +def test_compute_dynamic_initial_chunk_size(active, max_bs, max_ic, expected): + assert compute_dynamic_initial_chunk_size(active, max_bs, max_ic) == expected + + +@pytest.mark.parametrize( + "chunk_size,expected", + [ + (25, 16), + (50, 32), + (70, 64), + (8, 4), + (4, 2), + (2, 1), + (1, 1), + ], +) +def test_max_ic_for_chunk_size(chunk_size, expected): + assert max_ic_for_chunk_size(chunk_size) == expected def test_first_streaming_chunk_prepends_ref_code_context(): - tm = _tm(initial_chunk=10) + tm = _tm() rid = "r-ref" tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(10)] ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) @@ -177,7 +228,7 @@ def test_first_streaming_chunk_prepends_ref_code_context(): def test_ref_code_context_only_applies_to_first_streaming_chunk(): - tm = _tm(initial_chunk=10) + tm = _tm() rid = "r-ref2" tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)] tm.put_req_chunk[rid] = 1 @@ -196,7 +247,7 @@ def test_ref_code_context_only_applies_to_first_streaming_chunk(): def test_ref_code_context_can_be_buffered_before_first_emit(): - tm = _tm(initial_chunk=10) + tm = _tm() rid = "r-ref-buffered" ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index db8d5dbaa..d96ef74db 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -236,6 +236,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.put_req_chunk[request_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") + if is_finished: + self.code_prompt_token_ids.pop(request_id, None) + ######################################################################## # Cleanup ######################################################################## diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 1efab8ebe..871787173 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -58,7 +58,7 @@ class OpenAICreateSpeechRequest(BaseModel): initial_codec_chunk_frames: int | None = Field( default=None, ge=0, - description="Initial chunk size for reduced TTFA. Overrides stage config for this request.", + description="Per-request initial chunk size override. If null, computed dynamically based on server load.", ) @field_validator("stream_format") diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index de278edf7..5ff259193 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py @@ -29,7 +29,7 @@ class CUDAGraphDecoderWrapper: output = wrapper.decode(codes) # Automatically uses CUDA graph if possible """ - DEFAULT_CAPTURE_SIZES = [25, 50, 100, 150, 200, 250, 300] + DEFAULT_CAPTURE_SIZES = [2, 4, 8, 16, 25, 32, 50, 100, 150, 200, 250, 300] def __init__( self, diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index 19502add8..095aafe08 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -96,10 +96,6 @@ runtime: # Align with Omni: small chunks with sufficient context overlap. codec_chunk_frames: 25 codec_left_context_frames: 25 - # First chunk size for reduced TTFA (0 = disabled). - # When > 0, emits small chunks every N frames during the initial phase, - # then switches to codec_chunk_frames cadence. - initial_codec_chunk_frames: 0 edges: - from: 0 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml index 43df70fc2..d73739109 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml @@ -98,10 +98,6 @@ runtime: # Align with Omni: small chunks with sufficient context overlap. codec_chunk_frames: 25 codec_left_context_frames: 25 - # First chunk size for reduced TTFA (0 = disabled). - # When > 0, emits small chunks every N frames during the initial phase, - # then switches to codec_chunk_frames cadence. - initial_codec_chunk_frames: 0 edges: - from: 0 diff --git a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py new file mode 100644 index 000000000..7bdec0bb3 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def max_ic_for_chunk_size(chunk_size: int) -> int: + """Largest power of 2 strictly less than chunk_size.""" + if chunk_size <= 2: + return 1 + return 1 << ((chunk_size - 1).bit_length() - 1) + + +def compute_dynamic_initial_chunk_size( + active_requests: int, + max_batch_size: int, + max_ic: int, +) -> int: + """Select IC from power-of-2 steps [2, 4, ..., max_ic] based on load factor. + + - Low load: small IC (faster TTFA). + - High load: large IC (amortise decode cost). + """ + steps: list[int] = [] + v = 2 + while v <= max_ic: + steps.append(v) + v <<= 1 + if not steps: + return max(1, max_ic) + if max_batch_size <= 0: + return steps[0] + load_factor = min(active_requests / max_batch_size, 1.0) + idx = int(round(load_factor * (len(steps) - 1))) + return steps[idx] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index eb9753262..435e45f03 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -5,6 +5,11 @@ from typing import Any import torch from vllm.logger import init_logger +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + compute_dynamic_initial_chunk_size, + max_ic_for_chunk_size, +) + logger = init_logger(__name__) @@ -14,7 +19,7 @@ def talker2code2wav( prompt: Any = None, requires_multimodal_data: bool = False, ) -> list[Any]: - """Non-async processor: wait for talker to finish, then pass all codes to code2wav at once.""" + """Non-async: collect all talker codes, then pass to code2wav at once.""" from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs @@ -94,9 +99,12 @@ def talker2code2wav_async_chunk( cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} chunk_size = int(cfg.get("codec_chunk_frames", 25)) left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) - initial_chunk_size = int(cfg.get("initial_codec_chunk_frames", 0)) - # Per-request override (takes priority over stage config) + + # Per-request override takes priority over dynamic IC. + per_request_override = False + initial_chunk_size = 0 additional_information = getattr(request, "additional_information", None) + if ( additional_information is not None and hasattr(additional_information, "entries") @@ -105,12 +113,32 @@ def talker2code2wav_async_chunk( entry = additional_information.entries["initial_codec_chunk_frames"] if entry.list_data is not None and len(entry.list_data) == 1: initial_chunk_size = int(entry.list_data[0]) + per_request_override = True + + # Dynamic IC: always derived from chunk_size, recomputed every call to adapt to load. + if not per_request_override: + max_ic = max_ic_for_chunk_size(chunk_size) + active = sum(1 for v in transfer_manager.code_prompt_token_ids.values() if len(v) > 0) + capacity = getattr(transfer_manager, "scheduler_max_num_seqs", 1) + initial_chunk_size = compute_dynamic_initial_chunk_size(active, capacity, max_ic) + logger.debug( + "Dynamic IC: active=%d, capacity=%d, max_ic=%d, ic=%d, cs=%d, req=%s, keys=%d", + active, + capacity, + max_ic, + initial_chunk_size, + chunk_size, + request_id, + len(transfer_manager.code_prompt_token_ids), + ) + if chunk_size <= 0 or left_context_size_config < 0 or initial_chunk_size < 0: raise ValueError( f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " f"codec_left_context_frames={left_context_size_config}, " f"initial_codec_chunk_frames={initial_chunk_size}" ) + if initial_chunk_size > chunk_size: logger.warning( "initial_codec_chunk_frames=%d > codec_chunk_frames=%d, clamping to codec_chunk_frames.", @@ -120,8 +148,6 @@ def talker2code2wav_async_chunk( initial_chunk_size = chunk_size length = len(transfer_manager.code_prompt_token_ids[request_id]) - # Avoid emitting empty chunks during normal streaming. If the request is - # finished and nothing was produced, emit an EOF marker. if length <= 0: if finished: return { @@ -130,34 +156,32 @@ def talker2code2wav_async_chunk( } return None - in_initial_phase = initial_chunk_size > 0 and length <= chunk_size + in_initial_phase = initial_chunk_size > 0 and initial_chunk_size < chunk_size and length < chunk_size if in_initial_phase: - # Initial-chunk phase: emit every initial_chunk_size frames with full accumulated context. - already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size - pending = length - already_sent - if pending <= 0: + # IC phase: emit every initial_chunk_size frames with growing left context. + if not finished and length % initial_chunk_size != 0: return None - if pending < initial_chunk_size and not finished: - return None - context_length = min(pending, initial_chunk_size) - end_index = length - left_context_size = max(0, length - context_length) - window_frames = transfer_manager.code_prompt_token_ids[request_id][:length] + context_length = ( + length % initial_chunk_size if (finished and length % initial_chunk_size != 0) else initial_chunk_size + ) else: - # Normal phase: standard chunk_size cadence with left_context sliding window. - # Offset by initial_coverage so normal starts from where the initial-chunk phase left off. - initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0 + # Normal phase: offset so the first normal emit picks up after IC phase. + # IC is stateless (may change with load); any mismatch is absorbed by left_context. + initial_coverage = ( + ((chunk_size - 1) // initial_chunk_size) * initial_chunk_size if 0 < initial_chunk_size < chunk_size else 0 + ) adjusted = length - initial_coverage - chunk_length = adjusted % chunk_size - if chunk_length != 0 and not finished: + if not finished and adjusted % chunk_size != 0: return None + chunk_length = adjusted % chunk_size context_length = chunk_length if chunk_length != 0 else chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, int(end_index - context_length)) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] - # Match offline Base ICL decoding by prepending ref_code only once to the first decoder window. + end_index = min(length, left_context_size_config + context_length) + left_context_size = max(0, end_index - context_length) + window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + # Prepend ref_code to the first decoder window only. if transfer_manager.put_req_chunk[request_id] == 0: ref_code = request_payload.pop(request_id, None) if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: @@ -165,7 +189,6 @@ def talker2code2wav_async_chunk( window_frames = ref_frames + window_frames left_context_size += len(ref_frames) - # Pack context + chunk into codebook-major flat codes for adapter. code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist() return { diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml index 5ad0cc90e..60659a976 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml @@ -96,10 +96,6 @@ runtime: # Align with Omni: small chunks with sufficient context overlap. codec_chunk_frames: 25 codec_left_context_frames: 25 - # First chunk size for reduced TTFA (0 = disabled). - # When > 0, emits small chunks every N frames during the initial phase, - # then switches to codec_chunk_frames cadence. - initial_codec_chunk_frames: 0 edges: - from: 0