[5/N] Schedule chunk transfers as their own stage (#7406)

Signed-off-by: specture724 <specture724@gmail.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -2859,8 +2859,10 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
             self.kwargs = kwargs
             self.__class__.instances.append(self)
 
-        def push(self, frames):
+        def push(self, frames, *, on_consumed=None):
             self.pushes.append(np.array(frames, copy=True))
+            if on_consumed is not None:
+                on_consumed()
 
         def finish(self):
             return b"mp4"
@@ -2986,8 +2988,10 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
             self.kwargs = kwargs
             self.__class__.instances.append(self)
 
-        def push(self, frames):
+        def push(self, frames, *, on_consumed=None):
             del frames
+            if on_consumed is not None:
+                on_consumed()
 
         def finish(self):
             return b"mp4"

tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py

@@ -138,13 +138,13 @@ def test_request_batch_threshold_controls_transfers_without_losing_frames(monkey
     from vllm_omni.diffusion.utils import chunked_video
 
     transfers = []
-    quantize = chunked_video.chunk_to_uint8_frames
+    quantize = chunked_video.quantize_chunk
 
     def capture(chunk, value_range):
         transfers.append(chunk.shape[2])
         return quantize(chunk, value_range)
 
-    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
     params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value})
     vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32)
     videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params))

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -43,7 +44,7 @@ def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float |
 
 
 def _make_s2v_sampling(**overrides):
