[perf] Quantize frames to uint8 on-device before the post-decode D->H copy

[Mister-Raggs]

Summary

PostDecodeFrameProcessStage spends most of its time on a deferred fp32 device→host transfer mischarged to it, plus a single-threaded per-frame CPU *255/cast loop.

Fix: cast the decoded video to uint8 on the source device (typically CUDA) before the host copy. The transfer becomes 4× smaller (fp32 → uint8) and the elementwise work runs on the GPU instead of a Python per-frame CPU loop.

Mechanism proven and independently cross-validated. The per-path speedup (~−50–66%) is consistent; e2e impact scales with how much of total latency the post-decode path owns on a given model / resolution / GPU:

Config	PostDecodeStage	End-to-end	Source
Cosmos 2.5, A100, 720×1280×29f, 8 steps	5.991 s → 2.015 s (−66%)	29.33 s → 24.75 s (−15.6%)	author
Wan2.1-T2V-1.3B, H100, 720×1280×77f, 30 steps	0.580 s → 0.290 s (−50%)	202.42 s → 198.92 s (−1.7%)	independent A/B by @rich7420 (Modal, avg of 2 runs)

The post-decode path is ~25% of e2e on Cosmos/A100 but only ~0.3% on Wan/H100, so the same per-path halving lands very differently on e2e. Net-positive everywhere; magnitude is configuration-dependent. VAE and denoising stages are byte-untouched in both runs — the change is correctly scoped to the pixel path. Stacks with enable_torch_compile (independently ~−14%; different pipeline stages).

Root cause

samples is a pinned-CPU buffer; samples.copy_(output_batch.output) is non-blocking, and gen_time is logged immediately after it. The full fp32 video D→H (scales with resolution × frames × batch and can be large) therefore completes lazily and the first op in the post-decode block blocks on it — so a large transfer is attributed to PostDecodeFrameProcessStage. On top of that, the old code ran (x * 255).to(uint8) in a Python loop over frames on the CPU tensor.

Evidence:

• nsys (whole process tree, --trace-fork-before-exec): fragmented CPU tail + transfer, GPU idle during the stage.
• Synthetic microbench of the exact block: pure CPU compute is only ~0.46 s (T=29) / ~1.8 s (T=125) — far below the ~6 s observed on Cosmos/A100, i.e. the stage isn't compute-bound.
• Real-conditions microbench (pinned dest, non-blocking copy, no pre-sync): reproduces the deferred-transfer tail; the on-device-uint8 path is 2.056 s → 0.034 s (T=29) and 2.304 s → 0.132 s (T=125).
• Real pipeline + independent cross-model A/B: table above.

Correctness

Not bit-identical to the previous CPU cast: float → uint8 can differ by ≤1 LSB between CPU and GPU on boundary pixels, and clamp_(0, 255) additionally fixes a latent overflow (VAE output slightly outside [0, 1] previously wrapped mod 256 in the unclamped cast). Gated by SSIM rather than exact equality:

MS-SSIM (repo compute_video_ssim_torchvision, Cosmos 2.5, A100): 29/29 frames = 1.000000 exactly (verified all frames read and compared, mean = min = max = 1.0).

Skip the dead fp32 samples copy when return_frames=False (commit 2)

After commit 1, the post-decode frames path reads output_batch.output directly via the GPU vid_u8 cast, so samples is consumed in exactly one place — the result dict's samples field, gated on batch.return_frames ("samples": samples if batch.return_frames else None). When the caller doesn't ask for samples, the pinned ~50 MB fp32 alloc and its D→H copy (and the latent fallthrough .cpu()) are dead weight; the typical generate flow (save_video=True, return_frames=False) hits this on every call.

Commit 2 extends skip_pixel_prealloc to include not batch.return_frames and short-circuits the copy on the same condition — removing the residual ~2 s the PR text originally called out as a follow-up. Behavior is unchanged for callers with return_frames=True (they still receive the fp32 samples tensor) or for the latent fallthrough when return_frames=True.

