[test] gate compiled H3 fusion dispatcher contracts

fastvideo/models/dits/minimax_h3_fusions/modulation.py

@@ -380,6 +380,7 @@ def fused_rmsnorm_modulate(
     device-side bounds check would synchronize); callers are safe by
     construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
     """
+    _require_forward_only(x, weight, scale, shift)
     if torch.compiler.is_compiling():
         return torch.ops.fastvideo._minimax_h3_rmsnorm_modulate(x, weight, scale, shift, index, eps)
     return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)
@@ -400,6 +401,7 @@ def fused_residual_gate_rmsnorm_modulate(
     ``index`` values must lie in ``[0, table_rows)``; see
     :func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
     """
+    _require_forward_only(residual, branch, gate, weight, scale, shift)
     if torch.compiler.is_compiling():
         return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate(
             residual,

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

@@ -205,6 +205,9 @@ def fused_qknorm_rope(
     (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
     correctly.
     """
+    if torch.is_grad_enabled() and any(
+            isinstance(tensor, torch.Tensor) and tensor.requires_grad for tensor in (x, weight, cos, sin)):
+        raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
     if torch.compiler.is_compiling():
         return torch.ops.fastvideo._minimax_h3_qknorm_rope(x, weight, cos, sin, eps)
     return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)

fastvideo/models/dits/minimax_h3_fusions/swiglu.py

@@ -116,6 +116,8 @@ def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
     This is intentionally a strict kernel wrapper: callers own fallback policy and
     must only invoke it for a supported CUDA inference path.
     """
+    if torch.is_grad_enabled() and x.requires_grad:
+        raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd")
     if torch.compiler.is_compiling():
         return torch.ops.fastvideo._minimax_h3_swiglu(x)
     return _minimax_h3_swiglu_impl(x)

fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 from torch._subclasses.fake_tensor import FakeTensorMode
 
@@ -55,7 +56,7 @@ def test_minimax_h3_fusions_are_dynamic_fullgraph_traceable() -> None:
 
     compiled = torch.compile(_all_fusions, backend=capture_graph, fullgraph=True, dynamic=True)
 
-    with FakeTensorMode():
+    with torch.no_grad(), FakeTensorMode():
         x = torch.empty((2, 3, 8), device="cuda", dtype=torch.bfloat16)
         weight = torch.empty((8,), device="cuda", dtype=torch.bfloat16)
         table = torch.empty((6, 8), device="cuda", dtype=torch.bfloat16)
@@ -95,3 +96,64 @@ def test_minimax_h3_fusions_are_dynamic_fullgraph_traceable() -> None:
         torch.ops.fastvideo._minimax_h3_qknorm_rope,
         torch.ops.fastvideo._minimax_h3_swiglu,
     } <= call_targets
+
+
+@pytest.mark.gpu
+def test_minimax_h3_fusion_custom_ops_match_real_fullgraph_cuda() -> None:
+    """Cross-check fake metadata and execute the real CUDA dispatcher path."""
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        pytest.skip("BF16 CUDA is required")
+    pytest.importorskip("triton")
+
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(2026)
+    x = torch.randn((2, 3, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    branch = torch.randn_like(x)
+    weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16)
+    scale = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    shift = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    gate = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    index = torch.tensor((0, 3, 5), device=device, dtype=torch.int64)
+    q = torch.randn((2, 3, 4, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    q_weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16).contiguous()
+    rotary = torch.randn((3, 96), generator=generator, device=device, dtype=torch.bfloat16).contiguous()
+    packed_mlp = torch.randn((2, 3, 256), generator=generator, device=device, dtype=torch.bfloat16)
+
+    opchecks = (
+        (
+            torch.ops.fastvideo._minimax_h3_rmsnorm_modulate,
+            (x, weight, scale, shift, index, 1e-5),
+        ),
+        (
+            torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate,
+            (x, branch, gate, weight, scale, shift, index, 1e-5),
+        ),
+        (
+            torch.ops.fastvideo._minimax_h3_qknorm_rope,
+            (q, q_weight, rotary, rotary.clone(), 1e-5),
+        ),
+        (torch.ops.fastvideo._minimax_h3_swiglu, (packed_mlp,)),
+    )
+    for op, args in opchecks:
+        torch.library.opcheck(op, args)
+
+    compiled = torch.compile(_all_fusions, fullgraph=True, dynamic=True)
+    inputs = (
+        x,
+        branch,
+        weight,
+        scale,
+        shift,
+        gate,
+        index,
+        q,
+        q_weight,
+        rotary,
+        rotary.clone(),
+        packed_mlp,
+    )
+    with torch.no_grad():
+        expected = _all_fusions(*inputs)
+        actual = compiled(*inputs)
+    for actual_output, expected_output in zip(actual, expected, strict=True):
+        torch.testing.assert_close(actual_output, expected_output, atol=0.0, rtol=0.0)