-    values: dict[str, object] = {
+    values: dict[str, Any] = {
         "height": 16,
         "width": 16,
         "num_frames": 8,
@@ -749,13 +750,13 @@ def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch
     from vllm_omni.diffusion.utils import chunked_video
 
     transfers = []
-    quantize = chunked_video.chunk_to_uint8_frames
+    quantize = chunked_video.quantize_chunk
 
     def capture(chunk, value_range):
         transfers.append(chunk.shape[2])
         return quantize(chunk, value_range)
 
-    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
 
     with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
         platform.is_available.return_value = False

tests/diffusion/utils/test_chunked_transfer.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ownership and backpressure contract for chunk transfers to the host.
+
+The ring hands a chunk one reusable host slot for its whole trip. Getting the
+release point wrong is silent: a slot reused while its copy is still landing,
+or while an encoder still holds a view of it, corrupts frames rather than
+raising. These cover the release point, the backpressure it creates, and the
+escape path out of a wait.
+"""
+
+from __future__ import annotations
+
+import threading
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.utils.chunked_video import (
+    ChunkedVideoMP4Session,
+    ChunkLease,
+    PinnedChunkRing,
+    quantize_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _frames(value: int = 7, *, batch: int = 1) -> torch.Tensor:
+    """One quantized ``BTHWC`` chunk, the shape the ring transfers."""
+    return torch.full((batch, 2, 4, 6, 3), value, dtype=torch.uint8)
+
+
+def _cpu_ring(depth: int) -> PinnedChunkRing:
+    return PinnedChunkRing(depth=depth, device=torch.device("cpu"))
+
+
+def test_slot_returns_only_after_its_last_reader():
+    """Each batch entry reads the same buffer, so one release must not free it."""
+    ring = _cpu_ring(2)
+    lease = ring.transfer(_frames(batch=3), readers=3)
+
+    assert ring.slots_in_use == 1
+    view = lease.wait()
+    assert view.shape == (3, 2, 4, 6, 3)
+    # Waiting is not consuming: an encoder still holds a view of this buffer.
+    assert ring.slots_in_use == 1
+
+    lease.release()
+    lease.release()
+    assert ring.slots_in_use == 1
+
+    lease.release()
+    assert ring.slots_in_use == 0
+    # A late extra release must not free the slot a second time.
+    lease.release()
+    assert ring.slots_in_use == 0
+
+
+def test_transfer_waits_for_a_slot_instead_of_growing_the_pool():
+    """Depth is the bound on in-flight host buffers, so a full ring blocks."""
+    ring = _cpu_ring(1)
+    held = ring.transfer(_frames(1), readers=1)
+    second: list[object] = []
+
+    def take_a_slot() -> None:
+        second.append(ring.transfer(_frames(2), readers=1))
+
+    waiter = threading.Thread(target=take_a_slot)
+    waiter.start()
+    waiter.join(timeout=0.3)
+    assert waiter.is_alive(), "a full ring must not hand out a second slot"
+
+    held.wait()
+    held.release()
+    waiter.join(timeout=5)
+    assert not waiter.is_alive()
+    assert ring.slots_in_use == 1
+    assert ring.stats.slot_wait_seconds > 0
+
+
+def test_abort_wakes_a_transfer_waiting_for_a_slot():
+    """A dead consumer must surface as an error, not as a stuck producer."""
+    ring = _cpu_ring(1)
+    ring.transfer(_frames(), readers=1)
+    failures: list[BaseException] = []
+
+    def take_a_slot() -> None:
+        try:
+            ring.transfer(_frames(), readers=1)
+        except BaseException as exc:  # noqa: BLE001
+            failures.append(exc)
+
+    waiter = threading.Thread(target=take_a_slot)
+    waiter.start()
+    waiter.join(timeout=0.3)
+    assert waiter.is_alive()
+
+    ring.abort(RuntimeError("encoder died"))
+    waiter.join(timeout=5)
+    assert [str(exc) for exc in failures] == ["encoder died"]
+    # Later transfers stay failed rather than silently resuming.
+    with pytest.raises(RuntimeError, match="encoder died"):
+        ring.transfer(_frames(), readers=1)
+
+
+def test_encoder_releases_the_chunk_it_was_reading_when_muxing_fails(monkeypatch):
+    """A failed mux abandons the frame generator, which still owns a lent slot."""
+    from vllm_omni.diffusion.utils import media_utils
+
+    released: list[int] = []
+
+    def exploding_mux(frames, **kwargs):
+        # Read into the chunk, so the generator is suspended inside it.
+        next(iter(frames))
+        raise RuntimeError("mux failed")
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", exploding_mux)
+    encoder = media_utils.ChunkedMP4Encoder(width=6, height=4, fps=24)
+    encoder.push(np.zeros((2, 4, 6, 3), dtype=np.uint8), on_consumed=lambda: released.append(1))
+
+    with pytest.raises(RuntimeError, match="mux failed"):
+        encoder.finish()
+
+    assert released == [1], "the abandoned generator must not keep holding its chunk"
+
+
+class _RecordingEncoder:
+    """Stands in for one batch entry's encoder, honouring the release contract."""
+
+    instances: list[_RecordingEncoder] = []
+
+    def __init__(self, **kwargs):
+        self.pushes: list[int] = []
+        self.released = 0
+        self.fail_on_push = False
+        self.__class__.instances.append(self)
+
+    def push(self, frames, *, on_consumed=None):
+        if self.fail_on_push:
+            if on_consumed is not None:
+                on_consumed()
+                self.released += 1
+            raise RuntimeError("encoder rejected the chunk")
+        self.pushes.append(int(frames[0, 0, 0, 0]))
+        if on_consumed is not None:
+            on_consumed()
+            self.released += 1
+
+    def finish(self) -> bytes:
+        return b"mp4"
+
+    def abort(self) -> None:
+        return None
+
+
+@pytest.fixture
+def recording_encoders(monkeypatch):
+    _RecordingEncoder.instances = []
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder",
+        _RecordingEncoder,
+    )
+    return _RecordingEncoder.instances
+
+
+def _push_ramp(session: ChunkedVideoMP4Session, values: list[int]) -> None:
+    for value in values:
+        session.push(torch.full((1, 3, 1, 4, 6), value / 255.0))
+
+
+def test_session_feeds_chunks_in_producer_order(recording_encoders):
+    """Copies complete in issue order, and the encoder must see that order."""
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24)
+    _push_ramp(session, [10, 20, 30, 40])
+
+    assert session.finish() == [b"mp4"]
+    assert recording_encoders[0].pushes == [10, 20, 30, 40]
+
+
+def test_session_returns_the_slot_when_an_encoder_rejects_a_chunk(recording_encoders):
+    """A rejected chunk must not strand the buffer it was lent."""
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, transfer_slots=1)
+    session.push(torch.zeros(1, 3, 1, 4, 6))
+    recording_encoders[0].fail_on_push = True
+
+    with pytest.raises(RuntimeError, match="rejected"):
+        session.push(torch.zeros(1, 3, 1, 4, 6))
+
+    assert session.ring is not None
+    assert session.ring.slots_in_use == 0
+
+
+def test_session_abort_releases_every_slot_still_in_flight(recording_encoders):
+    """Teardown returns the pool even when nothing consumed the chunks."""
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, transfer_slots=2)
+    session.push(torch.zeros(1, 3, 1, 4, 6))
+
+    session.abort()
+
+    assert session.ring is not None
+    assert session.ring.slots_in_use == 0
+
+
+def test_session_gets_its_slot_back_after_a_real_encoder_failure():
+    """A real libx264 error must not strand the producer on the slot it lent."""
+    session = ChunkedVideoMP4Session(
+        value_range=(0.0, 1.0),
+        fps=24,
+        transfer_slots=1,
+        video_codec_options={"preset": "invalid-preset"},
+    )
+    failures: list[BaseException] = []
+
+    def produce() -> None:
+        try:
+            for _ in range(4):
+                session.push(torch.zeros(1, 3, 1, 16, 16))
+            session.finish()
+        except BaseException as exc:  # noqa: BLE001
+            failures.append(exc)
+
+    producer = threading.Thread(target=produce, daemon=True)
+    producer.start()
+    producer.join(timeout=30)
+    # With one slot, a lease the failed encoder never released leaves the
+    # producer waiting on the ring forever instead of seeing the codec error.
+    assert not producer.is_alive(), "the producer stalled on a slot the failed encoder kept"
+    assert len(failures) == 1
+
+    session.abort()
+    assert session.ring is not None
+    assert session.ring.slots_in_use == 0
+
+
+@hardware_test(res={"cuda": ["H100", "B200"]})
+def test_device_chunks_land_intact_through_the_pinned_ring():
+    """On an accelerator the copy is deferred, so its event must order the read."""
+    if not torch.accelerator.is_available():
+        pytest.skip("no accelerator available")
+
+    device = torch.device(torch.accelerator.current_accelerator().type, 0)
+    ring = PinnedChunkRing(depth=2, device=device)
+    pending: list[tuple[int, ChunkLease]] = []
+    for step in range(6):
+        value = 10 * step + 1
+        # Quantize on the compute stream immediately before handing the chunk
+        # over: without the source-ready event the copy would race this work.
+        chunk = torch.full((1, 3, 2, 4, 6), value / 255.0, device=device)
+        pending.append((value, ring.transfer(quantize_chunk(chunk, (0.0, 1.0)), readers=1)))
+        if len(pending) == 2:
+            for expected, lease in pending:
+                assert int(lease.wait()[0, 0, 0, 0, 0]) == expected
+                lease.release()
+            pending.clear()
+
+    assert ring.stats.transfers == 6
+    assert ring.stats.peak_slots_in_use == 2
+    assert ring.stats.peak_pending_bytes == 2 * 2 * 4 * 6 * 3

