[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 <noreply@anthropic.com>

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

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.
 

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):