[bugfix]: Fall back when large pinned output allocation fails (#1759)

fastvideo/entrypoints/video_generator.py

@@ -55,7 +55,7 @@ from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
 from fastvideo.logger import init_logger
 from fastvideo.pipelines import ForwardBatch
-from fastvideo.utils import align_to, shallow_asdict
+from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback, align_to, shallow_asdict
 from fastvideo.worker.executor import Executor
 
 fcntl: types.ModuleType | None
@@ -819,9 +819,8 @@ class VideoGenerator:
         if skip_pixel_prealloc:
             samples = torch.empty(0, device='cpu')
         else:
-            samples = torch.empty(
+            samples = allocate_cpu_tensor_with_pin_fallback(
                 (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
-                device='cpu',
                 pin_memory=fastvideo_args.pin_cpu_memory)
         thread.join()
 

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py

@@ -26,7 +26,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
-from fastvideo.utils import is_pin_memory_available
+from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback
 
 logger = init_logger(__name__)
 
@@ -135,11 +135,10 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
             output = None
             if is_output_rank:
-                output = torch.empty(
+                output = allocate_cpu_tensor_with_pin_fallback(
                     self.vae.decoded_pixel_shape(latents.shape),
-                    device="cpu",
                     dtype=torch.float32,
-                    pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
+                    pin_memory=fastvideo_args.pin_cpu_memory,
                 )
             # Attribute the streamed decoder computation while retaining
             # per-chunk device-to-host transfer and pinned-buffer reuse.

fastvideo/tests/entrypoints/test_video_generator.py

@@ -8,6 +8,7 @@ import torch
 from einops import rearrange
 
 import fastvideo.entrypoints.video_generator as video_generator_module
+import fastvideo.utils as fastvideo_utils
 from fastvideo.api import (
     ConfigValidationError,
     GenerationRequest,
@@ -273,6 +274,39 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path)
     assert result["video_path"] is None
 
 
+def test_generate_single_video_retries_pageable_output_when_pinned_allocation_fails(monkeypatch, tmp_path):
+    output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5
+    output_batch = _single_video_output_batch(output)
+    fastvideo_args = _single_video_args()
+    fastvideo_args.pin_cpu_memory = True
+    generator = _single_video_generator(output_batch, fastvideo_args)
+    real_empty = torch.empty
+    allocation_attempts = []
+    warnings = []
+
+    def fake_empty(size, **kwargs):
+        allocation_attempts.append(kwargs.get("pin_memory", False))
+        if kwargs.get("pin_memory"):
+            raise RuntimeError("CUDA error: out of memory")
+        return real_empty(size, **kwargs)
+
+    monkeypatch.setattr(fastvideo_utils, "is_pin_memory_available", lambda: True)
+    monkeypatch.setattr(fastvideo_utils.torch, "empty", fake_empty)
+    monkeypatch.setattr(fastvideo_utils.logger, "warning", lambda message, *args: warnings.append(message % args))
+
+    result = generator._generate_single_video(
+        prompt="pinned allocation fallback",
+        sampling_param=_small_sampling_param(save_video=False, return_frames=True),
+        fastvideo_args=fastvideo_args,
+        output_path=str(tmp_path / "unused.mp4"),
+    )
+
+    torch.testing.assert_close(result["samples"], output)
+    assert allocation_attempts[:2] == [True, False]
+    assert len(warnings) == 1
+    assert "retrying with pageable memory" in warnings[0]
+
+
 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

fastvideo/tests/stages/test_minimax_h3_vae_streaming.py

@@ -114,6 +114,49 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
     assert torch.all(result.output == 0.25)
 
 
+def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
+    latent_shape = (1, 4, 2, 4, 4)
+    latents = torch.randn(latent_shape)
+    rows = patchify_video_latents(latents, (1, 1, 1))
+    batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
+    batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape)
+    observed = {}
+
+    class VAE:
+
+        def to(self, device):
+            return self
+
+        def denormalize_latents(self, decoded_latents):
+            return decoded_latents
+
+        def decoded_pixel_shape(self, shape):
+            return (1, 3, 5, 16, 16)
+
+        def decode_to_pixels(self, decoded_latents, output):
+            output.zero_()
+
+    def fake_allocate(size, *, dtype=None, pin_memory=False):
+        observed["size"] = size
+        observed["dtype"] = dtype
+        observed["pin_memory"] = pin_memory
+        return torch.empty(size, dtype=dtype)
+
+    monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(minimax_h3_decoding, "allocate_cpu_tensor_with_pin_fallback", fake_allocate)
+
+    MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(
+        batch,
+        SimpleNamespace(output_type="pil", pin_cpu_memory=True, vae_cpu_offload=False, vae_parallel_decode=False),
+    )
+
+    assert observed == {
+        "size": (1, 3, 5, 16, 16),
+        "dtype": torch.float32,
+        "pin_memory": True,
+    }
+
+
 def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
     class VAE:
 

fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py

@@ -158,6 +158,22 @@ def test_decode_to_pixels_pinned_buffer_matches_dense_on_cuda() -> None:
     assert_close(pinned, expected, atol=0.0, rtol=0.0)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="pageable-buffer streaming requires CUDA")
+@torch.inference_mode()
+def test_decode_to_pixels_pageable_buffer_matches_dense_on_cuda() -> None:
+    """The allocation fallback's blocking chunk copies must equal dense decode."""
+    torch.manual_seed(20260814)
+    vae = _tiny_vae().to("cuda")
+    latents = torch.randn(1, 4, 12, 4, 4, device="cuda")
+    expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1).cpu()
+
+    pageable = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, pageable)
+
+    assert not pageable.is_pinned()
+    assert_close(pageable, expected, atol=0.0, rtol=0.0)
+
+
 def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None:
     vae = _tiny_vae()
     latents = torch.randn(1, 4, 2, 4, 4)

fastvideo/utils.py

@@ -1306,3 +1306,35 @@ def _cached_pin_memory_available(pid: int) -> bool:
 
 def is_pin_memory_available() -> bool:
     return _cached_pin_memory_available(os.getpid())
+
+
+def allocate_cpu_tensor_with_pin_fallback(
+    size: tuple[int, ...] | torch.Size,
+    *,
+    dtype: torch.dtype | None = None,
+    pin_memory: bool = False,
+) -> torch.Tensor:
+    """Allocate a CPU tensor, retrying pageable memory if pinning fails.
+
+    ``is_pin_memory_available`` intentionally uses a small capability probe.
+    Large decoded-video buffers can still exhaust the CUDA host allocator after
+    that probe succeeds, especially on unified-memory systems. Treat pinning as
+    a performance preference: preserve the requested shape and dtype by
+    retrying the allocation without pinning.
+    """
+    kwargs: dict[str, object] = {"device": "cpu"}
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+
+    if not pin_memory or not is_pin_memory_available():
+        return torch.empty(size, **kwargs)
+
+    try:
+        return torch.empty(size, pin_memory=True, **kwargs)
+    except RuntimeError as exc:
+        logger.warning(
+            "Pinned CPU allocation failed for shape %s; retrying with pageable memory: %s",
+            tuple(size),
+            exc,
+        )
+        return torch.empty(size, **kwargs)