vllm_omni/diffusion/utils/chunked_video.py

@@ -10,41 +10,293 @@ and lands in a bounded encoder. That consumer lives here so a model only has to
 publish finished chunks -- :class:`SupportsChunkedVAEDecode` for producers that
 can be driven, :class:`ChunkedVideoMP4Session` for producers that drive
 themselves.
+
+The move to the host runs as its own scheduled stage. A chunk is quantized on
+the accelerator, copied without blocking into one of a fixed number of reusable
+pinned host slots on a dedicated stream, and only read once its copy-completion
+event has fired. The producer therefore returns to decoding while the copy is
+still in flight, and the slot count bounds how many chunks can be in flight at
+once.
 """
 
 from __future__ import annotations
 
+import threading
+import time
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
 from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder
 
+logger = init_logger(__name__)
 
-def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray:
-    """Quantize a ``BCTHW`` chunk to ``BTHWC`` uint8, transferring once.
+# Re-checked while waiting for a slot so an abort recorded without a notify
+# still wakes the producer.
+_SLOT_POLL_SECONDS = 0.05
+
+_COPY_STREAMS: dict[tuple[str, int], torch.Stream] = {}
+_COPY_STREAM_LOCK = threading.Lock()
+
+
+def _copy_stream(device: torch.device) -> torch.Stream:
+    """Return the process-wide chunk-copy stream for ``device``.
+
+    One stream per device rather than one per request: copies are issued in
+    order on it anyway, and a stream per request would churn accelerator
+    resources for no ordering benefit.
+    """
+    with _COPY_STREAM_LOCK:
+        key = (device.type, device.index or 0)
+        stream = _COPY_STREAMS.get(key)
+        if stream is None:
+            stream = torch.Stream(device=device)
+            _COPY_STREAMS[key] = stream
+        return stream
+
+
+def quantize_chunk(chunk: torch.Tensor, value_range: tuple[float, float]) -> torch.Tensor:
+    """Quantize a ``BCTHW`` chunk to contiguous ``BTHWC`` uint8 on its own device.
 
     ``value_range`` is the interval the producer publishes, which differs per
     checkpoint (see ``SupportsChunkedVAEDecode.chunk_value_range``). Quantizing
-    on the accelerator means one transfer moves the final bytes rather than
-    float frames.
-
-    That transfer is synchronous: ``.cpu()`` blocks inside the producer's chunk
-    callback, so what overlaps the remaining decode today is the CPU H.264
-    encode of earlier chunks, not this copy. Splitting device preparation,
-    device-to-host copy, and encode into three independently scheduled stages --
-    a pool of reusable pinned host slots with non-blocking copies on a dedicated
-    stream, leases released by the encoder -- fits behind this same callback
-    contract, but is not implemented here.
+    on the accelerator means the transfer that follows moves the final bytes
+    rather than float frames, and making the result contiguous there keeps the
+    copy a single flat move instead of a strided host-side gather.
     """
     low, high = float(value_range[0]), float(value_range[1])
     if high <= low:
         raise ValueError(f"value_range must be increasing, got {value_range!r}")
     scale = 255.0 / (high - low)
     frames = chunk.clamp(low, high).sub(low).mul(scale).round().to(torch.uint8)
