[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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user