[perf] Add opt-in MiniMax-H3 Sol-Engine fusions (#1735)

fastvideo/envs.py

@@ -21,6 +21,7 @@ if TYPE_CHECKING:
     FASTVIDEO_TRACE_FUNCTION: int = 0
     FASTVIDEO_ATTENTION_BACKEND: str | None = None
     FASTVIDEO_FA4: bool = False
+    FASTVIDEO_MINIMAX_H3_FUSIONS: str = ""
     FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn"
     FASTVIDEO_TARGET_DEVICE: str = "cuda"
     MAX_JOBS: str | None = None
@@ -217,6 +218,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_FA4":
     lambda: os.getenv("FASTVIDEO_FA4", "0") != "0",
 
+    # Opt-in MiniMax-H3 inference-only Triton fusions adapted from the
+    # NVlabs/Sana Sol-Engine implementation. Accepts `all`, `1`, or a
+    # comma-separated subset of `modulate,qknorm_rope,swiglu`. An empty value
+    # (the default), `0`, or `none` keeps the eager implementation.
+    "FASTVIDEO_MINIMAX_H3_FUSIONS":
+    lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FUSIONS", ""),
+
     # Use dedicated multiprocess context for workers.
     "FASTVIDEO_WORKER_MULTIPROC_METHOD":
     lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"),

fastvideo/models/dits/minimax_h3.py

@@ -10,6 +10,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from fastvideo import envs
 from fastvideo.attention import DistributedAttention
 from fastvideo.attention.layer import DistributedAttention_VSA
 from fastvideo.attention.selector import get_attn_backend
@@ -23,12 +24,49 @@ 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,
+    minimax_h3_swiglu,
+)
 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"})
+
+
+def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]:
+    """Parse the independently switchable inference fusion set."""
+    raw = envs.FASTVIDEO_MINIMAX_H3_FUSIONS if value is None else value
+    normalized = raw.strip().lower()
+    if normalized in {"", "0", "none"}:
+        return frozenset()
+    if normalized in {"1", "all"}:
+        return _MINIMAX_H3_FUSION_NAMES
+    enabled = frozenset(item.strip() for item in normalized.split(",") if item.strip())
+    unknown = enabled - _MINIMAX_H3_FUSION_NAMES
+    if unknown:
+        supported = ",".join(sorted(_MINIMAX_H3_FUSION_NAMES))
+        raise ValueError(f"Unknown MiniMax H3 fusion(s) {sorted(unknown)}; expected a subset of {supported}.")
+    return enabled
+
+
+def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool:
+    """Triton kernels are inference-only and stay outside Dynamo capture.
+
+    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):
@@ -62,6 +100,7 @@ class MiniMaxH3FeedForward(nn.Module):
         ffn_dim: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        fuse_swiglu: bool = False,
     ) -> None:
         super().__init__()
         self.fc_in = ReplicatedLinear(
@@ -78,11 +117,15 @@ class MiniMaxH3FeedForward(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.fc_out",
         )
+        self.fuse_swiglu = fuse_swiglu
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states, _ = self.fc_in(hidden_states)
-        hidden_states, gate = hidden_states.chunk(2, dim=-1)
-        hidden_states = hidden_states * F.silu(gate)
+        if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states):
+            hidden_states = minimax_h3_swiglu(hidden_states)
+        else:
+            hidden_states, gate = hidden_states.chunk(2, dim=-1)
+            hidden_states = hidden_states * F.silu(gate)
         hidden_states, _ = self.fc_out(hidden_states)
         return hidden_states
 
@@ -99,6 +142,7 @@ class MiniMaxH3Attention(nn.Module):
         supported_attention_backends: tuple[AttentionBackendEnum, ...],
         quant_config: QuantizationConfig | None,
         prefix: str,
+        fuse_qknorm_rope: bool = False,
     ) -> None:
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -134,6 +178,7 @@ class MiniMaxH3Attention(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.to_out",
         )
+        self.fuse_qknorm_rope = fuse_qknorm_rope
         # VSA carries a learned gate on its pooled-compression branch. The H3
         # checkpoint has no such weight, so the loader zero-initializes it
         # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled
@@ -211,11 +256,18 @@ class MiniMaxH3Attention(nn.Module):
         query = query.unflatten(-1, (self.num_attention_heads, self.attention_head_dim))
         key = key.unflatten(-1, (self.num_attention_heads, self.attention_head_dim))
         value = value.unflatten(-1, (self.num_attention_heads, self.attention_head_dim))
-        query = self.norm_q(query)
-        key = self.norm_k(key)
-        if rotary_emb is not None:
-            query = self._apply_rotary_emb(query, rotary_emb)
-            key = self._apply_rotary_emb(key, rotary_emb)
+        if (self.fuse_qknorm_rope and rotary_emb is not None and _can_run_minimax_h3_fusion(query)):
+            cos, sin = rotary_emb
+            cos = cos.to(query.dtype)
+            sin = sin.to(query.dtype)
+            query = fused_qknorm_rope(query, self.norm_q.weight, cos, sin, self.norm_q.eps)
+            key = fused_qknorm_rope(key, self.norm_k.weight, cos, sin, self.norm_k.eps)
+        else:
+            query = self.norm_q(query)
+            key = self.norm_k(key)
+            if rotary_emb is not None:
+                query = self._apply_rotary_emb(query, rotary_emb)
+                key = self._apply_rotary_emb(key, rotary_emb)
 
         # H3 rotates only 96/128 channels, which the generic `freqs_cis`
         # branch cannot express. Apply it above, then pass no RoPE here.
@@ -397,6 +449,9 @@ class MiniMaxH3TransformerBlock(nn.Module):
         quant_config: QuantizationConfig | None,
         prefix: str,
         adaln_apply_silu: bool = True,
+        fuse_modulate: bool = False,
+        fuse_qknorm_rope: bool = False,
+        fuse_swiglu: bool = False,
     ) -> None:
         super().__init__()
         self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
@@ -408,6 +463,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             supported_attention_backends,
             quant_config,
             prefix=f"{prefix}.attn",
+            fuse_qknorm_rope=fuse_qknorm_rope,
         )
         self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
         self.ff = MiniMaxH3FeedForward(
@@ -415,6 +471,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             ffn_dim,
             quant_config=quant_config,
             prefix=f"{prefix}.ff",
+            fuse_swiglu=fuse_swiglu,
         )
         self.adaln_proj = MiniMaxH3AdaLayerNormModulation(
             time_embed_dim,
@@ -423,6 +480,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             prefix=f"{prefix}.adaln_proj",
             apply_silu=adaln_apply_silu,
         )
+        self.fuse_modulate = fuse_modulate
 
     def forward(
         self,
@@ -435,19 +493,39 @@ class MiniMaxH3TransformerBlock(nn.Module):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
             t.to(hidden_states.dtype) for t in self.adaln_proj(temb))
 
-        residual = hidden_states
-        norm_hidden_states = self.norm1(hidden_states)
-        norm_hidden_states = norm_hidden_states * (
-            1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
+        use_modulate_fusion = self.fuse_modulate and _can_run_minimax_h3_fusion(hidden_states)
+        if use_modulate_fusion:
+            norm_hidden_states = fused_rmsnorm_modulate(
+                hidden_states,
+                self.norm1.weight,
+                scale_msa,
+                shift_msa,
+                adaln_indices,
+                self.norm1.eps,
+            )
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+            norm_hidden_states = norm_hidden_states * (
+                1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
         attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len)
-        hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attention_output
-
-        residual = hidden_states
-        norm_hidden_states = self.norm2(hidden_states)
-        norm_hidden_states = norm_hidden_states * (
-            1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
+        if use_modulate_fusion:
+            hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate(
+                hidden_states,
+                attention_output,
+                gate_msa,
+                self.norm2.weight,
+                scale_mlp,
+                shift_mlp,
+                adaln_indices,
+                self.norm2.eps,
+            )
+        else:
+            hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (
+                1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
         feed_forward_output = self.ff(norm_hidden_states)
-        return residual + gate_mlp.index_select(0, adaln_indices) * feed_forward_output
+        return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output
 
 
 class MiniMaxH3Transformer3DModel(BaseDiT):
@@ -493,6 +571,17 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
     def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None:
         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 "
@@ -590,6 +679,9 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 config.quant_config,
                 prefix=f"{config.prefix}.transformer_blocks.{index}",
                 adaln_apply_silu=self.adaln_rank is None,
+                fuse_modulate="modulate" in self.enabled_fusions,
+                fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions,
+                fuse_swiglu="swiglu" in self.enabled_fusions,
             ) for index in range(arch.num_layers)
         ])
         self.norm_out = MiniMaxH3AdaLayerNormOut(
@@ -616,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

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference-only MiniMax H3 fusions adapted from NVlabs/Sana Sol-Engine.
+
+Source: https://github.com/NVlabs/Sana/tree/sol-engine/models/minimax_h3/GB200
+"""
+
+from .modulation import (
+    fused_residual_gate_rmsnorm_modulate,
+    fused_rmsnorm_modulate,
+)
+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",
+    "minimax_h3_swiglu",
+]

