[bugfix] FA4 JIT cache vs tuning env, MXFP8 import and scale dtype, H3 VSA gate double quantization (#2049)

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>

docs/contributing/env_vars.md

@@ -13,7 +13,9 @@ test in the same pull request.
    them directly with `os.environ.get("NAME")`, and the name must be in the external-variable allowlist
    (`EXTERNAL_ALLOWLIST` in the contract test). When FastVideo sets such a variable for the other tool, it calls
    `envs.set_external`, `envs.setdefault_external`, or `envs.unset_external`, and the name must be in
-   `EXTERNAL_WRITE_ALLOWLIST`. Variables that FastVideo's CI and CI tooling define (for example `TEST_SCOPE` and
+   `EXTERNAL_WRITE_ALLOWLIST`. Code that needs every variable of a family that another tool owns, such as FA4's
+   `FA4_*` tuning knobs, calls `envs.external_with_prefixes("FA4_")` with literal prefixes that the allowlist
+   covers. Variables that FastVideo's CI and CI tooling define (for example `TEST_SCOPE` and
    `PERF_RUN_SOURCE`) keep their names, and test code under `fastvideo/tests/` reads them directly; they are listed
    in `CI_ONLY_VARIABLES` in the contract test, together with the file that sets each one.
 2. **Read with `envs.NAME.get()`, write with `envs.NAME.set()`, and change a value in tests with
@@ -208,7 +210,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/attention/utils/fa4_env_cache.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Drop FA4's JIT kernel caches when its FA4_*/FA_* tuning environment changes.
+
+FA4 reads these knobs while tracing a kernel but does not include them in its compile-cache key.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+
+import fastvideo.envs as envs
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+FA4_ENV_PREFIXES = ("FA4_", "FA_")
+
+_JIT_CACHE_ATTRS = ("compile_cache", "compile_cache_pre", "compile_cache_post")
+
+_last_fingerprint: tuple[tuple[str, str], ...] | None = None
+
+
+def fa4_env_fingerprint(environ: Mapping[str, str] | None = None, ) -> tuple[tuple[str, str], ...]:
+    """Return the sorted ``(name, value)`` pairs of FA4 tuning env vars."""
+    source = envs.external_with_prefixes("FA4_", "FA_") if environ is None else environ
+    pairs = [(name, source[name]) for name in source if name.startswith(FA4_ENV_PREFIXES)]
+    pairs.sort()
+    return tuple(pairs)
+
+
+def clear_fa4_jit_caches(kernel_objects: Iterable[object]) -> int:
+    """Drop every FA4 JIT cache on ``kernel_objects``; return entries dropped."""
+    cleared = 0
+    for kernel_object in kernel_objects:
+        for attribute in _JIT_CACHE_ATTRS:
+            cache = getattr(kernel_object, attribute, None)
+            if isinstance(cache, dict) and cache:
+                cleared += len(cache)
+                cache.clear()
+    return cleared
+
+
+def reset_fa4_env_fingerprint() -> None:
+    """Forget the last observed fingerprint (tests / fresh-process semantics)."""
+    global _last_fingerprint
+    _last_fingerprint = None
+
+
+def invalidate_fa4_jit_caches_on_env_change(
+    kernel_objects: Iterable[object],
+    environ: Mapping[str, str] | None = None,
+) -> bool:
+    """Clear FA4's JIT caches if the FA4 tuning env changed since the last call (always on the first call)."""
+    global _last_fingerprint
+    fingerprint = fa4_env_fingerprint(environ)
+    if fingerprint == _last_fingerprint:
+        return False
+    previous, _last_fingerprint = _last_fingerprint, fingerprint
+    cleared = clear_fa4_jit_caches(kernel_objects)
+    if previous is not None:
+        changed = sorted(name for name in {k
+                                           for k, _ in previous} | {k
+                                                                    for k, _ in fingerprint}
+                         if dict(previous).get(name) != dict(fingerprint).get(name))
+        if cleared:
+            logger.info(
+                "FA4 tuning env changed (%s); dropped %d cached JIT kernel(s) "
+                "so the new setting is not served stale code", ", ".join(changed), cleared)
+        else:
+            logger.debug("FA4 tuning env changed (%s); FA4 JIT cache was empty", ", ".join(changed))
+    return True

fastvideo/attention/utils/flash_attn_cute.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import inspect
 from collections.abc import Callable
 
 import torch
 
+from fastvideo.attention.utils.fa4_env_cache import invalidate_fa4_jit_caches_on_env_change
 from fastvideo.logger import init_logger
 from fastvideo.platforms import current_platform
 
@@ -33,6 +35,17 @@ else:
     # This error will be caught in flash_attn.py or flash_attn_no_pad.py
     raise ImportError("flash_attn.cute is only available on CUDA devices; this error must be handled internally")
 
+# flash_attn.cute builds differ in which optional kwargs the varlen kernels take: some accept
+# max_seqlen_q/max_seqlen_k (and size the grid from them), older ones reject them. Forward an
+# optional kwarg only when the installed kernel declares it.
+_FWD_PARAMS = frozenset(inspect.signature(_flash_attn_fwd).parameters)
+_BWD_PARAMS = frozenset(inspect.signature(_flash_attn_bwd).parameters)
+
+
+def _accepted(params: frozenset[str], **kwargs) -> dict:
+    return {name: value for name, value in kwargs.items() if name in params}
+
+
 try:
     # FA2 serves the calls FA4 cute cannot on pre-sm90 GPUs (backward, GQA).
     # Optional so FA4-only installs can still import this module.
@@ -45,6 +58,13 @@ except ImportError:
 # Dynamo unwraps functools caches, so keep this cache inside the opaque helper.
 _SM90_OR_NEWER_BY_DEVICE: dict[int, bool] = {}
 
+# FA4's JIT cache key ignores its FA4_*/FA_* env knobs; the custom-op impls below run eagerly, so guard there.
+_FA4_KERNEL_OBJECTS = (_flash_attn_fwd, _flash_attn_bwd)
+
+
+def _guard_fa4_jit_cache() -> None:
+    invalidate_fa4_jit_caches_on_env_change(_FA4_KERNEL_OBJECTS)
+
 
 def _check_dropout(dropout_p: float) -> None:
     if dropout_p != 0.0:
@@ -94,6 +114,7 @@ def _flash_attn_cute_forward(
     causal: bool,
     deterministic: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    _guard_fa4_jit_cache()
     # _flash_attn_fwd returns (out, lse) on its empty-sequence early path but
     # (out, lse, p, row_max) on the main path at the pinned FA4 cute ref; take
     # the first two so both arities work.
@@ -235,14 +256,13 @@ def _flash_attn_cute_varlen_forward(
     causal: bool,
     deterministic: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    _guard_fa4_jit_cache()
     out, lse = _flash_attn_fwd(
         q,
         k,
         v,
         cu_seqlens_q=cu_seqlens_q,
         cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
         softmax_scale=softmax_scale,
         causal=causal,
         window_size_left=None,
@@ -251,6 +271,7 @@ def _flash_attn_cute_varlen_forward(
         num_splits=1,
         pack_gqa=None,
         return_lse=True,
+        **_accepted(_FWD_PARAMS, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k),
     )[:2]
     return out, lse
 
@@ -330,13 +351,14 @@ def _flash_attn_cute_varlen_backward_op(
         softmax_scale=softmax_scale,
         causal=causal,
         softcap=0.0,
-        window_size_left=None,
-        window_size_right=None,
         cu_seqlens_q=cu_seqlens_q,
         cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-        deterministic=deterministic,
+        **_accepted(_BWD_PARAMS,
+                    window_size_left=None,
+                    window_size_right=None,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    deterministic=deterministic),
     )
 
 
@@ -441,6 +463,7 @@ def _flash_attn_cute_fp4_forward(
     softmax_scale: float | None,
     causal: bool,
 ) -> torch.Tensor:
+    _guard_fa4_jit_cache()
     out = _flash_attn_fwd(
         q,
         k,

fastvideo/envs.py

@@ -222,6 +222,15 @@ def unset_external(name: str) -> None:
     os.environ.pop(name, None)
 
 
+def external_with_prefixes(*prefixes: str) -> dict[str, str]:
+    """Return the variables outside the registry whose names start with one of ``prefixes``.
+
+    For tools that read an open-ended family of variables, such as FA4's ``FA4_*`` tuning knobs. The contract test
+    checks each prefix against its read allowlist.
+    """
+    return {name: value for name, value in os.environ.items() if name.startswith(prefixes)}
+
+
 @contextmanager
 def override_external(name: str, value: str | None) -> Iterator[None]:
     """Set, or unset when ``value`` is None, a variable outside the registry, and restore it on exit.
@@ -427,7 +436,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/layers/mxfp8linear.py

@@ -7,7 +7,16 @@ import torch
 import torch.nn.functional as F
 import triton
 import triton.language as tl
-from quack.mx_utils import to_blocked, to_mx
+
+# Quack moved these helpers to quack.blockscaled.quantize, and some images ship neither.
+try:
+    from quack.mx_utils import to_blocked, to_mx
+except ImportError:
+    try:
+        from quack.blockscaled.quantize import to_blocked, to_mx
+    except ImportError:
+        to_blocked = None
+        to_mx = None
 
 MXFP8_BLOCK_SIZE = 32
 _BLOCKS_PER_PROGRAM = 16
@@ -31,6 +40,11 @@ def _quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor
 
 def quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Prequantize a weight with Quack and return hardware-blocked scales."""
+    if to_mx is None or to_blocked is None:
+        raise ImportError("MXFP8 weight quantization requires Quack's to_mx/to_blocked helpers "
+                          "(quack.mx_utils in older releases, quack.blockscaled.quantize in newer "
+                          "ones). Install a compatible quack to use the MXFP8 weight path; the "
+                          "Triton activation and SwiGLU quantizers do not need it.")
     _validate_mxfp8_matrix(matrix)
     return _quantize_mxfp8_weight_blockwise(matrix)
 
@@ -210,6 +224,13 @@ def swiglu_quantize_mxfp8_blockwise(preactivation: torch.Tensor) -> tuple[torch.
     return quantized, blocked_scales
 
 
+def _as_e8m0_scales(scales: torch.Tensor) -> torch.Tensor:
+    # F.scaled_mm rejects flashinfer's uint8 E8M0 scale buffers; they hold the same bytes as float8_e8m0fnu.
+    if scales.dtype == torch.uint8:
+        return scales.view(torch.float8_e8m0fnu)
+    return scales
+
+
 def mxfp8_scaled_mm(
     activation_values: torch.Tensor,
     activation_scales: torch.Tensor,
@@ -221,9 +242,9 @@ def mxfp8_scaled_mm(
     return F.scaled_mm(
         mat_a=activation_values,
         mat_b=weight_values.mT,
-        scale_a=activation_scales,
+        scale_a=_as_e8m0_scales(activation_scales),
         scale_recipe_a=F.ScalingType.BlockWise1x32,
-        scale_b=weight_scales,
+        scale_b=_as_e8m0_scales(weight_scales),
         scale_recipe_b=F.ScalingType.BlockWise1x32,
         swizzle_a=F.SwizzleType.SWIZZLE_32_4_4,
         swizzle_b=F.SwizzleType.SWIZZLE_32_4_4,

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
@@ -384,16 +403,18 @@ class MiniMaxH3Attention(nn.Module):
                 with STAGES.span("out_proj"):
                     hidden_states, _ = self.to_out(hidden_states)
                 return hidden_states
+        gate_linear = (self.to_gate_compress if self.to_gate_compress is not None and self._gate_active() else None)
         with STAGES.span("qkv_proj"):
-            # All three projections see the same activations. Reuse their FP8
-            # quantization when the loaded methods have identical granularity.
+            # q/k/v and the VSA gate see the same activations; share compatible input quantization.
             from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections
+            projections = (self.to_q, self.to_k, self.to_v) + ((gate_linear, ) if gate_linear is not None else ())
             if not torch.is_grad_enabled() and not torch.compiler.is_compiling():
-                query, key, value = _shared_input_projections((self.to_q, self.to_k, self.to_v), hidden_states)
+                projected = _shared_input_projections(projections, hidden_states)
             else:
-                query, _ = self.to_q(hidden_states)
-                key, _ = self.to_k(hidden_states)
-                value, _ = self.to_v(hidden_states)
+                projected = [projection(hidden_states)[0] for projection in projections]
+            query, key, value, *gate_out = projected
+            gate_compress = gate_out[0] if gate_out else None
+            del projected, gate_out
         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))
@@ -401,6 +422,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:
@@ -413,8 +450,7 @@ class MiniMaxH3Attention(nn.Module):
         # H3 rotates only 96/128 channels, which the generic `freqs_cis`
         # branch cannot express. Apply it above, then pass no RoPE here.
         extra_attention_kwargs = {}
-        if self.to_gate_compress is not None and self._gate_active():
-            gate_compress, _ = self.to_gate_compress(hidden_states)
+        if gate_compress is not None:
             extra_attention_kwargs["gate_compress"] = gate_compress.unflatten(
                 -1, (self.num_attention_heads, self.attention_head_dim))
         with STAGES.span("attention"):
@@ -658,6 +694,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 +710,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 +720,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 +731,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 +751,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 +780,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 +812,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 +874,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 +992,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,52 @@
+# 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 compute in fp32 and require the production H3 norm eps of ``1e-5``
+(enforced at the call sites). ``swiglu``, ``gated_residual`` and the qkv prologue
+round once at the store; ``rmsnorm_modulate`` also rounds its folded modulation
+tables and ``residual_gate_rmsnorm_modulate`` normalizes the stored residual, so
+those two match the default path to about 2.4e-3 rel-L2 rather than bitwise.
+"""
+
+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,192 @@
+# 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
+        vbase = 2 * plane.to(tl.int64)
+    else:
+        base = (batch * seq_len + tok) * num_heads * HD + h0 * HD
+        vbase = 2 * plane
+
+    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 + vbase + 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,239 @@
+# 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; the folded
+tables ``weight * (1 + scale)`` and ``shift`` are rounded to the activation dtype
+before the apply pass, so the output rounds twice::
+
+    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/models/dits/minimax_h3_vsa_fp4.py

@@ -186,25 +186,42 @@ def _layout_for(meta: MiniMaxH3VSAMetadata, rotary_emb: tuple[torch.Tensor, torc
     return layout
 
 
+def _prequantized_input_key(method: Any, linear: Any) -> tuple[Any, ...] | None:
+    """Key for linears that can share one pre-quantized input: unit-scale NVFP4, or FP8 by granularity."""
+    from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+    if type(method) is NVFP4QuantizeMethod:
+        if method.wants_prequantized_input() and method.uses_unit_activation_scale(linear):
+            return ("nvfp4", )
+        return None
+    if type(method) is FP8QuantizeMethod:
+        return ("fp8", method.granularity)
+    return None
+
+
 def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list[torch.Tensor]:
     """Share compatible FP8 preparation or unit-scale NVFP4 activation quantization.
 
     Calibrated or dynamic NVFP4 activation scales retain independent preparation.
     """
-    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
-    from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod
-
     methods = [getattr(linear, "quant_method", None) for linear in linears]
-    if any(m is None for m in methods):  # LoRA wrappers keep quant_method on base_layer.
-        return [linear(x)[0] for linear in linears]
-    same_nvfp4 = all(type(m) is NVFP4QuantizeMethod and m.wants_prequantized_input()
-                     and m.uses_unit_activation_scale(linear)
-                     for m, linear in zip(methods, linears, strict=True))
-    same_fp8 = all(type(m) is FP8QuantizeMethod and m.granularity == methods[0].granularity for m in methods)
-    if not (same_nvfp4 or same_fp8):
-        return [linear(x)[0] for linear in linears]
-    pre = methods[0].quantize_input(x)
-    return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)]
+    keys = [
+        _prequantized_input_key(m, linear) if m is not None else None
+        for m, linear in zip(methods, linears, strict=True)
+    ]
+    outputs: list[torch.Tensor | None] = [None] * len(linears)
+    groups: dict[tuple[Any, ...], list[int]] = {}
+    for index, key in enumerate(keys):
+        if key is None:
+            outputs[index] = linears[index](x)[0]
+        else:
+            groups.setdefault(key, []).append(index)
+    for indices in groups.values():
+        pre = methods[indices[0]].quantize_input(x)
+        for index in indices:
+            outputs[index] = methods[index].apply(linears[index], x, linears[index].bias, pre_quantized=pre)
+    return outputs  # type: ignore[return-value]
 
 
 @torch.compile(dynamic=True, fullgraph=True)
@@ -308,9 +325,11 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor,
     heads, dim = attn.num_attention_heads, attn.attention_head_dim
     with STAGES.span("tile_input"):
         x_tiles = layout.gather_in(hidden_states)[:, :logical]
+    gate_linear = attn.to_gate_compress if attn.to_gate_compress is not None and attn._gate_active() else None
+    projections = (attn.to_q, attn.to_k, attn.to_v) + ((gate_linear, ) if gate_linear is not None else ())
     with STAGES.span("qkv_proj"):
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles))
+        query, key, value, *gate_out = _shared_input_projections(projections, x_tiles)
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in (query, key, value))
     with STAGES.span("qknorm_rope"):
         cos, sin = layout.cos[:logical], layout.sin[:logical]
         if use_fused_rope:
@@ -320,11 +339,8 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor,
         else:
             query = attn._apply_rotary_emb(attn.norm_q(query), (cos, sin))
             key = attn._apply_rotary_emb(attn.norm_k(key), (cos, sin))
-    gate = None
-    if attn.to_gate_compress is not None and attn._gate_active():
-        with STAGES.span("gate_proj"):
-            gate, _ = attn.to_gate_compress(x_tiles)
-            gate = gate.unflatten(-1, (heads, dim))
+    gate = gate_out[0].unflatten(-1, (heads, dim)) if gate_out else None
+    del gate_out
     capture_root = envs.FASTVIDEO_H3_CAPTURE_QKV.get()
     if capture_root and attn._layer_idx in (0, 20, 41):
         from pathlib import Path
@@ -536,9 +552,11 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
         layout = _SPTileLayout(meta, rank, local_rows)
         meta._h3_fp4_sp_layout = layout  # type: ignore[attr-defined]
 
+    gate_linear = attn.to_gate_compress if attn.to_gate_compress is not None and attn._gate_active() else None
+    projections = (attn.to_q, attn.to_k, attn.to_v) + ((gate_linear, ) if gate_linear is not None else ())
     with STAGES.span("qkv_proj_rope"):
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states))
+        query, key, value, *gate_out = _shared_input_projections(projections, hidden_states)
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in (query, key, value))
         if use_fused_rope:
             from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
             cos, sin = rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype)
@@ -578,7 +596,7 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
     with STAGES.span("out_unpack"):
         out = _unpack_heads_fp8(payload, scale)  # [rows, H, D]
 
-    if attn.to_gate_compress is not None and attn._gate_active():
+    if gate_linear is not None:
         with STAGES.span("gate_compress"):
             v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
             out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(out.dtype)  # [Hs, n_tiles, D]
@@ -587,6 +605,5 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
             dist.all_gather_into_tensor(gathered, out_c.contiguous(), group=sp_group.device_group)
         with STAGES.span("gate_apply"):
             out_c_all = gathered.flatten(0, 1).transpose(0, 1)  # [n_tiles, H, D]
-            gate, _ = attn.to_gate_compress(hidden_states)
-            out = _apply_gate(out, out_c_all, layout.local_row_tile, gate[0].unflatten(-1, (heads, dim)))
+            out = _apply_gate(out, out_c_all, layout.local_row_tile, gate_out[0][0].unflatten(-1, (heads, dim)))
     return out.flatten(1, 2).unsqueeze(0)

fastvideo/tests/attention/test_fa4_env_cache.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FA4 JIT-cache invalidation on FA4_*/FA_* tuning-env changes.
+
+The CPU tests cover the fingerprint + invalidation logic with fake kernel
+objects, so they run on any machine. The last test needs a real FA4 install and
+a Blackwell/Hopper GPU: it shows that a knob change inside one process reuses
+the first arm's kernel without the guard (the bug) and compiles a second kernel
+with it (the fix).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.attention.utils import fa4_env_cache
+
+
+@pytest.fixture(autouse=True)
+def _fresh_fingerprint():
+    fa4_env_cache.reset_fa4_env_fingerprint()
+    yield
+    fa4_env_cache.reset_fa4_env_fingerprint()
+
+
+def _fake_kernels(*, fwd_entries: dict | None = None, bwd_entries: dict | None = None):
+    fwd = SimpleNamespace(compile_cache={} if fwd_entries is None else dict(fwd_entries))
+    bwd = SimpleNamespace(
+        compile_cache={} if bwd_entries is None else dict(bwd_entries),
+        compile_cache_pre={},
+        compile_cache_post={},
+    )
+    return fwd, bwd
+
+
+def test_fingerprint_selects_only_fa4_tuning_vars() -> None:
+    environ = {
+        "FA4_FP8_PV_USE_FUSED_PACK": "1",
+        "FA_DISABLE_2CTA": "0",
+        "FASTVIDEO_FA4": "1",  # a FastVideo flag, not an FA4 kernel knob
+        "PATH": "/usr/bin",
+    }
+    assert fa4_env_cache.fa4_env_fingerprint(environ) == (
+        ("FA4_FP8_PV_USE_FUSED_PACK", "1"),
+        ("FA_DISABLE_2CTA", "0"),
+    )
+
+
+def test_empty_environment_has_empty_fingerprint() -> None:
+    assert fa4_env_cache.fa4_env_fingerprint({}) == ()
+
+
+def test_first_call_drops_preexisting_kernels() -> None:
+    fwd, bwd = _fake_kernels(fwd_entries={"a": object()}, bwd_entries={"b": object()})
+
+    cleared = fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"})
+
+    assert cleared is True
+    assert fwd.compile_cache == {}
+    assert bwd.compile_cache == {}
+
+
+def test_unchanged_environment_keeps_the_compiled_kernel() -> None:
+    fwd, bwd = _fake_kernels()
+    environ = {"FA4_X": "1"}
+    assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), environ) is True
+    fwd.compile_cache["a"] = object()
+
+    assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), environ) is False
+    assert list(fwd.compile_cache) == ["a"]
+
+
+def test_changed_knob_drops_the_compiled_kernel() -> None:
+    fwd, bwd = _fake_kernels()
+    fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"})
+    fwd.compile_cache["a"] = object()
+
+    assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "0"}) is True
+    assert fwd.compile_cache == {}
+
+    fwd.compile_cache["a"] = object()
+    assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"}) is True
+    assert fwd.compile_cache == {}
+
+
+def test_unrelated_environment_changes_do_not_drop_kernels() -> None:
+    fwd, bwd = _fake_kernels()
+    fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {"FA4_X": "1"})
+    fwd.compile_cache["a"] = object()
+
+    assert fa4_env_cache.invalidate_fa4_jit_caches_on_env_change((fwd, bwd), {
+        "FA4_X": "1",
+        "CUDA_VISIBLE_DEVICES": "3",
+    }) is False
+    assert list(fwd.compile_cache) == ["a"]
+
+
+def test_missing_or_non_dict_caches_are_ignored() -> None:
+    bare = SimpleNamespace()
+    non_dict = SimpleNamespace(compile_cache=["not", "a", "dict"])
+
+    assert fa4_env_cache.clear_fa4_jit_caches((bare, non_dict)) == 0
+    assert non_dict.compile_cache == ["not", "a", "dict"]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FA4 needs a CUDA device")
+def test_real_fa4_kernel_is_recompiled_for_each_knob_value(monkeypatch) -> None:
+    """Regression test for the silent stale-kernel reuse.
+
+    Without the guard the second knob value is served the first value's kernel
+    (one ``cute.compile`` for two settings). With the guard each setting compiles
+    its own kernel.
+    """
+    interface = pytest.importorskip("flash_attn.cute.interface")
+    pytest.importorskip("flash_attn.cute.flash_fwd_sm100")
+    from fastvideo.attention.utils import flash_attn_cute
+
+    torch.manual_seed(0)
+    q = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(1, 512, 8, 128, device="cuda", dtype=torch.bfloat16)
+
+    compiles = {"n": 0}
+    original_compile = interface.cute.compile
+
+    def counting_compile(*args, **kwargs):
+        compiles["n"] += 1
+        return original_compile(*args, **kwargs)
+
+    monkeypatch.setattr(interface.cute, "compile", counting_compile)
+
+    def two_arms() -> int:
+        interface._flash_attn_fwd.compile_cache.clear()
+        fa4_env_cache.reset_fa4_env_fingerprint()
+        compiles["n"] = 0
+        for value in ("1", "0"):
+            with envs.override_external("FA_DISABLE_2CTA", value):
+                flash_attn_cute.flash_attn_func(q, k, v)
+        return compiles["n"]
+
+    monkeypatch.setattr(flash_attn_cute, "_guard_fa4_jit_cache", lambda: None)
+    assert two_arms() == 1
+
+    monkeypatch.undo()
+    monkeypatch.setattr(interface.cute, "compile", counting_compile)
+    assert two_arms() == 2

fastvideo/tests/attention/test_flash_attn_cute_varlen_kwargs.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FA4 varlen forwards optional kwargs only when the installed kernel declares them.
+
+``FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN=1`` routes H3 inference through the
+varlen custom ops. Some ``flash_attn.cute`` builds reject ``max_seqlen_q`` /
+``max_seqlen_k`` (and, in backward, ``window_size_*`` / ``deterministic``) with
+``TypeError: ... got an unexpected keyword argument``; the pinned build accepts
+them and sizes its grid from the max lengths, so they must still reach it.
+
+Each case installs a fake ``flash_attn.cute.interface`` with one of the two
+signatures, runs a packed-varlen forward and backward, and runs in a subprocess
+so the fake modules and the imported kernel references cannot leak into other
+tests.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+
+import pytest
+import torch
+
+_SKIP_EXIT_CODE = 77
+
+_SUBPROCESS_SCRIPT = r'''
+import importlib
+import importlib.machinery
+import json
+import sys
+import types
+
+import torch
+
+_SKIP_EXIT_CODE = 77
+
+import fastvideo  # noqa: F401  Import diffusers and the rest of the package before faking flash_attn.
+
+calls = {"fwd": [], "bwd": []}
+
+
+def _outputs(q):
+    return torch.empty_like(q), torch.empty((q.shape[1], q.shape[0]), dtype=torch.float32)
+
+
+def legacy_fwd(q, k, v, *, cu_seqlens_q=None, cu_seqlens_k=None, softmax_scale=None, causal=False,
+               window_size_left=None, window_size_right=None, softcap=0.0, num_splits=1, pack_gqa=None,
+               return_lse=True):
+    calls["fwd"].append({"cu_seqlens_q_is_input": cu_seqlens_q is cu_seqlens, "softmax_scale": softmax_scale})
+    return _outputs(q)
+
+
+def legacy_bwd(q, k, v, out, dout, lse, softmax_scale=None, causal=False, softcap=0.0, cu_seqlens_q=None,
+               cu_seqlens_k=None, seqused_q=None, seqused_k=None):
+    calls["bwd"].append({"cu_seqlens_q_is_input": cu_seqlens_q is cu_seqlens, "softmax_scale": softmax_scale})
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def pinned_fwd(q, k, v, *, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None,
+               softmax_scale=None, causal=False, window_size_left=None, window_size_right=None, softcap=0.0,
+               num_splits=1, pack_gqa=None, return_lse=True):
+    calls["fwd"].append({"cu_seqlens_q_is_input": cu_seqlens_q is cu_seqlens, "softmax_scale": softmax_scale,
+                         "max_seqlen_q": max_seqlen_q, "max_seqlen_k": max_seqlen_k})
+    return _outputs(q)
+
+
+def pinned_bwd(q, k, v, out, dout, lse, softmax_scale=None, causal=False, softcap=0.0, window_size_left=None,
+               window_size_right=None, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None,
+               max_seqlen_k=None, deterministic=False):
+    calls["bwd"].append({"cu_seqlens_q_is_input": cu_seqlens_q is cu_seqlens, "softmax_scale": softmax_scale,
+                         "max_seqlen_q": max_seqlen_q, "max_seqlen_k": max_seqlen_k,
+                         "deterministic": deterministic})
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+fake_fwd, fake_bwd = (pinned_fwd, pinned_bwd) if sys.argv[2] == "pinned" else (legacy_fwd, legacy_bwd)
+flash_attn = types.ModuleType("flash_attn")
+cute = types.ModuleType("flash_attn.cute")
+interface = types.ModuleType("flash_attn.cute.interface")
+for module, name in ((flash_attn, "flash_attn"), (cute, "flash_attn.cute"), (interface, "flash_attn.cute.interface")):
+    module.__spec__ = importlib.machinery.ModuleSpec(name, None)
+interface._flash_attn_fwd = fake_fwd
+interface._flash_attn_bwd = fake_bwd
+cute.interface = interface
+flash_attn.cute = cute
+sys.modules["flash_attn"] = flash_attn
+sys.modules["flash_attn.cute"] = cute
+sys.modules["flash_attn.cute.interface"] = interface
+
+try:
+    module = importlib.import_module("fastvideo.attention.utils.flash_attn_cute")
+except ImportError as exc:
+    print(exc, file=sys.stderr)
+    sys.exit(_SKIP_EXIT_CODE)
+module._flash_attn_fwd = fake_fwd
+module._flash_attn_bwd = fake_bwd
+
+q = torch.randn(7, 2, 8, device="cuda", requires_grad=True)
+k = torch.randn_like(q, requires_grad=True)
+v = torch.randn_like(q, requires_grad=True)
+cu_seqlens = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
+out, lse = torch.ops.fastvideo._flash_attn_cute_varlen_forward(q, k, v, cu_seqlens, cu_seqlens, 7, 7, 0.125, False,
+                                                              False)
+out.sum().backward()
+
+with open(sys.argv[1], "w") as handle:
+    json.dump({"out_shape": list(out.shape), "lse_shape": list(lse.shape), "q_grad": q.grad is not None,
+               "calls": calls}, handle)
+'''
+
+
+def _run(tmp_path, signature: str) -> dict:
+    if not torch.cuda.is_available():
+        pytest.skip("flash_attn.cute is only imported on CUDA builds")
+    result_path = tmp_path / f"varlen_{signature}.json"
+    completed = subprocess.run([sys.executable, "-c", _SUBPROCESS_SCRIPT,
+                                str(result_path), signature],
+                               capture_output=True,
+                               text=True,
+                               check=False)
+    if completed.returncode == _SKIP_EXIT_CODE:
+        pytest.skip(f"flash_attn.cute is not importable: {completed.stderr.strip()}")
+    assert completed.returncode == 0, completed.stderr
+    payload = json.loads(result_path.read_text())
+    assert payload["out_shape"] == [7, 2, 8]
+    assert payload["lse_shape"] == [2, 7]
+    assert payload["q_grad"]
+    assert len(payload["calls"]["fwd"]) == 1
+    assert len(payload["calls"]["bwd"]) == 1
+    for call in (payload["calls"]["fwd"][0], payload["calls"]["bwd"][0]):
+        assert call["cu_seqlens_q_is_input"]
+        assert call["softmax_scale"] == 0.125
+    return payload["calls"]
+
+
+def test_varlen_omits_kwargs_the_kernel_rejects(tmp_path) -> None:
+    calls = _run(tmp_path, "legacy")
+    assert "max_seqlen_q" not in calls["fwd"][0]
+    assert "max_seqlen_q" not in calls["bwd"][0]
+
+
+def test_varlen_forwards_max_seqlen_when_the_kernel_accepts_it(tmp_path) -> None:
+    calls = _run(tmp_path, "pinned")
+    assert calls["fwd"][0]["max_seqlen_q"] == 7
+    assert calls["fwd"][0]["max_seqlen_k"] == 7
+    assert calls["bwd"][0]["max_seqlen_q"] == 7
+    assert calls["bwd"][0]["max_seqlen_k"] == 7
+    assert calls["bwd"][0]["deterministic"] is False

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/contract/test_env_policy.py

@@ -6,7 +6,8 @@ The test parses every Python file under fastvideo/ (except fastvideo/third_party
 and the registry fastvideo/envs.py) with ``ast`` and reports code that:
 
 - reads the environment directly for a name outside EXTERNAL_ALLOWLIST (test
-  code under fastvideo/tests/ may also read CI_ONLY_VARIABLES);
+  code under fastvideo/tests/ may also read CI_ONLY_VARIABLES), or calls
+  envs.external_with_prefixes with a prefix that EXTERNAL_ALLOWLIST does not cover;
 - writes the environment directly (os.environ, os.putenv, monkeypatch.setenv),
   or through the envs.*_external helpers for a name outside
   EXTERNAL_WRITE_ALLOWLIST;
@@ -55,6 +56,8 @@ EXTERNAL_ALLOWLIST = {
     "NCCL_*": "NCCL communication library.",
     "TORCH_*": "PyTorch runtime settings.",
     "OMP_*": "OpenMP threading.",
+    "FA4_*": "FlashAttention-4 (flash_attn.cute) tuning knobs; fa4_env_cache.py fingerprints them.",
+    "FA_*": "FlashAttention-4 (flash_attn.cute) tuning knobs; fa4_env_cache.py fingerprints them.",
     "HF_HUB_*": "huggingface_hub settings.",
     "RANK": "Set by torchrun and other launchers.",
     "LOCAL_RANK": "Set by torchrun and other launchers.",
@@ -365,6 +368,12 @@ class EnvAccessScanner:
         elif (isinstance(func, ast.Attribute) and func.attr in EXTERNAL_WRITE_HELPERS
               and self._is_envs_module(func.value) and not is_external_write_allowed(func.attr, name, self.in_tests)):
             self.violations.append(("write", name, node.lineno))
+        elif (isinstance(func, ast.Attribute) and func.attr == "external_with_prefixes"
+              and self._is_envs_module(func.value)):
+            # A prefix read is allowed when an EXTERNAL_ALLOWLIST pattern covers the whole prefix.
+            for arg in node.args:
+                prefix = self._literal(arg)
+                self.violations.append(("read", prefix if prefix == "<dynamic>" else f"{prefix}*", node.lineno))
 
     def _check_registry_use(self, node: ast.Attribute, parent: ast.AST | None, in_function: bool) -> None:
         # Only a call counts: ``getter = envs.X.get`` neither reads the variable

fastvideo/tests/layers/test_mxfp8_optional_quack.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The in-tree MXFP8 module must import without Quack.
+
+``fastvideo.layers.mxfp8linear`` used to import ``to_blocked``/``to_mx`` from
+``quack.mx_utils`` at module scope. Quack is optional (and newer releases moved
+those helpers to ``quack.blockscaled.quantize``), so a missing or too-old Quack
+made the whole module -- and everything that imports the MXFP8 layer set --
+unimportable on an otherwise supported Blackwell image. Only weight
+prequantization needs Quack.
+"""
+
+from __future__ import annotations
+
+import builtins
+import importlib
+import sys
+
+import pytest
+import torch
+
+_MODULE_NAME = "fastvideo.layers.mxfp8linear"
+
+
+def _import_blocking_quack(monkeypatch: pytest.MonkeyPatch):
+    """Import the MXFP8 module while every ``quack`` import raises ImportError."""
+    real_import = builtins.__import__
+
+    def blocking_import(name, *args, **kwargs):
+        if name == "quack" or name.startswith("quack."):
+            raise ImportError(f"No module named {name!r} (blocked by the test)")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", blocking_import)
+    monkeypatch.delitem(sys.modules, _MODULE_NAME, raising=False)
+    try:
+        return importlib.import_module(_MODULE_NAME)
+    finally:
+        monkeypatch.delitem(sys.modules, _MODULE_NAME, raising=False)
+
+
+def test_mxfp8_module_imports_without_quack(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _import_blocking_quack(monkeypatch)
+    assert module.to_mx is None
+    assert module.to_blocked is None
+    assert callable(module.quantize_mxfp8_blockwise)
+    assert callable(module.swiglu_quantize_mxfp8_blockwise)
+    assert callable(module.mxfp8_scaled_mm)
+
+
+def test_weight_quantization_without_quack_raises_actionable_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _import_blocking_quack(monkeypatch)
+    with pytest.raises(ImportError, match="quack"):
+        module.quantize_mxfp8_weight_blockwise(torch.zeros(32, 64, dtype=torch.bfloat16))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Quack's MXFP8 helpers compile for CUDA tensors")
+def test_weight_quantization_with_quack_matches_blocked_layout() -> None:
+    pytest.importorskip("quack")
+    module = importlib.import_module(_MODULE_NAME)
+    if module.to_mx is None:
+        pytest.skip("the installed quack has neither mx_utils nor blockscaled.quantize")
+    values, scales = module.quantize_mxfp8_weight_blockwise(torch.randn(64, 128, dtype=torch.bfloat16, device="cuda"))
+    assert values.shape == (64, 128)
+    assert values.dtype == torch.float8_e4m3fn
+    assert scales.dtype == torch.float8_e8m0fnu
+    assert scales.numel() % 512 == 0

fastvideo/tests/layers/test_mxfp8_scaled_mm_scale_dtype.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MXFP8 GEMMs must accept flashinfer's ``uint8`` E8M0 scale buffers.
+
+``torch.nn.functional.scaled_mm`` requires ``torch.float8_e8m0fnu`` scales and
+rejects the ``uint8`` byte buffers that ``flashinfer.mxfp8_quantize`` returns
+with a generic "Invalid scaling configuration" ValueError. FastVideo's own
+Triton quantizers already return the e8m0 dtype, so the two paths disagreed on
+the same logical scale values.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from fastvideo.layers.mxfp8linear import (
+    _as_e8m0_scales,
+    mxfp8_scaled_mm,
+    quantize_mxfp8_blockwise,
+)
+
+
+def test_uint8_scales_are_presented_as_e8m0() -> None:
+    raw = torch.zeros(8, dtype=torch.uint8)
+    assert _as_e8m0_scales(raw).dtype == torch.float8_e8m0fnu
+    e8m0 = raw.view(torch.float8_e8m0fnu)
+    assert _as_e8m0_scales(e8m0) is e8m0
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+                    reason="the block-scaled MXFP8 GEMM needs a Blackwell GPU")
+def test_mxfp8_scaled_mm_accepts_uint8_and_e8m0_scales() -> None:
+    torch.manual_seed(0)
+    x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16)
+    w = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+    activation_values, activation_scales = quantize_mxfp8_blockwise(x)
+    weight_values, weight_scales = quantize_mxfp8_blockwise(w)
+
+    expected = mxfp8_scaled_mm(activation_values, activation_scales, weight_values, weight_scales, None)
+    from_uint8 = mxfp8_scaled_mm(
+        activation_values,
+        activation_scales.view(torch.uint8),
+        weight_values,
+        weight_scales.view(torch.uint8),
+        None,
+    )
+    torch.testing.assert_close(from_uint8, expected, atol=0, rtol=0)
+
+    reference = x.float() @ w.float().T
+    relative_error = (expected.float() - reference).norm() / reference.norm()
+    assert relative_error < 0.05, f"MXFP8 GEMM relative error {relative_error:.4f} is too large"

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))

fastvideo/tests/transformers/test_minimax_h3_mk_large_batch.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch
+
+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")
+
+
+def test_qkv_prologue_copies_v_above_int32_plane():
+    # B*S*H*D = 4 * 40000 * 56 * 128 = 1.14688e9 > 2**30, so ``2*plane``
+    # leaves int32 even though the base offsets are already int64.
+    if torch.cuda.get_device_properties(0).total_memory < 16 << 30:
+        pytest.skip("the shape needs ~13 GiB of device memory")
+    batch, seq, heads, head_dim, rope_dim = 4, 40000, 56, 128, 96
+    torch.manual_seed(0)
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    weight = torch.randn(head_dim, device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq, rope_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn(seq, rope_dim, device="cuda", dtype=torch.bfloat16)
+    out = mk.qkv_prologue(q, k, v, weight, weight, cos, sin)
+    torch.cuda.synchronize()
+    assert torch.equal(out[2 * batch:], v)

fastvideo/tests/transformers/test_minimax_h3_shared_input_quant.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+"""One activation quantization per compatible H3 attention projection group."""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+
+from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections
+
+
+class _FakeLinear(nn.Module):
+    """Minimal linear stand-in that records whether its own forward ran."""
+
+    def __init__(self, method: object) -> None:
+        super().__init__()
+        self.quant_method = method
+        self.bias = None
+        self.own_forwards = 0
+
+    def forward(self, x: torch.Tensor):
+        self.own_forwards += 1
+        return (x, None)
+
+
+class _FakeNVFP4Method:
+    """Version of ``NVFP4QuantizeMethod`` with a countable activation quantization."""
+
+    def __init__(self, unit_scale: bool = True) -> None:
+        self.unit_scale = unit_scale
+        self.quantize_calls = 0
+        self.apply_calls = 0
+
+    def wants_prequantized_input(self) -> bool:
+        return True
+
+    def uses_unit_activation_scale(self, layer: object) -> bool:
+        del layer
+        return self.unit_scale
+
+    def quantize_input(self, x: torch.Tensor):
+        self.quantize_calls += 1
+        return (x + 1, "scale", "global")
+
+    def apply(self, layer: object, x: torch.Tensor, bias: torch.Tensor | None, pre_quantized):
+        del x, bias
+        assert pre_quantized is not None
+        self.apply_calls += 1
+        return pre_quantized[0]
+
+
+def test_shared_input_projections_quantize_unit_scale_nvfp4_group_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    import fastvideo.layers.quantization.nvfp4_config as nvfp4
+
+    monkeypatch.setattr(nvfp4, "NVFP4QuantizeMethod", _FakeNVFP4Method)
+    methods = [_FakeNVFP4Method() for _ in range(4)]
+    linears = tuple(_FakeLinear(method) for method in methods)
+    x = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+    outputs = _shared_input_projections(linears, x)
+
+    assert [method.quantize_calls for method in methods] == [1, 0, 0, 0]
+    assert [method.apply_calls for method in methods] == [1, 1, 1, 1]
+    assert [linear.own_forwards for linear in linears] == [0, 0, 0, 0]
+    for output in outputs:
+        torch.testing.assert_close(output, x + 1)
+
+
+def test_shared_input_projections_leave_calibrated_nvfp4_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    import fastvideo.layers.quantization.nvfp4_config as nvfp4
+
+    monkeypatch.setattr(nvfp4, "NVFP4QuantizeMethod", _FakeNVFP4Method)
+    methods = [_FakeNVFP4Method(), _FakeNVFP4Method(), _FakeNVFP4Method(), _FakeNVFP4Method(unit_scale=False)]
+    linears = tuple(_FakeLinear(method) for method in methods)
+    x = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+    _shared_input_projections(linears, x)
+
+    assert [method.quantize_calls for method in methods] == [1, 0, 0, 0]
+    assert [method.apply_calls for method in methods] == [1, 1, 1, 0]
+    assert [linear.own_forwards for linear in linears] == [0, 0, 0, 1]
+
+
+def _install_fp8_buffers(layer, granularity: str, keep_weight: bool = False) -> None:
+    weight = layer.weight.data.float()
+    if granularity == "channel":
+        scale = (weight.abs().amax(dim=1, keepdim=True) / 448).clamp_min(1e-6)
+    else:
+        scale = (weight.abs().amax().reshape(1) / 448).clamp_min(1e-6)
+    layer.register_buffer("_fp8_weight", (weight / scale).to(torch.float8_e4m3fn))
+    layer.register_buffer("_fp8_weight_scale", scale)
+    if not keep_weight:
+        layer.register_parameter("weight", None)
+
+
+@pytest.mark.parametrize("granularity", ["tensor", "channel"])
+def test_shared_fp8_projections_cover_qkv_and_gate(granularity: str) -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9):
+        pytest.skip("sm89+ CUDA is required for FP8 GEMM")
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod
+
+    torch.manual_seed(23)
+    linears = tuple(
+        ReplicatedLinear(128, 256, bias=True, quant_config=FP8Config(granularity), prefix=f"block.attn.{name}")
+        for name in ("to_q", "to_k", "to_v", "to_gate_compress"))
+    linears[3].quant_method = FP8QuantizeMethod(granularity)
+    for layer in linears:
+        layer.to(device="cuda", dtype=torch.bfloat16)
+        layer.weight.data.normal_(std=0.1)
+        layer.bias.data.normal_(std=0.1)
+        _install_fp8_buffers(layer, granularity)
+    x = torch.randn(1, 272, 128, device="cuda", dtype=torch.bfloat16)
+    with torch.inference_mode():
+        reference = [layer(x)[0] for layer in linears]
+        with patch.object(FP8QuantizeMethod, "quantize_input", autospec=True,
+                          side_effect=FP8QuantizeMethod.quantize_input) as quant:
+            actual = _shared_input_projections(linears, x)
+        assert quant.call_count == 1
+        for expected, output in zip(reference, actual, strict=True):
+            torch.testing.assert_close(output, expected, atol=0, rtol=0)
+
+
+def test_shared_fp8_projections_keep_qkv_sharing_with_dense_gate() -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9):
+        pytest.skip("sm89+ CUDA is required for FP8 GEMM")
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod
+
+    torch.manual_seed(29)
+    linears = tuple(
+        ReplicatedLinear(128, 256, bias=True, quant_config=FP8Config("tensor"), prefix=f"block.attn.{name}")
+        for name in ("to_q", "to_k", "to_v"))
+    linears += (ReplicatedLinear(128, 256, bias=True, prefix="block.attn.to_gate_compress"), )
+    for layer in linears[:3]:
+        layer.to(device="cuda", dtype=torch.bfloat16)
+        layer.weight.data.normal_(std=0.1)
+        layer.bias.data.normal_(std=0.1)
+        _install_fp8_buffers(layer, "tensor")
+    linears[3].to(device="cuda", dtype=torch.bfloat16)
+    linears[3].weight.data.normal_(std=0.1)
+    linears[3].bias.data.normal_(std=0.1)
+    x = torch.randn(1, 272, 128, device="cuda", dtype=torch.bfloat16)
+    with torch.inference_mode():
+        reference = [layer(x)[0] for layer in linears]
+        with patch.object(FP8QuantizeMethod, "quantize_input", autospec=True,
+                          side_effect=FP8QuantizeMethod.quantize_input) as quant:
+            actual = _shared_input_projections(linears, x)
+        assert quant.call_count == 1
+        for expected, output in zip(reference, actual, strict=True):
+            torch.testing.assert_close(output, expected, atol=0, rtol=0)
+
+
+def test_generic_h3_attention_quantizes_qkv_and_gate_once(env_overrides, distributed_setup) -> None:
+    """The gate must reuse the q/k/v activation quantization, not prepare it again."""
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9):
+        pytest.skip("sm89+ CUDA is required for FP8 GEMM")
+    import fastvideo.layers.quantization.fp8_config as fp8_config
+    import fastvideo.envs as envs
+    from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.layers.quantization.fp8_config import FP8Config
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention
+    from fastvideo.platforms import AttentionBackendEnum
+
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3"))
+    env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1"))
+    env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(False))
+    env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False))
+    env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(False))
+    from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod
+
+    class _GateFP8Config(FP8Config):
+        """FP8 config that also tags the VSA compression gate (the NVFP4 profile's set)."""
+
+        def get_quant_method(self, layer, prefix):
+            if prefix.endswith("to_gate_compress"):
+                from fastvideo.layers.linear import LinearBase
+                if isinstance(layer, LinearBase):
+                    return FP8QuantizeMethod(self.granularity)
+            return super().get_quant_method(layer, prefix)
+
+    torch.manual_seed(31)
+    attn = MiniMaxH3Attention(256, 2, 128, 1e-5, (AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,),
+                              _GateFP8Config("channel"), "transformer_blocks.0.attn")
+    attn.to(device="cuda", dtype=torch.bfloat16)
+    for parameter in attn.parameters():
+        parameter.data.normal_(std=0.1)
+    for layer in (attn.to_q, attn.to_k, attn.to_v, attn.to_out):
+        _install_fp8_buffers(layer, "channel")
+    _install_fp8_buffers(attn.to_gate_compress, "channel", keep_weight=True)
+    meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=999, patch_size=(1, 1, 1),
+                                               VSA_sparsity=0.8, packed_segments=(65, 97, (4, 6, 10)),
+                                               device=torch.device("cuda"), tile_size=64)
+    length = meta.total_seq_length
+    x = torch.randn(1, length, 256, device="cuda", dtype=torch.bfloat16)
+    angles = torch.randn(length, 96, device="cuda")
+    rope = angles.cos(), angles.sin()
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta):
+        with patch.object(fp8_config, "_quantize_rowwise", side_effect=fp8_config._quantize_rowwise) as quant:
+            attn(x, rope, length)
+        assert quant.call_count == 2
+
+
+@pytest.mark.parametrize("tile_first", [False, True])
+def test_h3_attention_nvfp4_quantizes_qkv_and_gate_once(env_overrides, distributed_setup, tile_first: bool) -> None:
+    """The NVFP4 ``h3_dit_vsa`` profile prepares q/k/v/gate activations once."""
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("MiniMax-H3 NVFP4 inference requires an NVIDIA Blackwell GPU")
+    pytest.importorskip("flashinfer")
+    import fastvideo.envs as envs
+    import fastvideo.layers.quantization.nvfp4_config as nvfp4
+    from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention
+    from fastvideo.platforms import AttentionBackendEnum
+
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3"))
+    env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1"))
+    env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(False))
+    env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False))
+    env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(tile_first))
+    torch.manual_seed(37)
+    previous_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.bfloat16)
+    try:
+        attn = MiniMaxH3Attention(256, 2, 128, 1e-5, (AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,),
+                                  nvfp4.NVFP4Config(layer_profile="h3_dit_vsa"), "transformer_blocks.0.attn")
+    finally:
+        torch.set_default_dtype(previous_dtype)
+    attn = attn.cuda().eval()
+    for parameter in attn.parameters():
+        parameter.data.normal_(std=0.1)
+    nvfp4.convert_model_to_nvfp4(attn)
+    assert attn.to_q.weight is None
+    assert attn.to_gate_compress.weight is None
+    assert attn._gate_active()
+    meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=999, patch_size=(1, 1, 1),
+                                               VSA_sparsity=0.8, packed_segments=(65, 97, (4, 6, 10)),
+                                               device=torch.device("cuda"), tile_size=64)
+    length = meta.total_seq_length
+    x = torch.randn(1, length, 256, device="cuda", dtype=torch.bfloat16)
+    angles = torch.randn(length, 96, device="cuda")
+    rope = angles.cos(), angles.sin()
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta):
+        with patch.object(nvfp4, "_nvfp4_quantize", side_effect=nvfp4._nvfp4_quantize) as quant:
+            attn(x, rope, length)
+        assert quant.call_count == 2