[perf] MiniMax H3: opt-in batched VAE encoder tiles

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 "

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]

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))