[bugfix]: decode dynamic-scale NVFP4 H3 VAE tiles one at a time

fastvideo/envs.py

@@ -714,10 +714,14 @@ FASTVIDEO_H3_ENCODER_FUSED_DEQUANT = EnvBool(False,
                                              category="performance",
                                              doc="Expand the serialized NVFP4 MiniMax-H3 text encoder with one "
                                              "fused Triton pass on GPUs without FP4 GEMM.")
-FASTVIDEO_H3_VAE_TILE_BATCH = EnvInt(1,
+FASTVIDEO_H3_VAE_TILE_BATCH = EnvInt(8,
                                      category="performance",
-                                     doc="Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per "
-                                     "tile.")
+                                     doc="Spatial tiles per MiniMax-H3 video VAE decoder call. The ViT decoder "
+                                     "treats batch entries independently, so batching tiles matches per-tile "
+                                     "decoding, except for NVFP4 decoders with dynamic activation scales, which "
+                                     "always decode per tile. On GB200 at 768x1344, 8 tiles per call decodes 107 "
+                                     "frames in 3.7 s instead of 5.4 s and 362 frames in 12.8 s instead of 18.3 s "
+                                     "(+0.7 GiB peak). 1 decodes per tile.")
 FASTVIDEO_H3_REF2VA_MEMO_ENTRIES = EnvInt(0,
                                           category="performance",
                                           doc="Entries in each content-keyed MiniMax-H3 memo of reference encodes "
@@ -727,7 +731,8 @@ FASTVIDEO_H3_VAE_TILE_PARALLEL = EnvBool(False,
                                          category="performance",
                                          doc="Split MiniMax-H3 VAE spatial tiles (not only temporal chunks) across "
                                          "the sequence-parallel ranks in the parallel decode and keyframe encode; "
-                                         "bitwise equal to the serial path.")
+                                         "bitwise equal to the serial path. Requires "
+                                         "FASTVIDEO_H3_VAE_TILE_BATCH=1 (the split replays the per-tile path).")
 FASTVIDEO_H3_VAE_NVFP4_FUSED = EnvBool(True,
                                        category="performance",
                                        doc="Run NVFP4 MiniMax-H3 VAE decoder blocks through the bit-exact fused "

fastvideo/models/vaes/minimax_h3_parallel.py

@@ -252,7 +252,7 @@ def _tile_split_applies(vae: AutoencoderKLMiniMaxH3, x: torch.Tensor) -> bool:
     # The batched tile decode (FASTVIDEO_H3_VAE_TILE_BATCH) is a different
     # serial path; the tile split replays only the per-tile one.
     return (envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.get() and vae.use_tiling and not vae._tile_helpers_compiled
-            and x.shape[0] == 1 and _tile_batch_size() == 1 and not torch.is_grad_enabled())
+            and x.shape[0] == 1 and _tile_batch_size(vae.decoder) == 1 and not torch.is_grad_enabled())
 
 
 def _tile_grid(vae: AutoencoderKLMiniMaxH3, height: int, width: int):

fastvideo/models/vaes/minimax_h3_video.py

@@ -536,8 +536,15 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module):
         )
 
 
-def _tile_batch_size() -> int:
-    """Spatial tiles decoded per decoder call (``FASTVIDEO_H3_VAE_TILE_BATCH``, default 1 = per tile)."""
+def _has_dynamic_nvfp4_act_scale(module: nn.Module) -> bool:
+    from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import nvfp4_linears
+
+    return any(layer.act_scale == "dynamic" for layer in nvfp4_linears(module))
+
+
+def _tile_batch_size(decoder: nn.Module | None = None) -> int:
+    if decoder is not None and _has_dynamic_nvfp4_act_scale(decoder):
+        return 1
     return max(1, envs.FASTVIDEO_H3_VAE_TILE_BATCH.get())
 
 
@@ -838,7 +845,8 @@ class AutoencoderKLMiniMaxH3(nn.Module):
 
             ratio = self.spatial_compression_ratio
             rows = []
-            if _tile_batch_size() > 1 and len(set(y_lengths)) == 1 and len(set(x_lengths)) == 1:
+            per_call = _tile_batch_size(self.decoder)
+            if per_call > 1 and len(set(y_lengths)) == 1 and len(set(x_lengths)) == 1:
                 # Every tile of the grid has one shape, and the ViT decoder
                 # treats batch entries independently: decode the grid in a few
                 # large calls instead of one small call per tile.
