diff --git a/fastvideo/models/dits/minimax_h3_fusions/modulation.py b/fastvideo/models/dits/minimax_h3_fusions/modulation.py index a61950ec7..961250edb 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/modulation.py +++ b/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, diff --git a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py index 1ed24c689..1575cbf57 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py index 97773434d..9c703ce16 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py +++ b/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) diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py index 41926f2e0..d6e8c1679 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py +++ b/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)