-    return frames.permute(0, 2, 3, 4, 1).cpu().numpy()
+    return frames.permute(0, 2, 3, 4, 1).contiguous()
+
+
+def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray:
+    """Quantize a ``BCTHW`` chunk and block until it reaches the host.
+
+    The path a session takes when no pinned ring is available, and the one a
+    caller outside a session gets. A session on an accelerator uses
+    :class:`PinnedChunkRing` instead so the copy does not stall the producer.
+    """
+    return quantize_chunk(chunk, value_range).cpu().numpy()
+
+
+@dataclass
+class ChunkTransferStats:
+    """What the host transfer cost, for checking that it actually overlapped.
+
+    ``copy_wait_seconds`` staying near zero while ``transfers`` climbs is the
+    signal that copies completed behind the producer; it rising with them means
+    the producer is outrunning the copy stream. ``slot_wait_seconds`` rising
+    instead means the encoder is the bottleneck and backpressure has reached
+    the producer.
+    """
+
+    transfers: int = 0
+    slot_wait_seconds: float = 0.0
+    copy_wait_seconds: float = 0.0
+    peak_slots_in_use: int = 0
+    peak_pending_bytes: int = 0
+
+
+class _HostSlot:
+    """One reusable pinned host buffer.
+
+    The buffer is allocated on first use and never shrinks. ``batch_frames`` is
+    a flush threshold rather than a cap, so a later batch can be larger than
+    the first; the slot then keeps the larger buffer for the rest of the
+    request instead of reallocating per chunk.
+    """
+
+    def __init__(self, *, pinned: bool) -> None:
+        self._pinned = pinned
+        self._buffer: torch.Tensor | None = None
+        self.held_bytes = 0
+
+    def reserve(self, frames: torch.Tensor) -> torch.Tensor:
+        """Return a host view shaped like ``frames``, growing the buffer if needed."""
+        needed = frames.numel()
+        if self._buffer is None or self._buffer.numel() < needed:
+            try:
+                self._buffer = torch.empty(needed, dtype=frames.dtype, pin_memory=self._pinned)
+            except Exception:
+                # Pinning can be refused by the driver or unsupported on a
+                # platform. A pageable buffer still copies correctly, only
+                # synchronously, so degrade rather than fail the request.
+                logger.warning("Pinned host buffer unavailable; chunk copies fall back to pageable memory")
+                self._pinned = False
+                self._buffer = torch.empty(needed, dtype=frames.dtype)
+        self.held_bytes = needed * frames.element_size()
+        return self._buffer[:needed].view(frames.shape)
+
+
+class ChunkLease:
+    """One chunk in flight: a host slot, its copy event, and its readers.
+
+    The slot returns to the ring only after the last reader has finished with
+    it. Being handed to an encoder is not enough -- the encoder's queue holds a
+    view of this buffer until its muxing thread reads it -- so release is
+    driven by the encoder rather than by dequeuing.
+    """
+
+    def __init__(
+        self,
+        *,
+        ring: PinnedChunkRing,
+        slot: _HostSlot,
+        view: np.ndarray,
+        source: torch.Tensor,
+        landed: torch.Event | None,
+        readers: int,
+    ) -> None:
+        self._ring = ring
+        self._slot = slot
+        self._view = view
+        self._source: torch.Tensor | None = source
+        self._landed = landed
+        self._readers = readers
+        self._lock = threading.Lock()
+
+    def ready(self) -> bool:
+        """Whether the copy has landed, without blocking on it."""
+        return self._landed is None or self._landed.query()
+
+    def wait(self) -> np.ndarray:
+        """Block until the copy has landed, then hand over the host frames."""
+        if self._landed is not None:
+            started = time.perf_counter()
+            self._landed.synchronize()
+            self._ring.stats.copy_wait_seconds += time.perf_counter() - started
+            self._landed = None
+        # The source stays referenced until here so the allocator cannot hand
+        # its device memory to the next chunk while the copy is still reading.
+        self._source = None
+        return self._view
+
+    def release(self) -> None:
+        """Report one reader done; the last one returns the slot to the ring."""
+        with self._lock:
+            if self._readers <= 0:
+                return
+            self._readers -= 1
+            if self._readers > 0:
+                return
+        self._ring.reclaim(self._slot)
+
+    def discard(self) -> None:
+        """Drop an unread chunk during teardown, keeping the copy's memory alive.
+
+        The copy may still be writing into the slot and reading from the source
+        tensor, so this waits for it rather than letting either be reused
+        underneath an in-flight transfer.
+        """
+        if self._landed is not None:
+            self._landed.synchronize()
+            self._landed = None
+        self._source = None
+        with self._lock:
+            remaining, self._readers = self._readers, 0
+        if remaining > 0:
+            self._ring.reclaim(self._slot)
+
+
+class PinnedChunkRing:
+    """Fixed-depth pool of reusable host slots for chunk transfers.
+
+    A slot walks ``free -> copying -> ready -> encoding -> free`` and holds
+    exactly one chunk for that whole trip, so pending transfer bytes are
+    bounded by the depth times the largest chunk of the request. Waiting for a
+    free slot is where encoder backpressure reaches the producer.
+    """
+
+    def __init__(self, *, depth: int, device: torch.device) -> None:
+        if depth <= 0:
+            raise ValueError("transfer ring depth must be positive")
+        self.depth = depth
+        self.stats = ChunkTransferStats()
+        self._on_device = device.type != "cpu"
+        self._stream = _copy_stream(device) if self._on_device else None
+        self._lock = threading.Lock()
+        self._slot_freed = threading.Condition(self._lock)
+        self._free = [_HostSlot(pinned=self._on_device) for _ in range(depth)]
+        self._pending_bytes = 0
+        self._failure: BaseException | None = None
+
+    @property
+    def slots_in_use(self) -> int:
+        """Slots currently holding a chunk, whether copying, ready or encoding."""
+        with self._lock:
+            return self.depth - len(self._free)
+
+    def transfer(self, frames: torch.Tensor, *, readers: int) -> ChunkLease:
+        """Start one non-blocking copy of ``frames`` into a free slot."""
+        if readers <= 0:
+            raise ValueError("a transferred chunk needs at least one reader")
+        slot = self._acquire()
+        try:
+            host = slot.reserve(frames)
+            with self._lock:
+                self._pending_bytes += slot.held_bytes
+                self.stats.transfers += 1
+                self.stats.peak_pending_bytes = max(self.stats.peak_pending_bytes, self._pending_bytes)
+            landed = self._issue_copy(host, frames)
+        except BaseException:
+            self.reclaim(slot)
+            raise
+        return ChunkLease(ring=self, slot=slot, view=host.numpy(), source=frames, landed=landed, readers=readers)
+
+    def abort(self, error: BaseException | None = None) -> None:
+        """Fail every present and future wait for a slot."""
+        with self._slot_freed:
+            if self._failure is None:
+                self._failure = error or RuntimeError("chunk transfer ring aborted")
+            self._slot_freed.notify_all()
+
+    def reclaim(self, slot: _HostSlot) -> None:
+        """Return a slot whose last reader is done."""
+        with self._slot_freed:
+            self._pending_bytes -= slot.held_bytes
+            slot.held_bytes = 0
+            self._free.append(slot)
+            self._slot_freed.notify()
+
+    def _issue_copy(self, host: torch.Tensor, frames: torch.Tensor) -> torch.Event | None:
+        if self._stream is None:
+            host.copy_(frames)
+            return None
+        # The chunk was quantized on the compute stream, so the copy stream has
+        # to wait for that work before reading it.
+        source_ready = torch.Event()
+        source_ready.record()
+        self._stream.wait_event(source_ready)
+        with self._stream:
+            host.copy_(frames, non_blocking=True)
+            landed = torch.Event()
+            landed.record(self._stream)
+        return landed
+
+    def _acquire(self) -> _HostSlot:
+        started = time.perf_counter()
+        with self._slot_freed:
+            while not self._free and self._failure is None:
+                # ponytail: no request deadline yet, only an abort escape. A
+                # unified deadline across producer, slot, queue and encoder
+                # waits is RFC #6872 R4.
+                self._slot_freed.wait(timeout=_SLOT_POLL_SECONDS)
+            if self._failure is not None:
+                raise self._failure
+            slot = self._free.pop()
+            in_use = self.depth - len(self._free)
+            self.stats.peak_slots_in_use = max(self.stats.peak_slots_in_use, in_use)
+        self.stats.slot_wait_seconds += time.perf_counter() - started
+        return slot
 
 
 class ChunkedVideoMP4Session:
