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 bc6239e29..026103d1b 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 @@ -9,7 +9,11 @@ import torch from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( compute_dynamic_initial_chunk_size, + compute_ramp_emit, max_ic_for_chunk_size, + parse_chunk_ramp, + ramp_chunk_size, + ramp_cumulative, ) from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( talker2code2wav_async_chunk, @@ -35,21 +39,21 @@ 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): +def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frames=0, chunk_ramp=None): + extra = { + "codec_chunk_frames": chunk_frames, + "codec_left_context_frames": left_context, + "initial_codec_chunk_frames": initial_chunk_frames, + } + if chunk_ramp is not None: + extra["codec_chunk_ramp"] = chunk_ramp 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), request_payload={}, - connector=SimpleNamespace( - config={ - "extra": { - "codec_chunk_frames": chunk_frames, - "codec_left_context_frames": left_context, - "initial_codec_chunk_frames": initial_chunk_frames, - } - } - ), + connector=SimpleNamespace(config={"extra": extra}), ) @@ -549,3 +553,363 @@ def test_full_payload_emits_empty_finished_payload_on_degenerate_take(pooling_ou assert payload is not None assert payload["meta"]["finished"].item() is True assert payload["codes"]["audio"].numel() == 0 + + +_RAMP = [1, 4, 8, 16, 25] + + +class TestRampHelpers: + def test_parse_ramp_none_when_absent(self): + assert parse_chunk_ramp({}) is None + + def test_parse_ramp_valid(self): + assert parse_chunk_ramp({"codec_chunk_ramp": [1, 4, 8, 16, 25]}) == [1, 4, 8, 16, 25] + + def test_parse_ramp_from_string(self): + assert parse_chunk_ramp({"codec_chunk_ramp": "1, 4, 8, 16, 25"}) == [1, 4, 8, 16, 25] + + def test_parse_ramp_too_short(self): + assert parse_chunk_ramp({"codec_chunk_ramp": [25]}) is None + + def test_parse_ramp_non_positive(self): + assert parse_chunk_ramp({"codec_chunk_ramp": [1, 0, 8]}) is None + + def test_parse_ramp_bad_string(self): + assert parse_chunk_ramp({"codec_chunk_ramp": "a,b"}) is None + + def test_parse_ramp_int_returns_none(self): + assert parse_chunk_ramp({"codec_chunk_ramp": 4}) is None + + def test_parse_ramp_float_returns_none(self): + assert parse_chunk_ramp({"codec_chunk_ramp": 4.5}) is None + + def test_parse_ramp_mixed_list_returns_none(self): + assert parse_chunk_ramp({"codec_chunk_ramp": [4, "x"]}) is None + + def test_parse_ramp_warns_on_tail_mismatch(self, caplog): + import logging + + with caplog.at_level(logging.WARNING): + result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25) + assert result == [4, 4, 8] + assert any("reintroduces" in r.message for r in caplog.records) + + def test_parse_ramp_no_warn_on_tail_match(self, caplog): + import logging + + with caplog.at_level(logging.WARNING): + result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25) + assert result == [4, 4, 8, 16, 25] + assert not any("reintroduces" in r.message for r in caplog.records) + + @pytest.mark.parametrize( + "index,ramp,steady,expected", + [ + (0, _RAMP, 25, 1), + (1, _RAMP, 25, 4), + (2, _RAMP, 25, 8), + (3, _RAMP, 25, 16), + (4, _RAMP, 25, 25), + (5, _RAMP, 25, 25), + (100, _RAMP, 25, 25), + ], + ) + def test_ramp_chunk_size(self, index, ramp, steady, expected): + assert ramp_chunk_size(index, ramp, steady) == expected + + @pytest.mark.parametrize( + "index,ramp,steady,expected", + [ + (0, _RAMP, 25, 1), + (1, _RAMP, 25, 5), + (2, _RAMP, 25, 13), + (3, _RAMP, 25, 29), + (4, _RAMP, 25, 54), + (5, _RAMP, 25, 79), + (6, _RAMP, 25, 104), + (100, _RAMP, 25, 54 + 96 * 25), + ], + ) + def test_ramp_cumulative(self, index, ramp, steady, expected): + assert ramp_cumulative(index, ramp, steady) == expected + + @pytest.mark.parametrize( + "length,chunk_index,finished,expected_emit,expected_ctx", + [ + (0, 0, False, False, 0), + (1, 0, False, True, 1), + (3, 1, False, False, 0), + (5, 1, False, True, 4), + (12, 2, False, False, 0), + (13, 2, False, True, 8), + (28, 3, False, False, 0), + (29, 3, False, True, 16), + (53, 4, False, False, 0), + (54, 4, False, True, 25), + (78, 5, False, False, 0), + (79, 5, False, True, 25), + (1, 1, True, True, 0), + (3, 1, True, True, 2), + (5, 1, True, True, 4), + ], + ) + def test_compute_ramp_emit(self, length, chunk_index, finished, expected_emit, expected_ctx): + emit, ctx = compute_ramp_emit(length, chunk_index, _RAMP, 25, finished) + assert emit == expected_emit + assert ctx == expected_ctx + + +class TestChunkRampEmission: + RAMP = [1, 4, 8, 16, 25] + + 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_ramp_sequence_1_4_8_16_25(self): + tm = _tm(chunk_ramp=self.RAMP) + rid = "ramp-seq" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 1 + tm.ramp_chunk_count[rid] = 1 + + assert self._emit(tm, rid, 2) is None + assert self._emit(tm, rid, 3) is None + assert self._emit(tm, rid, 4) is None + p1 = self._emit(tm, rid, 5) + assert p1 is not None + assert len(p1.codes.audio) == _Q * 5 + assert p1.meta.left_context_size == 1 + tm.ramp_chunk_count[rid] = 2 + + for n in range(6, 13): + assert self._emit(tm, rid, n) is None + p2 = self._emit(tm, rid, 13) + assert p2 is not None + assert p2.meta.left_context_size == 5 + tm.ramp_chunk_count[rid] = 3 + + for n in range(14, 29): + assert self._emit(tm, rid, n) is None + p3 = self._emit(tm, rid, 29) + assert p3 is not None + assert p3.meta.left_context_size == 13 + tm.ramp_chunk_count[rid] = 4 + + for n in range(30, 54): + assert self._emit(tm, rid, n) is None + p4 = self._emit(tm, rid, 54) + assert p4 is not None + assert p4.meta.left_context_size == 25 + tm.ramp_chunk_count[rid] = 5 + + for n in range(55, 79): + assert self._emit(tm, rid, n) is None + p5 = self._emit(tm, rid, 79) + assert p5 is not None + assert p5.meta.left_context_size == 25 + + def test_ramp_finished_flush_mid_chunk(self): + tm = _tm(chunk_ramp=self.RAMP) + rid = "ramp-flush" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 3, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + assert len(p_fin.codes.audio) == _Q * (1 + 2) + assert p_fin.meta.left_context_size == 1 + + def test_ramp_finished_no_new_frames(self): + tm = _tm(chunk_ramp=self.RAMP) + rid = "ramp-no-new" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 1, 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_ramp_finished_at_exact_threshold(self): + tm = _tm(chunk_ramp=self.RAMP) + rid = "ramp-exact" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p1 = self._emit(tm, rid, 5, finished=True) + assert p1 is not None + assert p1.meta.finished.item() is True + + def test_ramp_backward_compat_without_config(self): + tm = _tm(initial_chunk_frames=1) + rid = "no-ramp" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 1 + tm.put_req_chunk[rid] = 1 + + for n in range(2, 26): + assert self._emit(tm, rid, n) is None + p1 = self._emit(tm, rid, 26) + assert p1 is not None + assert p1.meta.left_context_size == 1 + + def test_ramp_with_ref_code_first_chunk(self): + tm = _tm(chunk_ramp=self.RAMP, left_context=25) + rid = "ramp-ref" + tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(1)] + 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 + assert len(payload.codes.audio) == _Q * (2 + 1) + + def test_ramp_steady_state_after_ramp_exhausted(self): + tm = _tm(chunk_ramp=[1, 4], chunk_frames=25) + rid = "ramp-short" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p1 = self._emit(tm, rid, 5) + assert p1 is not None + tm.ramp_chunk_count[rid] = 2 + + for n in range(6, 30): + assert self._emit(tm, rid, n) is None + p2 = self._emit(tm, rid, 30) + assert p2 is not None + assert p2.meta.left_context_size == 5 + tm.ramp_chunk_count[rid] = 3 + + for n in range(31, 55): + assert self._emit(tm, rid, n) is None + p3 = self._emit(tm, rid, 55) + assert p3 is not None + assert p3.meta.left_context_size == 25 + + def test_ramp_profile_4_4_8_16_25(self): + """Ramp [4,4,8,16,25]: chunk 0=4 frames (320ms audio covers chunk 1 + gen time → no gap at chunk 0→1), gradual ramp to steady state.""" + tm = _tm(chunk_ramp=[4, 4, 8, 16, 25]) + rid = "ramp-4-4-8-16-25" + + for n in range(1, 4): + assert self._emit(tm, rid, n) is None + p0 = self._emit(tm, rid, 4) + assert p0 is not None + assert p0.meta.left_context_size == 0 + tm.ramp_chunk_count[rid] = 1 + + for n in range(5, 8): + assert self._emit(tm, rid, n) is None + p1 = self._emit(tm, rid, 8) + assert p1 is not None + assert p1.meta.left_context_size == 4 + tm.ramp_chunk_count[rid] = 2 + + for n in range(9, 16): + assert self._emit(tm, rid, n) is None + p2 = self._emit(tm, rid, 16) + assert p2 is not None + assert p2.meta.left_context_size == 8 + tm.ramp_chunk_count[rid] = 3 + + for n in range(17, 32): + assert self._emit(tm, rid, n) is None + p3 = self._emit(tm, rid, 32) + assert p3 is not None + assert p3.meta.left_context_size == 16 + tm.ramp_chunk_count[rid] = 4 + + for n in range(33, 57): + assert self._emit(tm, rid, n) is None + p4 = self._emit(tm, rid, 57) + assert p4 is not None + assert p4.meta.left_context_size == 25 + tm.ramp_chunk_count[rid] = 5 + + for n in range(58, 82): + assert self._emit(tm, rid, n) is None + p5 = self._emit(tm, rid, 82) + assert p5 is not None + assert p5.meta.left_context_size == 25 + + def test_ramp_resets_on_segment_boundary(self): + """After segment 1 emits chunks, a segment boundary clears + code_prompt_token_ids and ramp_chunk_count. Segment 2 must restart + from chunk index 0. + + Regression for: put_req_chunk is request-global and not reset at + segment boundaries. ramp_chunk_count is popped alongside + code_prompt_token_ids on is_segment_finished, so the ramp index + resets by construction.""" + tm = _tm(chunk_ramp=[4, 4, 8, 16, 25]) + rid = "ramp-segment" + + p0 = self._emit(tm, rid, 4) + assert p0 is not None + assert p0.meta.left_context_size == 0 + tm.ramp_chunk_count[rid] = 1 + tm.put_req_chunk[rid] = 1 + + p1 = self._emit(tm, rid, 8) + assert p1 is not None + tm.ramp_chunk_count[rid] = 2 + tm.put_req_chunk[rid] = 2 + + tm.code_prompt_token_ids[rid] = [] + tm.ramp_chunk_count.pop(rid, None) + + p_seg2_0 = self._emit(tm, rid, 4) + assert p_seg2_0 is not None + assert p_seg2_0.meta.left_context_size == 0 + + def test_ramp_resets_on_segment_boundary_one_chunk_history(self): + """Segment 1 emitted exactly 1 chunk (put_req_chunk=1). Segment 2 + must still restart from chunk index 0. + + Regression for: the hybrid put_req_chunk/derived scheme collided + when put_chunk=1 because length >= ramp_cumulative(0) turned true + at length=4, causing chunk_index to jump to 1 and dropping the + first 4 frames of segment 2.""" + tm = _tm(chunk_ramp=[4, 4, 8, 16, 25]) + rid = "ramp-segment-1chunk" + + p0 = self._emit(tm, rid, 4) + 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) + + p_seg2_0 = self._emit(tm, rid, 4) + assert p_seg2_0 is not None + assert p_seg2_0.meta.left_context_size == 0 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index a12c920ca..8f4c6440e 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -26,6 +26,13 @@ connectors: codec_left_context_frames: 72 # Emit only the first audio chunk early, then return to codec_chunk_frames. initial_codec_chunk_frames: 1 + # Chunk ramp (opt-in): gradual size increase for first N chunks to + # eliminate the IC→steady cliff (1→25). Uncomment to enable. + # Each entry is the size of that chunk index; indices past the table + # fall back to codec_chunk_frames. When enabled, ramp overrides + # initial_codec_chunk_frames (including the per-request API field). + # Example: + # codec_chunk_ramp: [4, 4, 8, 16, 25] # Common Stage1 decode buckets: # no-ref first/steady chunks: 25 / 97 frames # Base ref-context first/steady chunks: 73 / 169 frames diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 847383955..25bac72c0 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -31,6 +31,13 @@ connectors: ref_code_context_frames: 72 # Emit only the first audio chunk early, then return to codec_chunk_frames. initial_codec_chunk_frames: 1 + # Chunk ramp (opt-in): gradual size increase for first N chunks to + # eliminate the IC→steady cliff (1→25). Uncomment to enable. + # Each entry is the size of that chunk index; indices past the table + # fall back to codec_chunk_frames. When enabled, ramp overrides + # initial_codec_chunk_frames (including the per-request API field). + # Example: + # codec_chunk_ramp: [4, 4, 8, 16, 25] # Common Stage1 decode buckets: # no-ref first/steady chunks: 25 / 97 frames # Base ref-context first/steady chunks: 73 / 169 frames 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 ea3642c8d..1a41372a8 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 @@ -69,6 +69,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # mapping for request id and chunk id self.put_req_chunk: dict[str, int] = defaultdict(int) self.get_req_chunk: dict[str, int] = defaultdict(int) + # Segment-local chunk counter: incremented alongside put_req_chunk + # 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.finished_requests: set[str] = set() self.segment_finished_requests: set[str] = set() self.request_payload = {} @@ -339,6 +343,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if success: self.put_req_chunk[external_req_id] += 1 + self.ramp_chunk_count[external_req_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") # Sender uses struct attr access here; the receive path in # `_load_one_request` / `_update_request_payload` reads dict keys. @@ -360,6 +365,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: 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) cached_ic = getattr(self, "_cached_ic", None) if cached_ic is not None: cached_ic.pop(external_req_id, None) @@ -412,6 +418,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.request_payload.pop(external_req_id, None) 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._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 8b8956839..ad35c2349 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 @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging + +logger = logging.getLogger(__name__) + def max_ic_for_chunk_size(chunk_size: int) -> int: """Largest power of 2 strictly less than chunk_size.""" @@ -31,3 +35,95 @@ def compute_dynamic_initial_chunk_size( load_factor = min(active_requests / max_num_seqs, 1.0) idx = int(round(load_factor * (len(steps) - 1))) return steps[idx] + + +def parse_chunk_ramp(cfg: dict, steady: int | None = None) -> list[int] | None: + """Parse ``codec_chunk_ramp`` from connector extra config. + + Accepts a list of positive ints or a comma-separated string. + Returns ``None`` when the key is absent or invalid (ramp disabled). + Requires >= 2 entries; all entries must be > 0. + + When ``steady`` is provided, warns if ``ramp[-1] != steady`` (the last + ramp entry should match ``codec_chunk_frames`` to avoid reintroducing + the cliff the ramp is meant to eliminate). + """ + raw = cfg.get("codec_chunk_ramp") + if raw is None: + return None + try: + if isinstance(raw, str): + raw = [int(x.strip()) for x in raw.split(",")] + ramp = [int(x) for x in raw] + except (TypeError, ValueError): + logger.warning("codec_chunk_ramp parse error for %r; disabling ramp.", raw) + return None + if len(ramp) < 2: + logger.warning("codec_chunk_ramp needs >= 2 entries, got %d; disabling.", len(ramp)) + return None + if any(x <= 0 for x in ramp): + logger.warning("codec_chunk_ramp entries must be positive; disabling.") + return None + if steady is not None and ramp[-1] != steady: + logger.warning( + "codec_chunk_ramp[-1]=%d != codec_chunk_frames=%d; this reintroduces the chunk-size cliff at index %d.", + ramp[-1], + steady, + len(ramp) - 1, + ) + return ramp + + +def ramp_chunk_size(index: int, ramp: list[int], steady: int) -> int: + """Return the target chunk size for chunk ``index`` under the ramp schedule. + + Indices within the ramp table use ``ramp[index]``; indices past the table + fall back to ``steady`` (typically ``codec_chunk_frames``). + """ + if index < len(ramp): + return ramp[index] + return steady + + +def ramp_cumulative(index: int, ramp: list[int], steady: int) -> int: + """Cumulative frame count through chunk ``index`` (inclusive). + + O(len(ramp)) closed form: sum the ramp table once, then add + ``(index + 1 - len(ramp)) * steady`` for indices past the table. + """ + ramp_len = len(ramp) + if index < ramp_len: + return sum(ramp[: index + 1]) + return sum(ramp) + (index + 1 - ramp_len) * steady + + +def compute_ramp_emit( + length: int, + chunk_index: int, + ramp: list[int], + steady: int, + finished: bool, +) -> tuple[bool, int]: + """Decide whether to emit a chunk under the ramp schedule. + + Uses cumulative thresholds to determine emission boundaries. Each chunk + ``i`` covers frames ``[ramp_cumulative(i-1), ramp_cumulative(i))``. + + Returns: + ``(emit, context_length)`` + + - ``emit=False, context_length=0``: not enough frames yet, hold. + - ``emit=True, context_length>0``: emit with this many new frames. + - ``emit=True, context_length=0``: finished with no new frames + (caller should emit an empty-finished sentinel). + """ + threshold = ramp_cumulative(chunk_index, ramp, steady) + prev_threshold = ramp_cumulative(chunk_index - 1, ramp, steady) if chunk_index > 0 else 0 + + if not finished and length < threshold: + return False, 0 + + if finished and length <= prev_threshold: + return True, 0 + + return True, length - prev_threshold 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 597fc4db5..0d1b5ad74 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -15,7 +15,9 @@ from vllm_omni.data_entry_keys import ( ) from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( compute_dynamic_initial_chunk_size, + compute_ramp_emit, max_ic_for_chunk_size, + parse_chunk_ramp, ) from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, @@ -84,6 +86,12 @@ def talker2code2wav_async_chunk( configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0) ref_code_context_frames = int(cfg.get("ref_code_context_frames") or left_context_size_config) + # Ramp: parse once per transfer_manager (config is static for the adapter's + # lifetime). When active, ramp replaces IC/steady entirely — skip dynamic IC. + if not hasattr(transfer_manager, "_ramp_parsed"): + transfer_manager._ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_size) + ramp = transfer_manager._ramp_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 @@ -100,7 +108,8 @@ def talker2code2wav_async_chunk( fixed_initial_chunk_size = True # Dynamic IC: cache per request so boundaries stay stable for its lifetime. - if not fixed_initial_chunk_size: + # Skipped when ramp is active (ramp replaces IC/steady entirely). + if ramp is None and not fixed_initial_chunk_size: _ic_cache = getattr(transfer_manager, "_cached_ic", None) if _ic_cache is None: _ic_cache = {} @@ -143,21 +152,30 @@ def talker2code2wav_async_chunk( ) return None - 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: + if ramp is not None: + chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size, finished) + if not emit: return None - context_length = length if finished and length < initial_chunk_size else initial_chunk_size + if context_length == 0: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) else: - # The initial chunk is only for TTFA. After that, return to the normal - # codec chunk size so Code2Wav is not flooded by repeated tiny windows. - 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 + 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 end_index = min(length, left_context_size_config + context_length) left_context_size = max(0, end_index - 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 f9f13cbe7..b31b04688 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -134,6 +134,11 @@ class OmniConnectorModelRunnerMixin: # -- chunk index tracking (ported from OmniChunkTransferAdapter) -- self._put_req_chunk: dict[str, int] = defaultdict(int) self._get_req_chunk: dict[str, int] = defaultdict(int) + # Segment-local chunk counter: incremented alongside _put_req_chunk + # and popped at request cleanup. Note: the mixin path (uniproc mode) + # 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) # 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]] = {} @@ -283,6 +288,7 @@ class OmniConnectorModelRunnerMixin: self._send_side_request_payload.pop(k, None) self._code_prompt_token_ids.pop(k, None) self._cached_ic.pop(k, None) + self._ramp_chunk_count.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) @@ -998,6 +1004,7 @@ class OmniConnectorModelRunnerMixin: external_req_id = self._resolve_external_req_id(request, req_id) chunk_id = self._put_req_chunk[req_id] self._put_req_chunk[req_id] += 1 + self._ramp_chunk_count[req_id] += 1 connector_put_key = f"{external_req_id}_{self._stage_id}_{chunk_id}" logger.debug( @@ -1123,6 +1130,7 @@ class OmniConnectorModelRunnerMixin: return False self._put_req_chunk[request_id] += 1 + self._ramp_chunk_count[request_id] += 1 next_stage_id = self._next_stage_id connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" @@ -1582,6 +1590,10 @@ class OmniConnectorModelRunnerMixin: def put_req_chunk(self) -> dict[str, int]: return self._put_req_chunk + @property + def ramp_chunk_count(self) -> dict[str, int]: + return self._ramp_chunk_count + @property def request_payload(self) -> dict[str, dict[str, Any]]: return self._send_side_request_payload @@ -1964,6 +1976,7 @@ class OmniConnectorModelRunnerMixin: self._send_side_request_payload.pop(cleanup_req_id, None) 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) # ------------------------------------------------------------------ # # Payload accumulation (ported from OmniChunkTransferAdapter)