Further device-residency improvement — running make_grid/permute on-device with a single stacked transfer — is still a candidate follow-up, measured and SSIM-gated separately.

Scope

Pixel-video path only. output_type == "latent" and audio-only paths are byte-unchanged. Callers passing return_frames=True still receive the fp32 samples tensor as before; commit 2 only skips work when nothing reads it.

Test plan

• nsys + microbench root-cause attribution
• Real e2e on A100 (Cosmos 2.5 2B): PostDecode −66%, e2e −15.6%
• Independent cross-model A/B (Wan2.1-T2V-1.3B, H100): PostDecode −50%, e2e −1.7%, VAE/denoising untouched
• MS-SSIM gate (commit 1): 29/29 frames = 1.000000
• SSIM re-gate after commit 2 (Wan2.1-T2V-1.3B, RTX A5000, 480×832×29f / 20 steps, FlashAttention-2, return_frames=False): BASE (copy runs) vs CAND (copy skipped) = MS-SSIM mean/min/max = 1.000000 across all frames — bit-identical, confirming the skip is output-neutral
• CI (buildkite + fastcheck passing)
• Maintainer review

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR optimizes decoded-frame postprocessing by moving quantization/casting to uint8 onto the source device (typically CUDA) before transferring frames to CPU, while also clamping values to avoid wraparound artifacts.

Changes:

• Quantize/clamp video tensor on-device before CPU transfer to reduce transfer size and CPU work.
• Replace per-frame CPU scaling/casting loop with a list-comprehension over a pre-quantized tensor.
• Add explanatory comments about the performance rationale and overflow behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gemini-code-assist]

Code Review

This pull request optimizes video post-processing by performing quantization and clamping on the GPU before transferring data to the CPU, which reduces the transfer size and fixes a potential overflow bug. Feedback suggests further improving performance by moving the grid layout and permutation operations to the GPU and using a single stacked transfer to minimize CPU bottlenecks.

[gemini-code-assist]

While quantizing on-device significantly reduces the transfer size, the subsequent make_grid and permute operations are still being performed on the CPU within a single-threaded loop. Since the tensors are already on the device (typically CUDA), moving these operations to the GPU and then performing a single stacked transfer can further reduce the bottleneck in the PostDecodeFrameProcessStage.

Suggested change

	vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu()
	frames = [
	torchvision.utils.make_grid(x, nrow=6).permute(
	1, 2, 0).squeeze(-1).contiguous().numpy()
	for x in vid_u8
	]
	vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w")
	# Perform grid layout and permutation on device to avoid CPU bottlenecks.
	# Stacking before .cpu() ensures a single efficient D->H transfer.
	grids = torch.stack([
	torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1)
	for x in vid_u8
	])
	frames = [f for f in grids.cpu().numpy()]

[rich7420]

@Mister-Raggs Thanks for the patch!

A/B results — Wan2.1-T2V-1.3B, H100×1, 720×1280×77f, 30 steps
Same-container A/B on Modal (avg of 2 runs each):

Stage	Baseline	PR	Δ
PostDecodeFrameProcessStage	0.580 s	0.290 s	−50%
VideoSave	0.820 s	0.701 s	−14.5%
Decoding (VAE)	7.325 s	7.307 s	−0.2%
Denoising	192.701 s	189.658 s	−1.6%
End-to-end	202.42 s	198.92 s	−1.73%

Per-run PostDecode very stable: baseline 0.568 / 0.592 s, treatment 0.296 / 0.285 s.
The mechanism is real and confirmed. PostDecode halved exactly as described — the on-GPU uint8 cast + smaller D→H transfer is working as intended, and VAE/Denoising are untouched.
On Wan T2V, PostDecode was already only ~0.6 s (vs ~6 s on Cosmos 2.5 / A100 where the fix was measured). The savings scale with how much time was being lost in that path — on Wan T2V it was never the bottleneck, so e2e impact lands at −1.7% instead of −15%.
Still a clean net-positive fix. The impact just varies a lot by model/resolution/hardware depending on how large the post-decode buffer is and whether that D→H was already the hotspot.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease