From 3e26db40b0d66c1ad03c30bcf8e98fb8efba3f4e Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Thu, 17 Sep 2026 20:01:01 +0000 Subject: [PATCH] [perf] H3 decode: pin the frozen VAEs' host copies With the device-to-host copy gone, the remaining host-to-device copy is the stage's cost, and from pageable memory it is staged through a bounce buffer at roughly 2.6 GB/s instead of full PCIe speed. The host copies are made once, so pinning them costs nothing per request. It does cost the module's size in non-pageable host memory (10.4 GB for the H3 video VAE, 0.6 GB for the audio VAE, per rank) for as long as the pipeline is alive. It follows --pin-cpu-memory, which is on by default; --no-pin-cpu- memory keeps the copies pageable and only loses the bandwidth. Note that this flag did not already reach the VAE: its other use is FSDP's own CPUOffloadPolicy and the VAE does not go through FSDP. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01ATw7g6cxq9N2LtnMratsru --- fastvideo/models/pinned_offload.py | 20 ++++++++++++++----- .../minimax_h3/stages/minimax_h3_decoding.py | 4 ++-- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/fastvideo/models/pinned_offload.py b/fastvideo/models/pinned_offload.py index 480c3cf09..e48c3cb1a 100644 --- a/fastvideo/models/pinned_offload.py +++ b/fastvideo/models/pinned_offload.py @@ -14,6 +14,12 @@ and no device-to-host copy at all, with the module living on the host between requests exactly as before. Numerics are untouched -- the device tensors are byte copies of the same weights -- and peak host memory can only go down, because one buffer is reused instead of a new one being allocated per cycle. + +The host copies are pinned when ``pin`` is set (``--pin-cpu-memory``, on by +default), which is what lets the remaining host-to-device copy run at full +PCIe speed instead of being staged through a bounce buffer. That trades the +module's size in non-pageable host memory for the bandwidth; ``pin=False`` +keeps the copies pageable and everything else the same. """ from __future__ import annotations @@ -21,6 +27,7 @@ import torch from torch import nn from fastvideo.logger import init_logger +from fastvideo.utils import is_pin_memory_available logger = init_logger(__name__) @@ -33,31 +40,34 @@ def _tensors(module: nn.Module): yield "buffer:" + name, buf -def _host_copies(module: nn.Module) -> dict[str, torch.Tensor]: +def _host_copies(module: nn.Module, pin: bool) -> dict[str, torch.Tensor]: host = getattr(module, _HOST_ATTR, None) if host is not None: return host host, total = {}, 0 for name, t in _tensors(module): src = t.data if t.data.device.type == "cpu" else t.data.to("cpu") + if pin and not src.is_pinned(): + src = src.pin_memory() host[name] = src total += src.numel() * src.element_size() setattr(module, _HOST_ATTR, host) - logger.info("frozen offload: holding %.2f GB of %s on the host", total / 1e9, type(module).__name__) + logger.info("frozen offload: holding %.2f GB of %s on the host (%s)", total / 1e9, type(module).__name__, + "pinned" if pin else "pageable") return host -def load(module: nn.Module, device: torch.device) -> nn.Module: +def load(module: nn.Module, device: torch.device, pin: bool = True) -> nn.Module: """Put a frozen ``module`` on ``device``, copying from its host copies.""" if device.type != "cuda": return module.to(device) - host = _host_copies(module) + host = _host_copies(module, pin and is_pin_memory_available()) for name, t in _tensors(module): if t.data.device == device: continue src = host[name] dst = torch.empty_like(src, device=device) - dst.copy_(src) + dst.copy_(src, non_blocking=src.is_pinned()) t.data = dst return module diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 411e92a66..a12cca50d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -124,7 +124,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): if self.vae is None: raise RuntimeError("MiniMax-H3 full VAE decode requires a loaded video VAE.") - pinned_offload.load(self.vae, device) + pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory) try: latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": @@ -201,7 +201,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage): layout.num_audio_latents, ) device = get_local_torch_device() - pinned_offload.load(self.audio_vae, device) + pinned_offload.load(self.audio_vae, device, pin=fastvideo_args.pin_cpu_memory) try: latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent":