[4/N] Unify worker-side MP4 encoding and wire Wan into it

[specture724]

[4/4] Unify worker-side MP4 encoding and wire Wan into it

Part of the RFC #6872 implementation stack (4/4), stacked on #7016, #7017, and
#7018. GitHub can't set a fork branch as another PR's base, so this targets
main like the rest of the stack; the diff includes the earlier PRs' commits
until they merge.

Summary

[1/N] added the Wan VAE chunk callback but nothing called it. This wires it up
and, now that a second consumer exists, moves the consumer itself into the
framework.

Producing chunks is the model's business: Wan walks a causal frame-by-frame
loop, H3 walks overlapping clips, and Wan S2V decodes one clip per
autoregressive iteration. What happens after a chunk is committed is not --
every one of them quantizes to uint8, moves to the host once, and lands in a
bounded encoder. Each model carried its own copy of that, so a capability meant
to be model-agnostic had a per-backend consumer bolted onto it.

diffusion/utils/chunked_video.py now owns that half:

• ChunkedVideoMP4Session for producers that drive themselves (Wan S2V pushes
  finished clips out of its autoregressive loop).
• decode_to_mp4() for producers that can be driven through
  SupportsChunkedVAEDecode.

Batching, per-batch-entry encoders, audio muxing, the requested-size crop, and
quantization from the producer's declared chunk_value_range are parameters.
H3, Wan T2V/I2V, and Wan S2V share it: wan2_2/chunked_mp4.py drops from 219
to 77 lines (only request-option plumbing remains) and H3's decode_to_mp4
keeps only what is H3-specific -- decoding audio first, the crop, the fixed
rate. A model that wants worker-side encoding now implements one method and
declares one pixel range.

With preencode_mp4, each committed chunk leaves the accelerator while later
chunks are still decoding, so host transfer and H.264 encoding overlap the
remaining decode instead of following it. The pre-encoded path emits MP4 bytes,
so it rejects rather than silently downgrades what it cannot serve:
enable_frame_interpolation, which needs the frames it no longer materializes,
and non-array output types.

Validation

Two measurements, at different scopes. They are not comparable to each other.

VAE decode stage, T2V — real Wan2.1-T2V-1.3B VAE weights, L20X, latent
[1,16,65,60,104] → 257 frames at 480×832. Denoising is excluded, so this
isolates the stage the callback changes:

Path	Total (avg of 3)	Host transfers
Full decode + one-shot mux	~13.05 s	1
Wired pre-encode, per chunk	~9.68 s	65
Wired pre-encode, 17-frame batches	~9.81 s	13

~26% off the decode-plus-encode stage at an identical 257-frame count.
Consumer batching trades 65 host transfers for 13; wall clock is within noise of
the per-chunk path on this GPU, so the default is about transfer count, not
latency.

Full request, S2V — real Wan2.2-S2V-14B weights, L20X, TP=1, 448×448, 77
frames, 4 steps, warmed up, both orders run to expose ordering effects. The
baseline's MP4 encoding happens downstream of generate(), so it is measured
and added:

Path	Run 1	Run 2
Baseline generate()	15.95 s	15.86 s
Baseline downstream mux	0.48 s	0.46 s
Baseline end to end	16.44 s	16.31 s
Pre-encode end to end	15.80 s	15.86 s

~3% end to end. The win is small here because denoising dominates a full S2V
request and the mux is only ~0.47 s of ~16.4 s; it is the same ~0.5 s of
encoding, now absorbed into the decode loop. An earlier unwarmed run showed
~32%, which was entirely first-call overhead — that number is not real.

Output equivalence: 77 frames on both paths, audio stream present in the
pre-encoded container. The consumer move is behavior-preserving: real
Wan2.1-T2V-1.3B weights decode to the same 257 frames with an unchanged
pre-encode SHA. Decoded pixels differ by lossy H.264 re-encoding
(max_abs 37, mean_abs 1.70), which is expected; the bit-exact signal is the
pre-encode tensor comparison in [1/3].

Tests: 11 new cases for the framework consumer, covering per-batch-entry
containers, the crop, audio muxing, ranks that own no output, a VAE that does
not declare the capability, and the property the shared quantizer exists for --
a Wan [-1, 1] chunk and an H3 [0, 1] chunk of the same picture quantize to
the same pixels. Wan's own cases cover transfer coalescing and every rejected
combination. Two more drive the real S2V forward() over a two-clip loop on
CPU, checking that both clips reach the encoder, that every output decodes as a
playable container with audio, and that each request's waveform lands on its own
outputs.

Known limitation

With VAE patch parallelism, S2V broadcasts decoded frames to every rank because
the autoregressive loop needs them, so every rank encodes its own identical
copy. This matches the existing full-decode path, which likewise concatenates
the same tensor on every rank, and is left alone here rather than changing which
rank produces output.

🤖 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]

[P1] Abort the S2V MP4 session when clip generation fails. After an earlier clip starts encoders, a later denoise/VAE failure exits forward() without aborting the session. The encoder threads retain their containers/audio and wait for a sentinel that never arrives. An injected second-clip failure through the real S2V forward, using the CPU model fixture and real encoders, left all four encoders alive; session.abort() closed all four. Please wrap the session lifetime in cleanup covering every unsuccessful exit, including push/finish failures. pipeline_wan2_2_s2v.py:1492–1503.

[P2] Preserve serving codec defaults for Wan pre-encoding. With only extra_params={"preencode_mp4": true}, this helper returns None. Serving normally supplies preset=ultrafast, threads=0, but skips its encoder when the worker returns bytes. Enabling pre-encoding therefore changes the codec preset and encoding cost. Please apply the serving defaults when the key is absent, while preserving explicit null/empty/custom options. H3 already carries these defaults into worker encoding. chunked_mp4.py:63–66.

Reviewed fee2684b836f against ecdd762ff576, distinguishing this PR's changes from the inherited #7018 work. 210 targeted CPU tests passed across the shared consumer and Wan/S2V/H3 contracts. The peer-VAE-rank fix is present; the earlier #7018 metadata finding is inherited.

Local A/B used one reserved L20X, Python 3.12, torch 2.13.0, vLLM 0.28.0, diffusers 0.40.0, and the real Wan2.1-T2V-1.3B VAE (revision 0fad780a534b6463e45facd96134c9f345acfa5b). Controls: BF16, seed 7048 latent [1,16,65,60,104], 257 frames at 480×832/24fps, no audio/tiling/offload, identical H264 CRF18/ultrafast/threads=0. Baseline uses normal decode/postprocess/serving encode with eight converter workers; head uses decode_to_mp4, batch threshold 17. One warmup and two measured runs per side; loading and file writes excluded, shared caches retained.

	Measured decode-to-MP4 times	Peak allocated GPU memory
Base	7.288 / 7.290 s	8222.0 MiB
Head	6.708 / 6.661 s	7679.2 MiB

Mean stage latency fell 8.29%, but output equivalence did not hold: both pairs retained 257 frames, with decoded RGB MAE 3.426/255, maximum difference 56, and PSNR 34.56 dB. These two observations do not establish output-quality equivalence or a full-request speedup; denoising, HTTP and IPC were not measured.

[hsliuustc0106]

Re-reviewed af8aa1c66fa3 against 7e520f9cc47e. Both findings from the previous review remain unresolved:

[P1] Abort the S2V MP4 session when clip generation fails. A second-clip denoise failure still exits forward() without aborting the session. Reproducing this through the real S2V forward with the CPU model fixture and real encoders left all four encoder threads alive and unclosed, retaining their containers/audio while waiting for more input. Explicit session.abort() closed all four. Please guard the entire session lifetime and abort on unsuccessful exits, including push/finish failures. pipeline_wan2_2_s2v.py:1492–1503.

