[Core][Diffusion] Add byte backpressure to chunk transfers (#7916)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/serving/videos_api.md

@@ -335,6 +335,13 @@ window and output frame count stay unchanged. Wan S2V keeps its existing
 per-clip behavior by default. Zero, negative, fractional, boolean, string, and
 null values are rejected when pre-encoding is enabled.
 
+The transfer ring applies backpressure before another D2H copy when both of
+its two slots are occupied or the combined pending uint8 payload would exceed
+256 MiB. This byte bound keeps large resolutions and uneven final chunks from
+turning a fixed item count into unexpectedly large memory growth. A native VAE
+chunk larger than 256 MiB is admitted only when the ring is otherwise empty,
+so supported chunk shapes cannot deadlock.
+
 `preencode_mp4` applies to the complete-MP4 response paths only. The
 `/v1/realtime/video` WebSocket endpoint rejects it, because that path already
 overlaps encoding through its own incremental fragmented-MP4 encoder. Wan also

tests/diffusion/utils/test_chunked_transfer.py

@@ -33,8 +33,8 @@ def _frames(value: int = 7, *, batch: int = 1) -> torch.Tensor:
     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 _cpu_ring(depth: int, *, max_pending_bytes: int | None = None) -> PinnedChunkRing:
+    return PinnedChunkRing(depth=depth, device=torch.device("cpu"), max_pending_bytes=max_pending_bytes)
 
 
 def test_slot_returns_only_after_its_last_reader():
@@ -81,6 +81,44 @@ def test_transfer_waits_for_a_slot_instead_of_growing_the_pool():
     assert ring.stats.slot_wait_seconds > 0
 
 
+def test_transfer_waits_for_byte_budget_even_when_a_slot_is_free():
+    """Variable-sized chunks cannot bypass the memory bound via free slots."""
+    small = _frames()
+    large = _frames(batch=2)
+    large_bytes = large.numel() * large.element_size()
+    ring = _cpu_ring(3, max_pending_bytes=large_bytes)
+    held = ring.transfer(small, readers=1)
+    second: list[ChunkLease] = []
+
+    waiter = threading.Thread(target=lambda: second.append(ring.transfer(large, readers=1)))
+    waiter.start()
+    waiter.join(timeout=0.3)
+    assert waiter.is_alive(), "the byte budget must backpressure a producer even with free slots"
+
+    held.wait()
+    held.release()
+    waiter.join(timeout=5)
+    assert not waiter.is_alive()
+    assert ring.slots_in_use == 1
+    assert ring.pending_bytes == large_bytes
+    second[0].wait()
+    second[0].release()
+
+
+def test_one_oversized_chunk_is_admitted_without_deadlock():
+    """A native VAE chunk larger than the budget must make progress alone."""
+    chunk = _frames()
+    chunk_bytes = chunk.numel() * chunk.element_size()
+    ring = _cpu_ring(2, max_pending_bytes=chunk_bytes - 1)
+
+    lease = ring.transfer(chunk, readers=1)
+
+    assert ring.pending_bytes == chunk_bytes
+    assert ring.stats.peak_pending_bytes == chunk_bytes
+    lease.wait()
+    lease.release()
+
+
 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)
@@ -180,6 +218,26 @@ def test_session_feeds_chunks_in_producer_order(recording_encoders):
     assert recording_encoders[0].pushes == [10, 20, 30, 40]
 
 
+def test_session_drains_oldest_chunk_when_byte_budget_fills(recording_encoders, monkeypatch):
+    """Byte pressure must drain queued D2H work instead of deadlocking it."""
+    monkeypatch.setattr(ChunkLease, "ready", lambda self: False)
+    chunk_bytes = 1 * 1 * 4 * 6 * 3  # one quantized BTHWC chunk from _push_ramp
+    session = ChunkedVideoMP4Session(
+        value_range=(0.0, 1.0),
+        fps=24,
+        transfer_slots=2,
+        max_pending_bytes=chunk_bytes,
+    )
+
+    _push_ramp(session, [10, 20, 30])
+
+    assert session.finish() == [b"mp4"]
+    assert recording_encoders[0].pushes == [10, 20, 30]
+    assert session.ring is not None
+    assert session.ring.stats.peak_slots_in_use == 1
+    assert session.ring.stats.peak_pending_bytes == chunk_bytes
+
+
 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)

tests/diffusion/utils/test_chunked_video.py

@@ -116,3 +116,9 @@ def test_session_rejects_a_waveform_count_that_does_not_match_the_batch():
 def test_session_rejects_a_non_positive_batch():
     with pytest.raises(ValueError, match="batch_frames"):
         ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, batch_frames=0)
+
+
+@pytest.mark.parametrize("value", [0, -1, 1.5, True])
+def test_session_rejects_an_invalid_pending_byte_budget(value):
+    with pytest.raises(ValueError, match="max_pending_bytes"):
+        ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, max_pending_bytes=value)

vllm_omni/diffusion/utils/chunked_video.py

@@ -41,6 +41,11 @@ logger = init_logger(__name__)
 # still wakes the producer.
 _SLOT_POLL_SECONDS = 0.05
 
