diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 037762fa4..62a002089 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: FASTVIDEO_TRACE_FUNCTION: int = 0 FASTVIDEO_ATTENTION_BACKEND: str | None = None FASTVIDEO_FA4: bool = False + FASTVIDEO_MINIMAX_H3_FUSIONS: str = "" FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn" FASTVIDEO_TARGET_DEVICE: str = "cuda" MAX_JOBS: str | None = None @@ -217,6 +218,13 @@ environment_variables: dict[str, Callable[[], Any]] = { "FASTVIDEO_FA4": lambda: os.getenv("FASTVIDEO_FA4", "0") != "0", + # Opt-in MiniMax-H3 inference-only Triton fusions adapted from the + # NVlabs/Sana Sol-Engine implementation. Accepts `all`, `1`, or a + # comma-separated subset of `modulate,qknorm_rope,swiglu`. An empty value + # (the default), `0`, or `none` keeps the eager implementation. + "FASTVIDEO_MINIMAX_H3_FUSIONS": + lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FUSIONS", ""), + # Use dedicated multiprocess context for workers. "FASTVIDEO_WORKER_MULTIPROC_METHOD": lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"), diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index f05f16013..d45c30eb3 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from fastvideo import envs from fastvideo.attention import DistributedAttention from fastvideo.attention.layer import DistributedAttention_VSA from fastvideo.attention.selector import get_attn_backend @@ -24,11 +25,39 @@ from fastvideo.layers.mlp import MLP from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.minimax_h3_fusions import ( + fused_qknorm_rope, + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, + minimax_h3_swiglu, +) from fastvideo.platforms import AttentionBackendEnum from fastvideo.utils import get_compute_dtype MINIMAX_H3_MODALITY_NUM = 3 _CFG = MiniMaxH3Config() +_MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"}) + + +def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: + """Parse the independently switchable inference fusion set.""" + raw = envs.FASTVIDEO_MINIMAX_H3_FUSIONS if value is None else value + normalized = raw.strip().lower() + if normalized in {"", "0", "none"}: + return frozenset() + if normalized in {"1", "all"}: + return _MINIMAX_H3_FUSION_NAMES + enabled = frozenset(item.strip() for item in normalized.split(",") if item.strip()) + unknown = enabled - _MINIMAX_H3_FUSION_NAMES + if unknown: + supported = ",".join(sorted(_MINIMAX_H3_FUSION_NAMES)) + raise ValueError(f"Unknown MiniMax H3 fusion(s) {sorted(unknown)}; expected a subset of {supported}.") + return enabled + + +def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool: + """Triton kernels are inference-only and stay outside Dynamo capture.""" + return tensor.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling() class MiniMaxH3RotaryPosEmbed(nn.Module): @@ -62,6 +91,7 @@ class MiniMaxH3FeedForward(nn.Module): ffn_dim: int, quant_config: QuantizationConfig | None = None, prefix: str = "", + fuse_swiglu: bool = False, ) -> None: super().__init__() self.fc_in = ReplicatedLinear( @@ -78,11 +108,15 @@ class MiniMaxH3FeedForward(nn.Module): quant_config=quant_config, prefix=f"{prefix}.fc_out", ) + self.fuse_swiglu = fuse_swiglu def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.fc_in(hidden_states) - hidden_states, gate = hidden_states.chunk(2, dim=-1) - hidden_states = hidden_states * F.silu(gate) + if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states): + hidden_states = minimax_h3_swiglu(hidden_states) + else: + hidden_states, gate = hidden_states.chunk(2, dim=-1) + hidden_states = hidden_states * F.silu(gate) hidden_states, _ = self.fc_out(hidden_states) return hidden_states @@ -99,6 +133,7 @@ class MiniMaxH3Attention(nn.Module): supported_attention_backends: tuple[AttentionBackendEnum, ...], quant_config: QuantizationConfig | None, prefix: str, + fuse_qknorm_rope: bool = False, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -134,6 +169,7 @@ class MiniMaxH3Attention(nn.Module): quant_config=quant_config, prefix=f"{prefix}.to_out", ) + self.fuse_qknorm_rope = fuse_qknorm_rope # VSA carries a learned gate on its pooled-compression branch. The H3 # checkpoint has no such weight, so the loader zero-initializes it # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled @@ -211,11 +247,18 @@ class MiniMaxH3Attention(nn.Module): query = query.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) key = key.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) value = value.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) - query = self.norm_q(query) - key = self.norm_k(key) - if rotary_emb is not None: - query = self._apply_rotary_emb(query, rotary_emb) - key = self._apply_rotary_emb(key, rotary_emb) + if (self.fuse_qknorm_rope and rotary_emb is not None and _can_run_minimax_h3_fusion(query)): + cos, sin = rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = fused_qknorm_rope(query, self.norm_q.weight, cos, sin, self.norm_q.eps) + key = fused_qknorm_rope(key, self.norm_k.weight, cos, sin, self.norm_k.eps) + else: + query = self.norm_q(query) + key = self.norm_k(key) + if rotary_emb is not None: + query = self._apply_rotary_emb(query, rotary_emb) + key = self._apply_rotary_emb(key, rotary_emb) # H3 rotates only 96/128 channels, which the generic `freqs_cis` # branch cannot express. Apply it above, then pass no RoPE here. @@ -397,6 +440,9 @@ class MiniMaxH3TransformerBlock(nn.Module): quant_config: QuantizationConfig | None, prefix: str, adaln_apply_silu: bool = True, + fuse_modulate: bool = False, + fuse_qknorm_rope: bool = False, + fuse_swiglu: bool = False, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -408,6 +454,7 @@ class MiniMaxH3TransformerBlock(nn.Module): supported_attention_backends, quant_config, prefix=f"{prefix}.attn", + fuse_qknorm_rope=fuse_qknorm_rope, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -415,6 +462,7 @@ class MiniMaxH3TransformerBlock(nn.Module): ffn_dim, quant_config=quant_config, prefix=f"{prefix}.ff", + fuse_swiglu=fuse_swiglu, ) self.adaln_proj = MiniMaxH3AdaLayerNormModulation( time_embed_dim, @@ -423,6 +471,7 @@ class MiniMaxH3TransformerBlock(nn.Module): prefix=f"{prefix}.adaln_proj", apply_silu=adaln_apply_silu, ) + self.fuse_modulate = fuse_modulate def forward( self, @@ -435,19 +484,39 @@ class MiniMaxH3TransformerBlock(nn.Module): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( t.to(hidden_states.dtype) for t in self.adaln_proj(temb)) - residual = hidden_states - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices) + use_modulate_fusion = self.fuse_modulate and _can_run_minimax_h3_fusion(hidden_states) + if use_modulate_fusion: + norm_hidden_states = fused_rmsnorm_modulate( + hidden_states, + self.norm1.weight, + scale_msa, + shift_msa, + adaln_indices, + self.norm1.eps, + ) + else: + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * ( + 1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices) attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len) - hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attention_output - - residual = hidden_states - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices) + if use_modulate_fusion: + hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( + hidden_states, + attention_output, + gate_msa, + self.norm2.weight, + scale_mlp, + shift_mlp, + adaln_indices, + self.norm2.eps, + ) + else: + hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * ( + 1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices) feed_forward_output = self.ff(norm_hidden_states) - return residual + gate_mlp.index_select(0, adaln_indices) * feed_forward_output + return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output class MiniMaxH3Transformer3DModel(BaseDiT): @@ -493,6 +562,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: super().__init__(config, hf_config) arch = config.arch_config + self.enabled_fusions = _enabled_minimax_h3_fusions() sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 if arch.num_attention_heads % sp_world_size: raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " @@ -590,6 +660,9 @@ class MiniMaxH3Transformer3DModel(BaseDiT): config.quant_config, prefix=f"{config.prefix}.transformer_blocks.{index}", adaln_apply_silu=self.adaln_rank is None, + fuse_modulate="modulate" in self.enabled_fusions, + fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions, + fuse_swiglu="swiglu" in self.enabled_fusions, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/models/dits/minimax_h3_fusions/__init__.py b/fastvideo/models/dits/minimax_h3_fusions/__init__.py new file mode 100644 index 000000000..81637d218 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_fusions/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference-only MiniMax H3 fusions adapted from NVlabs/Sana Sol-Engine. + +Source: https://github.com/NVlabs/Sana/tree/sol-engine/models/minimax_h3/GB200 +""" + +from .modulation import ( + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, +) +from .qknorm_rope import fused_qknorm_rope +from .swiglu import minimax_h3_swiglu + +__all__ = [ + "fused_qknorm_rope", + "fused_residual_gate_rmsnorm_modulate", + "fused_rmsnorm_modulate", + "minimax_h3_swiglu", +] diff --git a/fastvideo/models/dits/minimax_h3_fusions/modulation.py b/fastvideo/models/dits/minimax_h3_fusions/modulation.py new file mode 100644 index 000000000..123a79b18 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_fusions/modulation.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 RMSNorm and row-indexed modulation fusions.""" + +from __future__ import annotations + +import math + +import torch + +try: + import triton + import triton.language as tl +except ImportError as exc: # pragma: no cover - depends on the runtime image + triton = None + tl = None + _TRITON_IMPORT_ERROR: ImportError | None = exc +else: + _TRITON_IMPORT_ERROR = None + + +__all__ = [ + "fused_residual_gate_rmsnorm_modulate", + "fused_rmsnorm_modulate", +] + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) +_rmsnorm_modulate_kernel = None +_residual_gate_rmsnorm_modulate_kernel = None + + +if triton is not None: + + @triton.jit + def _rmsnorm_modulate_kernel( + out_ptr, + x_ptr, + weight_ptr, + scale_ptr, + shift_ptr, + index_ptr, + n_cols, + n_index, + eps, + stride_x_row, + stride_scale_row, + stride_shift_row, + BLOCK: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < n_cols + x_offsets = row * stride_x_row + cols + table_row = tl.load(index_ptr + row % n_index).to(tl.int64) + + x = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / n_cols + normed = x * tl.math.rsqrt(variance + eps) + weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32) + scale = tl.load( + scale_ptr + table_row * stride_scale_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + shift = tl.load( + shift_ptr + table_row * stride_shift_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + output = normed * weight * (1.0 + scale) + shift + tl.store(out_ptr + x_offsets, output.to(out_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def _residual_gate_rmsnorm_modulate_kernel( + hidden_out_ptr, + normed_out_ptr, + residual_ptr, + branch_ptr, + gate_ptr, + weight_ptr, + scale_ptr, + shift_ptr, + index_ptr, + n_cols, + n_index, + eps, + stride_input_row, + stride_gate_row, + stride_scale_row, + stride_shift_row, + BLOCK: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < n_cols + input_offsets = row * stride_input_row + cols + table_row = tl.load(index_ptr + row % n_index).to(tl.int64) + + residual = tl.load(residual_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32) + branch = tl.load(branch_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load( + gate_ptr + table_row * stride_gate_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + hidden = residual + gate * branch + tl.store(hidden_out_ptr + input_offsets, hidden.to(hidden_out_ptr.dtype.element_ty), mask=mask) + + variance = tl.sum(hidden * hidden, axis=0) / n_cols + normed = hidden * tl.math.rsqrt(variance + eps) + weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32) + scale = tl.load( + scale_ptr + table_row * stride_scale_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + shift = tl.load( + shift_ptr + table_row * stride_shift_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + output = normed * weight * (1.0 + scale) + shift + tl.store(normed_out_ptr + input_offsets, output.to(normed_out_ptr.dtype.element_ty), mask=mask) + + +def _validate_contract( + x: torch.Tensor, + weight: torch.Tensor, + tables: tuple[torch.Tensor, ...], + index: torch.Tensor, + eps: float, +) -> None: + if x.ndim < 2: + raise ValueError(f"x must have shape (..., sequence_length, hidden_size), got {tuple(x.shape)}.") + if x.numel() == 0 or x.shape[-1] == 0: + raise ValueError("x must not be empty.") + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"x must use float16, bfloat16, or float32, got {x.dtype}.") + hidden_size = x.shape[-1] + sequence_length = x.shape[-2] + if weight.shape != (hidden_size, ): + raise ValueError(f"weight must have shape ({hidden_size},), got {tuple(weight.shape)}.") + if weight.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"weight must use float16, bfloat16, or float32, got {weight.dtype}.") + if index.ndim != 1 or index.numel() != sequence_length: + raise ValueError( + f"index must have shape ({sequence_length},) so it can wrap over batch rows, got {tuple(index.shape)}." + ) + if index.dtype not in (torch.int32, torch.int64): + raise TypeError(f"index must use int32 or int64, got {index.dtype}.") + if not isinstance(eps, (float, int)) or isinstance(eps, bool) or not math.isfinite(eps) or eps <= 0: + raise ValueError(f"eps must be a positive finite number, got {eps!r}.") + + table_rows = tables[0].shape[0] if tables and tables[0].ndim == 2 else None + for name, table in zip(("gate", "scale", "shift")[-len(tables):], tables, strict=True): + if table.ndim != 2 or table.shape[1] != hidden_size: + raise ValueError(f"{name} must have shape (table_rows, {hidden_size}), got {tuple(table.shape)}.") + if table.shape[0] == 0 or table.shape[0] != table_rows: + raise ValueError("all modulation tables must have the same non-zero row count.") + if table.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"{name} must use float16, bfloat16, or float32, got {table.dtype}.") + + tensors = (x, weight, *tables, index) + if any(tensor.device != x.device for tensor in tensors[1:]): + raise ValueError("x, weight, modulation tables, and index must be on the same device.") + + +def _validate_residual_branch(residual: torch.Tensor, branch: torch.Tensor) -> None: + if branch.shape != residual.shape: + raise ValueError(f"branch must match residual shape {tuple(residual.shape)}, got {tuple(branch.shape)}.") + if branch.dtype != residual.dtype: + raise TypeError(f"branch dtype must match residual dtype {residual.dtype}, got {branch.dtype}.") + if branch.device != residual.device: + raise ValueError("branch and residual must be on the same device.") + + +def _require_triton_cuda(x: torch.Tensor) -> None: + if triton is None: + detail = f": {_TRITON_IMPORT_ERROR}" if _TRITON_IMPORT_ERROR is not None else "" + raise RuntimeError(f"MiniMax H3 modulation fusion requires Triton{detail}.") + if x.device.type != "cuda": + raise RuntimeError(f"MiniMax H3 modulation fusion requires CUDA tensors, got device {x.device}.") + + +def _require_forward_only(*tensors: torch.Tensor) -> None: + if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in tensors): + raise RuntimeError("MiniMax H3 modulation fusion is forward-only and does not support autograd.") + + +def _next_power_of_two(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _num_warps(block_size: int) -> int: + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +def _row_addressable(table: torch.Tensor) -> torch.Tensor: + return table if table.stride(-1) == 1 else table.contiguous() + + +def fused_rmsnorm_modulate( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Run RMSNorm and row-indexed modulation in one strict Triton kernel.""" + _validate_contract(x, weight, (scale, shift), index, eps) + _require_forward_only(x, weight, scale, shift) + _require_triton_cuda(x) + + hidden_size = x.shape[-1] + flat_x = x.reshape(-1, hidden_size).contiguous() + weight = weight.contiguous() + scale = _row_addressable(scale) + shift = _row_addressable(shift) + index = index.contiguous() + output = torch.empty_like(flat_x) + block_size = _next_power_of_two(hidden_size) + _rmsnorm_modulate_kernel[(flat_x.shape[0], )]( + output, + flat_x, + weight, + scale, + shift, + index, + hidden_size, + index.numel(), + eps, + flat_x.stride(0), + scale.stride(0), + shift.stride(0), + BLOCK=block_size, + num_warps=_num_warps(block_size), + ) + return output.view_as(x) + + +def fused_residual_gate_rmsnorm_modulate( + residual: torch.Tensor, + branch: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse residual update, row-indexed gate, RMSNorm, and modulation.""" + _validate_residual_branch(residual, branch) + _validate_contract(residual, weight, (gate, scale, shift), index, eps) + _require_forward_only(residual, branch, gate, weight, scale, shift) + _require_triton_cuda(residual) + + hidden_size = residual.shape[-1] + flat_residual = residual.reshape(-1, hidden_size).contiguous() + flat_branch = branch.reshape(-1, hidden_size).contiguous() + weight = weight.contiguous() + gate = _row_addressable(gate) + scale = _row_addressable(scale) + shift = _row_addressable(shift) + index = index.contiguous() + hidden = torch.empty_like(flat_residual) + modulated = torch.empty_like(flat_residual) + block_size = _next_power_of_two(hidden_size) + _residual_gate_rmsnorm_modulate_kernel[(flat_residual.shape[0], )]( + hidden, + modulated, + flat_residual, + flat_branch, + gate, + weight, + scale, + shift, + index, + hidden_size, + index.numel(), + eps, + flat_residual.stride(0), + gate.stride(0), + scale.stride(0), + shift.stride(0), + BLOCK=block_size, + num_warps=_num_warps(block_size), + ) + return hidden.view_as(residual), modulated.view_as(residual) diff --git a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py new file mode 100644 index 000000000..fdaf1d047 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused per-head RMSNorm and partial rotary embedding for MiniMax H3.""" + +from __future__ import annotations + +import math + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: # pragma: no cover - exercised only in environments without Triton + triton = None + tl = None + HAVE_TRITON = False + + +if HAVE_TRITON: + + @triton.jit + def _qknorm_partial_rope_kernel( + out_ptr, + x_ptr, + weight_ptr, + cos_ptr, + sin_ptr, + head_dim, + rotary_dim, + half_rotary_dim, + num_heads, + seq_len, + eps, + BLOCK_SIZE: tl.constexpr, + ): + row = tl.program_id(0) + seq_index = (row // num_heads) % seq_len + cols = tl.arange(0, BLOCK_SIZE) + head_mask = cols < head_dim + row_offset = row * head_dim + + x = tl.load(x_ptr + row_offset + cols, mask=head_mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / head_dim + inv_rms = tl.math.rsqrt(variance + eps) + weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32) + normalized = x * inv_rms * weight + + rotary_mask = cols < rotary_dim + first_half = cols < half_rotary_dim + partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim) + partner_x = tl.load( + x_ptr + row_offset + partner_col, + mask=rotary_mask, + other=0.0, + ).to(tl.float32) + partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32) + partner_normalized = partner_x * inv_rms * partner_weight + rotated = tl.where(first_half, -partner_normalized, partner_normalized) + + table_offset = seq_index * rotary_dim + cols + cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32) + sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32) + rotary_output = normalized * cos + rotated * sin + output = tl.where(rotary_mask, rotary_output, normalized) + tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask) + + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _validate_inputs( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> tuple[int, int, int, int, int]: + for name, tensor in (("x", x), ("weight", weight), ("cos", cos), ("sin", sin)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + + if x.ndim != 4: + raise ValueError(f"x must have shape (batch, seq, heads, head_dim), got {tuple(x.shape)}") + batch, seq_len, num_heads, head_dim = x.shape + if min(batch, seq_len, num_heads, head_dim) <= 0: + raise ValueError(f"x dimensions must all be positive, got {tuple(x.shape)}") + if weight.shape != (head_dim, ): + raise ValueError(f"weight must have shape ({head_dim},), got {tuple(weight.shape)}") + if cos.ndim != 2: + raise ValueError(f"cos must have shape (seq, rotary_dim), got {tuple(cos.shape)}") + if sin.shape != cos.shape: + raise ValueError(f"sin must match cos shape {tuple(cos.shape)}, got {tuple(sin.shape)}") + if cos.shape[0] != seq_len: + raise ValueError(f"cos/sin sequence length must be {seq_len}, got {cos.shape[0]}") + + rotary_dim = cos.shape[1] + if rotary_dim <= 0: + raise ValueError(f"rotary_dim must be positive, got {rotary_dim}") + if rotary_dim > head_dim: + raise ValueError(f"rotary_dim must not exceed head_dim, got rotary_dim={rotary_dim}, head_dim={head_dim}") + if rotary_dim % 2: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}") + + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"x dtype must be float16, bfloat16, or float32, got {x.dtype}") + for name, tensor in (("weight", weight), ("cos", cos), ("sin", sin)): + if tensor.dtype != x.dtype: + raise TypeError(f"{name} dtype must match x dtype {x.dtype}, got {tensor.dtype}") + if tensor.device != x.device: + raise ValueError(f"{name} device must match x device {x.device}, got {tensor.device}") + + if not isinstance(eps, (float, int)) or not math.isfinite(float(eps)) or eps <= 0: + raise ValueError(f"eps must be a positive finite number, got {eps!r}") + return batch, seq_len, num_heads, head_dim, rotary_dim + + +def fused_qknorm_rope( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel. + + RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the + final store. Triton's reduction order and the absence of eager's BF16 + intermediate materializations can produce small, expected rounding drift. + """ + batch, seq_len, num_heads, head_dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps) + if not weight.is_contiguous(): + raise ValueError("weight must be contiguous") + if not cos.is_contiguous() or not sin.is_contiguous(): + raise ValueError("cos and sin must be contiguous (seq, rotary_dim) tables") + if not x.is_cuda: + raise RuntimeError("fused_qknorm_rope requires CUDA tensors") + if not HAVE_TRITON: + raise RuntimeError("fused_qknorm_rope requires Triton") + if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in (x, weight, cos, sin)): + raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd") + + flat_x = x.reshape(-1, head_dim).contiguous() + flat_out = torch.empty_like(flat_x) + block_size = 1 << (head_dim - 1).bit_length() + _qknorm_partial_rope_kernel[(flat_x.shape[0], )]( + flat_out, + flat_x, + weight, + cos, + sin, + head_dim, + rotary_dim, + rotary_dim // 2, + num_heads, + seq_len, + eps, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return flat_out.view(batch, seq_len, num_heads, head_dim) + + +__all__ = ["HAVE_TRITON", "fused_qknorm_rope"] diff --git a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py new file mode 100644 index 000000000..b7884639e --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3's value-first packed SwiGLU fusion.""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: # pragma: no cover - exercised only in environments without Triton + triton = None + tl = None + HAVE_TRITON = False + + +def _validate_input(x: torch.Tensor) -> int: + if x.ndim == 0: + raise ValueError("MiniMax H3 SwiGLU expects at least one dimension") + + packed_width = x.shape[-1] + if packed_width == 0 or packed_width % 2 != 0: + raise ValueError( + "MiniMax H3 SwiGLU expects a positive even last dimension containing packed (value, gate) halves, " + f"got {packed_width}" + ) + if not x.is_floating_point(): + raise TypeError(f"MiniMax H3 SwiGLU expects a floating-point tensor, got {x.dtype}") + return packed_width // 2 + + +if HAVE_TRITON: + + @triton.jit + def _minimax_h3_swiglu_kernel( + out_ptr, + x_ptr, + ffn_dim, + stride_in_row, + stride_out_row, + BLOCK_SIZE: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < ffn_dim + + value = tl.load(x_ptr + row * stride_in_row + cols, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(x_ptr + row * stride_in_row + ffn_dim + cols, mask=mask, other=0.0).to(tl.float32) + # Match Sol-Engine: keep the complete SwiGLU expression in FP32 and + # convert only the final output store. + out = value * (gate * tl.sigmoid(gate)) + tl.store(out_ptr + row * stride_out_row + cols, out.to(out_ptr.dtype.element_ty), mask=mask) + +else: + _minimax_h3_swiglu_kernel = None + + +def _num_warps(block_size: int) -> int: + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor: + """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input. + + This is intentionally a strict kernel wrapper: callers own fallback policy and + must only invoke it for a supported CUDA inference path. + """ + ffn_dim = _validate_input(x) + if not x.is_cuda: + raise ValueError("MiniMax H3 fused SwiGLU requires a CUDA tensor") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"MiniMax H3 fused SwiGLU supports float16, bfloat16, and float32, got {x.dtype}") + if torch.is_grad_enabled() and x.requires_grad: + raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd") + if _minimax_h3_swiglu_kernel is None: + raise RuntimeError("MiniMax H3 fused SwiGLU requires Triton") + + packed_width = x.shape[-1] + flat = x.reshape(-1, packed_width).contiguous() + output_shape = (*x.shape[:-1], ffn_dim) + if flat.shape[0] == 0: + return torch.empty(output_shape, dtype=x.dtype, device=x.device) + + out = torch.empty((flat.shape[0], ffn_dim), dtype=x.dtype, device=x.device) + block_size = triton.next_power_of_2(ffn_dim) + _minimax_h3_swiglu_kernel[(flat.shape[0],)]( + out, + flat, + ffn_dim, + flat.stride(0), + out.stride(0), + BLOCK_SIZE=block_size, + num_warps=_num_warps(block_size), + ) + return out.view(output_shape) + + +__all__ = ["HAVE_TRITON", "minimax_h3_swiglu"] diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py new file mode 100644 index 000000000..3bf544351 --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused routing and FA4 integration checks for MiniMax-H3 fusions.""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.platforms import AttentionBackendEnum + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("", frozenset()), + ("0", frozenset()), + ("none", frozenset()), + ("all", frozenset({"modulate", "qknorm_rope", "swiglu"})), + ("1", frozenset({"modulate", "qknorm_rope", "swiglu"})), + ("swiglu, modulate", frozenset({"swiglu", "modulate"})), + ], +) +def test_minimax_h3_fusion_selector(raw: str, expected: frozenset[str]) -> None: + from fastvideo.models.dits.minimax_h3 import _enabled_minimax_h3_fusions + + assert _enabled_minimax_h3_fusions(raw) == expected + + +def test_minimax_h3_fusion_selector_rejects_unknown_name() -> None: + from fastvideo.models.dits.minimax_h3 import _enabled_minimax_h3_fusions + + with pytest.raises(ValueError, match="Unknown MiniMax H3 fusion"): + _enabled_minimax_h3_fusions("swiglu,unknown") + + +def test_swiglu_fusion_stays_on_eager_path_with_grad(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.models.dits.minimax_h3 as h3 + + def unexpected_kernel(_: torch.Tensor) -> torch.Tensor: + raise AssertionError("inference-only fusion ran with grad enabled") + + monkeypatch.setattr(h3, "minimax_h3_swiglu", unexpected_kernel) + layer = h3.MiniMaxH3FeedForward(8, 16, fuse_swiglu=True) + inputs = torch.randn(2, 3, 8, requires_grad=True) + layer(inputs).sum().backward() + + assert inputs.grad is not None + + +def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exercise the real block wiring without loading any H3 checkpoint.""" + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 CUDA is required") + pytest.importorskip("triton") + flash_attn = pytest.importorskip("flash_attn") + if "fa4" not in getattr(flash_attn, "__version__", "").lower(): + pytest.skip("the focused integration test requires the FA4 environment") + + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") + monkeypatch.setenv("FASTVIDEO_FA4", "1") + monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") + monkeypatch.setenv("MASTER_PORT", "29573") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_RANK", "0") + + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.forward_context import set_forward_context + from fastvideo.models.dits.minimax_h3 import MiniMaxH3RotaryPosEmbed, MiniMaxH3TransformerBlock + + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + kwargs = dict( + hidden_size=128, + num_attention_heads=1, + attention_head_dim=128, + ffn_dim=256, + time_embed_dim=64, + norm_eps=1e-5, + qk_norm_eps=1e-5, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, ), + quant_config=None, + prefix="minimax_h3.test_block", + ) + previous_default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + eager = MiniMaxH3TransformerBlock(**kwargs) + fused = MiniMaxH3TransformerBlock( + **kwargs, + fuse_modulate=True, + fuse_qknorm_rope=True, + fuse_swiglu=True, + ) + finally: + torch.set_default_dtype(previous_default_dtype) + + with torch.no_grad(): + for name, parameter in eager.named_parameters(): + if "norm" in name and name.endswith("weight"): + parameter.fill_(1.0) + elif parameter.ndim > 1: + torch.nn.init.normal_(parameter, mean=0.0, std=0.02) + else: + parameter.zero_() + fused.load_state_dict(eager.state_dict(), strict=True) + + device = torch.device("cuda") + eager = eager.to(device=device, dtype=torch.bfloat16).eval() + fused = fused.to(device=device, dtype=torch.bfloat16).eval() + generator = torch.Generator(device=device).manual_seed(2026) + hidden_states = torch.randn(2, 12, 128, generator=generator, device=device, dtype=torch.bfloat16) + temb = torch.randn(2, 64, generator=generator, device=device, dtype=torch.bfloat16) + adaln_indices = torch.arange(12, device=device, dtype=torch.long).remainder(6) + position_ids = torch.zeros(12, 3, device=device, dtype=torch.float32) + position_ids[:, 0] = torch.arange(12, device=device) + rotary_emb = MiniMaxH3RotaryPosEmbed(rope_freq_dim=16, rope_theta=10000.0).to(device)(position_ids) + inputs = dict( + hidden_states=hidden_states, + temb=temb, + adaln_indices=adaln_indices, + rotary_emb=tuple(value.to(torch.bfloat16) for value in rotary_emb), + original_seq_len=12, + ) + + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + eager_output = eager(**inputs) + fused_output = fused(**inputs) + + # Sol-Engine keeps fused intermediates in FP32 registers until their + # final BF16 stores, so the opt-in path is close but not bit-identical. + torch.testing.assert_close(fused_output, eager_output, atol=3e-2, rtol=3e-2) + finally: + cleanup_dist_env_and_memory() diff --git a/fastvideo/tests/transformers/test_minimax_h3_modulation_fusion.py b/fastvideo/tests/transformers/test_minimax_h3_modulation_fusion.py new file mode 100644 index 000000000..21c0ac20e --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_modulation_fusion.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn.functional as F + +from fastvideo.models.dits.minimax_h3_fusions.modulation import ( + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, +) + + +EPS = 1e-6 +SOL_ENGINE_BF16_TOLERANCE = 3e-2 + + +def _chunk_tables(rows: int, hidden_size: int, *, device: torch.device | str = "cpu", dtype=torch.float32): + wide = torch.randn(rows, 6 * hidden_size, device=device, dtype=dtype) + tables = wide.chunk(6, dim=-1) + assert all(table.stride() == (6 * hidden_size, 1) for table in tables) + assert all(not table.is_contiguous() for table in tables) + return tables + + +def _eager_rmsnorm_modulate(x, weight, scale, shift, index): + normed = F.rms_norm(x, (x.shape[-1], ), weight, EPS) + return normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index) + + +def _eager_residual_gate_rmsnorm_modulate(residual, branch, gate, weight, scale, shift, index): + hidden = residual + gate.index_select(0, index) * branch + normed = F.rms_norm(hidden, (hidden.shape[-1], ), weight, EPS) + modulated = normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index) + return hidden, modulated + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU") +def test_bf16_sol_engine_fusions_match_production_eager_within_tolerance(): + """Sol-Engine keeps fused intermediates in FP32 until its output stores.""" + pytest.importorskip("triton") + torch.manual_seed(2) + device = torch.device("cuda") + batch, sequence_length, hidden_size, table_rows = 2, 9, 5376, 6 + residual = torch.randn(batch, sequence_length, hidden_size, device=device, dtype=torch.bfloat16) + branch = torch.randn_like(residual) + weight = torch.randn(hidden_size, device=device, dtype=torch.bfloat16) + shift, scale, gate, shift_mlp, scale_mlp, _ = _chunk_tables( + table_rows, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) + index = torch.tensor([5, 0, 4, 1, 3, 2, 5, 1, 0], device=device, dtype=torch.int64) + + expected_norm1 = _eager_rmsnorm_modulate(residual, weight, scale, shift, index) + actual_norm1 = fused_rmsnorm_modulate(residual, weight, scale, shift, index, EPS) + expected_hidden, expected_norm2 = _eager_residual_gate_rmsnorm_modulate( + residual, + branch, + gate, + weight, + scale_mlp, + shift_mlp, + index, + ) + actual_hidden, actual_norm2 = fused_residual_gate_rmsnorm_modulate( + residual, + branch, + gate, + weight, + scale_mlp, + shift_mlp, + index, + EPS, + ) + + # Production eager materializes BF16 after each PyTorch operator. The + # single-kernel Sol-Engine path deliberately removes those round points. + torch.testing.assert_close( + actual_norm1, + expected_norm1, + rtol=SOL_ENGINE_BF16_TOLERANCE, + atol=SOL_ENGINE_BF16_TOLERANCE, + ) + torch.testing.assert_close( + actual_hidden, + expected_hidden, + rtol=SOL_ENGINE_BF16_TOLERANCE, + atol=SOL_ENGINE_BF16_TOLERANCE, + ) + torch.testing.assert_close( + actual_norm2, + expected_norm2, + rtol=SOL_ENGINE_BF16_TOLERANCE, + atol=SOL_ENGINE_BF16_TOLERANCE, + ) + + +@pytest.mark.parametrize( + ("mutate", "error", "match"), + [ + (lambda args: args | {"x": args["x"][0, 0]}, ValueError, "shape"), + (lambda args: args | {"weight": args["weight"][:-1]}, ValueError, "weight"), + (lambda args: args | {"index": args["index"][:-1]}, ValueError, "index"), + (lambda args: args | {"index": args["index"].float()}, TypeError, "index"), + (lambda args: args | {"eps": 0.0}, ValueError, "eps"), + ], +) +def test_fusion_rejects_invalid_contracts(mutate, error, match): + args = { + "x": torch.randn(2, 3, 8), + "weight": torch.randn(8), + "scale": torch.randn(4, 8), + "shift": torch.randn(4, 8), + "index": torch.tensor([0, 3, 1]), + "eps": EPS, + } + with pytest.raises(error, match=match): + fused_rmsnorm_modulate(**mutate(args)) + + +def test_residual_fusion_rejects_mismatched_branch(): + residual = torch.randn(2, 3, 8) + with pytest.raises(ValueError, match="branch"): + fused_residual_gate_rmsnorm_modulate( + residual, + torch.randn(2, 2, 8), + torch.randn(4, 8), + torch.randn(8), + torch.randn(4, 8), + torch.randn(4, 8), + torch.tensor([0, 1, 2]), + EPS, + ) + + +def test_triton_wrappers_fail_explicitly_on_cpu(): + x = torch.randn(1, 2, 8) + branch = torch.randn_like(x) + weight = torch.randn(8) + gate = torch.randn(3, 8) + scale = torch.randn(3, 8) + shift = torch.randn(3, 8) + index = torch.tensor([0, 2]) + + with pytest.raises(RuntimeError, match="Triton|CUDA"): + fused_rmsnorm_modulate(x, weight, scale, shift, index, EPS) + with pytest.raises(RuntimeError, match="Triton|CUDA"): + fused_residual_gate_rmsnorm_modulate(x, branch, gate, weight, scale, shift, index, EPS) + + +def test_triton_wrappers_reject_autograd_before_backend_check(): + x = torch.randn(1, 2, 8) + branch = torch.randn_like(x, requires_grad=True) + weight = torch.randn(8, requires_grad=True) + gate = torch.randn(3, 8) + scale = torch.randn(3, 8) + shift = torch.randn(3, 8) + index = torch.tensor([0, 2]) + + with pytest.raises(RuntimeError, match="forward-only"): + fused_rmsnorm_modulate(x, weight, scale, shift, index, EPS) + with pytest.raises(RuntimeError, match="forward-only"): + fused_residual_gate_rmsnorm_modulate( + x, + branch, + gate, + weight.detach(), + scale, + shift, + index, + EPS, + ) diff --git a/fastvideo/tests/transformers/test_minimax_h3_qknorm_rope_fusion.py b/fastvideo/tests/transformers/test_minimax_h3_qknorm_rope_fusion.py new file mode 100644 index 000000000..402a56d20 --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_qknorm_rope_fusion.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn.functional as F + +from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention +from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import ( + HAVE_TRITON, + fused_qknorm_rope, +) + + +def _rotary_tables( + seq_len: int, + rotary_dim: int, + *, + dtype: torch.dtype, + device: torch.device | str, +) -> tuple[torch.Tensor, torch.Tensor]: + angles = torch.randn(seq_len, rotary_dim // 2, dtype=torch.float32, device=device) + angles = torch.cat((angles, angles), dim=-1) + return angles.cos().to(dtype), angles.sin().to(dtype) + + +def _eager_qknorm_rope( + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> torch.Tensor: + normalized = F.rms_norm(x, (x.shape[-1], ), weight, eps) + return MiniMaxH3Attention._apply_rotary_emb(normalized, (cos, sin)) + + +def test_fused_qknorm_rope_rejects_invalid_rotary_dim() -> None: + x = torch.randn(2, 3, 4, 128) + weight = torch.ones(128) + + cos, sin = _rotary_tables(3, 130, dtype=x.dtype, device=x.device) + with pytest.raises(ValueError, match="rotary_dim must not exceed head_dim"): + fused_qknorm_rope(x, weight, cos, sin, 1e-6) + + cos = torch.randn(3, 95) + sin = torch.randn_like(cos) + with pytest.raises(ValueError, match="rotary_dim must be even"): + fused_qknorm_rope(x, weight, cos, sin, 1e-6) + + +def test_fused_qknorm_rope_rejects_shape_mismatches() -> None: + x = torch.randn(2, 3, 4, 128) + weight = torch.ones(128) + cos, sin = _rotary_tables(3, 96, dtype=x.dtype, device=x.device) + + with pytest.raises(ValueError, match="x must have shape"): + fused_qknorm_rope(x[0], weight, cos, sin, 1e-6) + with pytest.raises(ValueError, match="weight must have shape"): + fused_qknorm_rope(x, weight[:-1], cos, sin, 1e-6) + with pytest.raises(ValueError, match="sequence length"): + fused_qknorm_rope(x, weight, cos[:-1], sin[:-1], 1e-6) + with pytest.raises(ValueError, match="sin must match cos shape"): + fused_qknorm_rope(x, weight, cos, sin[:, :-2], 1e-6) + + +@pytest.mark.parametrize("noncontiguous_input", ["weight", "cos", "sin"]) +def test_fused_qknorm_rope_rejects_noncontiguous_linear_inputs(noncontiguous_input: str) -> None: + x = torch.randn(2, 3, 4, 128) + weight = torch.ones(256)[::2] + cos = torch.randn(3, 192)[:, ::2] + sin = torch.randn(3, 192)[:, ::2] + assert not weight.is_contiguous() + assert not cos.is_contiguous() + assert not sin.is_contiguous() + + if noncontiguous_input != "weight": + weight = weight.contiguous() + if noncontiguous_input != "cos": + cos = cos.contiguous() + if noncontiguous_input != "sin": + sin = sin.contiguous() + + with pytest.raises(ValueError, match="must be contiguous"): + fused_qknorm_rope(x, weight, cos, sin, 1e-6) + + +def test_fused_qknorm_rope_requires_matching_precast_dtype() -> None: + x = torch.randn(2, 3, 4, 128, dtype=torch.float32) + weight = torch.ones(128, dtype=torch.float32) + cos, sin = _rotary_tables(3, 96, dtype=torch.bfloat16, device=x.device) + + with pytest.raises(TypeError, match="cos dtype must match x dtype"): + fused_qknorm_rope(x, weight, cos, sin, 1e-6) + + +def test_fused_qknorm_rope_does_not_accept_missing_rotary_tables() -> None: + x = torch.randn(2, 3, 4, 128) + weight = torch.ones(128) + + with pytest.raises(TypeError, match="cos must be a torch.Tensor"): + fused_qknorm_rope(x, weight, None, None, 1e-6) # type: ignore[arg-type] + + +def test_fused_qknorm_rope_requires_cuda() -> None: + x = torch.randn(2, 3, 4, 128) + weight = torch.ones(128) + cos, sin = _rotary_tables(3, 96, dtype=x.dtype, device=x.device) + + with pytest.raises(RuntimeError, match="requires CUDA"): + fused_qknorm_rope(x, weight, cos, sin, 1e-6) + + +@pytest.mark.parametrize( + "rotary_dim,shape,use_input_view", + [ + pytest.param(96, (2, 11, 5, 128), False, id="partial-96-batch2-seq11-heads5"), + pytest.param(128, (2, 7, 3, 128), True, id="full-128-noncontiguous-input-view"), + ], +) +def test_fused_qknorm_rope_matches_eager_bf16_cuda( + rotary_dim: int, + shape: tuple[int, ...], + use_input_view: bool, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Triton fusion") + if not HAVE_TRITON: + pytest.skip("Triton is required for the fusion") + + torch.manual_seed(1) + device = torch.device("cuda") + if use_input_view: + x = torch.randn(*shape[:-1], shape[-1] * 2, dtype=torch.bfloat16, device=device)[..., ::2] + assert not x.is_contiguous() + else: + x = torch.randn(shape, dtype=torch.bfloat16, device=device) + weight = (1.0 + 0.05 * torch.randn(shape[-1], dtype=torch.bfloat16, device=device)).contiguous() + cos, sin = _rotary_tables(shape[1], rotary_dim, dtype=x.dtype, device=device) + + with torch.inference_mode(): + actual = fused_qknorm_rope(x, weight, cos, sin, 1e-6) + expected = _eager_qknorm_rope(x, weight, cos, sin, 1e-6) + + # The fused kernel keeps RMSNorm and both RoPE products in FP32 registers + # until its final BF16 store. Eager materializes BF16 intermediates, and + # PyTorch/Triton reductions need not use the same summation order. + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert actual.shape == x.shape + assert actual.dtype == x.dtype + assert actual.is_contiguous() diff --git a/fastvideo/tests/transformers/test_minimax_h3_swiglu_fusion.py b/fastvideo/tests/transformers/test_minimax_h3_swiglu_fusion.py new file mode 100644 index 000000000..cd3af87b0 --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_swiglu_fusion.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused tests for MiniMax H3's value-first packed SwiGLU fusion.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from fastvideo.models.dits.minimax_h3_fusions.swiglu import minimax_h3_swiglu + + +def _require_bf16_triton_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("MiniMax H3 fused SwiGLU requires CUDA") + if not torch.cuda.is_bf16_supported(): + pytest.skip("MiniMax H3 fused SwiGLU parity requires BF16 support") + pytest.importorskip("triton", reason="MiniMax H3 fused SwiGLU requires Triton") + + +def _assert_bf16_parity(x: torch.Tensor) -> None: + value, gate = x.chunk(2, dim=-1) + expected = value * F.silu(gate) + actual = minimax_h3_swiglu(x) + assert actual.shape == (*x.shape[:-1], x.shape[-1] // 2) + assert actual.dtype == x.dtype + assert actual.device == x.device + # Sol-Engine keeps the full SwiGLU expression in FP32 until the output + # store, while eager F.silu materializes a BF16 intermediate. + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("last_dim", [0, 7]) +def test_minimax_h3_swiglu_rejects_nonpositive_or_odd_last_dimension(last_dim: int) -> None: + x = torch.empty((2, last_dim), dtype=torch.float32) + + with pytest.raises(ValueError, match="positive even last dimension"): + minimax_h3_swiglu(x) + + +def test_minimax_h3_swiglu_strict_wrapper_rejects_cpu() -> None: + with pytest.raises(ValueError, match="requires a CUDA tensor"): + minimax_h3_swiglu(torch.randn(2, 8)) + + +@pytest.mark.gpu +def test_minimax_h3_swiglu_multidimensional_bf16_gpu_parity() -> None: + _require_bf16_triton_cuda() + torch.manual_seed(0) + x = torch.randn((2, 3, 5, 66), device="cuda", dtype=torch.bfloat16) + + _assert_bf16_parity(x) + + +@pytest.mark.gpu +def test_minimax_h3_swiglu_noncontiguous_bf16_gpu_parity() -> None: + _require_bf16_triton_cuda() + torch.manual_seed(1) + storage = torch.randn((2, 3, 148), device="cuda", dtype=torch.bfloat16) + x = storage[..., ::2] + assert not x.is_contiguous() + + _assert_bf16_parity(x) + + +@pytest.mark.gpu +def test_minimax_h3_swiglu_real_ffn_dim_bf16_gpu() -> None: + _require_bf16_triton_cuda() + torch.manual_seed(2) + ffn_dim = 14336 + x = torch.randn((1, 2 * ffn_dim), device="cuda", dtype=torch.bfloat16) + + _assert_bf16_parity(x)