From 18b2a3510637a1941c57e32b976e55ea511eea8f Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:21:38 -0800 Subject: [PATCH 1/2] [bugfix] FA4: stop forwarding max_seqlen to the varlen kernel (#2043) Co-authored-by: Aryan Kumar --- fastvideo/attention/utils/flash_attn_cute.py | 26 ++- .../test_flash_attn_cute_varlen_kwargs.py | 149 ++++++++++++++++++ 2 files changed, 168 insertions(+), 7 deletions(-) create mode 100644 fastvideo/tests/attention/test_flash_attn_cute_varlen_kwargs.py diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index 00dafb06d..4b52da60b 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/fastvideo/attention/utils/flash_attn_cute.py @@ -1,5 +1,6 @@ from __future__ import annotations +import inspect from collections.abc import Callable import torch @@ -33,6 +34,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. @@ -241,8 +253,6 @@ def _flash_attn_cute_varlen_forward( 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 +261,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 +341,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), ) diff --git a/fastvideo/tests/attention/test_flash_attn_cute_varlen_kwargs.py b/fastvideo/tests/attention/test_flash_attn_cute_varlen_kwargs.py new file mode 100644 index 000000000..b064d2c67 --- /dev/null +++ b/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 From 0d1244a8699d3f2c2bd290e3bb08deed6b1e6e85 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:22:09 -0800 Subject: [PATCH 2/2] [perf] MiniMax H3: opt-in faster kernel tier for the DiT fused ops (#2042) Co-authored-by: Aryan Kumar --- docs/contributing/env_vars.md | 2 +- fastvideo/attention/layer.py | 36 ++- fastvideo/envs.py | 3 +- fastvideo/models/dits/minimax_h3.py | 99 ++++++-- .../dits/minimax_h3_fusions/mk/__init__.py | 52 ++++ .../minimax_h3_fusions/mk/gated_residual.py | 151 +++++++++++ .../minimax_h3_fusions/mk/qkv_prologue.py | 192 ++++++++++++++ .../mk/residual_gate_rmsnorm_modulate.py | 137 ++++++++++ .../minimax_h3_fusions/mk/rmsnorm_modulate.py | 239 ++++++++++++++++++ .../dits/minimax_h3_fusions/mk/swiglu.py | 123 +++++++++ .../tests/attention/test_forward_stacked.py | 40 +++ .../test_minimax_h3_fusion_modes.py | 35 +++ .../transformers/test_minimax_h3_mk_block.py | 120 +++++++++ .../test_minimax_h3_mk_kernels.py | 94 +++++++ .../test_minimax_h3_mk_large_batch.py | 26 ++ 15 files changed, 1322 insertions(+), 27 deletions(-) create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/__init__.py create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/residual_gate_rmsnorm_modulate.py create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/rmsnorm_modulate.py create mode 100644 fastvideo/models/dits/minimax_h3_fusions/mk/swiglu.py create mode 100644 fastvideo/tests/attention/test_forward_stacked.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_fusion_modes.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_mk_block.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_mk_kernels.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_mk_large_batch.py diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index e2cd6646a..80236691a 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -208,7 +208,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY` | str | unset | performance | Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather. | | `FASTVIDEO_H3_ENCODER_SPLIT` | bool | `0` | distributed | MiniMax-H3 runs the Qwen3-VL text encoder on a dedicated group of nodes, so the denoising ranks never load it. Same as FastVideoArgs.h3_encoder_split=True. | | `FASTVIDEO_H3_ENCODER_NODES` | int | `1` | distributed | Number of leading nodes that FASTVIDEO_H3_ENCODER_SPLIT dedicates to the MiniMax-H3 text encoder. Values below 1 count as 1. | -| `FASTVIDEO_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation. | +| `FASTVIDEO_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. The whole-tier value mk selects the MotionKernel kernels instead of the Sol-Engine ones. Empty, 0, or none keeps the eager implementation. | | `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS` | str | `""` | performance | MiniMax-H3 inference-only Triton kernels that reproduce the eager output bit for bit: all, 1, or a comma-separated subset of rope,modulate,swiglu. An op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keeps the fusion. Empty, 0, or none keeps the eager implementation. | | `FASTVIDEO_FSDP2_AUTOWRAP` | bool | `0` | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation. | | `FASTVIDEO_FSDP2_MIN_PARAMS` | int | `10000000` | performance | Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards. | diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index eef3c2e75..a64c6d604 100644 --- a/fastvideo/attention/layer.py +++ b/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) diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 56c9a8cb1..8268d95a7 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -427,7 +427,8 @@ FASTVIDEO_MINIMAX_H3_FUSIONS = EnvStr( "", category="performance", doc="MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of " - "modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.") + "modulate,qknorm_rope,swiglu. The whole-tier value mk selects the MotionKernel kernels instead of " + "the Sol-Engine ones. Empty, 0, or none keeps the eager implementation.") FASTVIDEO_MINIMAX_H3_EXACT_KERNELS = EnvStr( "", category="performance", diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index faa70ce60..3ce84cdd2 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -38,6 +38,7 @@ from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp, vsa_fp4_requested, vsa_tile_first_attention) from fastvideo.models.dits.minimax_h3_fusions import exact as exact_kernels +from fastvideo.models.dits.minimax_h3_fusions import mk as mk_kernels from fastvideo.models.dits.minimax_h3_fusions import ( HAVE_TRITON, fused_qknorm_rope, @@ -52,9 +53,13 @@ from fastvideo.utils import get_compute_dtype logger = init_logger(__name__) MINIMAX_H3_MODALITY_NUM = 3 +#: eps baked into the MotionKernel H3 kernels; their call sites require it. +_MK_NORM_EPS = 1e-5 _CFG = MiniMaxH3Config() _MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"}) _MINIMAX_H3_EXACT_KERNEL_NAMES = frozenset({"rope", "modulate", "swiglu"}) +#: Whole-tier opt-in that replaces the Sol kernels with the MotionKernel ones. +_MINIMAX_H3_MK_MODE = "mk" def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: @@ -65,7 +70,14 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: return frozenset() if normalized in {"1", "all"}: return _MINIMAX_H3_FUSION_NAMES + if normalized == _MINIMAX_H3_MK_MODE: + # ``mk`` is a whole-tier mode: every fusion is on, and each fused call + # site dispatches to the MotionKernel kernel instead of the Sol one. + return _MINIMAX_H3_FUSION_NAMES | {_MINIMAX_H3_MK_MODE} enabled = frozenset(item.strip() for item in normalized.split(",") if item.strip()) + if _MINIMAX_H3_MK_MODE in enabled: + raise ValueError(f"'{_MINIMAX_H3_MK_MODE}' is a whole-tier mode and cannot be combined with individual " + f"fusions: {sorted(enabled - {_MINIMAX_H3_MK_MODE})}") unknown = enabled - _MINIMAX_H3_FUSION_NAMES if unknown: supported = ",".join(sorted(_MINIMAX_H3_FUSION_NAMES)) @@ -139,6 +151,7 @@ class MiniMaxH3FeedForward(nn.Module): prefix: str = "", fuse_swiglu: bool = False, exact_swiglu: bool = False, + mk_swiglu: bool = False, ) -> None: super().__init__() self.fc_in = ReplicatedLinear( @@ -157,6 +170,7 @@ class MiniMaxH3FeedForward(nn.Module): ) self.fuse_swiglu = fuse_swiglu self.exact_swiglu = exact_swiglu + self.mk_swiglu = mk_swiglu self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance( self.fc_out.quant_method, MXFP8QuantizeMethod) # Inference-only token chunking: the 2 * ffn_dim intermediate is ~5.3x the block input @@ -181,7 +195,10 @@ class MiniMaxH3FeedForward(nn.Module): return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out) hidden_states, _ = self.fc_in(hidden_states) if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states): - hidden_states = minimax_h3_swiglu(hidden_states) + if self.mk_swiglu and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling(): + hidden_states = mk_kernels.swiglu(hidden_states) + else: + hidden_states = minimax_h3_swiglu(hidden_states) elif self.exact_swiglu and exact_kernels.supports_rowwise(hidden_states): hidden_states = exact_kernels.swiglu(hidden_states) else: @@ -206,6 +223,7 @@ class MiniMaxH3Attention(nn.Module): fuse_qknorm_rope: bool = False, fa4_packed_varlen: bool = False, exact_rope: bool = False, + mk_qkv: bool = False, vsa_fp4: dict[str, Any] | None = None, ) -> None: super().__init__() @@ -244,6 +262,7 @@ class MiniMaxH3Attention(nn.Module): ) self.fuse_qknorm_rope = fuse_qknorm_rope self.exact_rope = exact_rope + self.mk_qkv = mk_qkv # VSA carries a learned gate on its pooled-compression branch. The H3 # checkpoint has no such weight, so the loader zero-initializes it # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled @@ -401,6 +420,22 @@ class MiniMaxH3Attention(nn.Module): cos, sin = rotary_emb cos = cos.to(query.dtype) sin = sin.to(query.dtype) + if (self.mk_qkv and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling() + and self.norm_q.eps == _MK_NORM_EPS + and self.norm_k.eps == _MK_NORM_EPS and type(self.distributed_attention) is DistributedAttention + and not (self.to_gate_compress is not None and self._gate_active())): + # One launch replaces the two Sol QK-norm/RoPE kernels and the + # ``torch.cat`` stack: the prologue writes a contiguous + # ``[3*B, S, H, D]`` buffer that ``forward_stacked`` feeds to the + # backend directly, so q/k/v are never copied again. + qkv = mk_kernels.qkv_prologue(query, key, value, self.norm_q.weight, self.norm_k.weight, cos, sin) + with STAGES.span("attention"): + hidden_states, _ = self.distributed_attention.forward_stacked( + qkv, original_seq_len=original_seq_len, attn_metadata=attn_metadata) + hidden_states = hidden_states.flatten(2, 3).type_as(qkv) + with STAGES.span("out_proj"): + hidden_states, _ = self.to_out(hidden_states) + return hidden_states query = fused_qknorm_rope(query, self.norm_q.weight, cos, sin, self.norm_q.eps) key = fused_qknorm_rope(key, self.norm_k.weight, cos, sin, self.norm_k.eps) else: @@ -658,6 +693,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fa4_packed_varlen: bool = False, exact_kernels: frozenset[str] = frozenset(), vsa_fp4: dict[str, Any] | None = None, + mk_tier: bool = False, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -673,6 +709,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fa4_packed_varlen=fa4_packed_varlen, exact_rope="rope" in exact_kernels, vsa_fp4=vsa_fp4, + mk_qkv=mk_tier, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -682,6 +719,7 @@ class MiniMaxH3TransformerBlock(nn.Module): prefix=f"{prefix}.ff", fuse_swiglu=fuse_swiglu, exact_swiglu="swiglu" in exact_kernels, + mk_swiglu=mk_tier, ) self.adaln_proj = MiniMaxH3AdaLayerNormModulation( time_embed_dim, @@ -692,6 +730,7 @@ class MiniMaxH3TransformerBlock(nn.Module): ) self.fuse_modulate = fuse_modulate self.exact_modulate = "modulate" in exact_kernels + self.mk_tier = mk_tier def forward( self, @@ -711,15 +750,20 @@ class MiniMaxH3TransformerBlock(nn.Module): and exact_kernels.supports_rowwise(hidden_states, shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp)) if use_modulate_fusion: - with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): - norm_hidden_states = fused_rmsnorm_modulate( - hidden_states, - self.norm1.weight, - scale_msa, - shift_msa, - adaln_indices, - self.norm1.eps, - ) + if (self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling() + and self.norm1.eps == _MK_NORM_EPS): + norm_hidden_states = mk_kernels.rmsnorm_modulate(hidden_states, self.norm1.weight, scale_msa, + shift_msa, adaln_indices) + else: + with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): + norm_hidden_states = fused_rmsnorm_modulate( + hidden_states, + self.norm1.weight, + scale_msa, + shift_msa, + adaln_indices, + self.norm1.eps, + ) elif use_exact_modulate: with nvtx_range("minimax_h3.transformer_block.exact_modulate"): norm_hidden_states = exact_kernels.modulate(self.norm1(hidden_states), scale_msa, shift_msa, @@ -735,17 +779,23 @@ class MiniMaxH3TransformerBlock(nn.Module): # overlaps three full-width activations during the residual fusion. del norm_hidden_states if use_modulate_fusion: - with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): - hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( - hidden_states, - attention_output, - gate_msa, - self.norm2.weight, - scale_mlp, - shift_mlp, - adaln_indices, - self.norm2.eps, - ) + if (self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling() + and self.norm2.eps == _MK_NORM_EPS): + hidden_states, norm_hidden_states = mk_kernels.residual_gate_rmsnorm_modulate( + hidden_states, attention_output, gate_msa, self.norm2.weight, scale_mlp, shift_mlp, + adaln_indices) + else: + with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): + hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( + hidden_states, + attention_output, + gate_msa, + self.norm2.weight, + scale_mlp, + shift_mlp, + adaln_indices, + self.norm2.eps, + ) elif use_exact_modulate: with nvtx_range("minimax_h3.transformer_block.exact_modulate"): hidden_states = exact_kernels.gate_residual(hidden_states, gate_msa, attention_output, adaln_indices) @@ -761,6 +811,8 @@ class MiniMaxH3TransformerBlock(nn.Module): with nvtx_range("minimax_h3.transformer_block.feed_forward"), STAGES.span("feed_forward"): feed_forward_output = self.ff(norm_hidden_states) if use_modulate_fusion and not torch.compiler.is_compiling(): + if self.mk_tier and mk_kernels.HAVE_TRITON and not torch.compiler.is_compiling(): + return mk_kernels.gated_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices) return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output) if use_exact_modulate: return exact_kernels.gate_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices) @@ -821,6 +873,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT): logger.warning( "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; " "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions))) + self.use_mk_kernels = _MINIMAX_H3_MK_MODE in self.enabled_fusions + if self.use_mk_kernels and not mk_kernels.HAVE_TRITON: + logger.warning("FASTVIDEO_MINIMAX_H3_FUSIONS=mk requested but the MotionKernel kernels are " + "unavailable (Triton missing); every forward stays on the eager path.") self.enabled_exact_kernels = _enabled_minimax_h3_exact_kernels() if self.enabled_exact_kernels: if exact_kernels.HAVE_TRITON: @@ -935,6 +991,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), exact_kernels=self.enabled_exact_kernels, vsa_fp4=config.vsa_fp4, + mk_tier=self.use_mk_kernels, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/__init__.py b/fastvideo/models/dits/minimax_h3_fusions/mk/__init__.py new file mode 100644 index 000000000..bdd5b4ce1 --- /dev/null +++ b/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", +] diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py b/fastvideo/models/dits/minimax_h3_fusions/mk/gated_residual.py new file mode 100644 index 000000000..73b8a5c95 --- /dev/null +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py b/fastvideo/models/dits/minimax_h3_fusions/mk/qkv_prologue.py new file mode 100644 index 000000000..e8c38dae8 --- /dev/null +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/residual_gate_rmsnorm_modulate.py b/fastvideo/models/dits/minimax_h3_fusions/mk/residual_gate_rmsnorm_modulate.py new file mode 100644 index 000000000..1224c2494 --- /dev/null +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/rmsnorm_modulate.py b/fastvideo/models/dits/minimax_h3_fusions/mk/rmsnorm_modulate.py new file mode 100644 index 000000000..28ce6afd5 --- /dev/null +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/mk/swiglu.py b/fastvideo/models/dits/minimax_h3_fusions/mk/swiglu.py new file mode 100644 index 000000000..526063d3b --- /dev/null +++ b/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) diff --git a/fastvideo/tests/attention/test_forward_stacked.py b/fastvideo/tests/attention/test_forward_stacked.py new file mode 100644 index 000000000..37673cd22 --- /dev/null +++ b/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 diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_modes.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_modes.py new file mode 100644 index 000000000..7be6788ed --- /dev/null +++ b/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") diff --git a/fastvideo/tests/transformers/test_minimax_h3_mk_block.py b/fastvideo/tests/transformers/test_minimax_h3_mk_block.py new file mode 100644 index 000000000..ba82748bb --- /dev/null +++ b/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 == [] diff --git a/fastvideo/tests/transformers/test_minimax_h3_mk_kernels.py b/fastvideo/tests/transformers/test_minimax_h3_mk_kernels.py new file mode 100644 index 000000000..c72d21778 --- /dev/null +++ b/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)) diff --git a/fastvideo/tests/transformers/test_minimax_h3_mk_large_batch.py b/fastvideo/tests/transformers/test_minimax_h3_mk_large_batch.py new file mode 100644 index 000000000..86ddcac4f --- /dev/null +++ b/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)