[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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ATw7g6cxq9N2LtnMratsru

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
 

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":