From 5c05f448a96994e75972ed19bfbfaaf14a4688b1 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Sat, 31 Oct 2026 13:12:43 -0700 Subject: [PATCH] [perf] MiniMax H3: opt-in batched VAE encoder tiles --- fastvideo/envs.py | 5 +++++ fastvideo/models/vaes/minimax_h3_video.py | 15 ++++++++++++++- .../test_minimax_h3_video_vae_streaming.py | 18 ++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/fastvideo/envs.py b/fastvideo/envs.py index a1eabdcb4..c46897184 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -758,6 +758,11 @@ FASTVIDEO_H3_VAE_UINT8_PIXELS = EnvBool(True, 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_ENCODE_TILE_BATCH = EnvInt(1, + category="performance", + doc="Spatial tiles per MiniMax-H3 video VAE encoder call (reference and " + "keyframe encodes); 1 encodes per tile. Batched tiles differ from " + "per-tile encoding only in reduction order.") FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT = EnvBool(False, category="performance", doc="Fused dequantization epilogue for the MiniMax-H3 VAE INT8 " diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index 224719d88..24defbf11 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -792,8 +792,21 @@ class AutoencoderKLMiniMaxH3(nn.Module): self.tile_sample_min_width, self.tile_sample_min_overlap_width, ) + per_call = envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.get() rows = [] - for y_position, y_length in zip(y_indices, y_lengths): + if per_call > 1 and len(set(y_lengths)) == 1 and len(set(x_lengths)) == 1: + tiles = [ + x[..., y_position:y_position + y_length, x_position:x_position + x_length] + for y_position, y_length in zip(y_indices, y_lengths) + for x_position, x_length in zip(x_indices, x_lengths) + ] + encoded: list[torch.Tensor] = [] + for start in range(0, len(tiles), per_call): + batch = torch.cat(tiles[start:start + per_call], dim=0) + encoded.extend(self.quant_conv(self.encoder(batch)).split(x.shape[0], dim=0)) + columns = len(x_indices) + rows = [encoded[index:index + columns] for index in range(0, len(encoded), columns)] + for y_position, y_length in zip(y_indices, y_lengths) if not rows else (): row = [] for x_position, x_length in zip(x_indices, x_lengths): tile = x[..., y_position:y_position + y_length, x_position:x_position + x_length] diff --git a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py index 917f705cc..0bfefd74d 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -123,6 +123,24 @@ def test_decode_to_pixels_uint8_matches_fp32_quantization(device: str) -> None: assert torch.equal(streamed, fp32.mul_(255).clamp_(0, 255).to(torch.uint8)) +@torch.inference_mode() +def test_batched_tile_encode_matches_per_tile(monkeypatch) -> None: + torch.manual_seed(20260810) + vae = _tiny_vae() + vae.enable_tiling(tile_sample_min_height=8, + tile_sample_min_width=8, + tile_sample_min_overlap_height=2, + tile_sample_min_overlap_width=2) + pixels = torch.rand(1, 3, 1, 20, 20) + + monkeypatch.setenv("FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH", "1") + per_tile = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters + monkeypatch.setenv("FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH", "8") + batched = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters + + assert_close(batched, per_tile, atol=1e-5, rtol=1e-5) + + # 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))