[perf] MiniMax H3: quantize decoded pixels to uint8 on the GPU by default

fastvideo/entrypoints/video_generator.py

@@ -1041,10 +1041,16 @@ class VideoGenerator:
             # `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. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback.
-            vid_u8 = _quantize_video_frames_to_uint8(output_batch.output)
-            frames = [
-                torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8
-            ]
+            src_frames = output_batch.output
+            if (isinstance(src_frames, torch.Tensor) and src_frames.dim() == 5 and src_frames.dtype == torch.uint8
+                    and src_frames.shape[0] == 1 and src_frames.shape[2] > 1):
+                frames = list(src_frames[0].permute(1, 2, 3, 0).contiguous().cpu().numpy())
+            else:
+                vid_u8 = _quantize_video_frames_to_uint8(src_frames)
+                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:

fastvideo/envs.py

@@ -753,6 +753,11 @@ FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW = EnvBool(False,
                                                category="performance",
                                                doc="Use transposed weight views in the MiniMax-H3 VAE INT8 "
                                                "projections.")
+FASTVIDEO_H3_VAE_UINT8_PIXELS = EnvBool(True,
+                                        category="performance",
+                                        doc="Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per "
+                                        "chunk and stream uint8 into the host buffer; frames are bitwise identical "
+                                        "to the fp32 path. 0 streams fp32.")
 FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT = EnvBool(False,
                                               category="performance",
                                               doc="Fused dequantization epilogue for the MiniMax-H3 VAE INT8 "

fastvideo/models/vaes/minimax_h3_video.py

@@ -1016,7 +1016,7 @@ class AutoencoderKLMiniMaxH3(nn.Module):
 
     @staticmethod
     def _copy_chunk_pixels(pixels: torch.Tensor, output: torch.Tensor, frame_start: int, non_blocking: bool) -> None:
-        """Copy one finalized fp32 pixel chunk into the CPU ``output`` buffer.
+        """Copy one finalized fp32 pixel chunk into the CPU ``output`` buffer (fp32, or uint8 quantized on device).
 
         Device-to-host copies run per (batch, channel) plane: the temporal
         slice of ``output`` is strided across channels, but each plane is
@@ -1027,6 +1027,9 @@ class AutoencoderKLMiniMaxH3(nn.Module):
         before releasing the buffer.
         """
         target = output[:, :, frame_start:frame_start + pixels.shape[2]]
+        if output.dtype == torch.uint8:
+            # Same mul/clamp/truncating cast as the stage's fp32 host path, so frames are bitwise equal.
+            pixels = (pixels * 255).clamp_(0, 255).to(torch.uint8)
         if pixels.device.type == "cuda":
             pixels = pixels.contiguous()
             for batch_index in range(pixels.shape[0]):
@@ -1145,11 +1148,17 @@ class AutoencoderKLMiniMaxH3(nn.Module):
         return DecoderOutput(sample=decoded)
 
     def decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
-        """Stream decoded ``[0, 1]`` FP32 pixels into a caller-owned CPU buffer."""
+        """Stream decoded ``[0, 1]`` pixels into a caller-owned CPU buffer.
+
+        ``output`` is float32, or uint8 when the caller asks for the quantized
+        stream (``FASTVIDEO_H3_VAE_UINT8_PIXELS``); uint8 quantizes each chunk on
+        the GPU with the fp32 host path's exact mul/clamp/truncate arithmetic.
+        """
         expected_shape = self.decoded_pixel_shape(z.shape)
-        if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
+        if (output.device.type != "cpu" or output.dtype not in (torch.float32, torch.uint8)
+                or tuple(output.shape) != expected_shape):
             raise ValueError(
-                "`output` must be a CPU float32 tensor with shape "
+                "`output` must be a CPU float32 or uint8 tensor with shape "
                 f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
         try:
             if self.use_slicing and z.shape[0] > 1:

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

@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+import fastvideo.envs as envs
 from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -68,11 +69,16 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
     def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None:
         super().__init__()
         self.vae = vae
-        self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None
+        self._output_buffer: tuple[tuple[tuple[int, ...], bool, torch.dtype], torch.Tensor] | None = None
         self._unified_memory: bool | None = None
 
-    def _pixel_buffer(self, shape: tuple[int, ...], pin_memory: bool, device: torch.device) -> torch.Tensor:
-        """The CPU FP32 pixel buffer, reused across requests of the same geometry on unified memory.
+    def _pixel_buffer(self,
+                      shape: tuple[int, ...],
+                      pin_memory: bool,
+                      device: torch.device,
+                      dtype: torch.dtype = torch.float32) -> torch.Tensor:
+        """The CPU pixel buffer (FP32, or uint8 when chunks are quantized on the GPU), reused across requests of
+        the same geometry on unified memory.
 
         The stage returns a uint8 copy, so the buffer never escapes. Reuse matters on
         unified-memory GPUs (DGX Spark): device-to-host copies into never-touched pageable
@@ -84,11 +90,11 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
         if self._unified_memory is None:
             self._unified_memory = current_platform.has_unified_memory(device.index or 0)
         if not self._unified_memory:
-            return allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory)
-        key = (tuple(shape), bool(pin_memory))
+            return allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory)
+        key = (tuple(shape), bool(pin_memory), dtype)
         if self._output_buffer is None or self._output_buffer[0] != key:
             self._output_buffer = None
-            buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory)
+            buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory)
             if not buffer.is_pinned():
                 buffer.zero_()
             self._output_buffer = (key, buffer)
@@ -170,8 +176,9 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
             output = None
             if is_output_rank:
+                pixel_dtype = torch.uint8 if envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.get() and not parallel else torch.float32
                 output = self._pixel_buffer(tuple(self.vae.decoded_pixel_shape(latents.shape)),
-                                            fastvideo_args.pin_cpu_memory, device)
+                                            fastvideo_args.pin_cpu_memory, device, pixel_dtype)
             # Attribute the streamed decoder computation while retaining
             # per-chunk device-to-host transfer and pinned-buffer reuse.
             # fp16 for the release decoders; NVFP4 decoders were trained under bf16.

fastvideo/tests/stages/test_minimax_h3_vae_streaming.py

@@ -75,6 +75,7 @@ def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
 
 
 def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0")
     latent_shape = (1, 4, 2, 4, 4)
     latents = torch.randn(latent_shape)
     rows = patchify_video_latents(latents, (1, 1, 1))
@@ -129,6 +130,7 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
 
 
 def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0")
     latent_shape = (1, 4, 2, 4, 4)
     latents = torch.randn(latent_shape)
     rows = patchify_video_latents(latents, (1, 1, 1))
@@ -177,10 +179,50 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
     }
 
 
+def test_decode_stage_streams_uint8_by_default(monkeypatch) -> None:
+    latent_shape = (1, 4, 2, 4, 4)
+    rows = patchify_video_latents(torch.randn(latent_shape), (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):
+            observed["dtype"] = output.dtype
+            output.fill_(63)
+
+    monkeypatch.delenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", raising=False)
+    monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
+    result = MiniMaxH3VideoDecodingStage(VAE()).forward(
+        batch,
+        SimpleNamespace(
+            output_type="pil",
+            pin_cpu_memory=False,
+            vae_cpu_offload=False,
+            vae_parallel_decode=False,
+            pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))),
+        ),
+    )
+
+    assert observed["dtype"] == torch.uint8
+    assert torch.equal(result.output, torch.full((1, 3, 5, 16, 16), 63, dtype=torch.uint8))
+
+
 @pytest.mark.parametrize("unified_memory", [True, False])
 def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory) -> None:
     """On unified memory, same-geometry requests reuse one pre-touched CPU buffer; returned
     frames never alias it. Discrete GPUs allocate per request and hold nothing between them."""
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0")
     latent_shape = (1, 4, 2, 4, 4)
     rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1))
     allocations = []

fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py

@@ -107,6 +107,22 @@ def test_decode_to_pixels_matches_decode(latent_frames: int) -> None:
     assert_close(actual, expected, atol=0.0, rtol=0.0)
 
 
+@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="needs CUDA"))])
+@torch.inference_mode()
+def test_decode_to_pixels_uint8_matches_fp32_quantization(device: str) -> None:
+    torch.manual_seed(20260810)
+    vae = _tiny_vae().to(device)
+    latents = torch.randn(1, 4, 12, 4, 4, device=device)
+    fp32 = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    streamed = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.uint8)
+
+    vae.decode_to_pixels(latents, fp32)
+    vae.decode_to_pixels(latents, streamed)
+
+    assert torch.equal(streamed, fp32.mul_(255).clamp_(0, 255).to(torch.uint8))
+
+
 # 3: one chunk with pad tokens; 6: pad hitting the intra-clip tail;
 # 12: two blended chunks without padding; 13: three chunks plus pad trim.
 @pytest.mark.parametrize("latent_frames", (3, 6, 12, 13))