@@ -850,7 +858,6 @@ class AutoencoderKLMiniMaxH3(nn.Module):
                         for x_position, x_length in zip(x_indices, x_lengths)
                     ]
                     decoded: list[torch.Tensor] = []
-                    per_call = _tile_batch_size()
                     for start in range(0, len(latent_tiles), per_call):
                         batch = torch.cat(latent_tiles[start:start + per_call], dim=0)
                         # Clone before retaining: under the reduce-overhead compile the decoder output lives in

fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py

@@ -318,7 +318,9 @@ def test_fastvideo_args_strategy_literals_match_module() -> None:
 
 @pytest.fixture
 def tile_parallel():
-    with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(True):
+    # The tile split replays the per-tile serial forwards, so it stands down when the
+    # serial path batches tiles (FASTVIDEO_H3_VAE_TILE_BATCH=8 by default).
+    with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(True), envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(1):
         yield
 
 

fastvideo/tests/vaes/test_minimax_h3_vae_compile.py

@@ -350,8 +350,9 @@ def test_decode_clip_no_spatial_tiling_stage_ranges() -> None:
     ]
 
 
-def test_decode_clip_emits_tiled_stage_ranges() -> None:
+def test_decode_clip_emits_tiled_stage_ranges(monkeypatch) -> None:
     """Nest indexed decoder tiles between tile-splitting and stitching ranges."""
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", "1")
     vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
     vae.use_tiling = True
     vae.spatial_compression_ratio = 1

fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py

@@ -182,3 +182,84 @@ def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None:
 
     with pytest.raises(RuntimeError, match="wrote 0 frames"):
         vae.decode_to_pixels(latents, output)
+
+
+def test_tile_batch_defaults_to_eight(monkeypatch) -> None:
+    from fastvideo.models.vaes.minimax_h3_video import _tile_batch_size
+
+    monkeypatch.delenv("FASTVIDEO_H3_VAE_TILE_BATCH", raising=False)
+    assert _tile_batch_size() == 8
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", "3")
+    assert _tile_batch_size() == 3
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", "0")
+    assert _tile_batch_size() == 1
+
+
+def test_dynamic_nvfp4_act_scale_decodes_per_tile(monkeypatch) -> None:
+    from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import NVFP4DecoderLinear
+    from fastvideo.models.vaes.minimax_h3_video import _tile_batch_size
+
+    decoder = torch.nn.Module()
+    decoder.proj = NVFP4DecoderLinear(torch.zeros(16, 16), None, None, torch.float32, act_scale="dynamic")
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", "8")
+    assert _tile_batch_size(decoder) == 1
+    assert _tile_batch_size() == 8
+
+    decoder.proj.act_scale = "static"
+    assert _tile_batch_size(decoder) == 8
+    decoder.proj.act_scale = "unit"
+    assert _tile_batch_size(decoder) == 8
+    assert _tile_batch_size(torch.nn.Linear(16, 16)) == 8
+
+
+@torch.inference_mode()
+def test_dynamic_nvfp4_act_scale_decodes_one_tile_per_call(monkeypatch) -> None:
+    from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import NVFP4DecoderLinear
+
+    torch.manual_seed(20260814)
+    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)
+    latents = torch.randn(1, 4, 12, 4, 4)
+    monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", "8")
+    calls = []
+    decoder_forward = vae.decoder.forward
+
+    def record_batch(hidden_states):
+        calls.append(hidden_states.shape[0])
+        return decoder_forward(hidden_states)
+
+    vae.decoder.forward = record_batch
+    batched = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, batched)
+    batched_calls = len(calls)
+
+    calls.clear()
+    vae.decoder.dynamic_probe = NVFP4DecoderLinear(torch.zeros(16, 16), None, None, torch.float32, act_scale="dynamic")
+    per_tile = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, per_tile)
+
+    assert len(calls) > batched_calls
+    assert all(size == 1 for size in calls)
+    assert_close(per_tile, batched, atol=1e-6, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_batched_tile_decode_matches_per_tile(monkeypatch) -> None:
+    torch.manual_seed(20260814)
+    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)
+    latents = torch.randn(1, 4, 12, 4, 4)
+    outputs = {}
+    for tile_batch in (1, 8):
+        monkeypatch.setenv("FASTVIDEO_H3_VAE_TILE_BATCH", str(tile_batch))
+        output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+        vae.decode_to_pixels(latents, output)
+        outputs[tile_batch] = output
+
+    assert_close(outputs[1], outputs[8], atol=1e-6, rtol=0.0)