[perf]: quantize frames to uint8 on-device before the D->H copy

PostDecodeFrameProcessStage was charged ~6s/run (~25% of e2e on Cosmos
2.5, scaling with frames/resolution). Profiling (nsys + microbench)
showed the stage's own compute is only ~0.5-1.9s; the rest is the
non-blocking pinned-CPU samples.copy_(output) D->H of the full fp32
video (~0.3-1.4 GB) completing lazily and blocking the first
postprocess op, plus a single-threaded per-frame CPU *255/cast loop.

Cast to uint8 on the source device (typically CUDA) before the copy:
the transfer becomes 4x smaller (fp32 -> uint8) and the elementwise
work runs on the GPU. Microbench: 2.056s -> 0.034s (T=29), 2.304s ->
0.132s (T=125), 17-60x on the measurable cost.

clamp_(0, 255) additionally fixes a latent overflow: VAE output
slightly outside [0, 1] previously wrapped mod 256 in the unclamped
(x * 255).to(uint8) cast.

Output is not bit-identical to the old CPU cast (float->uint8 differs
<=1 LSB between CPU and GPU on boundary pixels); gate via SSIM rather
than exact equality. Scoped to the pixel-video path only; latent,
audio-only, and return-samples paths are unchanged.
This commit is contained in:
Raghav
2026-07-16 02:59:59 -07:00
parent 1c04ace573
commit 0af2e9e8ef
+17 -7
View File
@@ -815,13 +815,23 @@ class VideoGenerator:
elif not needs_frame_output:
frames = None
else:
videos = rearrange(samples, "b c t h w -> t b c h w")
frames = []
for x in videos:
x = torchvision.utils.make_grid(x, nrow=6)
x = x.permute(1, 2, 0).squeeze(-1)
x = (x * 255).to(torch.uint8)
frames.append(x.contiguous().cpu().numpy())
# Quantize on the source device (typically CUDA) BEFORE the
# device->host copy. The previous path read from the pinned-CPU
# `samples` buffer whose non-blocking D->H had not completed, so
# the first op here blocked on the full fp32 video transfer
# (~0.3-1.4 GB) and ran a single-threaded per-frame CPU *255/cast
# loop. Casting to uint8 on-device first makes the transfer 4x
# smaller and moves the elementwise work onto the GPU.
# clamp_() also fixes a latent overflow bug: VAE output slightly
# outside [0, 1] wrapped mod 256 in the old unclamped cast.
src = output_batch.output
vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8)
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
]
postprocess_time = time.perf_counter() - postprocess_start
logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
if logging_info is not None: