nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf] MiniMax H3: fuse VAE decoder QK RMSNorm and partial RoPE (#2041)

Browse Source 
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
2026-11-01 22:15:35 -08:00
committed by GitHub
co-authored by Aryan Kumar
parent 7e1ed5e28e
commit 166b9b8a07
6 changed files with 345 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/env_vars.md
+1
View File
@@ -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. |
fastvideo/envs.py
+6
View File
@@ -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 "
fastvideo/models/vaes/minimax_h3_qk_norm_rope.py
+118
View File
@@ -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,
)
fastvideo/models/vaes/minimax_h3_video.py
+57 -14
View File
@@ -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
fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope.py
+94
View File
@@ -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()
fastvideo/tests/vaes/test_minimax_h3_qk_norm_rope_fusion_gate.py
+69
View File
@@ -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))