@@ -58,7 +310,8 @@ class ChunkedVideoMP4Session:
     silent entry), so a caller with several outputs per prompt repeats a
     request's waveform across its entries. ``batch_frames`` coalesces transfers
     for producers that publish finer than a transfer is worth; ``crop`` trims
-    the decoder's padding to the requested output size.
+    the decoder's padding to the requested output size. ``transfer_slots`` sets
+    how many chunks may be in flight to the host at once.
     """
 
     def __init__(
@@ -72,6 +325,7 @@ class ChunkedVideoMP4Session:
         max_pending: int = 2,
         video_codec_options: dict[str, str] | None = None,
         crop: tuple[int, int] | None = None,
+        transfer_slots: int = 2,
     ) -> None:
         if batch_frames <= 0:
             raise ValueError("batch_frames must be positive")
@@ -83,9 +337,13 @@ class ChunkedVideoMP4Session:
         self._max_pending = max_pending
         self._video_codec_options = video_codec_options
         self._crop = crop
+        self._transfer_slots = transfer_slots
         self._encoders: list[ChunkedMP4Encoder] = []
         self._pending: list[torch.Tensor] = []
         self._pending_frames = 0
+        self._ring: PinnedChunkRing | None = None
+        self._ring_resolved = False
+        self._inflight: deque[ChunkLease] = deque()
 
     def push(self, chunk: torch.Tensor) -> None:
         """Queue one committed ``BCTHW`` chunk."""
@@ -100,33 +358,96 @@ class ChunkedVideoMP4Session:
     def finish(self) -> list[bytes]:
         """Flush what is pending and return one MP4 per batch entry."""
         self._flush()
+        while self._inflight:
+            self._feed_lease(self._inflight.popleft())
         return [encoder.finish() for encoder in self._encoders]
 
     def abort(self) -> None:
+        if self._ring is not None:
+            self._ring.abort()
+        while self._inflight:
+            self._inflight.popleft().discard()
         for encoder in self._encoders:
             encoder.abort()
 
+    def transfer_stats(self) -> ChunkTransferStats | None:
+        """Per-request transfer accounting, or ``None`` if no ring was used."""
+        return None if self._ring is None else self._ring.stats
+
+    @property
+    def ring(self) -> PinnedChunkRing | None:
+        """The transfer ring this session settled on, once a chunk has flushed."""
+        return self._ring
+
     def _flush(self) -> None:
         if not self._pending:
             return
-        frames = chunk_to_uint8_frames(torch.cat(self._pending, dim=2), self._value_range)
-        if not self._encoders:
-            self._encoders = [
-                ChunkedMP4Encoder(
-                    width=frames.shape[3],
-                    height=frames.shape[2],
-                    fps=self._fps,
-                    audio_waveform=self._waveform_for(index, frames.shape[0]),
-                    audio_sample_rate=self._audio_sample_rate,
-                    max_pending=self._max_pending,
-                    video_codec_options=self._video_codec_options,
-                )
-                for index in range(frames.shape[0])
-            ]
-        for index, encoder in enumerate(self._encoders):
-            encoder.push(np.ascontiguousarray(frames[index]))
+        frames = quantize_chunk(torch.cat(self._pending, dim=2), self._value_range)
         self._pending.clear()
         self._pending_frames = 0
+        self._ensure_encoders(frames)
+        ring = self._resolve_ring(frames.device)
+        if ring is None:
+            self._feed(frames.cpu().numpy(), release=None)
+            return
+        # Free a slot before asking for one, so the ring's own wait is only
+        # ever for the encoder to catch up rather than for this loop.
+        while len(self._inflight) >= ring.depth:
+            self._feed_lease(self._inflight.popleft())
+        self._inflight.append(ring.transfer(frames, readers=len(self._encoders)))
+        while self._inflight and self._inflight[0].ready():
+            self._feed_lease(self._inflight.popleft())
+
+    def _feed_lease(self, lease: ChunkLease) -> None:
+        try:
+            self._feed(lease.wait(), release=lease.release)
+        except BaseException:
+            lease.discard()
+            raise
+
+    def _feed(self, frames: np.ndarray, *, release: Callable[[], None] | None) -> None:
+        for index, encoder in enumerate(self._encoders):
+            try:
+                encoder.push(frames[index], on_consumed=release)
+            except BaseException:
+                # push() owns the callback for the entry it was given, so only
+                # the encoders this chunk never reached still hold a reader.
+                for _ in range(len(self._encoders) - index - 1):
+                    if release is not None:
+                        release()
+                raise
+
+    def _ensure_encoders(self, frames: torch.Tensor) -> None:
+        if self._encoders:
+            return
+        self._encoders = [
+            ChunkedMP4Encoder(
+                width=int(frames.shape[3]),
+                height=int(frames.shape[2]),
+                fps=self._fps,
+                audio_waveform=self._waveform_for(index, int(frames.shape[0])),
+                audio_sample_rate=self._audio_sample_rate,
+                max_pending=self._max_pending,
+                video_codec_options=self._video_codec_options,
+            )
+            for index in range(int(frames.shape[0]))
+        ]
+
+    def _resolve_ring(self, device: torch.device) -> PinnedChunkRing | None:
+        if self._ring_resolved:
+            return self._ring
+        self._ring_resolved = True
+        if self._transfer_slots <= 0:
+            return None
+        try:
+            self._ring = PinnedChunkRing(depth=self._transfer_slots, device=device)
+        except Exception:
+            logger.warning(
+                "Pinned chunk transfers unavailable on %s; falling back to blocking copies",
+                device,
+                exc_info=True,
+            )
+        return self._ring
 
     def _waveform_for(self, index: int, batch_size: int) -> np.ndarray | None:
         if self._audio_waveforms is None:
@@ -154,7 +475,18 @@ def decode_to_mp4(vae: Any, z: torch.Tensor, **session_kwargs: Any) -> list[byte
     session = ChunkedVideoMP4Session(value_range=vae.chunk_value_range, **session_kwargs)
     try:
         vae.decode_with_chunks(z, on_chunk=session.push)
-        return session.finish()
+        videos = session.finish()
     except BaseException:
         session.abort()
         raise
+    stats = session.transfer_stats()
+    if stats is not None:
+        logger.debug(
+            "Chunked MP4 transfers=%d slot_wait=%.3fs copy_wait=%.3fs peak_slots=%d peak_bytes=%d",
+            stats.transfers,
+            stats.slot_wait_seconds,
+            stats.copy_wait_seconds,
+            stats.peak_slots_in_use,
+            stats.peak_pending_bytes,
+        )
+    return videos

vllm_omni/diffusion/utils/media_utils.py

@@ -7,16 +7,36 @@ from __future__ import annotations
 import io
 import queue
 import threading
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from fractions import Fraction
-from typing import Any, cast
+from typing import Any, NamedTuple, cast
 
 import av
 import numpy as np
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 _CHUNKED_MP4_DONE = object()
 
 
+class _QueuedChunk(NamedTuple):
+    """One queued chunk plus the callback that returns its buffer to its owner."""
+
+    frames: np.ndarray
+    on_consumed: Callable[[], None] | None
+
+
+def _release_chunk(on_consumed: Callable[[], None] | None) -> None:
+    """Hand a chunk's buffer back, never letting the callback wedge the worker."""
+    if on_consumed is None:
+        return
+    try:
+        on_consumed()
+    except BaseException:  # noqa: BLE001
+        logger.exception("Chunked MP4 encoder failed to release a consumed chunk")
+
+
 def normalize_preencode_batch_frames(value: Any) -> int:
     """Validate the request's MP4 transfer/encoding batch threshold."""
     if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
