From c893b638ad06fa10fe6031c080c1cf675704fbd4 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 16:45:24 +0800 Subject: [PATCH] [Edge][WP5] Accumulate per-step codec delta rows (DELTA mode delivers only the newest row) Measured after the DELTA fix: 53 engine outputs for 52 rows, but the client still received everything at once because each per-step payload is the newest row only and the row alignment treated it as a cumulative snapshot (one yield at the first row, one at the end). accumulate_codec_rows folds delta rows and the final consolidated matrix into one growing matrix; unit test drives the generator with per-step deltas and asserts one yield per step. Co-Authored-By: Claude Fable 5.1 --- .../openai_api/test_codec_stream_protocol.py | 58 +++++++++++++++++++ vllm_omni/entrypoints/openai/codec_stream.py | 32 ++++++++++ .../entrypoints/openai/serving_speech.py | 21 ++++--- 3 files changed, 103 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/openai_api/test_codec_stream_protocol.py b/tests/entrypoints/openai_api/test_codec_stream_protocol.py index f4a11522f..ca0a57d24 100644 --- a/tests/entrypoints/openai_api/test_codec_stream_protocol.py +++ b/tests/entrypoints/openai_api/test_codec_stream_protocol.py @@ -222,3 +222,61 @@ def test_codec_snapshot_list_is_normalized(): assert out.shape == (4, 16) and list(out[:, 0]) == [0, 1, 2, 3] assert _coerce_codec_snapshots([None]) is None assert isinstance(_coerce_codec_snapshots([np.zeros(16)]), np.ndarray) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_accumulate_codec_rows_handles_deltas_and_snapshots(): + import numpy as np + + from vllm_omni.entrypoints.openai.codec_stream import accumulate_codec_rows + + r = lambda i: np.full((1, 16), i, dtype=np.int64) # noqa: E731 + acc = accumulate_codec_rows(None, r(0)) + acc = accumulate_codec_rows(acc, r(1)) + acc = accumulate_codec_rows(acc, r(2)) + assert acc.shape == (3, 16) and list(acc[:, 0]) == [0, 1, 2] + # duplicate of the last delta is not appended twice + assert accumulate_codec_rows(acc, acc).shape == (3, 16) + # final cumulative snapshot with a leading placeholder row replaces the matrix + final = np.concatenate([np.full((1, 16), -1), acc, r(3)], axis=0) + out = accumulate_codec_rows(acc, final) + assert out.shape == (5, 16) and out[0, 0] == -1 and out[-1, 0] == 3 + # 1-D row payload + assert accumulate_codec_rows(None, np.arange(16)).shape == (1, 16) + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.asyncio +async def test_generate_codec_chunks_yields_per_step(): + from types import SimpleNamespace + + import numpy as np + + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + + svc = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + svc._get_tts_adapter = lambda: SimpleNamespace( + codec_stream_spec={"codebook_size": 2048}, validates_generation=False + ) + + def out(token_ids, codes, finished): + return SimpleNamespace( + outputs=[SimpleNamespace(token_ids=token_ids, multimodal_output={"codes": {"audio": codes}})], + finished=finished, + ) + + rows = [np.full((1, 16), i + 10, dtype=np.int64) for i in range(4)] + + async def gen(): + # step k: token k and its codes row arrive together (DELTA mode: newest row only) + for k in range(4): + yield out([100 + k], rows[k], False) + # final: EOS token plus the consolidated cumulative matrix (EOS row is all zeros) + final = np.concatenate(rows + [np.zeros((1, 16), dtype=np.int64)], axis=0) + yield out([2150], final, True) + + got = [(None if c is None else c.shape[0], eos) async for c, eos in svc._generate_codec_chunks(gen(), "r1")] + assert got[:4] == [(1, False), (1, False), (1, False), (1, False)] + assert got[-1][1] is True and sum(n for n, _ in got if n) == 4 diff --git a/vllm_omni/entrypoints/openai/codec_stream.py b/vllm_omni/entrypoints/openai/codec_stream.py index 7ccb97221..ec0d0eaec 100644 --- a/vllm_omni/entrypoints/openai/codec_stream.py +++ b/vllm_omni/entrypoints/openai/codec_stream.py @@ -77,6 +77,38 @@ def unpack_codec_frame(buf: bytes, codebooks: int) -> tuple[int, int, int, np.nd return seq, n_frames, flags, arr +def accumulate_codec_rows(acc: np.ndarray | None, step: Any) -> np.ndarray: + """Fold a per-step ``codes.audio`` payload into the growing ``[rows, codebooks]`` matrix. + + DELTA-mode outputs carry only the newest row(s); the final consolidated + output (and CUMULATIVE mode) carries the whole matrix, possibly with a + leading placeholder row the deltas never had. A payload whose tail equals + the rows already accumulated (and that is longer) replaces the matrix; an + identical payload is a duplicate; anything else is appended. + """ + if hasattr(step, "detach"): + step = step.detach().cpu().numpy() + arr = np.asarray(step) + if arr.ndim == 1: + arr = arr.reshape(1, -1) + if arr.ndim != 2 or arr.size == 0: + return acc if acc is not None else np.zeros((0, arr.shape[-1] if arr.ndim >= 1 else 0), dtype=np.int64) + arr = arr.astype(np.int64, copy=False) + if acc is None or acc.size == 0: + return arr + n = acc.shape[0] + if arr.shape[0] > n and arr.shape[1] == acc.shape[1]: + # A longer payload that contains the accumulated rows as a contiguous + # block is a cumulative snapshot (it may add a placeholder row in + # front and/or newer rows behind): replace. + for i in range(arr.shape[0] - n + 1): + if np.array_equal(arr[i : i + n], acc): + return arr + if arr.shape == acc.shape and np.array_equal(arr, acc): + return acc + return np.concatenate([acc, arr], axis=0) + + def align_codec_rows(codes_cum: Any, token_ids_cum: list[int], codebook_size: int) -> np.ndarray: """Keep the codec rows that correspond to real codebook-0 tokens. diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index cdaa878ee..90e691ba5 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3316,19 +3316,23 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ): """Yield ``(codes[int64, n, codebooks], eos)`` deltas from a talker-only pipeline. - The AR stage publishes the *cumulative* ``codes.audio`` tensor on every - step (``omni_client_multimodal_output_keys``); this generator aligns the - rows to the output token ids (dropping the prefill placeholder and the - EOS row) and emits only the new rows. The final item always carries - ``eos=True`` (possibly with zero frames). + The AR stage publishes ``codes.audio`` on every step + (``omni_client_multimodal_output_keys``): in DELTA output mode only the + newest row(s), and a consolidated cumulative matrix on the final + output. ``accumulate_codec_rows`` folds both shapes into one growing + matrix; this generator aligns its rows to the output token ids + (dropping the prefill placeholder and the EOS row) and emits only the + new rows. The final item always carries ``eos=True`` (possibly with + zero frames). """ - from vllm_omni.entrypoints.openai.codec_stream import align_codec_rows + from vllm_omni.entrypoints.openai.codec_stream import accumulate_codec_rows, align_codec_rows adapter = self._get_tts_adapter() spec = getattr(adapter, "codec_stream_spec", None) or {} codebook_size = int(spec.get("codebook_size", 2048)) usage_acc = SpeechOutputTokenCounter() if (adapter is not None and adapter.validates_generation) else None token_ids: list[int] = [] + codes_cum: np.ndarray | None = None sent = 0 finished = False n_outputs = 0 @@ -3347,12 +3351,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): token_ids = list(new_ids) else: token_ids.extend(int(t) for t in new_ids) - codes_cum = self._extract_codec_output(res) + codes_step = self._extract_codec_output(res) done = bool(getattr(res, "finished", False)) - if codes_cum is not None: + if codes_step is not None: n_with_codes += 1 if first_codes_at < 0: first_codes_at = n_outputs + codes_cum = accumulate_codec_rows(codes_cum, codes_step) aligned = align_codec_rows(codes_cum, token_ids, codebook_size) delta = aligned[sent:] if delta.shape[0] >= min_frames or (done and delta.shape[0] > 0):