From 0af2e9e8ef2c7db48448738ebb0fbb8dcf5c0159 Mon Sep 17 00:00:00 2001 From: Raghav Date: Sat, 16 May 2026 10:29:28 -0700 Subject: [PATCH] [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. --- fastvideo/entrypoints/video_generator.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 40203dd00..255490434 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -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: