[1/N] Stream Wan VAE chunks to the media consumer

[specture724]

[1/3] Stream Wan VAE chunks to the media consumer

Part of the RFC #6872 implementation stack (1/4). The consumer for this seam is #7048, which wires Wan T2V, I2V, and S2V into worker-side MP4 encoding.

Summary

Add a callback seam to Wan VAE decode so committed temporal chunks can leave the GPU while later chunks are still decoding. Keep the existing full-tensor path unchanged when no callback is supplied.

The producer owns decode ordering; the consumer owns transfer, conversion, and encoding. Spatial-sharded decode publishes only complete output chunks from the owning rank, and callback failures still drain the decode path before being re-raised.

Validation

• Existing chunked-media tests pass, including lifecycle, callback failure, and audio-mux stability coverage.
• L20X, Wan2.1-T2V-1.3B, real VAE weights, latent [1,16,65,60,104], output 480×832, H.264.
  This isolates the VAE decode stage; it is not a full-request number:

Path	Total (avg of 3)	D2H transfers	Decoded frames
Full decode	~13.58 s	1	257
Chunk callback	~9.79 s	65	257
Chunk callback + consumer batching	~9.73 s	13	257

Consumer-side batching reduces total time by ~28% versus full decode while preserving the complete frame count. Encoded MP4 bytes are not used as a cross-run equality oracle because H.264 threading is not byte-deterministic — three repeated full-decode runs alone produce three different encoded SHA-256 sums. A direct comparison of the pre-encode float tensor (full decode vs. concatenated streamed chunks) is bit-exact: identical shape, max_abs_error=0, mean_abs_error=0, torch.equal=True, and identical SHA-256.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to be related to model: Wan.

Model owners: @bjf-frz

@specture724, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[hsliuustc0106]

CI lanes select tests with pytest -m 'core_model and cpu' (.buildkite/cuda/test-merge.yml:20) and this file has no pytestmark, so the new ChunkedMP4Encoder coverage never runs in any lane — please add pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion].

[hsliuustc0106]

Nothing in this stack passes on_chunk to Wan decode — 3/3 wires only MiniMax-H3 into worker-side encoding, so this seam plus the tiled_decode/spatial_shard_decode plumbing lands without a production caller. Is there a follow-up PR that wires Wan in? If not, consider landing the seam with its consumer; and the ~28% figure in the description comes from a standalone harness, not a wired serving path, so it's worth labeling it as seam validation only.

[hsliuustc0106]

fix DCO

[hsliuustc0106]

Local review against b568ccb; inline findings below.

[hsliuustc0106]

The callback branch bypasses the parent decode, including its use_slicing behavior. I reproduced this with the real Wan VAE at batch size 2 and use_slicing=True: normal decode sent batches [1, 1] through post_quant_conv, while callback decode sent [2]. Thus batched pre-encoding can silently lose the configured memory protection and OOM. Please either implement slice-aware callback delivery or explicitly reject callback decoding when slicing would be required.

[hsliuustc0106]

This is flaky with the default multithreaded H.264 encoder. Repeating this exact 8x16x24 workload locally produced 6 byte mismatches in 500 runs (and decoded-pixel differences from alternative valid lossy encodes); setting video_codec_options={"threads": "1"} on both paths produced 0/500. Please pin single-threaded encoding for a byte-stability test, or compare playable structure/frame count instead. The audio equality assertion below has the same video-side risk.

[hsliuustc0106]

After the first callback failure, callback_error is non-None, so every subsequent chunk reaches this else and is retained in decoded_chunks; because on_chunk is present, the list is never consumed. A six-chunk reproducer showed retained chunks growing [0, 1, 1, 2, 3, 4]. On high-resolution spatial-sharded decode this can rebuild the full output footprint and OOM before the original callback exception is re-raised. Append only when on_chunk is None; otherwise decode and discard after a recorded callback failure.

[hsliuustc0106]

The inherited AutoencoderKLWan.decode is wrapped by Diffusers' apply_forward_hook, but this direct callback path neither calls that wrapper nor applies the decorator itself. I verified an attached _hf_hook.pre_forward receives zero calls here. Accelerate-offloaded instances can therefore enter post_quant_conv while VAE modules are still off-device. Please preserve the decode pre-forward hook for the callback path (without double-invoking it on the existing super().decode path).

[hsliuustc0106]

Architecturally, the temporal callback is being introduced at the right layer: the Wan decoder loop is where complete causal temporal chunks become available, and restricting distributed support to spatial_shard matches the current output-ownership model. The stacked consumer in #7048 also keeps Wan-specific batching outside the generic media utility.

I think the public API should be adjusted, though. Rather than extending and partially reimplementing the Diffusers-compatible decode() method, it would be cleaner to expose an explicit optional capability, for example:

DecodedChunkConsumer = Callable[[torch.Tensor], None]

@runtime_checkable
class SupportsChunkedVAEDecode(Protocol):
    def decode_with_chunks(
        self,
        z: torch.Tensor,
        *,
        on_chunk: DecodedChunkConsumer,
    ) -> None: ...

The ordinary decode() contract could then remain unchanged, preserving slicing, forward hooks, and its normal tensor/DecoderOutput result. Both entry points can share a private temporal decoding primitive. The streaming method should document that all distributed ranks must invoke it, while the callback runs only on the output-owning rank.

I would prefer this callback-based method over returning a Python generator: distributed decoding must run the collective sequence to completion, whereas an abandoned or partially consumed generator could leave peer ranks blocked. A dedicated method also avoids using an empty DecoderOutput to represent both successful streaming and a non-owner distributed rank.

This follows the repository's existing Supports... protocol pattern and gives downstream pipelines an explicit, typed way to discover the capability without assuming that every VAE accepts an on_chunk argument.

[RuixiangMa]

If the worker fails between _raise_if_failed() and _queue.put(), _drain_until_done() discards the chunk. Error surfaces on next push() or finish(), not at the affected push(). No silent corruption, finish() always raises.

[RuixiangMa]

If queue is at max_pending and worker is inside PyAV, put blocks until worker consumes enough items. Not a deadlock — bounded by worker chunk processing time. Could use put_nowait + pre-drain for faster abort.

[RuixiangMa]

close() calls abort(), which doesn't return self._result. A user calling close() after push() loses encoded bytes without error. Consider making close call finish() or removing the alias.

[hsliuustc0106]

lgtm, please update the docs later