[P2] Preserve serving codec defaults for Wan pre-encoding. With only extra_params={"preencode_mp4": true}, the helper still returns None. The normal serving paths supply preset=ultrafast, threads=0, but bypass encoding once the worker returns bytes. Enabling pre-encoding therefore changes codec settings and their cost/compression behavior. Please apply the serving defaults when video_codec_options is absent, while preserving explicit null/empty/custom options. chunked_mp4.py:63–66.

The update is mostly a rebase onto merged #7018. The T2V adaptation preserves typed media on the output-owning rank and the legacy empty placeholder on non-owner ranks; I found no additional regression in that change.

Validation: 246 focused CPU tests passed across shared chunked video, Wan/S2V, H3, T2V forward and typed-media contracts. An additional T2V smoke check with model/VAE stubs and real CPU encoding passed for two requests × two outputs: the owner produced four playable MP4s; the simulated non-owner returned empty outputs without a typed-media batch error. Both earlier findings were reproduced on this head, and all reproducer-owned encoder threads were cleaned up. Python 3.12, torch 2.13.0, vLLM 0.29.0, PyAV 18.1.0.

The earlier GPU A/B was not rerun and is not current-head performance evidence. Full S2V denoising, HTTP end-to-end behavior and multi-GPU execution were not tested in this re-review.

[hsliuustc0106]

@MuweiZ any more comments?

[hsliuustc0106]

[P3] Encode only on the output-owning VAE-PP rank. After the broadcast every rank of _vae_pp_group holds the clip, so every rank pushes into its own session and runs quantize + D2H + a full H.264/AAC encode per clip — but the executor collects only rank 0's forward response (num_responses = 1 in multiproc_executor.py), so peer output is discarded. T2V/I2V/H3 all encode owner-only because their peer ranks receive no chunks. The broadcast itself must stay (the AR motion loop needs frames on all ranks); gating the push on the pp-group owner leaves peers with an empty session, and finish() == [] is already tolerated downstream (the T2V non-owner path).

[hsliuustc0106]

Re-reviewed 015a2a5e5d with an architecture/performance focus. Both prior findings are fixed:

• [P1 → fixed] Session abort. The clip loop and output assembly are wrapped in try/except BaseException → mp4_session.abort(). I traced abort() on every exit: idempotent after a failed finish() (already-closed encoders only join), a no-op before the first flush, and an encoder-worker failure drains its own queue so a blocked push cannot deadlock. The new test drives the second-clip failure through the real forward and asserts all four encoder threads exit.
• [P2 → fixed] Codec defaults. An absent video_codec_options now inherits the one-shot serving encoder's preset=ultrafast, threads=0; explicit None/{}/custom stay caller choices, all four covered by tests.

Architecture: the owner-only encoding convention holds for T2V/I2V/H3 — peers receive no chunks, decode_to_mp4 returns [], and H3's wrapper returns b"" (pipeline_minimax_h3.py:2213), which resolves the 09-09 non-owner-rank thread here. One gap remains on S2V (inline comment): it pushes on every VAE-PP rank while only rank 0's output is consumed.

Performance: the body's claims match the implementation. chunk_to_uint8_frames quantizes on-device and transfers once per flush batch; the synchronous D2H inside the producer callback and the chunk-count (not byte) bound on pending memory are now stated in the docstrings, consistent with the 09-09 threads — per-chunk timestamps and max_pending_bytes are reasonable future transport-layer work. S2V's per-clip .cpu() becomes an equivalent batched transfer, so there is no new critical-path stall, and encode overlaps the next clip's denoise.

CI on this head: main / AMD / NPU / build 3.11+3.12 / pre-commit / RTD green. vllm-omni-intel-ci failed in 54s; the log is not anonymously readable and the duration points at setup rather than tests — worth a rerun. I could not run the new tests locally (no torch+PyAV environment on the review host), so test evidence here is CI plus reading.

[lishunyang12]

Validate common.output_type here so preencoding rejects pil, pt, and latent requests.