+# Bound queued/in-flight uint8 video independently of resolution and chunk
+# length. A single larger chunk is still admitted when the ring is empty so a
+# legal native VAE chunk can never deadlock the producer.
+_DEFAULT_MAX_PENDING_BYTES = 256 * 1024**2
+
 _COPY_STREAMS: dict[tuple[str, int], torch.Stream] = {}
 _COPY_STREAM_LOCK = threading.Lock()
 
@@ -211,15 +216,24 @@ 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.
+    exactly one chunk for that whole trip. Admission is bounded by both slot
+    count and pending bytes; reaching either limit is where encoder
+    backpressure reaches the producer.
     """
 
-    def __init__(self, *, depth: int, device: torch.device) -> None:
+    def __init__(
+        self,
+        *,
+        depth: int,
+        device: torch.device,
+        max_pending_bytes: int | None = _DEFAULT_MAX_PENDING_BYTES,
+    ) -> None:
         if depth <= 0:
             raise ValueError("transfer ring depth must be positive")
+        if max_pending_bytes is not None and (type(max_pending_bytes) is not int or max_pending_bytes <= 0):
+            raise ValueError("max_pending_bytes must be a positive integer or None")
         self.depth = depth
+        self.max_pending_bytes = max_pending_bytes
         self.stats = ChunkTransferStats()
         self._on_device = device.type != "cpu"
         self._stream = _copy_stream(device) if self._on_device else None
@@ -235,15 +249,26 @@ class PinnedChunkRing:
         with self._lock:
             return self.depth - len(self._free)
 
+    @property
+    def pending_bytes(self) -> int:
+        """Bytes held by chunks currently copying, ready, or encoding."""
+        with self._lock:
+            return self._pending_bytes
+
+    def would_block(self, chunk_bytes: int) -> bool:
+        """Whether admitting ``chunk_bytes`` now would exceed either bound."""
+        with self._lock:
+            return not self._free or not self._bytes_available(chunk_bytes)
+
     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()
+        chunk_bytes = frames.numel() * frames.element_size()
+        slot = self._acquire(chunk_bytes)
         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)
@@ -282,10 +307,17 @@ class PinnedChunkRing:
             landed.record(self._stream)
         return landed
 
-    def _acquire(self) -> _HostSlot:
+    def _bytes_available(self, chunk_bytes: int) -> bool:
+        return (
+            self.max_pending_bytes is None
+            or self._pending_bytes == 0
+            or self._pending_bytes + chunk_bytes <= self.max_pending_bytes
+        )
+
+    def _acquire(self, chunk_bytes: int) -> _HostSlot:
         started = time.perf_counter()
         with self._slot_freed:
-            while not self._free and self._failure is None:
+            while (not self._free or not self._bytes_available(chunk_bytes)) 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.
@@ -293,6 +325,11 @@ class PinnedChunkRing:
             if self._failure is not None:
                 raise self._failure
             slot = self._free.pop()
+            # Reserve before dropping the lock so concurrent producers cannot
+            # both pass the byte budget. Setting held_bytes here also makes a
+            # failed host allocation reclaim the complete reservation.
+            slot.held_bytes = chunk_bytes
+            self._pending_bytes += chunk_bytes
             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
@@ -311,7 +348,8 @@ class ChunkedVideoMP4Session:
     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. ``transfer_slots`` sets
-    how many chunks may be in flight to the host at once.
+    how many chunks may be in flight to the host at once, while
+    ``max_pending_bytes`` bounds their combined uint8 payload size.
     """
 
     def __init__(
@@ -326,9 +364,12 @@ class ChunkedVideoMP4Session:
         video_codec_options: dict[str, str] | None = None,
         crop: tuple[int, int] | None = None,
         transfer_slots: int = 2,
+        max_pending_bytes: int | None = _DEFAULT_MAX_PENDING_BYTES,
     ) -> None:
         if batch_frames <= 0:
             raise ValueError("batch_frames must be positive")
+        if max_pending_bytes is not None and (type(max_pending_bytes) is not int or max_pending_bytes <= 0):
+            raise ValueError("max_pending_bytes must be a positive integer or None")
         self._value_range = value_range
         self._fps = fps
         self._audio_waveforms = audio_waveforms
@@ -338,6 +379,7 @@ class ChunkedVideoMP4Session:
         self._video_codec_options = video_codec_options
         self._crop = crop
         self._transfer_slots = transfer_slots
+        self._max_pending_bytes = max_pending_bytes
         self._encoders: list[ChunkedMP4Encoder] = []
         self._pending: list[torch.Tensor] = []
         self._pending_frames = 0
@@ -392,7 +434,8 @@ class ChunkedVideoMP4Session:
             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:
+        chunk_bytes = frames.numel() * frames.element_size()
+        while self._inflight and ring.would_block(chunk_bytes):
             self._feed_lease(self._inflight.popleft())
         self._inflight.append(ring.transfer(frames, readers=len(self._encoders)))
         while self._inflight and self._inflight[0].ready():
@@ -440,7 +483,11 @@ class ChunkedVideoMP4Session:
         if self._transfer_slots <= 0:
             return None
         try:
-            self._ring = PinnedChunkRing(depth=self._transfer_slots, device=device)
+            self._ring = PinnedChunkRing(
+                depth=self._transfer_slots,
+                device=device,
+                max_pending_bytes=self._max_pending_bytes,
+            )
         except Exception:
             logger.warning(
                 "Pinned chunk transfers unavailable on %s; falling back to blocking copies",