From 0af2e9e8ef2c7db48448738ebb0fbb8dcf5c0159 Mon Sep 17 00:00:00 2001 From: Raghav Date: Sat, 16 May 2026 10:29:28 -0700 Subject: [PATCH 1/6] [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: From 0399713e7b397b4d48f29130829c6d9ce5d6a0b5 Mon Sep 17 00:00:00 2001 From: Raghav Date: Sat, 16 May 2026 10:57:06 -0700 Subject: [PATCH 2/6] [perf]: clarify samples/output_batch.output equivalence; drop hardcoded transfer size Address review comments on #1362: - Document that `samples` is just the pinned-CPU mirror of `output_batch.output` with no intervening preprocessing, so sourcing from `output_batch.output` is the same data (Copilot). - Replace the hardcoded ~0.3-1.4 GB estimate with a description of the scaling relationship; it varies with resolution/frames/batch/dtype and would rot (Copilot). - Note the SSIM-gated (not bit-exact) equivalence inline. No behavior change. --- fastvideo/entrypoints/video_generator.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 255490434..57e533473 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -816,14 +816,20 @@ class VideoGenerator: frames = None else: # 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. + # device->host copy. `samples` above is just the pinned-CPU + # mirror of `output_batch.output` (`samples.copy_(output)` or + # `output.cpu()`) with no intervening preprocessing, so reading + # `output_batch.output` here is the same data — but its + # non-blocking D->H had not completed, so the first op that + # touched `samples` blocked on the full fp32 video transfer + # (which scales with resolution x frames x batch and can be + # large) and then 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. (Equivalence is SSIM-gated, not bit-exact: float->uint8 + # differs <=1 LSB CPU vs GPU.) 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() From 3c3da4d057f8677efa40d165e5236b8c326cd66e Mon Sep 17 00:00:00 2001 From: Raghav Date: Tue, 19 May 2026 17:32:49 -0700 Subject: [PATCH 3/6] [perf]: skip the fp32 samples D->H copy when return_frames=False MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit rebuilt the post-decode frames path to read `output_batch.output` directly via the GPU `vid_u8` cast, so `samples` is now consumed in exactly one place — the result dict's `samples` field, gated on `batch.return_frames`. 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. Extend `skip_pixel_prealloc` to include `not return_frames` so the pinned buffer is allocated only when needed, and short-circuit the copy/`.cpu()` on the same condition. No effect when `return_frames=True` or for latent callers that read `samples`; correctness is unchanged (SSIM-gated, same as the parent commit). Removes the residual ~2 s the PR text already calls out for the "only saving to disk" case. Co-Authored-By: Claude Opus 4.7 --- fastvideo/entrypoints/video_generator.py | 32 ++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 57e533473..e4f8bb46c 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -750,14 +750,20 @@ class VideoGenerator: latent_batch_size = _infer_latent_batch_size(batch) is_latent_output = fastvideo_args.output_type == "latent" needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output) - needs_samples_buffer = batch.return_frames or needs_frame_output - # When ``output_type == "latent"`` the forward output has latent - # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``) - # rather than the pre-allocation's pixel shape. Skip the pinned - # ~50 MB buffer entirely. Also skip it for metadata-only calls; - # neither the result nor save path will consume the decoded tensor. + # ``samples`` is consumed in exactly one place — the result dict + # (``"samples": samples if batch.return_frames else None``). + # Post-decode frame building reads ``output_batch.output`` + # directly (the GPU ``vid_u8`` path), not ``samples``. So when + # ``return_frames=False`` the pinned ~50 MB fp32 alloc + D->H + # copy are dead weight — the typical generate flow + # (``save_video=True``, ``return_frames=False``) hits this on + # every call. + # ``output_type == "latent"`` keeps its existing branch (shape + # mismatch falls through to ``.cpu()`` below) for callers that + # *do* ask for the latent samples via ``return_frames=True``. # ``skip_pixel_prealloc`` also gates the slow-path warning. - skip_pixel_prealloc = is_latent_output or not needs_samples_buffer + needs_samples_out = batch.return_frames + skip_pixel_prealloc = is_latent_output or not needs_samples_out if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: @@ -777,9 +783,11 @@ class VideoGenerator: "This usually means the executor/pipeline failed earlier.") audio_only = bool(output_batch.extra.get("audio_only")) - if not needs_samples_buffer or (audio_only and not batch.return_frames): - # Metadata-only/audio-only request: keep the empty placeholder and - # avoid the decoded tensor D->H copy. + if not needs_samples_out: + # Nothing downstream reads ``samples`` (the result dict + # returns None when ``return_frames=False``); keep the empty + # placeholder allocated above and skip the fp32 D->H copy + # entirely. pass elif audio_only: # Audio-only return-frames requests expose the small placeholder @@ -834,9 +842,7 @@ class VideoGenerator: 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 + 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) From 15a164a052b11707de7136457b80576cd8ccb40d Mon Sep 17 00:00:00 2001 From: Will Lin Date: Fri, 21 Aug 2026 02:34:16 +0000 Subject: [PATCH 4/6] [bugfix]: derive result size from the decoded output when the samples mirror is skipped PR #1362 commit 3 leaves `samples` as an empty placeholder when `return_frames=False`, but main picked up #1595's refiner size reporting in the meantime, and `_resolve_output_size(samples, ...)` silently fell back to the requested geometry in exactly the common save flow the PR optimizes. Read the geometry from `output_batch.output` instead (shape-only access, no D->H copy), gated on `needs_frame_output` so metadata-only and audio-only calls still never inspect the (possibly dropped) worker output. Co-Authored-By: Claude Fable 5 --- fastvideo/entrypoints/video_generator.py | 7 ++++- .../tests/entrypoints/test_video_generator.py | 31 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 04c704644..8b3e8d7d3 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -877,8 +877,13 @@ class VideoGenerator: # `GenerationResult.size` describes the produced media, not only the # base-stage request. Refiner pipelines can change the final pixel # dimensions, so derive this result metadata from the decoded output. + # Read the geometry from `output_batch.output` (a shape-only access, + # no D->H copy): when `return_frames=False` the `samples` mirror + # stays an empty placeholder and no longer carries the decoded + # shape. Metadata-only calls keep the request fallback and never + # inspect the (possibly dropped) worker output. output_size = _resolve_output_size( - samples, + output_batch.output if needs_frame_output else samples, (target_height, target_width, batch.num_frames), pixel_output=not is_latent_output and not audio_only, ) diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index d123f60a3..7ca9ba713 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -305,6 +305,37 @@ def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_p } +def test_generate_single_video_save_only_reports_refined_output_size(monkeypatch, tmp_path): + """`GenerationResult.size` must describe the decoded media even when the + fp32 `samples` mirror is skipped (`return_frames=False`, the CLI save + flow). Refiner pipelines can change the final pixel geometry, so the size + has to come from `output_batch.output`, not the base request. CPU-only.""" + # Refiner-style output: request asks for 2 frames of 16x16, pipeline + # produces 5 frames of 32x48. + output = torch.full((1, 3, 5, 32, 48), 0.5, dtype=torch.float32) + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + saved = {} + + def fake_mimsave(path, frames, *, fps, format): + saved["frame_count"] = len(frames) + + monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave) + + result = generator._generate_single_video( + prompt="refined save", + sampling_param=_small_sampling_param(save_video=True, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "refined.mp4"), + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["size"] == (32, 48, 5) + assert saved["frame_count"] == 5 + + def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path): audio = torch.zeros((16, ), dtype=torch.float32) output_batch = _single_video_output_batch( From f56f56704270e904c5981937b596732428b6aec9 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Fri, 21 Aug 2026 02:34:43 +0000 Subject: [PATCH 5/6] [misc]: correct the post-decode rationale comments The pre-#1362 `samples.copy_(output_batch.output)` never passed `non_blocking=True` (git log -S confirms), so drop the deferred non-blocking-transfer claim and describe the measured costs instead: a full fp32 D->H copy plus a single-threaded per-frame CPU loop. Also drop the reintroduced hardcoded "~50 MB" size estimate (same class of comment Copilot flagged and commit 0399713e7 removed elsewhere) and scope the "typical flow" claim to the CLI, since the SamplingParam API default is return_frames=True. Co-Authored-By: Claude Fable 5 --- fastvideo/entrypoints/video_generator.py | 30 +++++++++++------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 8b3e8d7d3..f81f338d2 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -808,14 +808,13 @@ class VideoGenerator: latent_batch_size = _infer_latent_batch_size(batch) is_latent_output = fastvideo_args.output_type == "latent" needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output) - # ``samples`` is consumed in exactly one place — the result dict - # (``"samples": samples if batch.return_frames else None``). + # A populated ``samples`` has exactly one consumer — the result + # dict (``"samples": samples if batch.return_frames else None``). # Post-decode frame building reads ``output_batch.output`` # directly (the GPU ``vid_u8`` path), not ``samples``. So when - # ``return_frames=False`` the pinned ~50 MB fp32 alloc + D->H - # copy are dead weight — the typical generate flow - # (``save_video=True``, ``return_frames=False``) hits this on - # every call. + # ``return_frames=False`` the pinned fp32 alloc + D->H copy are + # dead weight — the CLI generate flow (``save_video=True``, + # ``return_frames=False``) hits this on every call. # ``output_type == "latent"`` keeps its existing branch (shape # mismatch falls through to ``.cpu()`` below) for callers that # *do* ask for the latent samples via ``return_frames=True``. @@ -899,16 +898,15 @@ class VideoGenerator: # device->host copy. `samples` above is just the pinned-CPU # mirror of `output_batch.output` (`samples.copy_(output)` or # `output.cpu()`) with no intervening preprocessing, so reading - # `output_batch.output` here is the same data — but its - # non-blocking D->H had not completed, so the first op that - # touched `samples` blocked on the full fp32 video transfer - # (which scales with resolution x frames x batch and can be - # large) and then 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. (Equivalence is SSIM-gated, not bit-exact: float->uint8 + # `output_batch.output` here is the same data. The old path + # paid a full fp32 video D->H copy (which scales with + # resolution x frames x batch) and then a single-threaded + # per-frame CPU *255/cast loop. Casting to uint8 on-device + # makes the transfer 4x smaller, ships it in a single copy, + # 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. + # (Equivalence is SSIM-gated, not bit-exact: float->uint8 # differs <=1 LSB CPU vs GPU.) src = output_batch.output vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8) From aa95a4c18ea9c6847fb5e988d0de59970c6a5b03 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Fri, 21 Aug 2026 02:39:09 +0000 Subject: [PATCH 6/6] [misc]: add CPU regression tests for the on-device uint8 frame path Two tests certifying #1362's frame semantics without a GPU: - frames_match_legacy_cpu_loop: the quantize-then-grid path reproduces the legacy make_grid -> *255 -> uint8 per-frame loop bit-exactly for in-range fp32 pixels (batch>1 nrow=6 grid layout, odd frame count, uint8 HWC contract). Verified to pass against main's legacy loop too, so it pins both sides of the equivalence. - frames_clamp_out_of_range_pixels: out-of-[0,1] VAE output saturates at 0/255 instead of wrapping mod 256; fails on the pre-#1362 loop. Co-Authored-By: Claude Fable 5 --- .../tests/entrypoints/test_video_generator.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 7ca9ba713..0e7b940fa 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -2,8 +2,10 @@ import os from types import SimpleNamespace import warnings +import numpy as np import pytest import torch +from einops import rearrange import fastvideo.entrypoints.video_generator as video_generator_module from fastvideo.api import ( @@ -271,6 +273,70 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path) assert result["video_path"] is None +def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path): + """The on-device quantize path (#1362) must reproduce the legacy + per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly + for in-range fp32 pixels: same uint8 dtype, same HWC grid layout with + nrow=6 (batch>1), odd frame count. CPU-only: on CUDA the float->uint8 + cast may differ by <=1 LSB, but on CPU both orderings run identical + fp32 ops, so exact equality is required.""" + torch.manual_seed(0) + output = torch.rand((2, 3, 3, 16, 16), dtype=torch.float32) + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + sampling_param = _small_sampling_param(save_video=False, return_frames=True) + sampling_param.num_frames = 3 + sampling_param.num_videos_per_prompt = 2 + + result = generator._generate_single_video( + prompt="grid parity", + sampling_param=sampling_param, + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + legacy_frames = [] + for x in rearrange(output, "b c t h w -> t b c h w"): + grid = video_generator_module.torchvision.utils.make_grid(x, nrow=6) + grid = grid.permute(1, 2, 0).squeeze(-1) + legacy_frames.append((grid * 255).to(torch.uint8).contiguous().cpu().numpy()) + + torch.testing.assert_close(result["samples"], output) + assert len(result["frames"]) == 3 + for got, want in zip(result["frames"], legacy_frames, strict=True): + assert got.dtype == np.uint8 + assert got.shape == want.shape + np.testing.assert_array_equal(got, want) + + +def test_generate_single_video_frames_clamp_out_of_range_pixels(tmp_path): + """VAE output slightly outside [0, 1] must saturate at 0/255 in the + uint8 frames. The pre-#1362 unclamped cast wrapped mod 256 (e.g. + 1.5 -> 126). CPU-only.""" + output = torch.full((1, 3, 2, 16, 16), 1.5, dtype=torch.float32) + output[:, :, 1] = -0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="clamp", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + frames = result["frames"] + assert len(frames) == 2 + # make_grid passes a single image through without grid padding, so + # every pixel comes from the (clamped) output tensor. + assert frames[0].dtype == np.uint8 + assert frames[0].shape == (16, 16, 3) + assert (frames[0] == 255).all() + assert (frames[1] == 0).all() + + def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path): output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 output_batch = _single_video_output_batch(output)