diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 339756858..ae8d1d2f9 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -280,9 +280,9 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_H3_SPLICE_FROM_STEP` | int | `4` | eval | First denoising step run by FASTVIDEO_H3_SPLICE_TRANSFORMER. | | `FASTVIDEO_H3_ENCODER_LAYERWISE` | bool | `0` | performance | Stream MiniMax-H3 text-encoder language layers through exact-size pinned host memory (text-only prompts). | | `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT` | bool | `0` | performance | Expand the serialized NVFP4 MiniMax-H3 text encoder with one fused Triton pass on GPUs without FP4 GEMM. | -| `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per tile. | +| `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `8` | performance | 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` | int | `0` | performance | Entries in each content-keyed MiniMax-H3 memo of reference encodes (Qwen3-VL Ref2VA presentation, VAE keyframe latents); repeat references are reused exactly. 0 disables. | -| `FASTVIDEO_H3_VAE_TILE_PARALLEL` | bool | `0` | performance | 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. | +| `FASTVIDEO_H3_VAE_TILE_PARALLEL` | bool | `0` | performance | 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. Requires FASTVIDEO_H3_VAE_TILE_BATCH=1 (the split replays the per-tile path). | | `FASTVIDEO_H3_VAE_NVFP4_FUSED` | bool | `1` | performance | Run NVFP4 MiniMax-H3 VAE decoder blocks through the bit-exact fused Triton kernels at inference. | | `FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH` | bool | `1` | performance | Replay small fused NVFP4 MiniMax-H3 VAE decoder calls (at most four 1797-token tiles) as per-shape CUDA graphs. | | `FASTVIDEO_H3_VAE_NVFP4_DECODER` | path | unset | performance | Exported NVFP4 MiniMax-H3 video decoder (.pt, or .safetensors with bf16 masters or pre-packed NVFP4 weights) that replaces the VAE decoder at load; skips the INT8 ConvRot overlay. | diff --git a/fastvideo/envs.py b/fastvideo/envs.py index e4ca89460..d46a5cc4f 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -724,10 +724,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 " @@ -737,7 +741,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 " diff --git a/fastvideo/models/vaes/minimax_h3_parallel.py b/fastvideo/models/vaes/minimax_h3_parallel.py index 415e2239c..f39c875fc 100644 --- a/fastvideo/models/vaes/minimax_h3_parallel.py +++ b/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): diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index fa1ecbdbf..d24bec6cd 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/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 diff --git a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py index 90e91c744..06e9c8fbe 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py +++ b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py @@ -318,7 +318,7 @@ def test_fastvideo_args_strategy_literals_match_module() -> None: @pytest.fixture def tile_parallel(): - with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(True): + with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(True), envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(1): yield diff --git a/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py index dd573a7af..4517efe5c 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py +++ b/fastvideo/tests/vaes/test_minimax_h3_vae_compile.py @@ -15,6 +15,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +import fastvideo.envs as envs from fastvideo.models.vaes.minimax_h3_audio import ( MiniMaxH3AudioBigVGANDecoder, MiniMaxH3AudioVAE, @@ -350,8 +351,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(env_overrides) -> None: """Nest indexed decoder tiles between tile-splitting and stitching ranges.""" + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(1)) vae = _empty_typed_module(AutoencoderKLMiniMaxH3) vae.use_tiling = True vae.spatial_compression_ratio = 1 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 68e70da3e..d70f1ec70 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -3,6 +3,7 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.models.vaes.minimax_h3_video import ( MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig, @@ -182,3 +183,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(env_overrides) -> None: + from fastvideo.models.vaes.minimax_h3_video import _tile_batch_size + + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(None)) + assert _tile_batch_size() == 8 + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(3)) + assert _tile_batch_size() == 3 + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(0)) + assert _tile_batch_size() == 1 + + +def test_dynamic_nvfp4_act_scale_decodes_per_tile(env_overrides) -> 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") + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(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(env_overrides) -> 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) + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(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(env_overrides) -> 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): + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(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)