From 4ad455e61e2857450594189895fe860596f402dc Mon Sep 17 00:00:00 2001 From: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:11:05 +0800 Subject: [PATCH] [Feature][TTS] adaptive chunk ramp (Phase 2 buffer-feedback controller) (#6001) Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> --- .../test_qwen3_tts_async_chunk.py | 527 +++++++++++++++++- vllm_omni/deploy/qwen3_tts.yaml | 11 + .../deploy/qwen3_tts_high_concurrency.yaml | 11 + .../chunk_transfer_adapter.py | 3 + .../chunk_size_utils.py | 227 ++++++++ .../stage_input_processors/qwen3_tts.py | 97 +++- .../omni_connector_model_runner_mixin.py | 3 + 7 files changed, 871 insertions(+), 8 deletions(-) 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 013c0b99e..996486d49 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 @@ -8,9 +8,12 @@ import pytest import torch from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + AdaptiveChunkController, + compute_adaptive_emit, compute_dynamic_initial_chunk_size, compute_ramp_emit, max_ic_for_chunk_size, + parse_adaptive_config, parse_chunk_ramp, ramp_chunk_size, ramp_cumulative, @@ -41,7 +44,17 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None): ) -def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frames=0, chunk_ramp=None): +def _tm( + *, + chunk_frames=25, + left_context=25, + max_num_seqs=1, + initial_chunk_frames=0, + chunk_ramp=None, + adaptive=False, + adaptive_min=2, + adaptive_margin=200.0, +): extra = { "codec_chunk_frames": chunk_frames, "codec_left_context_frames": left_context, @@ -49,11 +62,16 @@ def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frame } if chunk_ramp is not None: extra["codec_chunk_ramp"] = chunk_ramp + if adaptive: + extra["codec_chunk_adaptive"] = True + extra["codec_chunk_min_frames"] = adaptive_min + extra["codec_chunk_safety_margin_ms"] = adaptive_margin return SimpleNamespace( code_prompt_token_ids=defaultdict(list), scheduler_max_num_seqs=max_num_seqs, put_req_chunk=defaultdict(int), ramp_chunk_count=defaultdict(int), + _adaptive_states={}, request_payload={}, connector=SimpleNamespace(config={"extra": extra}), ) @@ -605,18 +623,32 @@ class TestRampHelpers: def test_parse_ramp_warns_on_tail_mismatch(self, caplog): import logging - with caplog.at_level(logging.WARNING): + target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) assert result == [4, 4, 8] - assert any("reintroduces" in r.message for r in caplog.records) + assert "reintroduces" in caplog.text def test_parse_ramp_no_warn_on_tail_match(self, caplog): import logging - with caplog.at_level(logging.WARNING): + target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) assert result == [4, 4, 8, 16, 25] - assert not any("reintroduces" in r.message for r in caplog.records) + assert "reintroduces" not in caplog.text @pytest.mark.parametrize( "index,ramp,steady,expected", @@ -941,3 +973,488 @@ class TestChunkRampEmission: p_seg2_0 = self._emit(tm, rid, 4) assert p_seg2_0 is not None assert p_seg2_0.meta.left_context_size == 0 + + +class TestAdaptiveController: + def test_compute_next_chunk_size_basic(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=20.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=10, + ) + now = 0.5 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=100.0) + assert 2 <= size <= 25 + + def test_compute_next_chunk_size_clamp_min(self): + """min_frames floor holds when ramp_floor (last_target + delta) is + below min_frames. Uses a small EWMA so delta=2 and last_target=0, + i.e. ramp_floor=2 < min_frames=4 -> greedy=min_frames wins.""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=1, + ) + # buffer = 80 - 70 = 10ms (positive but < margin=200) + # greedy: available=0, max_n=0, greedy=min_frames=4 + # ramp_floor: last_target=0 + delta=max(2,int(10/15))=2 -> 2 < 4 + now = 0.07 + size = ctrl.compute_next_chunk_size(now, min_frames=4, max_frames=25, safety_margin_ms=200.0) + assert size == 4 # min_frames floor dominates ramp_floor=2 + + def test_compute_next_chunk_size_clamp_max(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=100, + ) + now = 1.0 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size == 25 + + def test_hysteresis_locks_steady(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=50, + ) + now = 0.1 + size1 = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert ctrl.steady_state_reached + assert size1 == 25 + size2 = ctrl.compute_next_chunk_size(now + 10.0, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size2 == 25 + + def test_ramp_floor_overrides_greedy_when_buffer_low(self): + """When buffer is negative (overloaded), ramp_floor still forces + gradual growth to avoid stagnating at min_frames (which would cause + too many Code2Wav decode passes and RTF explosion).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ) + # buffer = 160 - 200 = -40ms (negative) + # greedy: available=0, max_n=0, greedy=2 + # ramp_floor: last_target=4 + delta=max(2,int(50/15))=3 -> 7 + now = 0.2 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 7 # ramp_floor (unconditional) dominates greedy=2 + + def test_ramp_floor_dominates_when_last_target_set(self): + """Once last_target is set (post-emit), the unconditional ramp_floor + (last_target + ewma-scaled delta) dominates greedy whenever greedy is + smaller. This is intended: chunk size grows monotonically toward max + to minimize Code2Wav decode passes (pre-W5; see W5 coupling).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=4, + last_target=4, + ) + # buffer = 320 - 100 = 220ms (positive) + # greedy: available=170, max_n=int(170/50)=3, greedy=3 + # ramp_floor: 4 + delta=max(2,int(50/15))=3 -> 7 + now = 0.1 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size == 7 # ramp_floor (7) dominates greedy (3) + + def test_greedy_dominates_when_ramp_floor_low(self): + """When last_target=0 (fresh, no prior emit) and buffer is positive + but below the hysteresis threshold, greedy selects an intermediate + size larger than ramp_floor and greedy dominates. This is the only + regime where greedy (not ramp_floor or hysteresis) sets the size.""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=8, + last_target=0, + ) + # buffer = 640 - 100 = 540ms (positive, < hysteresis 1300ms) + # greedy: available=490, max_n=int(490/50)=9, greedy=9 + # ramp_floor: 0 + delta=max(2,int(50/15))=3 -> 3 + now = 0.1 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 9 # greedy (9) dominates ramp_floor (3) + + def test_record_emit_ewma_update(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=20.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=4, + ) + ctrl.record_emit(now=0.5, n_frames=8, accumulated=12, target_size=8) + assert ctrl.frames_emitted_this_segment == 12 + assert ctrl.accumulated_at_last_emit == 12 + assert ctrl.last_emit_time == 0.5 + assert ctrl.last_target == 8 # D6: stores target, not actual + measured = (0.5 - 0.0) * 1000.0 / 8 + expected = 0.3 * measured + 0.7 * 20.0 + assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01 + + def test_record_emit_first_then_second(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=0.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=0, + ) + ctrl.record_emit(now=0.4, n_frames=4, accumulated=4, target_size=4) + assert ctrl.frame_time_ewma_ms == 0.0 + assert ctrl.frames_emitted_this_segment == 4 + assert ctrl.last_target == 4 + + ctrl.record_emit(now=0.8, n_frames=4, accumulated=8, target_size=4) + expected = (0.8 - 0.4) * 1000.0 / 4 + assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01 + + def test_record_emit_collects_underrun(self): + """record_emit mirrors stream_client.py's playback-cursor underrun + model: per_chunk_gap_ms records each chunk's lateness, total_gap_s + accumulates gaps > 5ms. Deterministic given explicit `now`.""" + ctrl = AdaptiveChunkController( + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=0, + ) + # chunk0: now=0.0 (= first_emit_time), n=2 -> 0.16s audio + ctrl.record_emit(now=0.0, n_frames=2, accumulated=2, target_size=2) + assert ctrl.chunk_sizes == [2] + assert ctrl.per_chunk_gap_ms == [0.0] + assert ctrl.total_gap_s == 0.0 + assert abs(ctrl.playback_cursor_s - 0.16) < 1e-9 + # chunk1: now=0.2 (200ms wall) -> arrival 0.2, late by 0.2-0.16=0.04s + ctrl.record_emit(now=0.2, n_frames=2, accumulated=4, target_size=2) + assert ctrl.chunk_sizes == [2, 2] + assert abs(ctrl.per_chunk_gap_ms[1] - 40.0) < 1e-6 + assert abs(ctrl.total_gap_s - 0.04) < 1e-9 + assert abs(ctrl.playback_cursor_s - 0.36) < 1e-9 + + def test_ewma_zero_skips_hysteresis(self): + """When EWMA=0, hysteresis must not trigger (steady_gen_ms=0 + would make any positive buffer satisfy the condition).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=0.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=10, + ) + now = 0.001 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 2 # min_frames (EWMA=0 → max_n=min_frames) + + def test_ramp_divisor_configurable(self): + """ramp_divisor/ramp_delta_min are controller fields wired from config. + Defaults (15/2) preserve the original formula; changing them changes + delta. Lets a deploy tune the ramp via yaml without code changes, while + the unit tests (which use defaults) stay green.""" + # default divisor=15: ewma=50 -> delta=max(2,int(50/15))=3 -> floor=4+3=7 + c1 = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ) + assert c1.compute_next_chunk_size(0.2, 2, 25, 50.0) == 7 + # divisor=10: delta=max(2,int(50/10))=5 -> floor=4+5=9 + c2 = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ramp_divisor=10, + ) + assert c2.compute_next_chunk_size(0.2, 2, 25, 50.0) == 9 + # delta_min=4 + ewma=10 (int(10/15)=0): delta=max(4,0)=4 -> floor=4+4=8 + c3 = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ramp_delta_min=4, + ) + assert c3.compute_next_chunk_size(0.2, 2, 25, 50.0) == 8 + + +class TestAdaptiveEmitHelper: + @pytest.mark.parametrize( + "length,accumulated,target,finished,expected_emit,expected_ctx", + [ + (4, 0, 4, False, True, 4), + (3, 0, 4, False, False, 0), + (12, 4, 8, False, True, 8), + (10, 4, 8, False, False, 0), + # D3: finished flushes ALL remaining frames, not just target_size + (6, 4, 8, True, True, 2), + (4, 4, 8, True, True, 0), + (20, 4, 8, True, True, 16), # new_frames=16 > target=8 → flush 16 + (30, 0, 25, True, True, 30), # finished, new_frames=30 > target=25 → flush 30 + # B1: non-finished path also ships ALL accumulated frames (new_frames), + # not just target_size. When new_frames > target_size (e.g. scheduler + # stall dropped target below queued frames), returning target_size + # would cause the [-context_length:] slice to skip surplus frames. + (10, 0, 4, False, True, 10), # new_frames=10 > target=4 → ship 10 + (8, 2, 4, False, True, 6), # new_frames=6 > target=4 → ship 6 + (100, 0, 25, False, True, 100), # new_frames=100 > target=25 → ship 100 + ], + ) + def test_compute_adaptive_emit(self, length, accumulated, target, finished, expected_emit, expected_ctx): + emit, ctx = compute_adaptive_emit(length, accumulated, target, finished) + assert emit == expected_emit + assert ctx == expected_ctx + + +class TestAdaptiveConfigParsing: + def test_disabled_by_default(self): + enabled, min_f, margin, divisor, delta_min = parse_adaptive_config({}) + assert enabled is False + assert min_f == 2 + assert margin == 50.0 + assert divisor == 15 + assert delta_min == 2 + + def test_enabled(self): + cfg = {"codec_chunk_adaptive": True} + enabled, *_ = parse_adaptive_config(cfg) + assert enabled is True + + def test_custom_values(self): + cfg = { + "codec_chunk_adaptive": True, + "codec_chunk_min_frames": 4, + "codec_chunk_safety_margin_ms": 300.0, + "codec_chunk_ramp_divisor": 10, + "codec_chunk_ramp_delta_min": 3, + } + enabled, min_f, margin, divisor, delta_min = parse_adaptive_config(cfg) + assert enabled is True + assert min_f == 4 + assert margin == 300.0 + assert divisor == 10 + assert delta_min == 3 + + def test_min_frames_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_min_frames": 0} + _, min_f, _, _, _ = parse_adaptive_config(cfg) + assert min_f == 1 + + def test_ramp_divisor_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_divisor": 0} + _, _, _, divisor, _ = parse_adaptive_config(cfg) + assert divisor == 1 + + def test_ramp_delta_min_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_delta_min": 0} + _, _, _, _, delta_min = parse_adaptive_config(cfg) + assert delta_min == 1 + + +class TestAdaptiveEmission: + def _emit(self, tm, rid, n_frames, finished=False): + tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)] + return talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + + def test_adaptive_chunk0_uses_dynamic_ic(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-ic" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_chunk0_holds_until_ic(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-hold" + p = self._emit(tm, rid, 1) + assert p is None + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_finished_flush(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-flush" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 5, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + + def test_adaptive_finished_no_new_frames(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-no-new" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 2, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + assert p_fin.codes.audio.numel() == 0 + + def test_adaptive_precedence_over_fixed_ramp(self): + tm = _tm(chunk_ramp=[4, 4, 8, 16, 25], adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-precedence" + + p = self._emit(tm, rid, 1) + assert p is None + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_backward_compat(self): + tm = _tm(initial_chunk_frames=1) + rid = "no-adaptive" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 1 + + def test_adaptive_segment_reset(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-segment" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + tm.put_req_chunk[rid] = 1 + + tm.code_prompt_token_ids[rid] = [] + tm.ramp_chunk_count.pop(rid, None) + tm._adaptive_states.pop(rid, None) + + p_seg2 = self._emit(tm, rid, 2) + assert p_seg2 is not None + assert p_seg2.meta.left_context_size == 0 + + def test_adaptive_with_ref_code_first_chunk(self): + tm = _tm(adaptive=True, adaptive_min=2, left_context=25, max_num_seqs=8) + rid = "adaptive-ref" + tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(2)] + ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) + + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}}, + request=_req(rid, finished=False), + is_finished=False, + ) + + assert payload is not None + assert payload.meta.ref_context_size == 2 + assert payload.meta.ref_context_included is True + assert payload.meta.left_context_size == 2 + + +class TestAdaptiveAccumulationPath: + """Regression test for the frame-skipping bug (PR #6001 review B1). + + Appends one frame per call with a controlled clock, asserting that + the concatenation of every emitted chunk equals the accumulated + frame sequence, in order, with no gaps. Each frame has a distinct + value so the assertion can detect a skipped frame rather than just + a short total. + """ + + def test_no_frames_skipped_when_target_drops(self, monkeypatch): + """When greedy target_size is high (large buffer) and then drops + (clock jumps), new_frames > target_size at the emit point. + compute_adaptive_emit now returns new_frames (not target_size), + so all accumulated frames are shipped and the [-context_length:] + slice stays contiguous with the previous emit. last_target is set + to target_size (not n_frames) to prevent the ramp floor from + ratcheting up on transient overshoot. + + Clock schedule (max_num_seqs=4 gives IC=2 via dynamic IC): + t=0.0 chunk 0 (IC=2), emit 2. ewma=0. + t=0.05 emit target=4 (ramp_floor). ewma=12.5 (direct assign). + t=0.05 greedy dominates (target=18), 17 frames accumulate (holds). + t=100.0 clock jumps -> greedy collapses, target drops to + ramp_floor=6. new_frames=18 > 6 -> emit 18 (all new frames). + last_target=6 (not 18, preventing floor ratchet). + """ + tm = _tm( + adaptive=True, + adaptive_min=2, + max_num_seqs=4, + ) + rid = "adaptive-accum" + tm.code_prompt_token_ids[rid] = [] + + clock = [0.0] + monkeypatch.setattr( + "vllm_omni.model_executor.stage_input_processors.qwen3_tts.time.monotonic", + lambda: clock[0], + ) + + emitted_indices: list[int] = [] + total_frames = 40 + + def append_and_emit(frame_idx, finished=False): + tm.code_prompt_token_ids[rid].append([frame_idx] * _Q) + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + if p is not None: + if p.codes.audio.numel() > 0: + n = p.codes.audio.numel() // _Q + emitted_indices.extend(p.codes.audio[:n].tolist()) + tm.ramp_chunk_count[rid] += 1 + return p + + # Phase 1 (t=0.0): chunk 0, IC=2 + for i in range(2): + append_and_emit(i) + + # Phase 2 (t=0.05): build EWMA, emit target=4 at length=6 + clock[0] = 0.05 + for i in range(2, 6): + append_and_emit(i) + + # Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate + # 17 frames while holding (new_frames 1..17 < 18) + for i in range(6, 23): + append_and_emit(i) + + # Phase 4 (t=100.0): clock jumps, greedy collapses, target drops + # to ramp_floor=6. new_frames=18 > 6 -> bug trigger point. + clock[0] = 100.0 + for i in range(23, total_frames): + append_and_emit(i, finished=(i == total_frames - 1)) + + # Assert: all frames 0..39 emitted, in order, no gaps. + expected = list(range(total_frames)) + assert emitted_indices == expected, ( + f"Frame loss detected! " + f"Missing: {sorted(set(expected) - set(emitted_indices))}, " + f"Extra: {sorted(set(emitted_indices) - set(expected))}" + ) diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 58b191f77..747b68852 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -33,6 +33,17 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] + # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp): + # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback. + # Self-tunes across hardware (NPU/NVIDIA) and load conditions. + # codec_chunk_adaptive: true + # codec_chunk_min_frames: 2 + # Code2Wav decode + transfer margin + # codec_chunk_safety_margin_ms: 50 + # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp + # codec_chunk_ramp_divisor: 15 + # min chunk-size growth per emit + # codec_chunk_ramp_delta_min: 2 # Decoder CUDA Graph sizing: # * async_chunk=true derives all stateful ICL/xvec shapes from # codec_chunk_ramp when enabled, otherwise from diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index ff086beb5..7a73dd50a 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -38,6 +38,17 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] + # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp): + # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback. + # Self-tunes across hardware (NPU/NVIDIA) and load conditions. + # codec_chunk_adaptive: true + # codec_chunk_min_frames: 2 + # Code2Wav decode + transfer margin + # codec_chunk_safety_margin_ms: 50 + # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp + # codec_chunk_ramp_divisor: 15 + # min chunk-size growth per emit + # codec_chunk_ramp_delta_min: 2 # async_chunk derives stateful ICL/xvec graph shapes from the ramp when # enabled, otherwise from the fixed chunk settings above, and ignores # decode_cudagraph_capture_sizes. Keep B>1 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 34cc4851a..c57b0c7b2 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 @@ -94,6 +94,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # but popped at segment boundaries (unlike put_req_chunk which is # request-global for connector key continuity). self.ramp_chunk_count: dict[str, int] = defaultdict(int) + self._adaptive_states: dict[str, Any] = {} self.upstream_exhausted_requests: set[str] = set() self.segment_finished_requests: set[str] = set() self.request_payload = {} @@ -584,6 +585,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: self.code_prompt_token_ids.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) + self._adaptive_states.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) if cached_ic is not None: cached_ic.pop(external_req_id, None) @@ -677,6 +679,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.code_prompt_token_ids.pop(external_req_id, None) self.requests_num_chunks_sent.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) + self._adaptive_states.pop(external_req_id, None) self._pending_streaming_prefills.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) 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 index ad35c2349..9b9462ad9 100644 --- a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging +from dataclasses import dataclass, field +from typing import ClassVar logger = logging.getLogger(__name__) @@ -127,3 +129,228 @@ def compute_ramp_emit( return True, 0 return True, length - prev_threshold + + +@dataclass +class AdaptiveChunkController: + """Buffer-feedback controller for adaptive chunk sizing. + + Replaces the fixed ramp table with a greedy buffer-feedback rule. + Chunk 0 uses dynamic IC; chunk 1+ sizes are decided per-chunk from + live buffer-depth signals. + + State is per-segment: the controller is discarded and rebuilt at + segment boundaries via ``_adaptive_states.pop()`` in the adapter. + """ + + frame_time_ewma_ms: float = 0.0 + first_emit_time: float = 0.0 + last_emit_time: float = 0.0 + frames_emitted_this_segment: int = 0 + steady_state_reached: bool = False + accumulated_at_last_emit: int = 0 + last_target: int = 0 + # Per-request underrun telemetry, mirroring stream_client.py's playback + # cursor model so server-side numbers are comparable to client-reported + # AUDIO_UNDERRUN. Collected in record_emit; logged once at request end via + # log_summary(). total_gap_s accumulates only gaps > 5ms (matching + # stream_client's silence-insert threshold); per_chunk_gap_ms records every + # chunk's gap (including 0) for the full trajectory. + chunk_sizes: list = field(default_factory=list) + per_chunk_gap_ms: list = field(default_factory=list) + playback_cursor_s: float = 0.0 + total_gap_s: float = 0.0 + # Ramp-rate tuning, sourced from connector config (codec_chunk_ramp_divisor + # / codec_chunk_ramp_delta_min) so deploys can tune the ramp without code + # changes. Defaults preserve the original 15/2 behavior. + ramp_divisor: int = 15 + ramp_delta_min: int = 2 + + EWMA_ALPHA: ClassVar[float] = 0.3 + FRAME_DURATION_MS: ClassVar[float] = 80.0 + + def compute_next_chunk_size( + self, + now: float, + min_frames: int, + max_frames: int, + safety_margin_ms: float, + ) -> int: + if self.steady_state_reached: + return max_frames + + emitted_ms = self.frames_emitted_this_segment * self.FRAME_DURATION_MS + elapsed_ms = (now - self.first_emit_time) * 1000.0 + buffer_ms = emitted_ms - elapsed_ms + + # Hysteresis: only check when EWMA has a valid measurement. + if self.frame_time_ewma_ms > 0: + steady_gen_ms = max_frames * self.frame_time_ewma_ms + if buffer_ms >= steady_gen_ms + safety_margin_ms: + self.steady_state_reached = True + return max_frames + + # Greedy: select the largest n where buffer can absorb + # n * frame_time + safety_margin. + available_ms = max(0.0, buffer_ms - safety_margin_ms) + if self.frame_time_ewma_ms > 0: + max_n = int(available_ms / self.frame_time_ewma_ms) + else: + max_n = min_frames + greedy_result = max(min_frames, min(max_n, max_frames)) + + # Ramp floor: chunk size grows by an EWMA-scaled delta per emit. + # Unconditional (never gated on buffer sign): in the saturated + # regime (per-stream RTF > 1) buffer is structurally negative and + # shrinking chunks would multiply Code2Wav decode cost -- each chunk + # re-decodes the left_context, cost ~ (72+n)/n, so n=2 costs ~10x + # n=25. Growing toward max minimizes decode passes, the only lever + # pre-W5. delta scales with per-frame time: slow hardware (NPU, + # ewma~150) ramps fast (delta=10, ~3 emits to max); fast hardware + # (NVIDIA, ewma~16) ramps gently (delta=2) for lower TTFP. This + # cross-hardware self-tuning is why the adaptive controller exists. + # last_target is set to the controller's target_size in record_emit + # (on actual emit), not here -- compute_next_chunk_size may be called + # multiple times per emit as frames accumulate one-by-one. + # WARNING: do NOT gate ramp_floor on buffer>0. That collapses to + # greedy=min_frames under load and triggers a vocoder-cost death + # spiral (see #4855 linyueqian 2026-07-19 comment + W5 coupling). + # Downward chunk-size adaptation is deferred until W5 (streaming + # decoder state) lands -- see PR description. + if self.frame_time_ewma_ms > 0: + ramp_delta = max(self.ramp_delta_min, int(self.frame_time_ewma_ms / self.ramp_divisor)) + else: + ramp_delta = self.ramp_delta_min + ramp_floor = self.last_target + ramp_delta + result = max(ramp_floor, greedy_result) + result = min(result, max_frames) + + return result + + def record_emit( + self, + now: float, + n_frames: int, + accumulated: int, + target_size: int, + ) -> None: + if self.frames_emitted_this_segment > 0: + dt_ms = (now - self.last_emit_time) * 1000.0 + measured_ms = dt_ms / n_frames if n_frames > 0 else self.frame_time_ewma_ms + if self.frame_time_ewma_ms == 0: + self.frame_time_ewma_ms = measured_ms + else: + self.frame_time_ewma_ms = ( + self.EWMA_ALPHA * measured_ms + (1 - self.EWMA_ALPHA) * self.frame_time_ewma_ms + ) + # Per-chunk underrun telemetry: mirror stream_client.py's playback + # cursor so server-side total_gap is comparable to the client's + # AUDIO_UNDERRUN. arrival_s = wall since first emit; gap_s = how late + # this chunk is vs where playback had reached. + arrival_s = now - self.first_emit_time + audio_dur_s = n_frames * self.FRAME_DURATION_MS / 1000.0 + gap_s = max(0.0, arrival_s - self.playback_cursor_s) + if gap_s > 0.005: + self.total_gap_s += gap_s + self.per_chunk_gap_ms.append(gap_s * 1000.0) + self.chunk_sizes.append(n_frames) + self.playback_cursor_s = max(arrival_s, self.playback_cursor_s) + audio_dur_s + self.last_emit_time = now + self.frames_emitted_this_segment += n_frames + self.accumulated_at_last_emit = accumulated + # last_target tracks the controller's intended target_size, NOT the + # actual n_frames shipped. With the B1 fix (compute_adaptive_emit + # returns new_frames), n_frames can overshoot target_size during a + # latency spike. Using n_frames here would ratchet ramp_floor to the + # spike size and lock the controller toward max_frames prematurely. + self.last_target = target_size + + def log_summary(self, request_id: str) -> None: + """Emit one INFO line at segment/request end. + + Mirrors stream_client.py's underrun model so total_gap_s is + comparable to the client-reported AUDIO_UNDERRUN. Fires on + ``finished`` which is ``is_segment_finished or is_finished`` (see + chunk_transfer_adapter), so a resumable /v1/realtime session emits + one line per segment. INFO level so it is visible by default + without any env var or dedicated logger. Only emits fields vLLM + stats do NOT already report: chunk-size trajectory, per-chunk + underrun, and total underrun. (pure/wall/ewma were dropped: vLLM + already reports audio_duration, e2e/stage wall, and + tpot/time_per_output_unit respectively; a trailing ewma also could + not reflect its intra-request evolution.) + """ + logger.info( + "[ADAPTIVE] req=%s done: chunks=%s gaps_ms=%s total_gap=%.3fs", + request_id, + self.chunk_sizes, + [round(g, 1) for g in self.per_chunk_gap_ms], + self.total_gap_s, + ) + + +def compute_adaptive_emit( + length: int, + accumulated_at_last_emit: int, + target_size: int, + finished: bool, +) -> tuple[bool, int]: + """Decide whether to emit under the adaptive controller. + + Same return contract as ``compute_ramp_emit``: + + - ``(False, 0)``: hold. + - ``(True, >0)``: emit with this many new frames. + - ``(True, 0)``: finished with no new frames (empty-finished sentinel). + """ + new_frames = length - accumulated_at_last_emit + + if finished and new_frames <= 0: + return True, 0 + + # When finished, flush ALL remaining frames to avoid data loss. + # Truncating to target_size would silently drop new_frames - target_size frames. + if finished: + return True, new_frames + + if not finished and new_frames < target_size: + return False, 0 + + # Ship ALL accumulated frames, not just target_size. When new_frames + # exceeds target_size (e.g. a scheduler stall dropped target below the + # already-queued frames), returning target_size would cause the caller's + # [-context_length:] slice to take the newest frames, silently skipping + # the surplus. Shipping new_frames keeps the slice contiguous with the + # previous emit and preserves Code2Wav's per-request decoder cache. + return True, new_frames + + +def parse_adaptive_config( + cfg: dict, +) -> tuple[bool, int, float, int, int]: + """Parse adaptive chunk settings from connector extra config. + + The adaptive max is ``codec_chunk_frames`` (applied at the call site + as ``max_frames``); there is no separate max config -- capping below + steady would undermine the controller's ramp-to-large design (more + Code2Wav re-decode tax) and is a saturated-regime death-spiral risk, + so it is intentionally not exposed. + + Returns: + ``(enabled, min_frames, safety_margin_ms, ramp_divisor, ramp_delta_min)`` + """ + enabled = bool(cfg.get("codec_chunk_adaptive", False)) + min_frames = int(cfg.get("codec_chunk_min_frames", 2)) + margin_ms = float(cfg.get("codec_chunk_safety_margin_ms", 50.0)) + ramp_divisor = int(cfg.get("codec_chunk_ramp_divisor", 15)) + ramp_delta_min = int(cfg.get("codec_chunk_ramp_delta_min", 2)) + if min_frames < 1: + logger.warning("codec_chunk_min_frames must be >= 1, clamping to 1.") + min_frames = 1 + if ramp_divisor < 1: + logger.warning("codec_chunk_ramp_divisor must be >= 1, clamping to 1.") + ramp_divisor = 1 + if ramp_delta_min < 1: + logger.warning("codec_chunk_ramp_delta_min must be >= 1, clamping to 1.") + ramp_delta_min = 1 + return enabled, min_frames, margin_ms, ramp_divisor, ramp_delta_min 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 da5fa840d..423d288fe 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -1,5 +1,6 @@ """Stage input processor for Qwen3-TTS: Talker -> Code2Wav.""" +import time from collections.abc import Mapping from typing import Any @@ -14,9 +15,12 @@ from vllm_omni.data_entry_keys import ( to_dict, ) from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + AdaptiveChunkController, + compute_adaptive_emit, compute_dynamic_initial_chunk_size, compute_ramp_emit, max_ic_for_chunk_size, + parse_adaptive_config, parse_chunk_ramp, ) from vllm_omni.model_executor.stage_input_processors.tts_utils import ( @@ -108,6 +112,18 @@ def talker2code2wav_async_chunk( transfer_manager._ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_size) ramp = transfer_manager._ramp_parsed + # Adaptive: parse once per transfer_manager. + # Takes precedence over fixed ramp when both are configured. + 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 + # Per-request override takes priority over dynamic IC. fixed_initial_chunk_size = configured_initial_chunk_size > 0 initial_chunk_size = configured_initial_chunk_size @@ -124,8 +140,12 @@ def talker2code2wav_async_chunk( fixed_initial_chunk_size = True # Dynamic IC: cache per request so boundaries stay stable for its lifetime. - # Skipped when ramp is active (ramp replaces IC/steady entirely). - if ramp is None and not fixed_initial_chunk_size: + # Skipped when fixed ramp is active (ramp replaces IC/steady entirely) or + # a fixed IC is configured — unless adaptive is active, in which case + # dynamic IC always runs (used for chunk 0), even if + # initial_codec_chunk_frames is configured. + skip_dynamic_ic = (ramp is not None or fixed_initial_chunk_size) and not adaptive_enabled + if not skip_dynamic_ic: _ic_cache = getattr(transfer_manager, "_cached_ic", None) if _ic_cache is None: _ic_cache = {} @@ -172,7 +192,71 @@ def talker2code2wav_async_chunk( ) return None - if ramp is not None: + 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(request_id) + chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + first_chunk = chunk_index <= 0 + + if ctrl is None or chunk_index == 0: + use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size + + if use_first_chunk and length <= initial_chunk_size: + if not finished and length < initial_chunk_size: + return None + context_length = length if finished and length < initial_chunk_size else initial_chunk_size + else: + initial_coverage = initial_chunk_size if use_first_chunk else 0 + adjusted = length - initial_coverage + 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 + + # Chunk 0: target_size equals context_length (no overshoot possible + # — IC logic emits exactly what's accumulated up to the IC boundary). + target_size = context_length + + 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[request_id] = ctrl + else: + now = time.monotonic() + target_size = ctrl.compute_next_chunk_size( + now, + adaptive_min, + chunk_size, + adaptive_margin, + ) + emit, context_length = compute_adaptive_emit( + length, + ctrl.accumulated_at_last_emit, + target_size, + finished, + ) + if not emit: + return None + if context_length == 0: + ctrl.log_summary(request_id) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + request_id=request_id, + left_context_size=0, + finished=torch.tensor(True, dtype=torch.bool), + ), + ) + + elif ramp is not None: chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) first_chunk = chunk_index <= 0 emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size, finished) @@ -206,6 +290,13 @@ def talker2code2wav_async_chunk( if finished and first_chunk: context_length = length + if adaptive_enabled: + ctrl = transfer_manager._adaptive_states.get(request_id) + if ctrl is not None: + ctrl.record_emit(time.monotonic(), context_length, length, target_size) + if finished: + ctrl.log_summary(request_id) + # Code2Wav keeps the quantizer/conv/Transformer context per request. Ship # only the newly completed codec frames after the first chunk. window_frames = transfer_manager.code_prompt_token_ids[request_id][-context_length:] diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 0f233b1da..1c8fe5b9f 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -179,6 +179,7 @@ class OmniConnectorModelRunnerMixin: # does not have segment boundary infrastructure; multi-segment support # is only available via chunk_transfer_adapter (distributed path). self._ramp_chunk_count: dict[str, int] = defaultdict(int) + self._adaptive_states: dict[str, Any] = {} # Send-side async accumulation / staging buffer. Receive-side payload # ownership lives in ``_local_stage_payload_cache``. self._send_side_request_payload: dict[str, dict[str, Any]] = {} @@ -329,6 +330,7 @@ class OmniConnectorModelRunnerMixin: self._code_prompt_token_ids.pop(k, None) self._cached_ic.pop(k, None) self._ramp_chunk_count.pop(k, None) + self._adaptive_states.pop(k, None) self._kv_pending_transfers.pop(req_id, None) self._kv_active_transfers.discard(req_id) self._kv_completed_transfers.discard(req_id) @@ -2027,6 +2029,7 @@ class OmniConnectorModelRunnerMixin: self._code_prompt_token_ids.pop(cleanup_req_id, None) self._cached_ic.pop(cleanup_req_id, None) self._ramp_chunk_count.pop(cleanup_req_id, None) + self._adaptive_states.pop(cleanup_req_id, None) # ------------------------------------------------------------------ # # Payload accumulation (ported from OmniChunkTransferAdapter)