fastvideo/models/dits/minimax_h3_fusions/modulation.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 RMSNorm and row-indexed modulation fusions."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError as exc:  # pragma: no cover - depends on the runtime image
+    triton = None
+    tl = None
+    _TRITON_IMPORT_ERROR: ImportError | None = exc
+else:
+    _TRITON_IMPORT_ERROR = None
+
+
+__all__ = [
+    "fused_residual_gate_rmsnorm_modulate",
+    "fused_rmsnorm_modulate",
+]
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+_rmsnorm_modulate_kernel = None
+_residual_gate_rmsnorm_modulate_kernel = None
+
+
+if triton is not None:
+
+    @triton.jit
+    def _rmsnorm_modulate_kernel(
+        out_ptr,
+        x_ptr,
+        weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        index_ptr,
+        n_cols,
+        n_index,
+        eps,
+        stride_x_row,
+        stride_scale_row,
+        stride_shift_row,
+        BLOCK: tl.constexpr,
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK)
+        mask = cols < n_cols
+        x_offsets = row * stride_x_row + cols
+        table_row = tl.load(index_ptr + row % n_index).to(tl.int64)
+
+        x = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32)
+        variance = tl.sum(x * x, axis=0) / n_cols
+        normed = x * tl.math.rsqrt(variance + eps)
+        weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(
+            scale_ptr + table_row * stride_scale_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        shift = tl.load(
+            shift_ptr + table_row * stride_shift_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        output = normed * weight * (1.0 + scale) + shift
+        tl.store(out_ptr + x_offsets, output.to(out_ptr.dtype.element_ty), mask=mask)
+
+    @triton.jit
+    def _residual_gate_rmsnorm_modulate_kernel(
+        hidden_out_ptr,
+        normed_out_ptr,
+        residual_ptr,
+        branch_ptr,
+        gate_ptr,
+        weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        index_ptr,
+        n_cols,
+        n_index,
+        eps,
+        stride_input_row,
+        stride_gate_row,
+        stride_scale_row,
+        stride_shift_row,
+        BLOCK: tl.constexpr,
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK)
+        mask = cols < n_cols
+        input_offsets = row * stride_input_row + cols
+        table_row = tl.load(index_ptr + row % n_index).to(tl.int64)
+
+        residual = tl.load(residual_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+        branch = tl.load(branch_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(
+            gate_ptr + table_row * stride_gate_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        hidden = residual + gate * branch
+        tl.store(hidden_out_ptr + input_offsets, hidden.to(hidden_out_ptr.dtype.element_ty), mask=mask)
+
+        variance = tl.sum(hidden * hidden, axis=0) / n_cols
+        normed = hidden * tl.math.rsqrt(variance + eps)
+        weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(
+            scale_ptr + table_row * stride_scale_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        shift = tl.load(
+            shift_ptr + table_row * stride_shift_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        output = normed * weight * (1.0 + scale) + shift
+        tl.store(normed_out_ptr + input_offsets, output.to(normed_out_ptr.dtype.element_ty), mask=mask)
+
+
+def _validate_contract(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    tables: tuple[torch.Tensor, ...],
+    index: torch.Tensor,
+    eps: float,
+) -> None:
+    if x.ndim < 2:
+        raise ValueError(f"x must have shape (..., sequence_length, hidden_size), got {tuple(x.shape)}.")
+    if x.numel() == 0 or x.shape[-1] == 0:
+        raise ValueError("x must not be empty.")
+    if x.dtype not in _SUPPORTED_DTYPES:
+        raise TypeError(f"x must use float16, bfloat16, or float32, got {x.dtype}.")
+    hidden_size = x.shape[-1]
+    sequence_length = x.shape[-2]
+    if weight.shape != (hidden_size, ):
+        raise ValueError(f"weight must have shape ({hidden_size},), got {tuple(weight.shape)}.")
+    if weight.dtype not in _SUPPORTED_DTYPES:
+        raise TypeError(f"weight must use float16, bfloat16, or float32, got {weight.dtype}.")
+    if index.ndim != 1 or index.numel() != sequence_length:
+        raise ValueError(
+            f"index must have shape ({sequence_length},) so it can wrap over batch rows, got {tuple(index.shape)}."
+        )
+    if index.dtype not in (torch.int32, torch.int64):
+        raise TypeError(f"index must use int32 or int64, got {index.dtype}.")
+    if not isinstance(eps, (float, int)) or isinstance(eps, bool) or not math.isfinite(eps) or eps <= 0:
+        raise ValueError(f"eps must be a positive finite number, got {eps!r}.")
+
+    table_rows = tables[0].shape[0] if tables and tables[0].ndim == 2 else None
+    for name, table in zip(("gate", "scale", "shift")[-len(tables):], tables, strict=True):
+        if table.ndim != 2 or table.shape[1] != hidden_size:
+            raise ValueError(f"{name} must have shape (table_rows, {hidden_size}), got {tuple(table.shape)}.")
+        if table.shape[0] == 0 or table.shape[0] != table_rows:
+            raise ValueError("all modulation tables must have the same non-zero row count.")
+        if table.dtype not in _SUPPORTED_DTYPES:
+            raise TypeError(f"{name} must use float16, bfloat16, or float32, got {table.dtype}.")
+
+    tensors = (x, weight, *tables, index)
+    if any(tensor.device != x.device for tensor in tensors[1:]):
+        raise ValueError("x, weight, modulation tables, and index must be on the same device.")
+
+
+def _validate_residual_branch(residual: torch.Tensor, branch: torch.Tensor) -> None:
+    if branch.shape != residual.shape:
+        raise ValueError(f"branch must match residual shape {tuple(residual.shape)}, got {tuple(branch.shape)}.")
+    if branch.dtype != residual.dtype:
+        raise TypeError(f"branch dtype must match residual dtype {residual.dtype}, got {branch.dtype}.")
+    if branch.device != residual.device:
+        raise ValueError("branch and residual must be on the same device.")
+
+
+def _require_triton_cuda(x: torch.Tensor) -> None:
+    if triton is None:
+        detail = f": {_TRITON_IMPORT_ERROR}" if _TRITON_IMPORT_ERROR is not None else ""
+        raise RuntimeError(f"MiniMax H3 modulation fusion requires Triton{detail}.")
+    if x.device.type != "cuda":
+        raise RuntimeError(f"MiniMax H3 modulation fusion requires CUDA tensors, got device {x.device}.")
+
+
+def _require_forward_only(*tensors: torch.Tensor) -> None:
+    if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in tensors):
+        raise RuntimeError("MiniMax H3 modulation fusion is forward-only and does not support autograd.")
+
+
+def _next_power_of_two(value: int) -> int:
+    return 1 << (value - 1).bit_length()
+
+
+def _num_warps(block_size: int) -> int:
+    if block_size >= 8192:
+        return 16
+    if block_size >= 2048:
+        return 8
+    return 4
+
+
+def _row_addressable(table: torch.Tensor) -> torch.Tensor:
+    return table if table.stride(-1) == 1 else table.contiguous()
+
+
+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).
+    """
+    _validate_contract(x, weight, (scale, shift), index, eps)
+    _require_forward_only(x, weight, scale, shift)
+    _require_triton_cuda(x)
+
+    hidden_size = x.shape[-1]
+    flat_x = x.reshape(-1, hidden_size).contiguous()
+    weight = weight.contiguous()
+    scale = _row_addressable(scale)
+    shift = _row_addressable(shift)
+    index = index.contiguous()
+    output = torch.empty_like(flat_x)
+    block_size = _next_power_of_two(hidden_size)
+    _rmsnorm_modulate_kernel[(flat_x.shape[0], )](
+        output,
+        flat_x,
+        weight,
+        scale,
+        shift,
+        index,
+        hidden_size,
+        index.numel(),
+        eps,
+        flat_x.stride(0),
+        scale.stride(0),
+        shift.stride(0),
+        BLOCK=block_size,
+        num_warps=_num_warps(block_size),
+    )
+    return output.view_as(x)
+
+
+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.
+    """
+    _validate_residual_branch(residual, branch)
+    _validate_contract(residual, weight, (gate, scale, shift), index, eps)
+    _require_forward_only(residual, branch, gate, weight, scale, shift)
+    _require_triton_cuda(residual)
+
+    hidden_size = residual.shape[-1]
+    flat_residual = residual.reshape(-1, hidden_size).contiguous()
+    flat_branch = branch.reshape(-1, hidden_size).contiguous()
+    weight = weight.contiguous()
+    gate = _row_addressable(gate)
+    scale = _row_addressable(scale)
+    shift = _row_addressable(shift)
+    index = index.contiguous()
+    hidden = torch.empty_like(flat_residual)
+    modulated = torch.empty_like(flat_residual)
+    block_size = _next_power_of_two(hidden_size)
+    _residual_gate_rmsnorm_modulate_kernel[(flat_residual.shape[0], )](
+        hidden,
+        modulated,
+        flat_residual,
+        flat_branch,
+        gate,
+        weight,
+        scale,
+        shift,
+        index,
+        hidden_size,
+        index.numel(),
+        eps,
+        flat_residual.stride(0),
+        gate.stride(0),
+        scale.stride(0),
+        shift.stride(0),
+        BLOCK=block_size,
+        num_warps=_num_warps(block_size),
+    )
+    return hidden.view_as(residual), modulated.view_as(residual)

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused per-head RMSNorm and partial rotary embedding for MiniMax H3."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:  # pragma: no cover - exercised only in environments without Triton
+    triton = None
+    tl = None
+    HAVE_TRITON = False
+
+
+if HAVE_TRITON:
+
+    @triton.jit
+    def _qknorm_partial_rope_kernel(
+        out_ptr,
+        x_ptr,
+        weight_ptr,
+        cos_ptr,
+        sin_ptr,
+        head_dim,
+        rotary_dim,
+        half_rotary_dim,
+        num_heads,
+        seq_len,
+        eps,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        # int64, like the sibling kernels: with int32 program ids,
+        # ``row * head_dim`` wraps once the flattened input reaches 2**31
+        # elements (H3's 56 heads x 128 head_dim crosses that at
+        # batch*seq >= 299_593 tokens per rank) and the loads/stores below
+        # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
+        row = tl.program_id(0).to(tl.int64)
+        seq_index = (row // num_heads) % seq_len
+        cols = tl.arange(0, BLOCK_SIZE)
+        head_mask = cols < head_dim
+        row_offset = row * head_dim
+
+        x = tl.load(x_ptr + row_offset + cols, mask=head_mask, other=0.0).to(tl.float32)
+        variance = tl.sum(x * x, axis=0) / head_dim
+        inv_rms = tl.math.rsqrt(variance + eps)
+        weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+        normalized = x * inv_rms * weight
+
+        rotary_mask = cols < rotary_dim
+        first_half = cols < half_rotary_dim
+        partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+        partner_x = tl.load(
+            x_ptr + row_offset + partner_col,
+            mask=rotary_mask,
+            other=0.0,
+        ).to(tl.float32)
+        partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32)
+        partner_normalized = partner_x * inv_rms * partner_weight
+        rotated = tl.where(first_half, -partner_normalized, partner_normalized)
+
+        table_offset = seq_index * rotary_dim + cols
+        cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+        sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+        rotary_output = normalized * cos + rotated * sin
+        output = tl.where(rotary_mask, rotary_output, normalized)
+        tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+def _validate_inputs(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+) -> tuple[int, int, int, int, int]:
+    for name, tensor in (("x", x), ("weight", weight), ("cos", cos), ("sin", sin)):
+        if not isinstance(tensor, torch.Tensor):
+            raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}")
+
+    if x.ndim != 4:
+        raise ValueError(f"x must have shape (batch, seq, heads, head_dim), got {tuple(x.shape)}")
+    batch, seq_len, num_heads, head_dim = x.shape
+    if min(batch, seq_len, num_heads, head_dim) <= 0:
+        raise ValueError(f"x dimensions must all be positive, got {tuple(x.shape)}")
+    if weight.shape != (head_dim, ):
+        raise ValueError(f"weight must have shape ({head_dim},), got {tuple(weight.shape)}")
+    if cos.ndim != 2:
+        raise ValueError(f"cos must have shape (seq, rotary_dim), got {tuple(cos.shape)}")
+    if sin.shape != cos.shape:
+        raise ValueError(f"sin must match cos shape {tuple(cos.shape)}, got {tuple(sin.shape)}")
+    if cos.shape[0] != seq_len:
+        raise ValueError(f"cos/sin sequence length must be {seq_len}, got {cos.shape[0]}")
+
+    rotary_dim = cos.shape[1]
+    if rotary_dim <= 0:
+        raise ValueError(f"rotary_dim must be positive, got {rotary_dim}")
+    if rotary_dim > head_dim:
+        raise ValueError(f"rotary_dim must not exceed head_dim, got rotary_dim={rotary_dim}, head_dim={head_dim}")
+    if rotary_dim % 2:
+        raise ValueError(f"rotary_dim must be even, got {rotary_dim}")
+
+    if x.dtype not in _SUPPORTED_DTYPES:
+        raise TypeError(f"x dtype must be float16, bfloat16, or float32, got {x.dtype}")
+    for name, tensor in (("weight", weight), ("cos", cos), ("sin", sin)):
+        if tensor.dtype != x.dtype:
+            raise TypeError(f"{name} dtype must match x dtype {x.dtype}, got {tensor.dtype}")
+        if tensor.device != x.device:
+            raise ValueError(f"{name} device must match x device {x.device}, got {tensor.device}")
+
+    if not isinstance(eps, (float, int)) or not math.isfinite(float(eps)) or eps <= 0:
+        raise ValueError(f"eps must be a positive finite number, got {eps!r}")
+    return batch, seq_len, num_heads, head_dim, rotary_dim
+
+
+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.
+    """
+    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")
+    if not cos.is_contiguous() or not sin.is_contiguous():
+        raise ValueError("cos and sin must be contiguous (seq, rotary_dim) tables")
+    if not x.is_cuda:
+        raise RuntimeError("fused_qknorm_rope requires CUDA tensors")
+    if not HAVE_TRITON:
+        raise RuntimeError("fused_qknorm_rope requires Triton")
+    if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in (x, weight, cos, sin)):
+        raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
+
+    flat_x = x.reshape(-1, head_dim).contiguous()
+    flat_out = torch.empty_like(flat_x)
+    block_size = 1 << (head_dim - 1).bit_length()
+    _qknorm_partial_rope_kernel[(flat_x.shape[0], )](
+        flat_out,
+        flat_x,
+        weight,
+        cos,
+        sin,
+        head_dim,
+        rotary_dim,
+        rotary_dim // 2,
+        num_heads,
+        seq_len,
+        eps,
+        BLOCK_SIZE=block_size,
+        num_warps=4,
+    )
+    return flat_out.view(batch, seq_len, num_heads, head_dim)
+
+
+__all__ = ["HAVE_TRITON", "fused_qknorm_rope"]

fastvideo/models/dits/minimax_h3_fusions/swiglu.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3's value-first packed SwiGLU fusion."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAVE_TRITON = True
+except ImportError:  # pragma: no cover - exercised only in environments without Triton
+    triton = None
+    tl = None
+    HAVE_TRITON = False
+
+
+def _validate_input(x: torch.Tensor) -> int:
+    if x.ndim == 0:
+        raise ValueError("MiniMax H3 SwiGLU expects at least one dimension")
+
+    packed_width = x.shape[-1]
+    if packed_width == 0 or packed_width % 2 != 0:
+        raise ValueError(
+            "MiniMax H3 SwiGLU expects a positive even last dimension containing packed (value, gate) halves, "
+            f"got {packed_width}"
+        )
+    if not x.is_floating_point():
+        raise TypeError(f"MiniMax H3 SwiGLU expects a floating-point tensor, got {x.dtype}")
+    return packed_width // 2
+
+
+if HAVE_TRITON:
+
+    @triton.jit
+    def _minimax_h3_swiglu_kernel(
+        out_ptr,
+        x_ptr,
+        ffn_dim,
+        stride_in_row,
+        stride_out_row,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        cols = tl.arange(0, BLOCK_SIZE)
+        mask = cols < ffn_dim
+
+        value = tl.load(x_ptr + row * stride_in_row + cols, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(x_ptr + row * stride_in_row + ffn_dim + cols, mask=mask, other=0.0).to(tl.float32)
+        # Match Sol-Engine: keep the complete SwiGLU expression in FP32 and
+        # convert only the final output store.
+        out = value * (gate * tl.sigmoid(gate))
+        tl.store(out_ptr + row * stride_out_row + cols, out.to(out_ptr.dtype.element_ty), mask=mask)
+
+else:
+    _minimax_h3_swiglu_kernel = None
+
+
+def _num_warps(block_size: int) -> int:
+    if block_size >= 8192:
+        return 16
+    if block_size >= 2048:
+        return 8
+    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.
+    """
+    ffn_dim = _validate_input(x)
+    if not x.is_cuda:
+        raise ValueError("MiniMax H3 fused SwiGLU requires a CUDA tensor")
+    if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise TypeError(f"MiniMax H3 fused SwiGLU supports float16, bfloat16, and float32, got {x.dtype}")
+    if torch.is_grad_enabled() and x.requires_grad:
+        raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd")
+    if _minimax_h3_swiglu_kernel is None:
+        raise RuntimeError("MiniMax H3 fused SwiGLU requires Triton")
+
+    packed_width = x.shape[-1]
+    flat = x.reshape(-1, packed_width).contiguous()
+    output_shape = (*x.shape[:-1], ffn_dim)
+    if flat.shape[0] == 0:
+        return torch.empty(output_shape, dtype=x.dtype, device=x.device)
+
+    out = torch.empty((flat.shape[0], ffn_dim), dtype=x.dtype, device=x.device)
+    block_size = triton.next_power_of_2(ffn_dim)
+    _minimax_h3_swiglu_kernel[(flat.shape[0],)](
+        out,
+        flat,
+        ffn_dim,
+        flat.stride(0),
+        out.stride(0),
+        BLOCK_SIZE=block_size,
+        num_warps=_num_warps(block_size),
+    )
+    return out.view(output_shape)
+
+
+__all__ = ["HAVE_TRITON", "minimax_h3_swiglu"]

fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Focused routing and FA4 integration checks for MiniMax-H3 fusions."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@pytest.mark.parametrize(
+    ("raw", "expected"),
+    [
+        ("", frozenset()),
+        ("0", frozenset()),
+        ("none", frozenset()),
+        ("all", frozenset({"modulate", "qknorm_rope", "swiglu"})),
+        ("1", frozenset({"modulate", "qknorm_rope", "swiglu"})),
+        ("swiglu, modulate", frozenset({"swiglu", "modulate"})),
+    ],
+)
+def test_minimax_h3_fusion_selector(raw: str, expected: frozenset[str]) -> None:
+    from fastvideo.models.dits.minimax_h3 import _enabled_minimax_h3_fusions
+
+    assert _enabled_minimax_h3_fusions(raw) == expected
+
+
+def test_minimax_h3_fusion_selector_rejects_unknown_name() -> None:
+    from fastvideo.models.dits.minimax_h3 import _enabled_minimax_h3_fusions
+
+    with pytest.raises(ValueError, match="Unknown MiniMax H3 fusion"):
+        _enabled_minimax_h3_fusions("swiglu,unknown")
+
+
+def test_swiglu_fusion_stays_on_eager_path_with_grad(monkeypatch: pytest.MonkeyPatch) -> None:
+    import fastvideo.models.dits.minimax_h3 as h3
+
+    def unexpected_kernel(_: torch.Tensor) -> torch.Tensor:
+        raise AssertionError("inference-only fusion ran with grad enabled")
+
+    monkeypatch.setattr(h3, "minimax_h3_swiglu", unexpected_kernel)
+    layer = h3.MiniMaxH3FeedForward(8, 16, fuse_swiglu=True)
+    inputs = torch.randn(2, 3, 8, requires_grad=True)
+    layer(inputs).sum().backward()
+
+    assert inputs.grad is not None
+
+
+def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Exercise the real block wiring without loading any H3 checkpoint."""
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        pytest.skip("BF16 CUDA is required")
+    pytest.importorskip("triton")
+    flash_attn = pytest.importorskip("flash_attn")
+    if "fa4" not in getattr(flash_attn, "__version__", "").lower():
+        pytest.skip("the focused integration test requires the FA4 environment")
+
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
+    monkeypatch.setenv("FASTVIDEO_FA4", "1")
+    monkeypatch.setenv("MASTER_ADDR", "127.0.0.1")
+    monkeypatch.setenv("MASTER_PORT", "29573")
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3RotaryPosEmbed, MiniMaxH3TransformerBlock
+
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    try:
+        kwargs = dict(
+            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.FLASH_ATTN, ),
+            quant_config=None,
+            prefix="minimax_h3.test_block",
+        )
+        previous_default_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(torch.bfloat16)
+        try:
+            eager = MiniMaxH3TransformerBlock(**kwargs)
+            fused = MiniMaxH3TransformerBlock(
+                **kwargs,
+                fuse_modulate=True,
+                fuse_qknorm_rope=True,
+                fuse_swiglu=True,
+            )
+        finally:
+            torch.set_default_dtype(previous_default_dtype)
+
+        with torch.no_grad():
+            for name, parameter in eager.named_parameters():
+                if "norm" in name and name.endswith("weight"):
+                    parameter.fill_(1.0)
+                elif parameter.ndim > 1:
+                    torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
+                else:
+                    parameter.zero_()
+        fused.load_state_dict(eager.state_dict(), strict=True)
+
+        device = torch.device("cuda")
+        eager = eager.to(device=device, dtype=torch.bfloat16).eval()
+        fused = fused.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 = 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):
+            eager_output = eager(**inputs)
+            fused_output = fused(**inputs)
+
+        # Sol-Engine keeps fused intermediates in FP32 registers until their
+        # final BF16 stores, so the opt-in path is close but not bit-identical.
+        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()

fastvideo/tests/transformers/test_minimax_h3_modulation_fusion.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fastvideo.models.dits.minimax_h3_fusions.modulation import (
+    fused_residual_gate_rmsnorm_modulate,
+    fused_rmsnorm_modulate,
+)
+
+
+EPS = 1e-6
+SOL_ENGINE_BF16_TOLERANCE = 3e-2
+
+
+def _chunk_tables(rows: int, hidden_size: int, *, device: torch.device | str = "cpu", dtype=torch.float32):
+    wide = torch.randn(rows, 6 * hidden_size, device=device, dtype=dtype)
+    tables = wide.chunk(6, dim=-1)
+    assert all(table.stride() == (6 * hidden_size, 1) for table in tables)
+    assert all(not table.is_contiguous() for table in tables)
+    return tables
+
+
+def _eager_rmsnorm_modulate(x, weight, scale, shift, index):
+    normed = F.rms_norm(x, (x.shape[-1], ), weight, EPS)
+    return normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index)
+
+
+def _eager_residual_gate_rmsnorm_modulate(residual, branch, gate, weight, scale, shift, index):
+    hidden = residual + gate.index_select(0, index) * branch
+    normed = F.rms_norm(hidden, (hidden.shape[-1], ), weight, EPS)
+    modulated = normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index)
+    return hidden, modulated
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU")
+def test_bf16_sol_engine_fusions_match_production_eager_within_tolerance():
+    """Sol-Engine keeps fused intermediates in FP32 until its output stores."""
+    pytest.importorskip("triton")
+    torch.manual_seed(2)
+    device = torch.device("cuda")
+    batch, sequence_length, hidden_size, table_rows = 2, 9, 5376, 6
+    residual = torch.randn(batch, sequence_length, hidden_size, device=device, dtype=torch.bfloat16)
+    branch = torch.randn_like(residual)
+    weight = torch.randn(hidden_size, device=device, dtype=torch.bfloat16)
+    shift, scale, gate, shift_mlp, scale_mlp, _ = _chunk_tables(
+        table_rows,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    index = torch.tensor([5, 0, 4, 1, 3, 2, 5, 1, 0], device=device, dtype=torch.int64)
+
+    expected_norm1 = _eager_rmsnorm_modulate(residual, weight, scale, shift, index)
+    actual_norm1 = fused_rmsnorm_modulate(residual, weight, scale, shift, index, EPS)
+    expected_hidden, expected_norm2 = _eager_residual_gate_rmsnorm_modulate(
+        residual,
+        branch,
+        gate,
+        weight,
+        scale_mlp,
+        shift_mlp,
+        index,
+    )
+    actual_hidden, actual_norm2 = fused_residual_gate_rmsnorm_modulate(
+        residual,
+        branch,
+        gate,
+        weight,
+        scale_mlp,
+        shift_mlp,
+        index,
+        EPS,
+    )
+
+    # Production eager materializes BF16 after each PyTorch operator. The
+    # single-kernel Sol-Engine path deliberately removes those round points.
+    torch.testing.assert_close(
+        actual_norm1,
+        expected_norm1,
+        rtol=SOL_ENGINE_BF16_TOLERANCE,
+        atol=SOL_ENGINE_BF16_TOLERANCE,
+    )
+    torch.testing.assert_close(
+        actual_hidden,
+        expected_hidden,
+        rtol=SOL_ENGINE_BF16_TOLERANCE,
+        atol=SOL_ENGINE_BF16_TOLERANCE,
+    )
+    torch.testing.assert_close(
+        actual_norm2,
+        expected_norm2,
+        rtol=SOL_ENGINE_BF16_TOLERANCE,
+        atol=SOL_ENGINE_BF16_TOLERANCE,
+    )
+
+
+@pytest.mark.parametrize(
+    ("mutate", "error", "match"),
+    [
+        (lambda args: args | {"x": args["x"][0, 0]}, ValueError, "shape"),
+        (lambda args: args | {"weight": args["weight"][:-1]}, ValueError, "weight"),
+        (lambda args: args | {"index": args["index"][:-1]}, ValueError, "index"),
+        (lambda args: args | {"index": args["index"].float()}, TypeError, "index"),
+        (lambda args: args | {"eps": 0.0}, ValueError, "eps"),
+    ],
+)
+def test_fusion_rejects_invalid_contracts(mutate, error, match):
+    args = {
+        "x": torch.randn(2, 3, 8),
+        "weight": torch.randn(8),
+        "scale": torch.randn(4, 8),
+        "shift": torch.randn(4, 8),
+        "index": torch.tensor([0, 3, 1]),
+        "eps": EPS,
+    }
+    with pytest.raises(error, match=match):
+        fused_rmsnorm_modulate(**mutate(args))
+
+
+def test_residual_fusion_rejects_mismatched_branch():
+    residual = torch.randn(2, 3, 8)
+    with pytest.raises(ValueError, match="branch"):
+        fused_residual_gate_rmsnorm_modulate(
+            residual,
+            torch.randn(2, 2, 8),
+            torch.randn(4, 8),
+            torch.randn(8),
+            torch.randn(4, 8),
+            torch.randn(4, 8),
+            torch.tensor([0, 1, 2]),
+            EPS,
+        )
+
+
+def test_triton_wrappers_fail_explicitly_on_cpu():
+    x = torch.randn(1, 2, 8)
+    branch = torch.randn_like(x)
+    weight = torch.randn(8)
+    gate = torch.randn(3, 8)
+    scale = torch.randn(3, 8)
+    shift = torch.randn(3, 8)
+    index = torch.tensor([0, 2])
+
+    with pytest.raises(RuntimeError, match="Triton|CUDA"):
+        fused_rmsnorm_modulate(x, weight, scale, shift, index, EPS)
+    with pytest.raises(RuntimeError, match="Triton|CUDA"):
+        fused_residual_gate_rmsnorm_modulate(x, branch, gate, weight, scale, shift, index, EPS)
+
+
+def test_triton_wrappers_reject_autograd_before_backend_check():
+    x = torch.randn(1, 2, 8)
+    branch = torch.randn_like(x, requires_grad=True)
+    weight = torch.randn(8, requires_grad=True)
+    gate = torch.randn(3, 8)
+    scale = torch.randn(3, 8)
+    shift = torch.randn(3, 8)
+    index = torch.tensor([0, 2])
+
+    with pytest.raises(RuntimeError, match="forward-only"):
+        fused_rmsnorm_modulate(x, weight, scale, shift, index, EPS)
+    with pytest.raises(RuntimeError, match="forward-only"):
+        fused_residual_gate_rmsnorm_modulate(
+            x,
+            branch,
+            gate,
+            weight.detach(),
+            scale,
+            shift,
+            index,
+            EPS,
+        )

fastvideo/tests/transformers/test_minimax_h3_qknorm_rope_fusion.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention
+from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import (
+    HAVE_TRITON,
+    fused_qknorm_rope,
+)
+
+
+def _rotary_tables(
+    seq_len: int,
+    rotary_dim: int,
+    *,
+    dtype: torch.dtype,
+    device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    angles = torch.randn(seq_len, rotary_dim // 2, dtype=torch.float32, device=device)
+    angles = torch.cat((angles, angles), dim=-1)
+    return angles.cos().to(dtype), angles.sin().to(dtype)
+
+
+def _eager_qknorm_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    normalized = F.rms_norm(x, (x.shape[-1], ), weight, eps)
+    return MiniMaxH3Attention._apply_rotary_emb(normalized, (cos, sin))
+
+
+def test_fused_qknorm_rope_rejects_invalid_rotary_dim() -> None:
+    x = torch.randn(2, 3, 4, 128)
+    weight = torch.ones(128)
+
+    cos, sin = _rotary_tables(3, 130, dtype=x.dtype, device=x.device)
+    with pytest.raises(ValueError, match="rotary_dim must not exceed head_dim"):
+        fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+    cos = torch.randn(3, 95)
+    sin = torch.randn_like(cos)
+    with pytest.raises(ValueError, match="rotary_dim must be even"):
+        fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+
+def test_fused_qknorm_rope_rejects_shape_mismatches() -> None:
+    x = torch.randn(2, 3, 4, 128)
+    weight = torch.ones(128)
+    cos, sin = _rotary_tables(3, 96, dtype=x.dtype, device=x.device)
+
+    with pytest.raises(ValueError, match="x must have shape"):
+        fused_qknorm_rope(x[0], weight, cos, sin, 1e-6)
+    with pytest.raises(ValueError, match="weight must have shape"):
+        fused_qknorm_rope(x, weight[:-1], cos, sin, 1e-6)
+    with pytest.raises(ValueError, match="sequence length"):
+        fused_qknorm_rope(x, weight, cos[:-1], sin[:-1], 1e-6)
+    with pytest.raises(ValueError, match="sin must match cos shape"):
+        fused_qknorm_rope(x, weight, cos, sin[:, :-2], 1e-6)
+
+
+@pytest.mark.parametrize("noncontiguous_input", ["weight", "cos", "sin"])
+def test_fused_qknorm_rope_rejects_noncontiguous_linear_inputs(noncontiguous_input: str) -> None:
+    x = torch.randn(2, 3, 4, 128)
+    weight = torch.ones(256)[::2]
+    cos = torch.randn(3, 192)[:, ::2]
+    sin = torch.randn(3, 192)[:, ::2]
+    assert not weight.is_contiguous()
+    assert not cos.is_contiguous()
+    assert not sin.is_contiguous()
+
+    if noncontiguous_input != "weight":
+        weight = weight.contiguous()
+    if noncontiguous_input != "cos":
+        cos = cos.contiguous()
+    if noncontiguous_input != "sin":
+        sin = sin.contiguous()
+
+    with pytest.raises(ValueError, match="must be contiguous"):
+        fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+
+def test_fused_qknorm_rope_requires_matching_precast_dtype() -> None:
+    x = torch.randn(2, 3, 4, 128, dtype=torch.float32)
+    weight = torch.ones(128, dtype=torch.float32)
+    cos, sin = _rotary_tables(3, 96, dtype=torch.bfloat16, device=x.device)
+
+    with pytest.raises(TypeError, match="cos dtype must match x dtype"):
+        fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+
+def test_fused_qknorm_rope_does_not_accept_missing_rotary_tables() -> None:
+    x = torch.randn(2, 3, 4, 128)
+    weight = torch.ones(128)
+
+    with pytest.raises(TypeError, match="cos must be a torch.Tensor"):
+        fused_qknorm_rope(x, weight, None, None, 1e-6)  # type: ignore[arg-type]
+
+
+def test_fused_qknorm_rope_requires_cuda() -> None:
+    x = torch.randn(2, 3, 4, 128)
+    weight = torch.ones(128)
+    cos, sin = _rotary_tables(3, 96, dtype=x.dtype, device=x.device)
+
+    with pytest.raises(RuntimeError, match="requires CUDA"):
+        fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+
+@pytest.mark.parametrize(
+    "rotary_dim,shape,use_input_view",
+    [
+        pytest.param(96, (2, 11, 5, 128), False, id="partial-96-batch2-seq11-heads5"),
+        pytest.param(128, (2, 7, 3, 128), True, id="full-128-noncontiguous-input-view"),
+    ],
+)
+def test_fused_qknorm_rope_matches_eager_bf16_cuda(
+    rotary_dim: int,
+    shape: tuple[int, ...],
+    use_input_view: bool,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the Triton fusion")
+    if not HAVE_TRITON:
+        pytest.skip("Triton is required for the fusion")
+
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    if use_input_view:
+        x = torch.randn(*shape[:-1], shape[-1] * 2, dtype=torch.bfloat16, device=device)[..., ::2]
+        assert not x.is_contiguous()
+    else:
+        x = torch.randn(shape, dtype=torch.bfloat16, device=device)
+    weight = (1.0 + 0.05 * torch.randn(shape[-1], dtype=torch.bfloat16, device=device)).contiguous()
+    cos, sin = _rotary_tables(shape[1], rotary_dim, dtype=x.dtype, device=device)
+
+    with torch.inference_mode():
+        actual = fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+        expected = _eager_qknorm_rope(x, weight, cos, sin, 1e-6)
+
+    # The fused kernel keeps RMSNorm and both RoPE products in FP32 registers
+    # until its final BF16 store. Eager materializes BF16 intermediates, and
+    # PyTorch/Triton reductions need not use the same summation order.
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+    assert actual.shape == x.shape
+    assert actual.dtype == x.dtype
+    assert actual.is_contiguous()
+
+
+def test_fused_qknorm_rope_matches_eager_beyond_int32_element_count() -> None:
+    """Regression: kernel row offsets must be int64.
+
+    With int32 offsets, ``row * head_dim`` wraps once the flattened input
+    crosses 2**31 elements and the kernel reads/writes out of bounds (CUDA
+    illegal memory access). ``(1, 8_500_000, 2, 128)`` is 2.176e9 elements,
+    just past the boundary; for H3's 56 heads x 128 head_dim the equivalent
+    is ``batch*seq >= 299_593`` tokens per rank, reachable at SP=1.
+
+    GPU assumption: needs ~16 GiB free CUDA memory (input + output at
+    bf16 plus the fp32 rotary-table construction); skips below 20 GiB.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the Triton fusion")
+    if not HAVE_TRITON:
+        pytest.skip("Triton is required for the fusion")
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < 20 * 1024**3:
+        pytest.skip("needs ~20 GiB free GPU memory for a >2**31-element input")
+
+    heads, head_dim, rotary_dim = 2, 128, 96
+    seq_len = 8_500_000
+    assert seq_len * heads * head_dim > 2**31
+
+    torch.manual_seed(9)
+    device = torch.device("cuda")
+    x = torch.randn(1, seq_len, heads, head_dim, dtype=torch.bfloat16, device=device)
+    weight = (1.0 + 0.05 * torch.randn(head_dim, dtype=torch.bfloat16, device=device)).contiguous()
+    cos, sin = _rotary_tables(seq_len, rotary_dim, dtype=x.dtype, device=device)
+
+    with torch.inference_mode():
+        fused = fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+        torch.cuda.synchronize()
+        # Compare only head/tail slices against eager: a full-tensor eager
+        # reference would double peak memory for no extra coverage, and the
+        # tail rows are exactly the ones an int32 wrap corrupts first.
+        expected_head = _eager_qknorm_rope(x[:, :8], weight, cos[:8], sin[:8], 1e-6)
+        expected_tail = _eager_qknorm_rope(x[:, -8:], weight, cos[-8:], sin[-8:], 1e-6)
+        torch.testing.assert_close(fused[:, :8], expected_head, atol=2e-2, rtol=2e-2)
+        torch.testing.assert_close(fused[:, -8:], expected_tail, atol=2e-2, rtol=2e-2)

fastvideo/tests/transformers/test_minimax_h3_swiglu_fusion.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Focused tests for MiniMax H3's value-first packed SwiGLU fusion."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from fastvideo.models.dits.minimax_h3_fusions.swiglu import minimax_h3_swiglu
+
+
+def _require_bf16_triton_cuda() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MiniMax H3 fused SwiGLU requires CUDA")
+    if not torch.cuda.is_bf16_supported():
+        pytest.skip("MiniMax H3 fused SwiGLU parity requires BF16 support")
+    pytest.importorskip("triton", reason="MiniMax H3 fused SwiGLU requires Triton")
+
+
+def _assert_bf16_parity(x: torch.Tensor) -> None:
+    value, gate = x.chunk(2, dim=-1)
+    expected = value * F.silu(gate)
+    actual = minimax_h3_swiglu(x)
+    assert actual.shape == (*x.shape[:-1], x.shape[-1] // 2)
+    assert actual.dtype == x.dtype
+    assert actual.device == x.device
+    # Sol-Engine keeps the full SwiGLU expression in FP32 until the output
+    # store, while eager F.silu materializes a BF16 intermediate.
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.parametrize("last_dim", [0, 7])
+def test_minimax_h3_swiglu_rejects_nonpositive_or_odd_last_dimension(last_dim: int) -> None:
+    x = torch.empty((2, last_dim), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="positive even last dimension"):
+        minimax_h3_swiglu(x)
+
+
+def test_minimax_h3_swiglu_strict_wrapper_rejects_cpu() -> None:
+    with pytest.raises(ValueError, match="requires a CUDA tensor"):
+        minimax_h3_swiglu(torch.randn(2, 8))
+
+
+@pytest.mark.gpu
+def test_minimax_h3_swiglu_multidimensional_bf16_gpu_parity() -> None:
+    _require_bf16_triton_cuda()
+    torch.manual_seed(0)
+    x = torch.randn((2, 3, 5, 66), device="cuda", dtype=torch.bfloat16)
+
+    _assert_bf16_parity(x)
+
+
+@pytest.mark.gpu
+def test_minimax_h3_swiglu_noncontiguous_bf16_gpu_parity() -> None:
+    _require_bf16_triton_cuda()
+    torch.manual_seed(1)
+    storage = torch.randn((2, 3, 148), device="cuda", dtype=torch.bfloat16)
+    x = storage[..., ::2]
+    assert not x.is_contiguous()
+
+    _assert_bf16_parity(x)
+
+
+@pytest.mark.gpu
+def test_minimax_h3_swiglu_real_ffn_dim_bf16_gpu() -> None:
+    _require_bf16_triton_cuda()
+    torch.manual_seed(2)
+    ffn_dim = 14336
+    x = torch.randn((1, 2 * ffn_dim), device="cuda", dtype=torch.bfloat16)
+
+    _assert_bf16_parity(x)