nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf] integrate compile-safe H3 fusions for matched benchmarking

Browse Source
This commit is contained in:
Will Lin
2026-08-22 03:09:42 +00:00
parent a72cfddade 9709cc037e
commit 84766e0174
5 changed files with 317 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/dits/minimax_h3.py
+9 -10
View File
@@ -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,
fastvideo/models/dits/minimax_h3_fusions/modulation.py
+135 -14
View File
@@ -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,
)
fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py
+50 -11
View File
@@ -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"]
fastvideo/models/dits/minimax_h3_fusions/swiglu.py
+26 -6
View File
@@ -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"]
fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py
+97
View File
@@ -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