From 9709cc037e0d6e37361022a7de67f4878c18a7c7 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Sat, 22 Aug 2026 03:09:14 +0000 Subject: [PATCH] [perf] MiniMax-H3: keep Triton fusions active under regional compile --- fastvideo/models/dits/minimax_h3.py | 19 ++- .../dits/minimax_h3_fusions/modulation.py | 149 ++++++++++++++++-- .../dits/minimax_h3_fusions/qknorm_rope.py | 61 +++++-- .../models/dits/minimax_h3_fusions/swiglu.py | 32 +++- .../test_minimax_h3_fusion_compile.py | 97 ++++++++++++ 5 files changed, 317 insertions(+), 41 deletions(-) create mode 100644 fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 84949fa59..355ea00f8 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -61,13 +61,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 trace as opaque custom ops. 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()) + return HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() class MiniMaxH3RotaryPosEmbed(nn.Module): @@ -712,16 +712,15 @@ class MiniMaxH3Transformer3DModel(BaseDiT): 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. + The inference-only Triton fusions expose fake-backed custom operators, + so Dynamo can keep them active as opaque nodes inside each fullgraph + block instead of tracing into their launcher implementation. """ 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))) + logger.info( + "MiniMax H3 inference fusions remain active under torch.compile through custom-op boundaries: %s", + ",".join(sorted(self.enabled_fusions)), + ) def materialize_non_persistent_buffers( self, diff --git a/fastvideo/models/dits/minimax_h3_fusions/modulation.py b/fastvideo/models/dits/minimax_h3_fusions/modulation.py index 5cc3213b3..a61950ec7 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/modulation.py +++ b/fastvideo/models/dits/minimax_h3_fusions/modulation.py @@ -202,7 +202,7 @@ def _row_addressable(table: torch.Tensor) -> torch.Tensor: return table if table.stride(-1) == 1 else table.contiguous() -def fused_rmsnorm_modulate( +def _fused_rmsnorm_modulate_impl( x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, @@ -210,13 +210,6 @@ def fused_rmsnorm_modulate( index: torch.Tensor, eps: float, ) -> torch.Tensor: - """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) @@ -248,7 +241,7 @@ def fused_rmsnorm_modulate( return output.view_as(x) -def fused_residual_gate_rmsnorm_modulate( +def _fused_residual_gate_rmsnorm_modulate_impl( residual: torch.Tensor, branch: torch.Tensor, gate: torch.Tensor, @@ -258,11 +251,6 @@ 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. - - ``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) @@ -300,3 +288,136 @@ def fused_residual_gate_rmsnorm_modulate( num_warps=_num_warps(block_size), ) return hidden.view_as(residual), modulated.view_as(residual) + + +# Dynamo must see the Triton launches as opaque nodes. In eager mode the +# public wrappers below keep their strict, actionable input validation; while +# compiling, these custom-op boundaries avoid tracing into Triton's launcher +# and let the surrounding H3 block remain one full graph. +@torch.library.custom_op( + "fastvideo::_minimax_h3_rmsnorm_modulate", + mutates_args=(), + device_types="cuda", +) +def _fused_rmsnorm_modulate_op( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> torch.Tensor: + return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps) + + +@torch.library.register_fake("fastvideo::_minimax_h3_rmsnorm_modulate") +def _fused_rmsnorm_modulate_fake( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> torch.Tensor: + del weight, scale, shift, index, eps + return x.new_empty(x.shape) + + +@torch.library.custom_op( + "fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate", + mutates_args=(), + device_types="cuda", +) +def _fused_residual_gate_rmsnorm_modulate_op( + residual: torch.Tensor, + branch: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return _fused_residual_gate_rmsnorm_modulate_impl( + residual, + branch, + gate, + weight, + scale, + shift, + index, + eps, + ) + + +@torch.library.register_fake("fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate") +def _fused_residual_gate_rmsnorm_modulate_fake( + residual: torch.Tensor, + branch: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del branch, gate, weight, scale, shift, index, eps + return residual.new_empty(residual.shape), residual.new_empty(residual.shape) + + +def fused_rmsnorm_modulate( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> torch.Tensor: + """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). + """ + 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) + + +def fused_residual_gate_rmsnorm_modulate( + residual: torch.Tensor, + branch: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """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. + """ + if torch.compiler.is_compiling(): + return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate( + residual, + branch, + gate, + weight, + scale, + shift, + index, + eps, + ) + return _fused_residual_gate_rmsnorm_modulate_impl( + residual, + branch, + gate, + weight, + scale, + shift, + index, + eps, + ) diff --git a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py index d8aad9bef..1ed24c689 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py +++ b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py @@ -121,23 +121,13 @@ def _validate_inputs( return batch, seq_len, num_heads, head_dim, rotary_dim -def fused_qknorm_rope( +def _fused_qknorm_rope_impl( x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float, ) -> torch.Tensor: - """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel. - - RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the - final store. Triton's reduction order and the absence of eager's BF16 - intermediate materializations can produce small, expected rounding drift. - - Row offsets are computed in int64, so inputs beyond 2**31 total elements - (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address - correctly. - """ batch, seq_len, num_heads, head_dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps) if not weight.is_contiguous(): raise ValueError("weight must be contiguous") @@ -171,4 +161,53 @@ def fused_qknorm_rope( return flat_out.view(batch, seq_len, num_heads, head_dim) +@torch.library.custom_op( + "fastvideo::_minimax_h3_qknorm_rope", + mutates_args=(), + device_types="cuda", +) +def _fused_qknorm_rope_op( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> torch.Tensor: + return _fused_qknorm_rope_impl(x, weight, cos, sin, eps) + + +@torch.library.register_fake("fastvideo::_minimax_h3_qknorm_rope") +def _fused_qknorm_rope_fake( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> torch.Tensor: + del weight, cos, sin, eps + return x.new_empty(x.shape) + + +def fused_qknorm_rope( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel. + + RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the + final store. Triton's reduction order and the absence of eager's BF16 + intermediate materializations can produce small, expected rounding drift. + + Row offsets are computed in int64, so inputs beyond 2**31 total elements + (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address + correctly. + """ + 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) + + __all__ = ["HAVE_TRITON", "fused_qknorm_rope"] diff --git a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py index b7884639e..97773434d 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py +++ b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py @@ -65,12 +65,7 @@ def _num_warps(block_size: int) -> int: return 4 -def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor: - """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input. - - This is intentionally a strict kernel wrapper: callers own fallback policy and - must only invoke it for a supported CUDA inference path. - """ +def _minimax_h3_swiglu_impl(x: torch.Tensor) -> torch.Tensor: ffn_dim = _validate_input(x) if not x.is_cuda: raise ValueError("MiniMax H3 fused SwiGLU requires a CUDA tensor") @@ -101,4 +96,29 @@ def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor: return out.view(output_shape) +@torch.library.custom_op( + "fastvideo::_minimax_h3_swiglu", + mutates_args=(), + device_types="cuda", +) +def _minimax_h3_swiglu_op(x: torch.Tensor) -> torch.Tensor: + return _minimax_h3_swiglu_impl(x) + + +@torch.library.register_fake("fastvideo::_minimax_h3_swiglu") +def _minimax_h3_swiglu_fake(x: torch.Tensor) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], x.shape[-1] // 2)) + + +def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor: + """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input. + + This is intentionally a strict kernel wrapper: callers own fallback policy and + must only invoke it for a supported CUDA inference path. + """ + if torch.compiler.is_compiling(): + return torch.ops.fastvideo._minimax_h3_swiglu(x) + return _minimax_h3_swiglu_impl(x) + + __all__ = ["HAVE_TRITON", "minimax_h3_swiglu"] diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py new file mode 100644 index 000000000..41926f2e0 --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fullgraph trace contract for the MiniMax-H3 inference fusions.""" + +from __future__ import annotations + +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from fastvideo.models.dits.minimax_h3_fusions import ( + fused_qknorm_rope, + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, + minimax_h3_swiglu, +) + + +def _all_fusions( + x: torch.Tensor, + branch: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + q: torch.Tensor, + q_weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + packed_mlp: torch.Tensor, +) -> tuple[torch.Tensor, ...]: + norm1 = fused_rmsnorm_modulate(x, weight, scale, shift, index, 1e-5) + hidden, norm2 = fused_residual_gate_rmsnorm_modulate( + x, + branch, + gate, + weight, + scale, + shift, + index, + 1e-5, + ) + q_out = fused_qknorm_rope(q, q_weight, cos, sin, 1e-5) + mlp_out = minimax_h3_swiglu(packed_mlp) + return norm1, hidden, norm2, q_out, mlp_out + + +def test_minimax_h3_fusions_are_dynamic_fullgraph_traceable() -> None: + """Dynamo sees four opaque ops without requiring a physical CUDA GPU.""" + graphs: list[torch.fx.GraphModule] = [] + + def capture_graph(graph: torch.fx.GraphModule, example_inputs: list[torch.Tensor]): + del example_inputs + graphs.append(graph) + return graph.forward + + compiled = torch.compile(_all_fusions, backend=capture_graph, fullgraph=True, dynamic=True) + + with 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) + index = torch.empty((3,), device="cuda", dtype=torch.int64) + q = torch.empty((2, 3, 4, 8), device="cuda", dtype=torch.bfloat16) + rotary = torch.empty((3, 8), device="cuda", dtype=torch.bfloat16) + packed_mlp = torch.empty((2, 3, 16), device="cuda", dtype=torch.bfloat16) + + outputs = compiled( + x, + x, + weight, + table, + table, + table, + index, + q, + weight, + rotary, + rotary, + packed_mlp, + ) + + assert tuple(output.shape for output in outputs) == ( + (2, 3, 8), + (2, 3, 8), + (2, 3, 8), + (2, 3, 4, 8), + (2, 3, 8), + ) + assert all(output.is_contiguous() for output in outputs) + assert len(graphs) == 1 + call_targets = {node.target for node in graphs[0].graph.nodes if node.op == "call_function"} + assert { + torch.ops.fastvideo._minimax_h3_rmsnorm_modulate, + torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate, + torch.ops.fastvideo._minimax_h3_qknorm_rope, + torch.ops.fastvideo._minimax_h3_swiglu, + } <= call_targets