From 166b9b8a077690228df408a7611996c631a68ca5 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Sun, 1 Nov 2026 22:15:35 -0800 Subject: [PATCH] [perf] MiniMax H3: fuse VAE decoder QK RMSNorm and partial RoPE (#2041) Co-authored-by: Aryan Kumar --- docs/contributing/env_vars.md | 1 + fastvideo/envs.py | 6 + .../models/vaes/minimax_h3_qk_norm_rope.py | 118 ++++++++++++++++++ fastvideo/models/vaes/minimax_h3_video.py | 71 ++++++++--- .../vaes/test_minimax_h3_qk_norm_rope.py | 94 ++++++++++++++ ...est_minimax_h3_qk_norm_rope_fusion_gate.py | 69 ++++++++++ 6 files changed, 345 insertions(+), 14 deletions(-) create mode 100644 fastvideo/models/vaes/minimax_h3_qk_norm_rope.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope_fusion_gate.py diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index ae8d1d2f9..a6ef71769 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -281,6 +281,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_H3_ENCODER_LAYERWISE` | bool | `0` | performance | Stream MiniMax-H3 text-encoder language layers through exact-size pinned host memory (text-only prompts). | | `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT` | bool | `0` | performance | Expand the serialized NVFP4 MiniMax-H3 text encoder with one fused Triton pass on GPUs without FP4 GEMM. | | `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `8` | performance | Spatial tiles per MiniMax-H3 video VAE decoder call. The ViT decoder treats batch entries independently, so batching tiles matches per-tile decoding, except for NVFP4 decoders with dynamic activation scales, which always decode per tile. On GB200 at 768x1344, 8 tiles per call decodes 107 frames in 3.7 s instead of 5.4 s and 362 frames in 12.8 s instead of 18.3 s (+0.7 GiB peak). 1 decodes per tile. | +| `FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE` | bool | `0` | performance | Fuse the MiniMax-H3 video VAE decoder per-head QK RMSNorm and partial RoPE into one Triton pass per attention layer (same fp32/inference-dtype arithmetic; the RMSNorm reduction order differs, like the DiT FASTVIDEO_MINIMAX_H3_FUSIONS qknorm_rope). | | `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES` | int | `0` | performance | Entries in each content-keyed MiniMax-H3 memo of reference encodes (Qwen3-VL Ref2VA presentation, VAE keyframe latents); repeat references are reused exactly. 0 disables. | | `FASTVIDEO_H3_VAE_TILE_PARALLEL` | bool | `0` | performance | Split MiniMax-H3 VAE spatial tiles (not only temporal chunks) across the sequence-parallel ranks in the parallel decode and keyframe encode; bitwise equal to the serial path. Requires FASTVIDEO_H3_VAE_TILE_BATCH=1 (the split replays the per-tile path). | | `FASTVIDEO_H3_VAE_NVFP4_FUSED` | bool | `1` | performance | Run NVFP4 MiniMax-H3 VAE decoder blocks through the bit-exact fused Triton kernels at inference. | diff --git a/fastvideo/envs.py b/fastvideo/envs.py index d46a5cc4f..6c09305c9 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -732,6 +732,12 @@ FASTVIDEO_H3_VAE_TILE_BATCH = EnvInt(8, "always decode per tile. On GB200 at 768x1344, 8 tiles per call decodes 107 " "frames in 3.7 s instead of 5.4 s and 362 frames in 12.8 s instead of 18.3 s " "(+0.7 GiB peak). 1 decodes per tile.") +FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE = EnvBool(False, + category="performance", + doc="Fuse the MiniMax-H3 video VAE decoder per-head QK RMSNorm and " + "partial RoPE into one Triton pass per attention layer (same " + "fp32/inference-dtype arithmetic; the RMSNorm reduction order differs, " + "like the DiT FASTVIDEO_MINIMAX_H3_FUSIONS qknorm_rope).") FASTVIDEO_H3_REF2VA_MEMO_ENTRIES = EnvInt(0, category="performance", doc="Entries in each content-keyed MiniMax-H3 memo of reference encodes " diff --git a/fastvideo/models/vaes/minimax_h3_qk_norm_rope.py b/fastvideo/models/vaes/minimax_h3_qk_norm_rope.py new file mode 100644 index 000000000..2f7fc4a08 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_qk_norm_rope.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused QK RMSNorm + partial RoPE for the MiniMax-H3 video VAE decoder. + +One Triton pass per decoder attention layer replaces the eager prologue of +``MiniMaxH3VideoAttention.forward``: + + query = norm_q(query.float()).to(query.dtype) # fp32 RMSNorm, no affine + key = norm_k(key.float()).to(key.dtype) + q = cat([q_rot * cos + rotate_half(q_rot) * sin, q_pass]) + k = cat([k_rot * cos + rotate_half(k_rot) * sin, k_pass]) + +The arithmetic is the eager arithmetic (fp32 RMSNorm, RoPE on the activation dtype); +only the fp32 RMSNorm reduction order differs, because ``tl.sum`` reduces the 64 +per-head values with a different tree than ``at::native`` layer norm. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _qk_norm_rope_kernel( + q_ptr, + k_ptr, + cos_ptr, + sin_ptr, + seq_len, + num_heads, + eps, + HD: tl.constexpr, + RD: tl.constexpr, + HB: tl.constexpr, + NHG: tl.constexpr, +): + """Normalize and rope ``HB`` heads of one (batch, token) row of Q and K. + + ``q``/``k`` are contiguous ``[B, S, H, HD]``; ``cos``/``sin`` are contiguous + ``[B, S, 1, RD]``. ``HB`` heads share one program; ``HD`` and ``RD`` must be + powers of two and ``RD <= HD``. + """ + 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 + base = bt * (num_heads * HD) + h0 * HD + + idx = tl.arange(0, HD) + # rotate_half: out[i] = -x[i + RD/2] for i < RD/2, +x[i - RD/2] otherwise. + 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)) + + cb = (batch * seq_len + tok) * RD + rmask = idx < RD + cos = tl.load(cos_ptr + cb + idx, mask=rmask, other=1.0) + sin = tl.load(sin_ptr + cb + idx, mask=rmask, other=0.0) + dt = q_ptr.dtype.element_ty + # The eager path casts cos/sin to the activation dtype before the multiply. + cos = tl.where(rmask, cos.to(dt).to(tl.float32), 1.0) + sin = tl.where(rmask, sin.to(dt).to(tl.float32), 0.0) + sin = tl.where(idx < RD // 2, -sin, sin) + cos2 = tl.broadcast_to(cos[None, :], (HB, HD)) + sin2 = tl.broadcast_to(sin[None, :], (HB, HD)) + + for which in tl.static_range(2): + ptr = q_ptr if which == 0 else k_ptr + x = tl.load(ptr + base + off, mask=mask, other=0.0).to(tl.float32) + x = tl.reshape(x, (HB, HD)) + inv = tl.math.rsqrt(tl.sum(x * x, 1) / HD + eps)[:, None] + y = (x * inv).to(dt).to(tl.float32) + rotated = tl.gather(y, gidx, 1) + y = (y * cos2).to(dt).to(tl.float32) + (rotated * sin2).to(dt).to(tl.float32) + tl.store(ptr + base + off, tl.reshape(y, (NB, )).to(dt), mask=mask) + + +def _heads_per_program(num_heads: int) -> int: + """Largest power-of-two head group (<= 8) that divides ``num_heads``.""" + for group in (8, 4, 2): + if num_heads % group == 0: + return group + return 1 + + +def qk_norm_rope_( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + num_heads: int, + head_dim: int, + rope_dim: int, + eps: float = 1e-5, +) -> None: + """Fuse QK RMSNorm + partial RoPE in place on contiguous ``[B, S, H, D]`` Q/K.""" + batch, seq_len, heads, _ = query.shape + group = _heads_per_program(heads) + _qk_norm_rope_kernel[(batch * seq_len * triton.cdiv(heads, group), )]( + query, + key, + cos, + sin, + seq_len, + heads, + eps, + HD=head_dim, + RD=rope_dim, + HB=group, + NHG=triton.cdiv(heads, group), + num_warps=1, + ) diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index d24bec6cd..ff5b97528 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -294,6 +294,55 @@ class MiniMaxH3VideoRotaryPosEmbed(nn.Module): return angles.cos(), angles.sin() +def _apply_rotary( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the rotate-half partial RoPE to activation-dtype ``[B, S, H, D]`` Q/K.""" + rotary_dim = cos.shape[-1] + query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] + key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] + query_first, query_second = query_rotary.chunk(2, dim=-1) + key_first, key_second = key_rotary.chunk(2, dim=-1) + query_rotated = torch.cat([-query_second, query_first], dim=-1) + key_rotated = torch.cat([-key_second, key_first], dim=-1) + query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1) + key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1) + return query, key + + +def _rotary_tables_fit(query: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> bool: + batch, seq_len, head_dim = query.shape[0], query.shape[1], query.shape[-1] + if cos.dim() != 4 or cos.shape != sin.shape: + return False + if cos.shape[0] != batch or cos.shape[1] != seq_len or cos.shape[2] != 1: + return False + if cos.device != query.device or sin.device != query.device: + return False + if not cos.is_contiguous() or not sin.is_contiguous(): + return False + rotary_dim = cos.shape[-1] + return 0 < rotary_dim <= head_dim and rotary_dim % 2 == 0 + + +def _can_fuse_qk_norm_rope( + query: torch.Tensor, + key: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, +) -> bool: + """The Triton prologue needs CUDA, contiguous Q/K, packed per-token rotary tables, and a power-of-two head dim.""" + if rotary_emb is None or not query.is_cuda or not query.is_contiguous(): + return False + if not key.is_contiguous() or key.shape != query.shape or key.dtype != query.dtype or not query.is_floating_point(): + return False + if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad): + return False + head_dim = query.shape[-1] + return head_dim & (head_dim - 1) == 0 and _rotary_tables_fit(query, rotary_emb[0], rotary_emb[1]) + + class MiniMaxH3VideoAttention(nn.Module): def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-5, bias: bool = True) -> None: @@ -350,21 +399,15 @@ class MiniMaxH3VideoAttention(nn.Module): query, key, value = (projection.unflatten(2, (self.heads, -1)) for projection in projections) del projections - query = self.norm_q(query.float()).to(query.dtype) - key = self.norm_k(key.float()).to(key.dtype) - if rotary_emb is not None: + if envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.get() and _can_fuse_qk_norm_rope(query, key, rotary_emb): + from fastvideo.models.vaes.minimax_h3_qk_norm_rope import qk_norm_rope_ cos, sin = rotary_emb - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - rotary_dim = cos.shape[-1] - query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] - key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] - query_first, query_second = query_rotary.chunk(2, dim=-1) - key_first, key_second = key_rotary.chunk(2, dim=-1) - query_rotated = torch.cat([-query_second, query_first], dim=-1) - key_rotated = torch.cat([-key_second, key_first], dim=-1) - query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1) - key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1) + qk_norm_rope_(query, key, cos, sin, self.heads, query.shape[-1], cos.shape[-1], float(self.norm_q.eps)) + else: + query = self.norm_q(query.float()).to(query.dtype) + key = self.norm_k(key.float()).to(key.dtype) + if rotary_emb is not None: + query, key = _apply_rotary(query, key, rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype)) if self.attn_impl is not None and query.device.type != "cpu": # VAE decoding has no diffusion-step metadata, so call the selected diff --git a/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope.py b/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope.py new file mode 100644 index 000000000..f252a0351 --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE: fused VAE-decoder QK RMSNorm + partial RoPE.""" + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.models.vaes.minimax_h3_qk_norm_rope import qk_norm_rope_ +from fastvideo.models.vaes.minimax_h3_video import MiniMaxH3VideoAttention + + +def _eager_prologue(query, key, cos, sin, eps): + """The VAE decoder's eager QK prologue (the fused kernel's reference).""" + def rms_norm(x): + return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps) + + query = rms_norm(query).to(query.dtype) + key = rms_norm(key).to(key.dtype) + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + rotary_dim = cos.shape[-1] + + def rotate(x): + rotary, passthrough = x[..., :rotary_dim], x[..., rotary_dim:] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat([-second, first], dim=-1) + return torch.cat([rotary * cos + rotated * sin, passthrough], dim=-1) + + return rotate(query), rotate(key) + + +def test_fused_prologue_is_a_cpu_no_op(env_overrides) -> None: + """Without CUDA the env flag must leave the eager path (and its output) untouched.""" + torch.manual_seed(20260819) + attention = MiniMaxH3VideoAttention(dim=32, heads=2, dim_head=16, bias=True).eval() + hidden_states = torch.randn(2, 7, 32) + rotary_emb = (torch.randn(2, 7, 1, 12), torch.randn(2, 7, 1, 12)) + + with torch.no_grad(): + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.override(None)) + eager = attention(hidden_states, rotary_emb) + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.override(True)) + flagged = attention(hidden_states, rotary_emb) + + assert torch.equal(eager, flagged) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the fused kernel requires CUDA") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_fused_prologue_matches_the_eager_arithmetic(dtype: torch.dtype) -> None: + """Production geometry: 32 heads x 64 dims, partial RoPE on 48 dims.""" + torch.manual_seed(20260820) + batch, seq, heads, head_dim, rope_dim = 8, 337, 32, 64, 48 + query = (torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=dtype) * 0.02) + key = (torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=dtype) * 0.02) + cos = torch.rand(seq, rope_dim, device="cuda", dtype=torch.float32) + sin = torch.rand(seq, rope_dim, device="cuda", dtype=torch.float32) + reference_query, reference_key = _eager_prologue(query, key, cos[None, :, None], sin[None, :, None], 1e-5) + + fused_query, fused_key = query.clone(), key.clone() + qk_norm_rope_(fused_query, fused_key, cos[None, :, None].expand(batch, -1, -1, -1).contiguous(), + sin[None, :, None].expand(batch, -1, -1, -1).contiguous(), heads, head_dim, rope_dim) + + for name, fused, reference in (("q", fused_query, reference_query), ("k", fused_key, reference_key)): + difference = (fused.float() - reference.float()).abs() + bit_equal = (fused == reference).float().mean().item() + relative_l2 = (difference.norm() / reference.float().norm()).item() + # Only the fp32 RMSNorm reduction order differs; RoPE rounds exactly like eager. + assert bit_equal > 0.9995, name + assert relative_l2 < max(1e-5, 0.01 * torch.finfo(dtype).eps), name + assert difference.max().item() <= 4 * torch.finfo(dtype).eps * reference.float().abs().max().item() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the fused kernel requires CUDA") +@torch.inference_mode() +def test_attention_forward_matches_eager_within_one_ulp(env_overrides) -> None: + """A 32-head attention block produces the same output with the fused prologue. + + The prologue's RMSNorm reduction order is the only difference, so the + attention output lands within one activation-dtype ULP of the eager path. + """ + torch.manual_seed(20260821) + attention = MiniMaxH3VideoAttention(dim=2048, heads=32, dim_head=64, bias=True).cuda().eval() + hidden_states = torch.randn(1, 129, 2048, device="cuda") + rotary_emb = (torch.randn(1, 129, 1, 48, device="cuda"), torch.randn(1, 129, 1, 48, device="cuda")) + with torch.autocast("cuda", dtype=torch.float16): + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.override(None)) + eager = attention(hidden_states, rotary_emb) + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.override(True)) + fused = attention(hidden_states, rotary_emb) + + difference = (fused.float() - eager.float()).abs().max().item() + assert difference <= 2 * torch.finfo(torch.float16).eps * eager.float().abs().max().item() diff --git a/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope_fusion_gate.py b/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope_fusion_gate.py new file mode 100644 index 000000000..0afddc4a5 --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope_fusion_gate.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +from fastvideo.models.vaes.minimax_h3_video import _can_fuse_qk_norm_rope, _rotary_tables_fit + + +@pytest.fixture +def cuda_tensors(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + + +def _query_key(batch: int = 2, seq_len: int = 5, heads: int = 4, head_dim: int = 16): + return (torch.randn(batch, seq_len, heads, head_dim), torch.randn(batch, seq_len, heads, head_dim)) + + +def test_packed_rotary_tables_fuse(cuda_tensors) -> None: + query, key = _query_key() + cos = torch.randn(2, 5, 1, 8) + sin = torch.randn(2, 5, 1, 8) + + assert _can_fuse_qk_norm_rope(query, key, (cos, sin)) + + +def test_broadcast_rotary_tables_stay_eager(cuda_tensors) -> None: + query, key = _query_key() + cos = torch.randn(1, 5, 1, 8) + sin = torch.randn(1, 5, 1, 8) + + assert not _can_fuse_qk_norm_rope(query, key, (cos, sin)) + assert not _can_fuse_qk_norm_rope(query, key, (cos.expand(2, -1, -1, -1), sin.expand(2, -1, -1, -1))) + + +def test_non_contiguous_rotary_tables_stay_eager(cuda_tensors) -> None: + query, key = _query_key() + cos = torch.randn(2, 5, 1, 16)[..., :8] + sin = torch.randn(2, 5, 1, 16)[..., :8] + + assert not cos.is_contiguous() and not sin.is_contiguous() + assert not _can_fuse_qk_norm_rope(query, key, (cos, sin)) + assert not _rotary_tables_fit(query, cos, sin) + + +def test_rotary_tables_must_match_query_tokens(cuda_tensors) -> None: + query, key = _query_key() + cos = torch.randn(2, 6, 1, 8) + sin = torch.randn(2, 6, 1, 8) + + assert not _can_fuse_qk_norm_rope(query, key, (cos, sin)) + + +def test_grad_required_stays_eager(cuda_tensors) -> None: + query, key = _query_key() + query.requires_grad_(True) + cos = torch.randn(2, 5, 1, 8) + sin = torch.randn(2, 5, 1, 8) + + assert torch.is_grad_enabled() + assert not _can_fuse_qk_norm_rope(query, key, (cos, sin)) + with torch.no_grad(): + assert _can_fuse_qk_norm_rope(query, key, (cos, sin)) + + +def test_mismatched_key_dtype_stays_eager(cuda_tensors) -> None: + query, key = _query_key() + cos = torch.randn(2, 5, 1, 8) + sin = torch.randn(2, 5, 1, 8) + + assert not _can_fuse_qk_norm_rope(query, key.double(), (cos, sin))