[misc]: H3 MK tier: drop narration comments, move tests under the collected transformers lane

docs/contributing/env_vars.md

@@ -208,7 +208,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY`                   | str                                                                | unset                                                 | performance | Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather.                                                                                                                                                                                                                                                                                                                                                |
 | `FASTVIDEO_H3_ENCODER_SPLIT`                               | bool                                                               | `0`                                                   | distributed | MiniMax-H3 runs the Qwen3-VL text encoder on a dedicated group of nodes, so the denoising ranks never load it. Same as FastVideoArgs.h3_encoder_split=True.                                                                                                                                                                                                                                                                                 |
 | `FASTVIDEO_H3_ENCODER_NODES`                               | int                                                                | `1`                                                   | distributed | Number of leading nodes that FASTVIDEO_H3_ENCODER_SPLIT dedicates to the MiniMax-H3 text encoder. Values below 1 count as 1.                                                                                                                                                                                                                                                                                                                |
-| `FASTVIDEO_MINIMAX_H3_FUSIONS`                             | str                                                                | `""`                                                  | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.                                                                                                                                                                                                                                                                             |
+| `FASTVIDEO_MINIMAX_H3_FUSIONS`                             | str                                                                | `""`                                                  | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. The whole-tier value mk selects the MotionKernel kernels instead of the Sol-Engine ones. Empty, 0, or none keeps the eager implementation.                                                                                                                                                                                    |
 | `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS`                       | str                                                                | `""`                                                  | performance | MiniMax-H3 inference-only Triton kernels that reproduce the eager output bit for bit: all, 1, or a comma-separated subset of rope,modulate,swiglu. An op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keeps the fusion. Empty, 0, or none keeps the eager implementation.                                                                                                                                                                |
 | `FASTVIDEO_FSDP2_AUTOWRAP`                                 | bool                                                               | `0`                                                   | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation.                                                                                                                                                                                                                                                                                                                |
 | `FASTVIDEO_FSDP2_MIN_PARAMS`                               | int                                                                | `10000000`                                            | performance | Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards.                                                                                                                                                                                                                                                                                                                                                                   |

fastvideo/attention/layer.py

@@ -136,16 +136,44 @@ class DistributedAttention(nn.Module):
         """
         # Check input shapes
         assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
-        batch_size, _, num_heads, _ = q.shape
+        # Stack QKV
+        qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]
+        return self.forward_stacked(
+            qkv,
+            original_seq_len=original_seq_len,
+            replicated_q=replicated_q,
+            replicated_k=replicated_k,
+            replicated_v=replicated_v,
+            freqs_cis=freqs_cis,
+            attn_metadata=attn_metadata,
+        )
+
+    def forward_stacked(
+        self,
+        qkv: torch.Tensor,
+        original_seq_len: int | None = None,
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Attention over a pre-stacked ``[3*batch, seq_len, num_heads, head_dim]`` QKV buffer.
+
+        Identical to :meth:`forward` after its ``torch.cat``, so a fused QKV
+        prologue that already wrote the stack (e.g. MiniMax H3's MotionKernel
+        prologue) can skip that copy. Sequence-parallel redistribution,
+        ``freqs_cis``, replicated tokens and backend pre/post-processing are
+        handled exactly as in :meth:`forward`.
+        """
+        stacked, _, num_heads, _ = qkv.shape
+        batch_size = stacked // 3
         local_rank = get_sp_parallel_rank()
         world_size = get_sp_world_size()
 
         forward_context: ForwardContext = get_forward_context()
         ctx_attn_metadata = attn_metadata if attn_metadata is not None else forward_context.attn_metadata
 
-        # Stack QKV
-        qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]
-
         # Redistribute heads across sequence dimension
         qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
 

fastvideo/envs.py

@@ -427,7 +427,8 @@ FASTVIDEO_MINIMAX_H3_FUSIONS = EnvStr(
     "",
     category="performance",
     doc="MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of "
-    "modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.")
+    "modulate,qknorm_rope,swiglu. The whole-tier value mk selects the MotionKernel kernels instead of "
+    "the Sol-Engine ones. Empty, 0, or none keeps the eager implementation.")
 FASTVIDEO_MINIMAX_H3_EXACT_KERNELS = EnvStr(
     "",
     category="performance",

fastvideo/models/dits/minimax_h3.py

@@ -38,6 +38,7 @@ from fastvideo.models.dits.base import BaseDiT
 from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp,
                                                        vsa_fp4_requested, vsa_tile_first_attention)
 from fastvideo.models.dits.minimax_h3_fusions import exact as exact_kernels
+from fastvideo.models.dits.minimax_h3_fusions import mk as mk_kernels
 from fastvideo.models.dits.minimax_h3_fusions import (
     HAVE_TRITON,
     fused_qknorm_rope,
@@ -52,9 +53,13 @@ from fastvideo.utils import get_compute_dtype
 logger = init_logger(__name__)
 
 MINIMAX_H3_MODALITY_NUM = 3
+#: eps baked into the MotionKernel H3 kernels; their call sites require it.
+_MK_NORM_EPS = 1e-5
 _CFG = MiniMaxH3Config()
 _MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"})
 _MINIMAX_H3_EXACT_KERNEL_NAMES = frozenset({"rope", "modulate", "swiglu"})
+#: Whole-tier opt-in that replaces the Sol kernels with the MotionKernel ones.
+_MINIMAX_H3_MK_MODE = "mk"
 
 
 def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]:
@@ -65,7 +70,14 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]:
         return frozenset()
     if normalized in {"1", "all"}:
         return _MINIMAX_H3_FUSION_NAMES
+    if normalized == _MINIMAX_H3_MK_MODE:
+        # ``mk`` is a whole-tier mode: every fusion is on, and each fused call
+        # site dispatches to the MotionKernel kernel instead of the Sol one.
+        return _MINIMAX_H3_FUSION_NAMES | {_MINIMAX_H3_MK_MODE}
     enabled = frozenset(item.strip() for item in normalized.split(",") if item.strip())
+    if _MINIMAX_H3_MK_MODE in enabled:
+        raise ValueError(f"'{_MINIMAX_H3_MK_MODE}' is a whole-tier mode and cannot be combined with individual "
+                         f"fusions: {sorted(enabled - {_MINIMAX_H3_MK_MODE})}")
     unknown = enabled - _MINIMAX_H3_FUSION_NAMES
     if unknown:
         supported = ",".join(sorted(_MINIMAX_H3_FUSION_NAMES))
@@ -139,6 +151,7 @@ class MiniMaxH3FeedForward(nn.Module):
         prefix: str = "",
         fuse_swiglu: bool = False,
         exact_swiglu: bool = False,
+        mk_swiglu: bool = False,
     ) -> None:
         super().__init__()
         self.fc_in = ReplicatedLinear(
@@ -157,6 +170,7 @@ class MiniMaxH3FeedForward(nn.Module):
         )
         self.fuse_swiglu = fuse_swiglu
         self.exact_swiglu = exact_swiglu
+        self.mk_swiglu = mk_swiglu
         self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance(
             self.fc_out.quant_method, MXFP8QuantizeMethod)
         # Inference-only token chunking: the 2 * ffn_dim intermediate is ~5.3x the block input
@@ -181,7 +195,10 @@ class MiniMaxH3FeedForward(nn.Module):
             return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out)
         hidden_states, _ = self.fc_in(hidden_states)
         if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states):
-            hidden_states = minimax_h3_swiglu(hidden_states)
+            if self.mk_swiglu and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling():
+                hidden_states = mk_kernels.swiglu(hidden_states)
+            else:
+                hidden_states = minimax_h3_swiglu(hidden_states)
         elif self.exact_swiglu and exact_kernels.supports_rowwise(hidden_states):
             hidden_states = exact_kernels.swiglu(hidden_states)
         else:
@@ -206,6 +223,7 @@ class MiniMaxH3Attention(nn.Module):
         fuse_qknorm_rope: bool = False,
         fa4_packed_varlen: bool = False,
         exact_rope: bool = False,
+        mk_qkv: bool = False,
         vsa_fp4: dict[str, Any] | None = None,
     ) -> None:
         super().__init__()
@@ -244,6 +262,7 @@ class MiniMaxH3Attention(nn.Module):
         )
         self.fuse_qknorm_rope = fuse_qknorm_rope
         self.exact_rope = exact_rope
+        self.mk_qkv = mk_qkv
         # 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
@@ -401,6 +420,22 @@ class MiniMaxH3Attention(nn.Module):
             cos, sin = rotary_emb
             cos = cos.to(query.dtype)
             sin = sin.to(query.dtype)
+            if (self.mk_qkv and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling()
+                    and self.norm_q.eps == _MK_NORM_EPS
+                    and self.norm_k.eps == _MK_NORM_EPS and type(self.distributed_attention) is DistributedAttention
+                    and not (self.to_gate_compress is not None and self._gate_active())):
+                # One launch replaces the two Sol QK-norm/RoPE kernels and the
+                # ``torch.cat`` stack: the prologue writes a contiguous
+                # ``[3*B, S, H, D]`` buffer that ``forward_stacked`` feeds to the
+                # backend directly, so q/k/v are never copied again.
+                qkv = mk_kernels.qkv_prologue(query, key, value, self.norm_q.weight, self.norm_k.weight, cos, sin)
+                with STAGES.span("attention"):
+                    hidden_states, _ = self.distributed_attention.forward_stacked(
+                        qkv, original_seq_len=original_seq_len, attn_metadata=attn_metadata)
+                hidden_states = hidden_states.flatten(2, 3).type_as(qkv)
+                with STAGES.span("out_proj"):
+                    hidden_states, _ = self.to_out(hidden_states)
+                return hidden_states
             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:
@@ -658,6 +693,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
         fa4_packed_varlen: bool = False,
         exact_kernels: frozenset[str] = frozenset(),
         vsa_fp4: dict[str, Any] | None = None,
+        mk_tier: bool = False,
     ) -> None:
         super().__init__()
         self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
@@ -673,6 +709,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             fa4_packed_varlen=fa4_packed_varlen,
             exact_rope="rope" in exact_kernels,
             vsa_fp4=vsa_fp4,
+            mk_qkv=mk_tier,
         )
         self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
         self.ff = MiniMaxH3FeedForward(
@@ -682,6 +719,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             prefix=f"{prefix}.ff",
             fuse_swiglu=fuse_swiglu,
             exact_swiglu="swiglu" in exact_kernels,
+            mk_swiglu=mk_tier,
         )
         self.adaln_proj = MiniMaxH3AdaLayerNormModulation(
             time_embed_dim,
@@ -692,6 +730,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
         )
         self.fuse_modulate = fuse_modulate
         self.exact_modulate = "modulate" in exact_kernels
+        self.mk_tier = mk_tier
 
     def forward(
         self,
@@ -711,15 +750,20 @@ class MiniMaxH3TransformerBlock(nn.Module):
                               and exact_kernels.supports_rowwise(hidden_states, shift_msa, scale_msa, gate_msa,
                                                                  shift_mlp, scale_mlp, gate_mlp))
         if use_modulate_fusion:
-            with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
-                norm_hidden_states = fused_rmsnorm_modulate(
-                    hidden_states,
-                    self.norm1.weight,
-                    scale_msa,
-                    shift_msa,
-                    adaln_indices,
-                    self.norm1.eps,
-                )
+            if (self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling()
+                    and self.norm1.eps == _MK_NORM_EPS):
+                norm_hidden_states = mk_kernels.rmsnorm_modulate(hidden_states, self.norm1.weight, scale_msa,
+                                                                 shift_msa, adaln_indices)
+            else:
+                with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
+                    norm_hidden_states = fused_rmsnorm_modulate(
+                        hidden_states,
+                        self.norm1.weight,
+                        scale_msa,
+                        shift_msa,
+                        adaln_indices,
+                        self.norm1.eps,
+                    )
         elif use_exact_modulate:
             with nvtx_range("minimax_h3.transformer_block.exact_modulate"):
                 norm_hidden_states = exact_kernels.modulate(self.norm1(hidden_states), scale_msa, shift_msa,
@@ -735,17 +779,23 @@ class MiniMaxH3TransformerBlock(nn.Module):
         # overlaps three full-width activations during the residual fusion.
         del norm_hidden_states
         if use_modulate_fusion:
-            with nvtx_range("minimax_h3.transformer_block.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,
-                )
+            if (self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling()
+                    and self.norm2.eps == _MK_NORM_EPS):
+                hidden_states, norm_hidden_states = mk_kernels.residual_gate_rmsnorm_modulate(
+                    hidden_states, attention_output, gate_msa, self.norm2.weight, scale_mlp, shift_mlp,
+                    adaln_indices)
+            else:
+                with nvtx_range("minimax_h3.transformer_block.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,
+                    )
         elif use_exact_modulate:
             with nvtx_range("minimax_h3.transformer_block.exact_modulate"):
                 hidden_states = exact_kernels.gate_residual(hidden_states, gate_msa, attention_output, adaln_indices)
@@ -761,6 +811,8 @@ class MiniMaxH3TransformerBlock(nn.Module):
         with nvtx_range("minimax_h3.transformer_block.feed_forward"), STAGES.span("feed_forward"):
             feed_forward_output = self.ff(norm_hidden_states)
         if use_modulate_fusion and not torch.compiler.is_compiling():
+            if self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling():
+                return mk_kernels.gated_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices)
             return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output)
         if use_exact_modulate:
             return exact_kernels.gate_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices)
@@ -821,6 +873,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 logger.warning(
                     "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; "
                     "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions)))
+        self.use_mk_kernels = _MINIMAX_H3_MK_MODE in self.enabled_fusions
+        if self.use_mk_kernels and not mk_kernels.HAVE_TRITON:
+            logger.warning("FASTVIDEO_MINIMAX_H3_FUSIONS=mk requested but the MotionKernel kernels are "
+                           "unavailable (Triton missing); every forward stays on the eager path.")
         self.enabled_exact_kernels = _enabled_minimax_h3_exact_kernels()
         if self.enabled_exact_kernels:
             if exact_kernels.HAVE_TRITON:
@@ -935,6 +991,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(),
                 exact_kernels=self.enabled_exact_kernels,
                 vsa_fp4=config.vsa_fp4,
+                mk_tier=self.use_mk_kernels,
             ) for index in range(arch.num_layers)
         ])
         self.norm_out = MiniMaxH3AdaLayerNormOut(

fastvideo/models/dits/minimax_h3_fusions/mk/__init__.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernels (opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``).
+
+These are the MotionKernel agent-discovered Triton kernels for the MiniMax-H3
+inference blocks, kept separate from the Sol-Engine kernels in the parent
+package so the default path stays byte-for-byte identical. Each kernel is a
+drop-in replacement for one Sol fusion or eager expression:
+
+============================================  ======================================
+MotionKernel                                   replaces
+============================================  ======================================
+``qkv_prologue``                              2x ``fused_qknorm_rope`` + the
+                                              ``torch.cat`` QKV stack
+``rmsnorm_modulate``                          ``fused_rmsnorm_modulate``
+``residual_gate_rmsnorm_modulate``            ``fused_residual_gate_rmsnorm_modulate``
+``swiglu``                                    ``minimax_h3_swiglu``
+``gated_residual``                            the compiled ``_gated_residual``
+============================================  ======================================
+
+All five run in fp32 internally with a single rounding at the bf16/fp16 store
+and require the production H3 norm eps of ``1e-5`` (enforced at the call sites).
+"""
+
+try:
+    import triton  # noqa: F401  (each module imports it directly)
+
+    from .gated_residual import gated_residual
+    from .qkv_prologue import qkv_prologue
+    from .residual_gate_rmsnorm_modulate import residual_gate_rmsnorm_modulate
+    from .rmsnorm_modulate import rmsnorm_modulate
+    from .swiglu import swiglu
+
+    HAVE_TRITON = True
+except ImportError:  # pragma: no cover - environments without Triton
+    gated_residual = None
+    qkv_prologue = None
+    residual_gate_rmsnorm_modulate = None
+    rmsnorm_modulate = None
+    swiglu = None
+    HAVE_TRITON = False
+
+__all__ = [
+    "HAVE_TRITON",
+    "gated_residual",
+    "qkv_prologue",
+    "residual_gate_rmsnorm_modulate",
+    "rmsnorm_modulate",
+    "swiglu",
+]

fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernel: Fused H3 row-indexed gated residual (hidden + gate[index] * branch).
+
+Found with MotionKernel (agent-driven kernel search). Tuned on GB200 / sm_100 for the
+MiniMax-H3 inference shapes; opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``.
+
+Triton kernel for MiniMax H3's post-FFN row-indexed gated residual.
+
+``out = hidden + gate[index] * branch`` with the math in fp32 and a single
+bf16 rounding at the store.
+
+Design notes (GB200, sm_100a):
+  * The op moves 3 activation-sized streams (2 reads + 1 write); it is purely
+    bandwidth bound.  A flat 1-D grid over ``tokens * hidden`` gives fully
+    coalesced, 16-byte-vectorized accesses on all three streams with no lane
+    padding (the 5376-wide rows are not padded to a power of two).
+  * Split full-block / tail code paths.  Every non-tail block runs completely
+    unmasked, so the compiler emits unpredicated 128-bit loads/stores; only the
+    final block pays for predicates.  Measured 175.6 -> 173.6 us weighted over
+    the five production shapes (~1.1%), consistently on every shape.
+  * The gathered AdaLN gate row is tiny (<= 12 rows) so it stays resident in
+    L2; the per-element row/column split needed to address it costs only ALU,
+    of which this kernel uses ~1% of peak.
+  * 32-bit offsets (all production shapes fit) halve index register pressure.
+  * ``evict_last`` on the three loads is worth ~2%: it keeps the streamed
+    lines in L2 long enough for the DRAM scheduler to coalesce the
+    read/write turnaround (measured 269 -> 266 us at 59k tokens).
+  * Roofline: a minimal unmasked 2-read/1-write Triton ``c = a + b`` kernel
+    (same bytes, no gather) runs at 267.7 us / 7.11 TB/s at the largest shape;
+    this kernel measures 263.2 us / 7.24 TB/s, i.e. it is at (slightly above)
+    the measured 3-stream streaming roofline, so the body is not tunable
+    further.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _h3_gated_residual_kernel(
+    out_ptr,
+    hidden_ptr,
+    branch_ptr,
+    gate_ptr,
+    index_ptr,
+    n_elems,
+    n_index,
+    D: tl.constexpr,
+    GS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    NFULL: tl.constexpr,
+    WIDE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if WIDE:
+        off = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        mask = off < n_elems
+        row = off // D
+        col = off - row * D
+        table_row = tl.load(index_ptr + row % n_index, mask=mask, other=0).to(tl.int32)
+        hidden = tl.load(hidden_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32)
+        branch = tl.load(branch_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32)
+        gate = tl.load(gate_ptr + table_row * GS + col, mask=mask, eviction_policy="evict_last").to(tl.float32)
+        tl.store(
+            out_ptr + off,
+            (hidden + gate * branch).to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
+    else:
+        base = pid * BLOCK
+        if base + BLOCK <= NFULL:
+            off = base + tl.arange(0, BLOCK)
+            row = off // D
+            col = off - row * D
+            table_row = tl.load(index_ptr + row % n_index).to(tl.int32)
+            hidden = tl.load(hidden_ptr + off, eviction_policy="evict_last").to(tl.float32)
+            branch = tl.load(branch_ptr + off, eviction_policy="evict_last").to(tl.float32)
+            gate = tl.load(gate_ptr + table_row * GS + col, eviction_policy="evict_last").to(tl.float32)
+            tl.store(out_ptr + off, (hidden + gate * branch).to(out_ptr.dtype.element_ty))
+        else:
+            off = base + tl.arange(0, BLOCK)
+            mask = off < n_elems
+            row = off // D
+            col = off - row * D
+            table_row = tl.load(index_ptr + row % n_index, mask=mask, other=0).to(tl.int32)
+            hidden = tl.load(hidden_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32)
+            branch = tl.load(branch_ptr + off, mask=mask, eviction_policy="evict_last").to(tl.float32)
+            gate = tl.load(
+                gate_ptr + table_row * GS + col,
+                mask=mask,
+                eviction_policy="evict_last",
+            ).to(tl.float32)
+            tl.store(
+                out_ptr + off,
+                (hidden + gate * branch).to(out_ptr.dtype.element_ty),
+                mask=mask,
+            )
+
+
+def gated_residual(
+    hidden: torch.Tensor,
+    gate: torch.Tensor,
+    branch: torch.Tensor,
+    index: torch.Tensor,
+) -> torch.Tensor:
+    """``hidden + gate[index] * branch`` in one pass."""
+    if not hidden.is_cuda:
+        raise ValueError("h3_gated_residual requires CUDA tensors")
+    if branch.shape != hidden.shape:
+        raise ValueError("branch must match hidden")
+    hidden_size = hidden.shape[-1]
+    if index.ndim != 1 or index.numel() != hidden.shape[-2]:
+        raise ValueError("index must hold one table row per token")
+    if gate.shape[-1] != hidden_size:
+        raise ValueError("gate table must match the hidden width")
+    flat_hidden = hidden.reshape(-1, hidden_size)
+    flat_branch = branch.reshape(-1, hidden_size)
+    if not flat_hidden.is_contiguous():
+        flat_hidden = flat_hidden.contiguous()
+    if not flat_branch.is_contiguous():
+        flat_branch = flat_branch.contiguous()
+    if gate.stride(-1) != 1:
+        gate = gate.contiguous()
+    gate_row_stride = gate.stride(0)
+    index = index.contiguous()
+    out = torch.empty_like(flat_hidden)
+
+    n_elems = flat_hidden.numel()
+    if n_elems == 0:
+        return out.view_as(hidden)
+    block = 4096
+    wide = n_elems >= 2**31
+    nfull = (n_elems // block) * block
+    grid = (triton.cdiv(n_elems, block), )
+    _h3_gated_residual_kernel[grid](
+        out,
+        flat_hidden,
+        flat_branch,
+        gate,
+        index,
+        n_elems,
+        index.numel(),
+        D=hidden_size,
+        GS=gate_row_stride,
+        BLOCK=block,
+        NFULL=nfull,
+        WIDE=wide,
+        num_warps=8,
+    )
+    return out.view_as(hidden)

fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernel: Fused H3 QKV prologue: per-head QK RMSNorm + partial RoPE + stacked QKV output.
+
+Found with MotionKernel (agent-driven kernel search). Tuned on GB200 / sm_100 for the
+MiniMax-H3 inference shapes; opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``.
+
+Fused QKV prologue: per-head QK RMSNorm + partial RoPE on Q and K plus a straight V copy.
+
+Production runs two Sol-Engine ``_qknorm_partial_rope_kernel`` launches (one Triton program per
+128-wide (token, head) row of Q and of K) plus ``torch.cat([q, k, v])``. Here a single launch
+produces Q, K and V together, so every byte is read and written exactly once.
+
+Each program owns ``HB`` consecutive head rows of one token, i.e. a *contiguous flat* chunk of
+``HB * head_dim`` elements. Because the tile is flat and contiguous, each tensor costs exactly
+one wide (128-bit) load and one wide store per thread instead of a pile of per-sub-block
+addressing. The grid is ordered so that the head-group index is the *fastest* varying dimension:
+consecutive CTAs then cover one token's whole ``H * D`` slab in address order, which turns the
+six DRAM streams (q, k, v in; q, k, v out) into pure sequential scans (~7.1 vs ~6.6 TB/s).
+
+All three input tiles are issued as loads *before* any store (q, k and v in flight together,
+then the stores), so the six streams overlap instead of being serialised by the store that
+follows each load. With ``HB=4`` (512 elements, one warp) this reaches the measured streaming
+roofline of this GPU: 7.08 TB/s versus 7.10 TB/s for a pure three-stream copy of the same bytes
+(the previous load/store/load/store schedule measured 7.00 TB/s).
+
+The rotate-half partner ``x[i +/- rotary_dim/2]`` is materialised with an in-register
+``tl.gather`` along the head dim (warp-shuffle lowering for a single warp), so the whole row is
+loaded once: no gather loads from memory, no lane padding. All arithmetic is fp32 with one
+rounding at the bf16 store; V is copied verbatim.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+H3_NORM_EPS = 1e-5
+
+
+@triton.jit
+def _qkv_norm_rope_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    qw_ptr,
+    kw_ptr,
+    cos_ptr,
+    sin_ptr,
+    out_ptr,
+    num_heads,
+    seq_len,
+    plane,  # elements per output plane = B*S*H*D
+    eps,
+    HD: tl.constexpr,
+    RD: tl.constexpr,
+    HB: tl.constexpr,
+    NHG: tl.constexpr,
+    I64: tl.constexpr,
+):
+    # 1-D grid: gridDim.y is capped at 65535, below the 73k-token shape.
+    pid = tl.program_id(0)
+    h0 = (pid % NHG) * HB
+    bt = pid // NHG
+    tok = bt % seq_len
+    batch = bt // seq_len
+
+    NB: tl.constexpr = HB * HD
+    off = tl.arange(0, NB)
+    hrow = h0 + off // HD
+    mask = hrow < num_heads
+
+    if I64:
+        base = (batch.to(tl.int64) * seq_len + tok) * num_heads * HD + h0 * HD
+    else:
+        base = (batch * seq_len + tok) * num_heads * HD + h0 * HD
+
+    idx = tl.arange(0, HD)
+    pidx = tl.where(idx < RD, tl.where(idx < RD // 2, idx + RD // 2, idx - RD // 2), idx)
+    gidx = tl.broadcast_to(pidx[None, :], (HB, HD))
+    c = tl.load(cos_ptr + tok * RD + idx, mask=idx < RD, other=1.0).to(tl.float32)[None, :]
+    s = tl.load(sin_ptr + tok * RD + idx, mask=idx < RD, other=0.0).to(tl.float32)
+    s = tl.where(idx < RD // 2, -s, s)[None, :]
+
+    wq = tl.load(qw_ptr + idx).to(tl.float32)[None, :]
+    wk = tl.load(kw_ptr + idx).to(tl.float32)[None, :]
+
+    xq = tl.load(q_ptr + base + off, mask=mask, other=0.0).to(tl.float32)
+    xk = tl.load(k_ptr + base + off, mask=mask, other=0.0).to(tl.float32)
+    xv = tl.load(v_ptr + base + off, mask=mask, other=0.0)
+
+    tl.store(out_ptr + 2 * plane + base + off, xv, mask=mask)
+
+    y = tl.reshape(xq, (HB, HD))
+    inv = tl.math.rsqrt(tl.sum(y * y, 1) / HD + eps)[:, None]
+    y = y * inv * wq
+    p = tl.gather(y, gidx, 1)
+    y = y * c + p * s
+    tl.store(out_ptr + base + off, tl.reshape(y, (NB, )).to(out_ptr.dtype.element_ty), mask=mask)
+
+    y = tl.reshape(xk, (HB, HD))
+    inv = tl.math.rsqrt(tl.sum(y * y, 1) / HD + eps)[:, None]
+    y = y * inv * wk
+    p = tl.gather(y, gidx, 1)
+    y = y * c + p * s
+    tl.store(out_ptr + plane + base + off, tl.reshape(y, (NB, )).to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def _pick_config(num_heads: int):
+    """(HB, num_warps) -- consecutive head rows per program.
+
+    One warp per program with 4 rows gives 16 elements (two 128-bit accesses) per thread per
+    tensor, the measured streaming optimum for this GPU: every row stays inside one warp so the
+    RMSNorm reduction and the rotate-half gather are pure warp shuffles, and 32 such CTAs per SM
+    keep enough loads in flight to run at the copy roofline.
+    """
+    if num_heads >= 4:
+        return 4, 1
+    hb = max(triton.next_power_of_2(max(num_heads, 1)), 1)
+    return hb, 1
+
+
+def _launch(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    b, s, h, d = query.shape
+    out = torch.empty((3 * b, s, h, d), device=query.device, dtype=query.dtype)
+    hb, warps = _pick_config(h)
+    nhg = triton.cdiv(h, hb)
+    _qkv_norm_rope_kernel[(b * s * nhg, )](
+        query,
+        key,
+        value,
+        q_weight,
+        k_weight,
+        cos,
+        sin,
+        out,
+        h,
+        s,
+        b * s * h * d,
+        H3_NORM_EPS,
+        HD=d,
+        RD=cos.shape[1],
+        HB=hb,
+        NHG=nhg,
+        I64=out.numel() >= 2**31,
+        num_warps=warps,
+        num_stages=1,
+    )
+    return out
+
+
+def qkv_prologue(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """Returns the [3*B, S, H, D] stack (rotated Q, rotated K, V), as FastVideo's attention path builds it."""
+    if query.ndim != 4 or key.shape != query.shape or value.shape != query.shape:
+        raise ValueError("query/key/value must share a [B, S, H, D] shape")
+    if cos.shape != sin.shape or cos.shape[0] != query.shape[1]:
+        raise ValueError("cos/sin must be [S, rotary_dim] tables")
+    if not query.is_contiguous():
+        query = query.contiguous()
+    if not key.is_contiguous():
+        key = key.contiguous()
+    if not value.is_contiguous():
+        value = value.contiguous()
+    if not q_weight.is_contiguous():
+        q_weight = q_weight.contiguous()
+    if not k_weight.is_contiguous():
+        k_weight = k_weight.contiguous()
+    if cos.dtype != query.dtype:
+        cos = cos.to(query.dtype)
+    if sin.dtype != query.dtype:
+        sin = sin.to(query.dtype)
+    if not cos.is_contiguous():
+        cos = cos.contiguous()
+    if not sin.is_contiguous():
+        sin = sin.contiguous()
+    return _launch(query, key, value, q_weight, k_weight, cos, sin)

fastvideo/models/dits/minimax_h3_fusions/mk/residual_gate_rmsnorm_modulate.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernel: Fused H3 residual gate + RMSNorm + AdaLN modulation.
+
+Found with MotionKernel (agent-driven kernel search). Tuned on GB200 / sm_100 for the
+MiniMax-H3 inference shapes; opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``.
+
+Chunked gated residual + RMSNorm + AdaLN modulation (H3 post-attention).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+H3_NORM_EPS = 1e-5
+
+_BLOCK = 1024
+_NUM_WARPS = 4
+_NUM_STAGES = 3
+
+
+@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_gate,
+    stride_scale,
+    stride_shift,
+    BLOCK: tl.constexpr,
+    IDENTITY_INDEX: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK)
+    base = row * n_cols
+    if IDENTITY_INDEX:
+        table_row = tl.load(index_ptr + row).to(tl.int64)
+    else:
+        table_row = tl.load(index_ptr + (row % n_index)).to(tl.int64)
+    gbase = table_row * stride_gate
+    sbase = table_row * stride_scale
+    hbase = table_row * stride_shift
+
+    acc = tl.zeros((BLOCK, ), tl.float32)
+    for c in range(0, n_cols, BLOCK):
+        cur = c + cols
+        m = cur < n_cols
+        off = base + cur
+        residual = tl.load(residual_ptr + off, mask=m, other=0.0).to(tl.float32)
+        branch = tl.load(branch_ptr + off, mask=m, other=0.0).to(tl.float32)
+        gate = tl.load(gate_ptr + gbase + cur, mask=m, other=0.0).to(tl.float32)
+        hidden = residual + gate * branch
+        tl.store(
+            hidden_out_ptr + off,
+            hidden.to(hidden_out_ptr.dtype.element_ty),
+            mask=m,
+        )
+        acc += hidden * hidden
+
+    rstd = tl.math.rsqrt(tl.sum(acc, axis=0) / n_cols + eps)
+
+    for c in range(0, n_cols, BLOCK):
+        cur = c + cols
+        m = cur < n_cols
+        off = base + cur
+        hidden = tl.load(hidden_out_ptr + off, mask=m, other=0.0, cache_modifier=".cg").to(tl.float32)
+        weight = tl.load(weight_ptr + cur, mask=m, other=0.0).to(tl.float32)
+        scale = tl.load(scale_ptr + sbase + cur, mask=m, other=0.0).to(tl.float32)
+        shift = tl.load(shift_ptr + hbase + cur, mask=m, other=0.0).to(tl.float32)
+        output = hidden * rstd * weight * (1.0 + scale) + shift
+        tl.store(
+            normed_out_ptr + off,
+            output.to(normed_out_ptr.dtype.element_ty),
+            mask=m,
+        )
+
+
+def _row_addressable(table: torch.Tensor) -> torch.Tensor:
+    return table if table.stride(-1) == 1 else table.contiguous()
+
+
+def 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,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return ``(hidden, normed)`` exactly as FastVideo's optimized profile does."""
+    if not residual.is_cuda:
+        raise ValueError("h3_residual_gate_rmsnorm_modulate requires CUDA tensors")
+    if branch.shape != residual.shape or branch.dtype != residual.dtype:
+        raise ValueError("branch must match residual shape and dtype")
+    hidden_size = residual.shape[-1]
+    if index.ndim != 1 or index.numel() != residual.shape[-2]:
+        raise ValueError("index must hold one table row per token")
+    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)
+    _residual_gate_rmsnorm_modulate_kernel[(flat_residual.shape[0], )](
+        hidden,
+        modulated,
+        flat_residual,
+        flat_branch,
+        gate,
+        weight,
+        scale,
+        shift,
+        index,
+        hidden_size,
+        index.numel(),
+        H3_NORM_EPS,
+        gate.stride(0),
+        scale.stride(0),
+        shift.stride(0),
+        IDENTITY_INDEX=index.numel() == flat_residual.shape[0],
+        BLOCK=_BLOCK,
+        num_warps=_NUM_WARPS,
+        num_stages=_NUM_STAGES,
+    )
+    return hidden.view_as(residual), modulated.view_as(residual)

fastvideo/models/dits/minimax_h3_fusions/mk/rmsnorm_modulate.py

@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernel: Fused H3 RMSNorm + row-indexed AdaLN modulation.
+
+Found with MotionKernel (agent-driven kernel search). Tuned on GB200 / sm_100 for the
+MiniMax-H3 inference shapes; opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``.
+
+MiniMax H3 RMSNorm + row-indexed AdaLN modulation (two-kernel fast path).
+
+Math (FastVideo ``minimax_h3.py`` lines 729-731), FP32 internally with one
+rounding at the store::
+
+    y = rms_norm(x, weight, eps) * (1 + scale[idx]) + shift[idx]
+
+Shape notes that drive the design (GB200, hidden D = 5376 = 2**8 * 21, so D is
+*not* a power of two):
+
+* The vendored Sol-Engine starter runs one program per row with ``BLOCK =
+  next_pow2(D) = 8192`` and ``num_warps=16``; 34% of its lanes are masked off
+  and the whole 8192-lane tile is live in registers.
+* The AdaLN tables are row-strided views (row stride ``6 * D``) with only
+  ``table_rows`` (<= 12) distinct rows, so every token re-reads the same
+  ``weight``, ``scale`` and ``shift`` vectors: 3 * D * 2 bytes of gather per
+  row on top of the D * 2 bytes of activation.
+
+This candidate keeps the public interface and splits the work:
+
+1. ``_fuse_adaln_kernel`` (tiny: ``table_rows`` rows) folds the per-table-row
+   modulation into ``a[t] = weight * (1 + scale[t])`` and ``b[t] = shift[t]``,
+   stored in the activation dtype. The main pass then gathers two vectors per
+   token instead of three, and they are read at the activation width.
+2. ``_rmsnorm_modulate_kernel`` is a chunked single-row program with no lane
+   padding: ``ceil(D / BLOCK)`` chunks of ``BLOCK`` lanes (BLOCK = 2048 for
+   D = 5376, i.e. three chunks with only the last one partially masked). It
+   makes a first pass over the row to accumulate the sum of squares, then a
+   second pass that re-reads the row (an L1/L2 hit, the row is 10.5 KiB) and
+   applies the store. That keeps the live register tile at one chunk instead
+   of a whole padded row, which is what the starter loses to.
+
+Performance notes (GB200, measured against a minimal Triton streaming copy of
+the same byte count on the same shapes, i.e. the true roofline):
+
+* copy roofline per production shape: 58.9 / 98.6 / 129.1 / 140.8 / 183.9 us
+  (6.2-6.9 TB/s), mean 122.3 us.
+* this kernel: mean of five shapes ~126.3 us, i.e. ~94.5% of the copy
+  roofline re-measured for this attempt (55.6 / 95.1 / 124.7 / 138.0 / 184.6 us,
+  mean 119.6 us); on the 59k shape 188.2 us vs 184.6 us (98%).
+
+The step that mattered most over the starter (134.8 us -> 128.5 us) was making
+the hidden width a ``tl.constexpr`` and unrolling the chunk loops with
+``tl.static_range``: the lane masks, the row offset and the reduction become
+static, so every activation/table/store access compiles to an aligned 16-byte
+vector access instead of scalar traffic on an unknown-alignment base pointer.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+H3_NORM_EPS = 1e-5
+
+_MAIN_BLOCK = 2048
+_TABLE_BLOCK = 1024
+
+
+@triton.jit
+def _fuse_adaln_kernel(
+    a_ptr,
+    b_ptr,
+    weight_ptr,
+    scale_ptr,
+    shift_ptr,
+    stride_scale_row,
+    stride_shift_row,
+    NC: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """One program per (table row, column chunk): a = w * (1 + scale), b = shift."""
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = cols < NC
+    w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    s = tl.load(scale_ptr + row * stride_scale_row + cols, mask=mask, other=0.0).to(tl.float32)
+    h = tl.load(shift_ptr + row * stride_shift_row + cols, mask=mask, other=0.0).to(tl.float32)
+    out = row * NC + cols
+    tl.store(a_ptr + out, (w * (1.0 + s)).to(a_ptr.dtype.element_ty), mask=mask)
+    tl.store(b_ptr + out, h.to(b_ptr.dtype.element_ty), mask=mask)
+
+
+@triton.jit
+def _rmsnorm_modulate_kernel(
+    out_ptr,
+    x_ptr,
+    a_ptr,
+    b_ptr,
+    index_ptr,
+    n_index,
+    eps,
+    stride_a_row,
+    stride_b_row,
+    NC: tl.constexpr,
+    BLOCK: tl.constexpr,
+    NCHUNK: tl.constexpr,
+):
+    """One program per row: sumsq pass over the row, then an apply pass.
+
+    ``NC`` (the hidden width) is a ``constexpr`` so that lane masks, row
+    offsets and the reduction are fully static: every activation, table and
+    store access is then a 16-byte vector access on an aligned base, which is
+    what lifts this pass to the streaming-copy roofline. The row is 10.5 KiB,
+    so the second pass re-reads it from L1 instead of DRAM.
+    """
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK)
+    table_row = tl.load(index_ptr + row % n_index).to(tl.int64)
+    base = row * NC
+
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for chunk in tl.static_range(NCHUNK):
+        off = chunk * BLOCK + cols
+        x = tl.load(x_ptr + base + off, mask=off < NC, other=0.0, eviction_policy="evict_last").to(tl.float32)
+        acc += x * x
+    rstd = 1.0 / tl.sqrt(tl.sum(acc, axis=0) / NC + eps)
+
+    chunk = NCHUNK - 1
+    off = chunk * BLOCK + cols
+    mask = off < NC
+    x = tl.load(x_ptr + base + off, mask=mask, other=0.0, eviction_policy="evict_last").to(tl.float32)
+    a = tl.load(
+        a_ptr + table_row * stride_a_row + off,
+        mask=mask,
+        other=0.0,
+        eviction_policy="evict_last",
+    ).to(tl.float32)
+    b = tl.load(
+        b_ptr + table_row * stride_b_row + off,
+        mask=mask,
+        other=0.0,
+        eviction_policy="evict_last",
+    ).to(tl.float32)
+    y = x * rstd * a + b
+    tl.store(
+        out_ptr + base + off,
+        y.to(out_ptr.dtype.element_ty),
+        mask=mask,
+        cache_modifier=".cs",
+    )
+
+    for chunk in tl.static_range(NCHUNK - 1):
+        off = chunk * BLOCK + cols
+        mask = off < NC
+        x = tl.load(x_ptr + base + off, mask=mask, other=0.0).to(tl.float32)
+        a = tl.load(a_ptr + table_row * stride_a_row + off, mask=mask, other=0.0).to(tl.float32)
+        b = tl.load(b_ptr + table_row * stride_b_row + off, mask=mask, other=0.0).to(tl.float32)
+        y = x * rstd * a + b
+        tl.store(
+            out_ptr + base + off,
+            y.to(out_ptr.dtype.element_ty),
+            mask=mask,
+            cache_modifier=".cs",
+        )
+
+
+def _next_power_of_two(value: int) -> int:
+    return 1 << (value - 1).bit_length()
+
+
+def _main_block(hidden: int) -> int:
+    return min(_MAIN_BLOCK, _next_power_of_two(hidden))
+
+
+def _num_warps(block: int) -> int:
+    if block >= 512:
+        return 4
+    return 2
+
+
+def _row_addressable(table: torch.Tensor) -> torch.Tensor:
+    return table if table.stride(-1) == 1 else table.contiguous()
+
+
+def rmsnorm_modulate(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+) -> torch.Tensor:
+    """RMSNorm + per-token AdaLN modulation with one activation read and write."""
+    if not x.is_cuda:
+        raise ValueError("h3_rmsnorm_modulate requires CUDA tensors")
+    hidden = x.shape[-1]
+    if index.ndim != 1 or index.numel() != x.shape[-2]:
+        raise ValueError("index must hold one table row per token")
+    if scale.shape[-1] != hidden or shift.shape[-1] != hidden:
+        raise ValueError("modulation tables must match the hidden width")
+    flat_x = x.reshape(-1, hidden).contiguous()
+    weight = weight.contiguous()
+    scale = _row_addressable(scale)
+    shift = _row_addressable(shift)
+    index = index.contiguous()
+    output = torch.empty_like(flat_x)
+
+    table_rows = scale.shape[0]
+    tab_dtype = flat_x.dtype if flat_x.dtype in (torch.bfloat16, torch.float16) else torch.float32
+    fused_a = torch.empty((table_rows, hidden), dtype=tab_dtype, device=flat_x.device)
+    fused_b = torch.empty((table_rows, hidden), dtype=tab_dtype, device=flat_x.device)
+
+    table_block = min(_TABLE_BLOCK, _next_power_of_two(hidden))
+    _fuse_adaln_kernel[(table_rows, triton.cdiv(hidden, table_block))](
+        fused_a,
+        fused_b,
+        weight,
+        scale,
+        shift,
+        scale.stride(0),
+        shift.stride(0),
+        NC=hidden,
+        BLOCK=table_block,
+        num_warps=_num_warps(table_block),
+    )
+
+    block = _main_block(hidden)
+    _rmsnorm_modulate_kernel[(flat_x.shape[0], )](
+        output,
+        flat_x,
+        fused_a,
+        fused_b,
+        index,
+        index.numel(),
+        H3_NORM_EPS,
+        fused_a.stride(0),
+        fused_b.stride(0),
+        NC=hidden,
+        BLOCK=block,
+        NCHUNK=triton.cdiv(hidden, block),
+        num_warps=_num_warps(block),
+    )
+    return output.view_as(x)

fastvideo/models/dits/minimax_h3_fusions/mk/swiglu.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernel: Fused H3 SwiGLU over the packed gate/value halves.
+
+Found with MotionKernel (agent-driven kernel search). Tuned on GB200 / sm_100 for the
+MiniMax-H3 inference shapes; opt-in through ``FASTVIDEO_MINIMAX_H3_FUSIONS=mk``.
+
+MiniMax H3 value-first packed SwiGLU -- chunked, unpadded, streaming.
+
+Math (FastVideo ``minimax_h3.py`` @ 7203b4976, lines 188-189)::
+
+    value, gate = hidden_states.chunk(2, dim=-1)   # value first, gate second
+    out = value * F.silu(gate)
+
+Kept in FP32 throughout with one rounding at the store (matches Sol-Engine).
+
+Why this beats the vendored Sol-Engine kernel: that kernel runs one program per
+row with ``BLOCK = next_pow2(ffn) = 16384`` lanes, so 12.5% of every lane group
+is idle and each program keeps two 16384-wide FP32 tiles live (heavy register
+pressure, low occupancy). F=14336 = 14 x 1024, so a chunked grid needs no
+padding at all. Each program handles one 1024-element column chunk of one row:
+one 128-bit vector load per thread per half, FP32 math, one store.
+
+Launch-shape notes measured on GB200 (Pika shapes):
+  * A flat 1D grid whose column chunk is ``pid % NTILES`` (so concurrently
+    resident CTAs stream adjacent 2 KiB blocks of the same row) is the fastest
+    enumeration: ~443 us at 37718 tokens = 7.33 TB/s, versus a measured
+    ~7.4 TB/s pure-read / 7.8 TB/s pure-write streaming roofline. Enumerating
+    rows fastest costs ~8% (478 us); BLOCK=512 or BLOCK=2048 with more warps
+    cost 2-28%.
+  * The flat grid also removes the ``grid.y <= 65535`` limit that the 768p
+    243-frame shape (73394 rows) exceeds: adding a row-group loop there cost
+    4% (888 us vs 852 us).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+_BLOCK = 1024
+_NUM_WARPS = 4
+_I32_LIMIT = 0x7FFFFFFF
+
+
+@triton.jit
+def _h3_swiglu_kernel(
+    out_ptr,
+    x_ptr,
+    ffn_dim,
+    n_rows,
+    stride_in_row,
+    stride_out_row,
+    NTILES: tl.constexpr,
+    BLOCK: tl.constexpr,
+    EXACT: tl.constexpr,
+    GUARD: tl.constexpr,
+    I64: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    col = pid % NTILES
+    row = pid // NTILES
+    cols = col * BLOCK + tl.arange(0, BLOCK)
+    if I64:
+        row = row.to(tl.int64)
+    if GUARD:
+        if row < n_rows:
+            _body(out_ptr, x_ptr, ffn_dim, row, cols, stride_in_row, stride_out_row, BLOCK, EXACT)
+    else:
+        _body(out_ptr, x_ptr, ffn_dim, row, cols, stride_in_row, stride_out_row, BLOCK, EXACT)
+
+
+@triton.jit
+def _body(out_ptr, x_ptr, ffn_dim, row, cols, stride_in_row, stride_out_row, BLOCK: tl.constexpr, EXACT: tl.constexpr):
+    in_off = row * stride_in_row + cols
+    if EXACT:
+        value = tl.load(x_ptr + in_off).to(tl.float32)
+        gate = tl.load(x_ptr + in_off + ffn_dim).to(tl.float32)
+        out = value * (gate * tl.sigmoid(gate))
+        tl.store(out_ptr + row * stride_out_row + cols, out.to(out_ptr.dtype.element_ty))
+    else:
+        mask = cols < ffn_dim
+        value = tl.load(x_ptr + in_off, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(x_ptr + in_off + ffn_dim, mask=mask, other=0.0).to(tl.float32)
+        out = value * (gate * tl.sigmoid(gate))
+        tl.store(out_ptr + row * stride_out_row + cols, out.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def swiglu(x: torch.Tensor) -> torch.Tensor:
+    """Launch the chunked unpadded SwiGLU."""
+    if not x.is_cuda:
+        raise ValueError("h3_swiglu requires a CUDA tensor")
+    packed_width = x.shape[-1]
+    if packed_width == 0 or packed_width % 2:
+        raise ValueError("h3_swiglu needs a positive even last dimension")
+    ffn_dim = packed_width // 2
+    flat = x.reshape(-1, packed_width)
+    if not flat.is_contiguous():
+        flat = flat.contiguous()
+    output_shape = (*x.shape[:-1], ffn_dim)
+    rows = flat.shape[0]
+    if rows == 0:
+        return torch.empty(output_shape, dtype=x.dtype, device=x.device)
+    out = torch.empty((rows, ffn_dim), dtype=x.dtype, device=x.device)
+
+    block = min(_BLOCK, triton.next_power_of_2(ffn_dim))
+    exact = (ffn_dim % block) == 0
+    ntiles = (ffn_dim + block - 1) // block
+    i64 = rows * packed_width > _I32_LIMIT
+    num_warps = _NUM_WARPS if block >= 1024 else 2
+    _h3_swiglu_kernel[(ntiles * rows, )](
+        out,
+        flat,
+        ffn_dim,
+        rows,
+        flat.stride(0),
+        out.stride(0),
+        NTILES=ntiles,
+        BLOCK=block,
+        EXACT=exact,
+        GUARD=False,
+        I64=i64,
+        num_warps=num_warps,
+    )
+    return out.view(output_shape)

fastvideo/tests/attention/test_forward_stacked.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``DistributedAttention.forward`` is the stacked path plus a ``torch.cat``."""
+import torch
+
+from fastvideo.attention.layer import DistributedAttention
+
+
+class _Recorder:
+
+    def __init__(self):
+        self.qkv = None
+        self.kwargs = None
+
+
+def test_forward_cats_then_delegates_to_forward_stacked():
+    rec = _Recorder()
+
+    def fake_forward_stacked(self, qkv, **kwargs):
+        rec.qkv = qkv
+        rec.kwargs = kwargs
+        return torch.zeros(2, 3, 4, 5), None
+
+    original = DistributedAttention.forward_stacked
+    DistributedAttention.forward_stacked = fake_forward_stacked
+    try:
+        q = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5)
+        k = q + 1
+        v = q + 2
+        # Bypass __init__ (it builds a real backend); ``forward`` only needs the
+        # compiler-policy flag and the patched ``forward_stacked``.
+        dummy = object.__new__(DistributedAttention)
+        dummy._compile_forward_enabled = False
+        out, replicated = DistributedAttention.forward(dummy, q, k, v, original_seq_len=3)
+    finally:
+        DistributedAttention.forward_stacked = original
+
+    assert replicated is None
+    assert rec.qkv.shape == (6, 3, 4, 5)
+    assert torch.equal(rec.qkv[:2], q) and torch.equal(rec.qkv[2:4], k) and torch.equal(rec.qkv[4:], v)
+    assert rec.kwargs["original_seq_len"] == 3

fastvideo/tests/transformers/test_minimax_h3_fusion_modes.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The MotionKernel ``mk`` tier is opt-in and never changes the default fusion set."""
+import pytest
+
+from fastvideo.models.dits.minimax_h3 import (_MINIMAX_H3_FUSION_NAMES, _MINIMAX_H3_MK_MODE,
+                                              _enabled_minimax_h3_fusions)
+
+
+def test_default_is_unchanged():
+    for value in ("", "0", "none", None):
+        assert _enabled_minimax_h3_fusions(value) == frozenset()
+
+
+def test_all_is_unchanged():
+    assert _enabled_minimax_h3_fusions("all") == _MINIMAX_H3_FUSION_NAMES
+    assert _enabled_minimax_h3_fusions("1") == _MINIMAX_H3_FUSION_NAMES
+
+
+def test_subset_is_unchanged():
+    assert _enabled_minimax_h3_fusions("modulate,swiglu") == frozenset({"modulate", "swiglu"})
+
+
+def test_mk_enables_every_fusion_plus_the_tier_marker():
+    assert _enabled_minimax_h3_fusions("mk") == _MINIMAX_H3_FUSION_NAMES | {_MINIMAX_H3_MK_MODE}
+    assert _enabled_minimax_h3_fusions(" MK ") == _MINIMAX_H3_FUSION_NAMES | {_MINIMAX_H3_MK_MODE}
+
+
+def test_mk_cannot_be_combined_with_individual_fusions():
+    with pytest.raises(ValueError, match="whole-tier mode"):
+        _enabled_minimax_h3_fusions("mk,swiglu")
+
+
+def test_unknown_fusion_still_raises():
+    with pytest.raises(ValueError, match="Unknown MiniMax H3 fusion"):
+        _enabled_minimax_h3_fusions("modulate,nope")

fastvideo/tests/transformers/test_minimax_h3_mk_block.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""A real H3 transformer block runs the MotionKernel tier and matches the Sol tier.
+
+Small random-initialised block, SDPA backend, no checkpoint: this exercises the
+``mk_tier`` dispatch of every fused call site (including the stacked QKV
+prologue) on the actual module wiring rather than on standalone kernels.
+"""
+import pytest
+import torch
+
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.dits.minimax_h3 import MiniMaxH3TransformerBlock
+from fastvideo.models.dits.minimax_h3_fusions import mk
+from fastvideo.platforms import AttentionBackendEnum
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available() or not mk.HAVE_TRITON,
+                                reason="MotionKernel H3 kernels need CUDA + Triton")
+
+HIDDEN, HEADS, HEAD_DIM, ROTARY_DIM, FFN, TABLE_ROWS, TOKENS = 512, 4, 128, 96, 1024, 6, 256
+
+
+def _block(mk_tier: bool, fused: bool = True) -> MiniMaxH3TransformerBlock:
+    torch.manual_seed(0)
+    torch.set_default_dtype(torch.bfloat16)
+    try:
+        return MiniMaxH3TransformerBlock(
+            hidden_size=HIDDEN,
+            num_attention_heads=HEADS,
+            attention_head_dim=HEAD_DIM,
+            ffn_dim=FFN,
+            time_embed_dim=128,
+            norm_eps=1e-5,
+            qk_norm_eps=1e-5,
+            supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ),
+            quant_config=None,
+            prefix="test",
+            fuse_modulate=fused,
+            fuse_qknorm_rope=fused,
+            fuse_swiglu=fused,
+            mk_tier=mk_tier,
+        ).cuda()
+    finally:
+        torch.set_default_dtype(torch.float32)
+
+
+def _inputs():
+    torch.manual_seed(1)
+    return (torch.randn(1, TOKENS, HIDDEN, device="cuda",
+                        dtype=torch.bfloat16), torch.randn(TABLE_ROWS, 128, device="cuda", dtype=torch.bfloat16),
+            torch.randint(0, TABLE_ROWS, (TOKENS, ),
+                          device="cuda", dtype=torch.int64), (torch.randn(TOKENS,
+                                                                          ROTARY_DIM,
+                                                                          device="cuda",
+                                                                          dtype=torch.bfloat16),
+                                                              torch.randn(TOKENS,
+                                                                          ROTARY_DIM,
+                                                                          device="cuda",
+                                                                          dtype=torch.bfloat16)), TOKENS)
+
+
+def test_mk_tier_dispatches_and_matches_sol(monkeypatch, distributed_setup):
+    calls = {
+        name: 0
+        for name in ("qkv_prologue", "rmsnorm_modulate", "residual_gate_rmsnorm_modulate", "swiglu", "gated_residual")
+    }
+    for name in calls:
+        target = getattr(mk, name)
+
+        def counted(*args, _target=target, _name=name, **kwargs):
+            calls[_name] += 1
+            return _target(*args, **kwargs)
+
+        monkeypatch.setattr(mk, name, counted)
+
+    sol, mk_block, eager = _block(False), _block(True), _block(False, fused=False)
+    # The loader zero-initialises ReplicatedLinear (weights arrive from the
+    # checkpoint) and that would make the block an identity, hiding every kernel
+    # difference. Give all three copies the same non-trivial weights instead.
+    torch.manual_seed(2)
+    for name, param in sol.named_parameters():
+        if name.endswith(".weight") and param.dim() == 2 or name.endswith(".bias"):
+            param.data.normal_(0.0, 0.02)
+        elif name.endswith(".weight"):
+            param.data.fill_(1.0)
+    mk_block.load_state_dict(sol.state_dict())
+    eager.load_state_dict(sol.state_dict())
+    inputs = _inputs()
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+        expected = sol(*inputs)
+        actual = mk_block(*inputs)
+        unfused = eager(*inputs)
+
+    assert all(count > 0 for count in calls.values()), calls
+    base = torch.linalg.vector_norm(expected.float())
+
+    def rel(other):
+        return (torch.linalg.vector_norm(other.float() - expected.float()) / base).item()
+
+    floor, rel = rel(unfused), rel(actual)
+    print(f"eager-vs-sol block rel-L2 = {floor:.3e}; mk-vs-sol block rel-L2 = {rel:.3e}")
+    assert rel <= 3 * floor, (rel, floor)
+
+
+def test_mk_tier_defers_to_sol_while_compiling(monkeypatch, distributed_setup):
+    """A traced block must keep using the compile-safe Sol custom ops."""
+    calls = []
+    for name in ("qkv_prologue", "rmsnorm_modulate", "residual_gate_rmsnorm_modulate", "swiglu", "gated_residual"):
+        target = getattr(mk, name)
+
+        def counted(*args, _target=target, _name=name, **kwargs):
+            calls.append(_name)
+            return _target(*args, **kwargs)
+
+        monkeypatch.setattr(mk, name, counted)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    block = _block(True)
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+        block(*_inputs())
+    assert calls == []

fastvideo/tests/transformers/test_minimax_h3_mk_kernels.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MotionKernel H3 kernels match the in-tree Sol/eager paths on the production shapes.
+
+GPU-only (Triton on CUDA); the shapes use H3's real hidden width, head layout and
+AdaLN table stride, only the token count is reduced to keep the test quick.
+"""
+import pytest
+import torch
+
+from fastvideo.models.dits.minimax_h3_fusions import (fused_qknorm_rope, fused_residual_gate_rmsnorm_modulate,
+                                                      fused_rmsnorm_modulate, minimax_h3_swiglu)
+from fastvideo.models.dits.minimax_h3_fusions import mk
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available() or not mk.HAVE_TRITON,
+                                reason="MotionKernel H3 kernels need CUDA + Triton")
+
+HIDDEN, HEADS, HEAD_DIM, ROTARY_DIM, FFN_DIM, EPS = 5376, 56, 128, 96, 14336, 1e-5
+TOKENS = 512
+
+
+def _rel_l2(actual, expected):
+    return (torch.linalg.vector_norm(actual.float() - expected.float()) /
+            torch.linalg.vector_norm(expected.float())).item()
+
+
+@pytest.fixture(scope="module")
+def production_inputs():
+    torch.manual_seed(0)
+    device = "cuda"
+    mod = torch.randn(6, 6, HIDDEN, device=device, dtype=torch.bfloat16)
+    # The production AdaLN tables are row-strided views of this [rows, 6*D] buffer.
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (mod[:, i, :] for i in range(6))
+    return {
+        "x": torch.randn(TOKENS, HIDDEN, device=device, dtype=torch.bfloat16),
+        "branch": torch.randn(TOKENS, HIDDEN, device=device, dtype=torch.bfloat16),
+        "weight": torch.randn(HIDDEN, device=device, dtype=torch.bfloat16),
+        "qk_weight": torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16),
+        "scale_msa": scale_msa,
+        "shift_msa": shift_msa,
+        "gate_msa": gate_msa,
+        "scale_mlp": scale_mlp,
+        "shift_mlp": shift_mlp,
+        "gate_mlp": gate_mlp,
+        "index": torch.randint(0, 6, (TOKENS, ), device=device, dtype=torch.int64),
+    }
+
+
+def test_rmsnorm_modulate_matches_sol(production_inputs):
+    i = production_inputs
+    expected = fused_rmsnorm_modulate(i["x"], i["weight"], i["scale_msa"], i["shift_msa"], i["index"], EPS)
+    actual = mk.rmsnorm_modulate(i["x"], i["weight"], i["scale_msa"], i["shift_msa"], i["index"])
+    assert actual.shape == expected.shape and actual.dtype == expected.dtype
+    assert _rel_l2(actual, expected) < 5e-3
+
+
+def test_residual_gate_rmsnorm_modulate_matches_sol(production_inputs):
+    i = production_inputs
+    args = (i["branch"], i["gate_msa"], i["weight"], i["scale_mlp"], i["shift_mlp"], i["index"])
+    exp_hidden, exp_normed = fused_residual_gate_rmsnorm_modulate(i["x"], *args, EPS)
+    act_hidden, act_normed = mk.residual_gate_rmsnorm_modulate(i["x"], *args)
+    assert _rel_l2(act_hidden, exp_hidden) < 1e-3
+    assert _rel_l2(act_normed, exp_normed) < 5e-3
+
+
+def test_swiglu_matches_sol(production_inputs):
+    packed = torch.randn(TOKENS, 2 * FFN_DIM, device="cuda", dtype=torch.bfloat16)
+    assert _rel_l2(mk.swiglu(packed), minimax_h3_swiglu(packed)) < 1e-3
+
+
+def test_gated_residual_matches_eager(production_inputs):
+    i = production_inputs
+    expected = i["x"] + i["gate_mlp"].index_select(0, i["index"]) * i["branch"]
+    assert _rel_l2(mk.gated_residual(i["x"], i["gate_mlp"], i["branch"], i["index"]), expected) < 5e-3
+
+
+def test_qkv_prologue_matches_two_sol_launches(production_inputs):
+    i = production_inputs
+    q = torch.randn(1, TOKENS, HEADS, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    cos = torch.randn(TOKENS, ROTARY_DIM, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn(TOKENS, ROTARY_DIM, device="cuda", dtype=torch.bfloat16)
+    expected = torch.cat(
+        [fused_qknorm_rope(q, i["qk_weight"], cos, sin, EPS),
+         fused_qknorm_rope(k, i["qk_weight"], cos, sin, EPS), v],
+        dim=0)
+    actual = mk.qkv_prologue(q, k, v, i["qk_weight"], i["qk_weight"], cos, sin)
+    assert actual.shape == (3, TOKENS, HEADS, HEAD_DIM)
+    assert _rel_l2(actual, expected) < 1e-4
+
+
+def test_kernels_reject_cpu_tensors():
+    with pytest.raises(ValueError):
+        mk.swiglu(torch.randn(4, 8))