From cbab605efff2377a4600a978903464cd3c899f30 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Fri, 21 Aug 2026 10:06:23 +0000 Subject: [PATCH] [misc] MiniMax-H3 fusions: exact eager fallback, enable/inert logging, engagement test - _can_run_minimax_h3_fusion now also requires Triton availability, so an enabled fusion on a CUDA build without a working Triton falls back to eager instead of hitting the strict wrappers' RuntimeError mid-forward. - One-time logger.info of the resolved fusion set at model init (and a warning when the set is requested without Triton), plus a prepare_for_compile hook warning that torch.compile capture makes the fusions inert inside compiled block forwards. - Positive-engagement routing test: counts 1/1/2/1 fused-kernel calls in one CUDA inference block forward and asserts a grad-enabled forward leaves the counters unchanged, so a guard regression to always-eager can no longer pass silently. - Document the index-bounds contract ([0, table_rows)) on the modulation wrappers; a device-side check would synchronize, and in-model callers are safe by construction. Co-Authored-By: Claude Fable 5 --- fastvideo/models/dits/minimax_h3.py | 37 +++++++- .../dits/minimax_h3_fusions/__init__.py | 3 +- .../dits/minimax_h3_fusions/modulation.py | 14 ++- .../test_minimax_h3_fusion_routing.py | 95 +++++++++++++++++++ 4 files changed, 144 insertions(+), 5 deletions(-) diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index d45c30eb3..a36d2bc5f 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -24,8 +24,10 @@ from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps +from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.minimax_h3_fusions import ( + HAVE_TRITON, fused_qknorm_rope, fused_residual_gate_rmsnorm_modulate, fused_rmsnorm_modulate, @@ -34,6 +36,8 @@ from fastvideo.models.dits.minimax_h3_fusions import ( from fastvideo.platforms import AttentionBackendEnum from fastvideo.utils import get_compute_dtype +logger = init_logger(__name__) + MINIMAX_H3_MODALITY_NUM = 3 _CFG = MiniMaxH3Config() _MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"}) @@ -56,8 +60,13 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool: - """Triton kernels are inference-only and stay outside Dynamo capture.""" - return tensor.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling() + """Triton kernels are inference-only and stay outside Dynamo capture. + + The ``HAVE_TRITON`` check makes the eager fallback exact: on a CUDA build + whose Triton failed to import, an enabled fusion falls back instead of + hitting the strict wrappers' hard RuntimeError mid-forward. + """ + return (HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling()) class MiniMaxH3RotaryPosEmbed(nn.Module): @@ -563,6 +572,16 @@ class MiniMaxH3Transformer3DModel(BaseDiT): super().__init__(config, hf_config) arch = config.arch_config self.enabled_fusions = _enabled_minimax_h3_fusions() + if self.enabled_fusions: + if HAVE_TRITON: + logger.info( + "MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled and " + "torch.compile-captured forwards fall back to eager).", + ",".join(sorted(self.enabled_fusions))) + else: + logger.warning( + "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; " + "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions))) sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 if arch.num_attention_heads % sp_world_size: raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " @@ -689,6 +708,20 @@ class MiniMaxH3Transformer3DModel(BaseDiT): ) self.__post_init__() + def prepare_for_compile(self) -> None: + """Pipeline hook, called once right before torch.compile wraps the blocks. + + Dynamo capture traces the eager branch of every fusion guard, so an + enabled ``FASTVIDEO_MINIMAX_H3_FUSIONS`` set is silently inert inside + compiled block forwards (H3 compiles per-block by default). Say so + once instead of leaving the flag looking active. + """ + if self.enabled_fusions: + logger.warning( + "torch.compile is enabled for MiniMax H3, so the requested inference fusions (%s) are " + "inert inside compiled block forwards; the compiled eager path runs instead.", + ",".join(sorted(self.enabled_fusions))) + def materialize_non_persistent_buffers( self, device: torch.device, diff --git a/fastvideo/models/dits/minimax_h3_fusions/__init__.py b/fastvideo/models/dits/minimax_h3_fusions/__init__.py index 81637d218..d534e01df 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/__init__.py +++ b/fastvideo/models/dits/minimax_h3_fusions/__init__.py @@ -8,10 +8,11 @@ from .modulation import ( fused_residual_gate_rmsnorm_modulate, fused_rmsnorm_modulate, ) -from .qknorm_rope import fused_qknorm_rope +from .qknorm_rope import HAVE_TRITON, fused_qknorm_rope from .swiglu import minimax_h3_swiglu __all__ = [ + "HAVE_TRITON", "fused_qknorm_rope", "fused_residual_gate_rmsnorm_modulate", "fused_rmsnorm_modulate", diff --git a/fastvideo/models/dits/minimax_h3_fusions/modulation.py b/fastvideo/models/dits/minimax_h3_fusions/modulation.py index 123a79b18..5cc3213b3 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/modulation.py +++ b/fastvideo/models/dits/minimax_h3_fusions/modulation.py @@ -210,7 +210,13 @@ def fused_rmsnorm_modulate( index: torch.Tensor, eps: float, ) -> torch.Tensor: - """Run RMSNorm and row-indexed modulation in one strict Triton kernel.""" + """Run RMSNorm and row-indexed modulation in one strict Triton kernel. + + ``index`` values must lie in ``[0, table_rows)``. Unlike eager + ``index_select``, the kernel does not raise on out-of-range values (a + device-side bounds check would synchronize); callers are safe by + construction (``timestep_indices * 3 + token_tags``, SP pads with 0). + """ _validate_contract(x, weight, (scale, shift), index, eps) _require_forward_only(x, weight, scale, shift) _require_triton_cuda(x) @@ -252,7 +258,11 @@ def fused_residual_gate_rmsnorm_modulate( index: torch.Tensor, eps: float, ) -> tuple[torch.Tensor, torch.Tensor]: - """Fuse residual update, row-indexed gate, RMSNorm, and modulation.""" + """Fuse residual update, row-indexed gate, RMSNorm, and modulation. + + ``index`` values must lie in ``[0, table_rows)``; see + :func:`fused_rmsnorm_modulate` for why the wrapper does not check them. + """ _validate_residual_branch(residual, branch) _validate_contract(residual, weight, (gate, scale, shift), index, eps) _require_forward_only(residual, branch, gate, weight, scale, shift) diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py index 3bf544351..176974108 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py +++ b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py @@ -134,3 +134,98 @@ def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4( torch.testing.assert_close(fused_output, eager_output, atol=3e-2, rtol=3e-2) finally: cleanup_dist_env_and_memory() + + +def test_minimax_h3_fusions_engage_on_cuda_inference(monkeypatch: pytest.MonkeyPatch) -> None: + """Pin the positive side of the routing guard. + + The parity test above still passes if ``_can_run_minimax_h3_fusion`` + silently degrades to always-False (both blocks then run the identical + eager path), so count the fused-kernel calls: one CUDA inference forward + through a fully fused block must hit ``fused_rmsnorm_modulate`` once, + ``fused_residual_gate_rmsnorm_modulate`` once, ``fused_qknorm_rope`` + twice (q and k), and ``minimax_h3_swiglu`` once -- and a grad-enabled + forward must leave every counter unchanged. + """ + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 CUDA is required") + pytest.importorskip("triton") + + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") + monkeypatch.setenv("MASTER_PORT", "29574") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_RANK", "0") + + import fastvideo.models.dits.minimax_h3 as h3 + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.forward_context import set_forward_context + + calls = dict.fromkeys(("rmsnorm_modulate", "residual_gate_rmsnorm_modulate", "qknorm_rope", "swiglu"), 0) + + def _counting(name: str, real): + + def wrapper(*args, **kwargs): + calls[name] += 1 + return real(*args, **kwargs) + + return wrapper + + monkeypatch.setattr(h3, "fused_rmsnorm_modulate", _counting("rmsnorm_modulate", h3.fused_rmsnorm_modulate)) + monkeypatch.setattr(h3, "fused_residual_gate_rmsnorm_modulate", + _counting("residual_gate_rmsnorm_modulate", h3.fused_residual_gate_rmsnorm_modulate)) + monkeypatch.setattr(h3, "fused_qknorm_rope", _counting("qknorm_rope", h3.fused_qknorm_rope)) + monkeypatch.setattr(h3, "minimax_h3_swiglu", _counting("swiglu", h3.minimax_h3_swiglu)) + + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + block = h3.MiniMaxH3TransformerBlock( + hidden_size=128, + num_attention_heads=1, + attention_head_dim=128, + ffn_dim=256, + time_embed_dim=64, + norm_eps=1e-5, + qk_norm_eps=1e-5, + supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ), + quant_config=None, + prefix="minimax_h3.engagement_block", + fuse_modulate=True, + fuse_qknorm_rope=True, + fuse_swiglu=True, + ) + device = torch.device("cuda") + block = block.to(device=device, dtype=torch.bfloat16).eval() + + generator = torch.Generator(device=device).manual_seed(2026) + hidden_states = torch.randn(2, 12, 128, generator=generator, device=device, dtype=torch.bfloat16) + temb = torch.randn(2, 64, generator=generator, device=device, dtype=torch.bfloat16) + adaln_indices = torch.arange(12, device=device, dtype=torch.long).remainder(6) + position_ids = torch.zeros(12, 3, device=device, dtype=torch.float32) + position_ids[:, 0] = torch.arange(12, device=device) + rotary_emb = h3.MiniMaxH3RotaryPosEmbed(rope_freq_dim=16, rope_theta=10000.0).to(device)(position_ids) + inputs = dict( + hidden_states=hidden_states, + temb=temb, + adaln_indices=adaln_indices, + rotary_emb=tuple(value.to(torch.bfloat16) for value in rotary_emb), + original_seq_len=12, + ) + + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + block(**inputs) + engaged = dict(calls) + assert engaged == { + "rmsnorm_modulate": 1, + "residual_gate_rmsnorm_modulate": 1, + "qknorm_rope": 2, + "swiglu": 1, + }, engaged + + grad_inputs = {**inputs, "hidden_states": hidden_states.clone().requires_grad_(True)} + with set_forward_context(current_timestep=0, attn_metadata=None): + block(**grad_inputs) + assert dict(calls) == engaged, f"a fusion ran under grad: {calls} vs {engaged}" + finally: + cleanup_dist_env_and_memory()