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.
Raghav
0af2e9e8ef