@@ -88,9 +108,10 @@ class ChunkedMP4Encoder:
         self._state_lock = threading.Lock()
 
         def run() -> None:
+            frames = self._frames()
             try:
                 self._result = mux_av_video_audio_bytes(
-                    self._frames(),
+                    frames,
                     width=self.width,
                     height=self.height,
                     audio_waveform=audio_waveform,
@@ -103,6 +124,11 @@ class ChunkedMP4Encoder:
                 )
             except BaseException as exc:
                 self._error = exc
+                # Closing is what runs the generator's cleanup and releases the
+                # chunk it was mid-way through. Waiting for collection would
+                # never get there: the traceback just stored on this encoder
+                # keeps the abandoned generator alive for the process's life.
+                frames.close()
                 self._drain_until_done()
 
         self._thread = threading.Thread(target=run, name="chunked-mp4", daemon=True)
@@ -110,19 +136,29 @@ class ChunkedMP4Encoder:
 
     def _drain_until_done(self) -> None:
         # A flush/mux failure can occur after _frames consumed the sentinel.
-        if not self._input_done:
-            while self._queue.get() is not _CHUNKED_MP4_DONE:
-                pass
+        if self._input_done:
+            return
+        while True:
+            entry = self._queue.get()
+            if entry is _CHUNKED_MP4_DONE:
+                return
+            _release_chunk(cast(_QueuedChunk, entry).on_consumed)
 
     def _frames(self):
         while True:
