[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 <noreply@anthropic.com>

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."""
+    """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()
+
+    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,

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",

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)

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