-            chunk = self._queue.get()
-            if chunk is _CHUNKED_MP4_DONE:
+            entry = self._queue.get()
+            if entry is _CHUNKED_MP4_DONE:
                 self._input_done = True
                 return
-            assert isinstance(chunk, np.ndarray)
-            for frame_data in chunk:
-                yield av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+            assert isinstance(entry, _QueuedChunk)
+            try:
+                for frame_data in entry.frames:
+                    yield av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+            finally:
+                # Also runs when the muxer raises part-way through this chunk
+                # and abandons the generator, so a lent buffer is never
+                # stranded in the encoder.
+                _release_chunk(entry.on_consumed)
 
     def _validate_chunk(self, chunk: np.ndarray) -> None:
         _validate_video_chunk(chunk, width=self.width, height=self.height)
@@ -133,11 +169,22 @@ class ChunkedMP4Encoder:
         if self._closed:
             raise RuntimeError("ChunkedMP4Encoder is already closed")
 
-    def push(self, chunk: np.ndarray) -> None:
-        """Queue a chunk; asynchronous failures surface here or at finish()."""
-        self._validate_chunk(chunk)
-        self._raise_if_failed()
-        self._queue.put(chunk)
+    def push(self, chunk: np.ndarray, *, on_consumed: Callable[[], None] | None = None) -> None:
+        """Queue a chunk; asynchronous failures surface here or at finish().
+
+        ``on_consumed`` fires exactly once after this encoder has read
+        ``chunk``, including when the chunk is discarded by an abort or by a
+        failed encode. Ownership transfers on call: whether this returns or
+        raises, the callback is this encoder's responsibility, so a caller
+        lending a pooled host buffer never compensates on the error path.
+        """
+        try:
+            self._validate_chunk(chunk)
+            self._raise_if_failed()
+        except BaseException:
+            _release_chunk(on_consumed)
+            raise
+        self._queue.put(_QueuedChunk(chunk, on_consumed))
         self._raise_if_failed()
 
     def _send_done(self) -> None:
@@ -168,9 +215,11 @@ class ChunkedMP4Encoder:
                 self._aborted = True
                 while True:
                     try:
-                        self._queue.get_nowait()
+                        entry = self._queue.get_nowait()
                     except queue.Empty:
                         break
+                    if entry is not _CHUNKED_MP4_DONE:
+                        _release_chunk(cast(_QueuedChunk, entry).on_consumed)
                 self._send_done()
         self._thread.join()