minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion][HunyuanVideo] Fuse eager QKV packing and high-quality QKNorm (#34617)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-13 19:54:00 +08:00
committed by GitHub
parent 993e24df75
commit 3c1791a7df
5 changed files with 588 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/diffusion/hunyuan_qknorm.py
+99
View File
@@ -0,0 +1,99 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
from functools import cache
import torch
import torch.nn as nn
from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
logger = logging.getLogger(__name__)
_FUSION = QualityGatedFusion(
name="HunyuanVideo strided QK RMSNorm",
marker_attr="_sgl_hunyuan_qknorm_site",
enabled_attr="_sgl_hunyuan_qknorm_enabled",
)
@cache
def _get_qk_rmsnorm_cute():
try:
# Use FlashInfer's re-exported CuTe entry point directly: the public
# ``rmsnorm`` wrapper adds custom-op dispatch to every Hunyuan block.
from flashinfer.norm import qk_rmsnorm_cute
except ImportError:
return None
return qk_rmsnorm_cute
def mark_hunyuan_qknorm_site(module: nn.Module) -> None:
_FUSION.mark(module)
def _site_reject_reason(_site: nn.Module) -> str | None:
if _get_qk_rmsnorm_cute() is None:
return "FlashInfer CuTe QK RMSNorm unavailable"
return None
def mount_hunyuan_qknorm(root: nn.Module) -> bool:
return _FUSION.mount(
root,
reject_reason=_site_reject_reason,
logger=logger,
)
def unmount_hunyuan_qknorm(root: nn.Module) -> None:
_FUSION.unmount(root)
def try_hunyuan_qknorm(
site: nn.Module,
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor] | None:
"""Normalize strided Hunyuan Q/K views without materializing inputs."""
if not (
_FUSION.is_enabled(site)
and not torch.compiler.is_compiling()
and q.is_cuda
and q.dtype == torch.bfloat16
and k.dtype == q.dtype
and q_weight.dtype == q.dtype
and k_weight.dtype == q.dtype
and q.stride(-1) == 1
and k.stride(-1) == 1
):
return None
qk_rmsnorm_cute = _get_qk_rmsnorm_cute()
if qk_rmsnorm_cute is None:
return None
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
q_shape = q.shape
k_shape = k.shape
qk_rmsnorm_cute(
q.reshape(-1, q_shape[-2], q_shape[-1]),
q_weight,
q_out.reshape(-1, q_shape[-2], q_shape[-1]),
eps,
enable_pdl=True,
)
qk_rmsnorm_cute(
k.reshape(-1, k_shape[-2], k_shape[-1]),
k_weight,
k_out.reshape(-1, k_shape[-2], k_shape[-1]),
eps,
enable_pdl=True,
)
return q_out, k_out
python/sglang/kernels/ops/diffusion/triton/hunyuan_qkv_pack.py
+214
View File
@@ -0,0 +1,214 @@
# SPDX-License-Identifier: Apache-2.0
import torch
import triton
import triton.language as tl
@triton.autotune(
configs=[
triton.Config({"BLOCK_HEADS": 1, "BLOCK_HALF": 64}, num_warps=2),
triton.Config({"BLOCK_HEADS": 2, "BLOCK_HALF": 64}, num_warps=4),
triton.Config({"BLOCK_HEADS": 4, "BLOCK_HALF": 64}, num_warps=4),
triton.Config({"BLOCK_HEADS": 8, "BLOCK_HALF": 64}, num_warps=8),
],
key=["num_heads", "head_dim"],
)
@triton.jit
def _hunyuan_qkv_rope_pack_kernel(
img_q_ptr,
img_k_ptr,
img_v_ptr,
txt_q_ptr,
txt_k_ptr,
txt_v_ptr,
cos_ptr,
sin_ptr,
output_ptr,
img_tokens,
txt_tokens,
num_heads,
head_dim,
stride_iqb,
stride_iqs,
stride_iqh,
stride_ikb,
stride_iks,
stride_ikh,
stride_ivb,
stride_ivs,
stride_ivh,
stride_tqb,
stride_tqs,
stride_tqh,
stride_tkb,
stride_tks,
stride_tkh,
stride_tvb,
stride_tvs,
stride_tvh,
stride_cos,
stride_sin,
BLOCK_HEADS: tl.constexpr,
BLOCK_HALF: tl.constexpr,
):
token = tl.program_id(0)
head_block = tl.program_id(1)
total_tokens = img_tokens + txt_tokens
batch = token // total_tokens
seq = token - batch * total_tokens
heads = head_block * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
head_mask = heads < num_heads
half = tl.arange(0, BLOCK_HALF)
half_mask = half < head_dim // 2
mask = head_mask[:, None] & half_mask[None, :]
even = 2 * half
odd = even + 1
output_row = (
batch * total_tokens * num_heads * head_dim
+ seq * num_heads * head_dim
+ heads[:, None] * head_dim
)
plane_stride = tl.num_programs(0) * num_heads * head_dim
if seq < img_tokens:
q_row = (
img_q_ptr
+ batch * stride_iqb
+ seq * stride_iqs
+ heads[:, None] * stride_iqh
)
k_row = (
img_k_ptr
+ batch * stride_ikb
+ seq * stride_iks
+ heads[:, None] * stride_ikh
)
v_row = (
img_v_ptr
+ batch * stride_ivb
+ seq * stride_ivs
+ heads[:, None] * stride_ivh
)
cos_row = cos_ptr + seq * stride_cos + half
sin_row = sin_ptr + seq * stride_sin + half
cos = tl.load(cos_row, mask=half_mask, other=0.0).to(tl.float32)[None, :]
sin = tl.load(sin_row, mask=half_mask, other=0.0).to(tl.float32)[None, :]
q0 = tl.load(q_row + even[None, :], mask=mask, other=0.0)
q1 = tl.load(q_row + odd[None, :], mask=mask, other=0.0)
k0 = tl.load(k_row + even[None, :], mask=mask, other=0.0)
k1 = tl.load(k_row + odd[None, :], mask=mask, other=0.0)
v0 = tl.load(v_row + even[None, :], mask=mask, other=0.0)
v1 = tl.load(v_row + odd[None, :], mask=mask, other=0.0)
q0f, q1f = q0.to(tl.float32), q1.to(tl.float32)
k0f, k1f = k0.to(tl.float32), k1.to(tl.float32)
oq0 = tl.fma(-q1f, sin, q0f * cos)
oq1 = tl.fma(q0f, sin, q1f * cos)
ok0 = tl.fma(-k1f, sin, k0f * cos)
ok1 = tl.fma(k0f, sin, k1f * cos)
else:
txt_seq = seq - img_tokens
q_row = (
txt_q_ptr
+ batch * stride_tqb
+ txt_seq * stride_tqs
+ heads[:, None] * stride_tqh
)
k_row = (
txt_k_ptr
+ batch * stride_tkb
+ txt_seq * stride_tks
+ heads[:, None] * stride_tkh
)
v_row = (
txt_v_ptr
+ batch * stride_tvb
+ txt_seq * stride_tvs
+ heads[:, None] * stride_tvh
)
oq0 = tl.load(q_row + even[None, :], mask=mask, other=0.0).to(tl.float32)
oq1 = tl.load(q_row + odd[None, :], mask=mask, other=0.0).to(tl.float32)
ok0 = tl.load(k_row + even[None, :], mask=mask, other=0.0).to(tl.float32)
ok1 = tl.load(k_row + odd[None, :], mask=mask, other=0.0).to(tl.float32)
v0 = tl.load(v_row + even[None, :], mask=mask, other=0.0)
v1 = tl.load(v_row + odd[None, :], mask=mask, other=0.0)
tl.store(output_ptr + output_row + even[None, :], oq0, mask=mask)
tl.store(output_ptr + output_row + odd[None, :], oq1, mask=mask)
tl.store(output_ptr + plane_stride + output_row + even[None, :], ok0, mask=mask)
tl.store(output_ptr + plane_stride + output_row + odd[None, :], ok1, mask=mask)
tl.store(output_ptr + 2 * plane_stride + output_row + even[None, :], v0, mask=mask)
tl.store(output_ptr + 2 * plane_stride + output_row + odd[None, :], v1, mask=mask)
def hunyuan_qkv_rope_pack(
img_q: torch.Tensor,
img_k: torch.Tensor,
img_v: torch.Tensor,
txt_q: torch.Tensor,
txt_k: torch.Tensor,
txt_v: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
tensors = (img_q, img_k, img_v, txt_q, txt_k, txt_v)
if any(x.ndim != 4 for x in tensors):
raise ValueError("QKV tensors must have shape [B, S, H, D]")
if any(not x.is_cuda or x.dtype != torch.bfloat16 for x in tensors):
raise ValueError("QKV tensors must be CUDA bfloat16 tensors")
if any(x.device != img_q.device for x in tensors):
raise ValueError("QKV tensors must be on the same CUDA device")
batch, img_tokens, num_heads, head_dim = img_q.shape
txt_tokens = txt_q.shape[1]
expected_img = (batch, img_tokens, num_heads, head_dim)
expected_txt = (batch, txt_tokens, num_heads, head_dim)
if any(tuple(x.shape) != expected_img for x in (img_q, img_k, img_v)):
raise ValueError("image QKV shapes must match")
if any(tuple(x.shape) != expected_txt for x in (txt_q, txt_k, txt_v)):
raise ValueError("text QKV shapes must match")
if any(x.stride(-1) != 1 for x in tensors):
raise ValueError("QKV last dimensions must be contiguous")
if head_dim <= 0 or head_dim > 128 or head_dim % 2:
raise ValueError("head_dim must be positive, even, and <= 128")
if cos.ndim != 2 or sin.ndim != 2 or cos.shape != sin.shape:
raise ValueError("cos and sin must have matching [S, D/2] shapes")
if cos.shape[0] < img_tokens or cos.shape[1] != head_dim // 2:
raise ValueError("cos/sin shape does not cover image tokens and head_dim")
if not cos.is_cuda or not sin.is_cuda or cos.stride(-1) != 1 or sin.stride(-1) != 1:
raise ValueError("cos and sin must be CUDA and last-dim contiguous")
if cos.device != img_q.device or sin.device != img_q.device:
raise ValueError("QKV and cos/sin tensors must be on the same CUDA device")
total_tokens = img_tokens + txt_tokens
storage = torch.empty(
(3, batch, total_tokens, num_heads, head_dim),
device=img_q.device,
dtype=img_q.dtype,
)
args = []
for x in tensors:
args.extend((x.stride(0), x.stride(1), x.stride(2)))
with torch.cuda.device(img_q.device):
_hunyuan_qkv_rope_pack_kernel[
lambda meta: (
batch * total_tokens,
triton.cdiv(num_heads, meta["BLOCK_HEADS"]),
)
](
*tensors,
cos,
sin,
storage,
img_tokens,
txt_tokens,
num_heads,
head_dim,
*args,
cos.stride(0),
sin.stride(0),
)
return tuple(storage.unbind(dim=0))
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+172 -31
View File
@@ -8,6 +8,23 @@ import numpy as np
import torch
import torch.nn as nn
from sglang.kernels.ops.diffusion.bitexact_gate import (
BitExactFusionGate,
tensors_equal,
)
from sglang.kernels.ops.diffusion.fused_linear_gelu import (
can_fuse_linear_gelu,
fused_gelu_active,
fused_linear_gelu_tanh,
mark_fused_gelu_site,
)
from sglang.kernels.ops.diffusion.hunyuan_qknorm import (
mark_hunyuan_qknorm_site,
try_hunyuan_qknorm,
)
from sglang.kernels.ops.diffusion.triton.hunyuan_qkv_pack import (
hunyuan_qkv_rope_pack,
)
from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.fsdp import (
is_double_block,
@@ -61,6 +78,134 @@ from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
_HUNYUAN_QKV_PACK = BitExactFusionGate("HunyuanVideo QKV RoPE pack", per_signature=True)
_HUNYUAN_QKV_PACK_SIGS = _HUNYUAN_QKV_PACK.verified_sigs
assert _HUNYUAN_QKV_PACK_SIGS is not None
def _hunyuan_qknorm(
site: nn.Module,
q: torch.Tensor,
k: torch.Tensor,
q_norm: RMSNorm,
k_norm: RMSNorm,
) -> tuple[torch.Tensor, torch.Tensor]:
fused = try_hunyuan_qknorm(
site,
q,
k,
q_norm.weight,
k_norm.weight,
q_norm.variance_epsilon,
)
if fused is not None:
return fused
return q_norm(q.contiguous()).to(q), k_norm(k.contiguous()).to(k)
class HunyuanMLP(MLP):
"""Hunyuan MLP with a quality-gated cublasLt GELU epilogue."""
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
mark_fused_gelu_site(self, "fc_in")
def forward(self, x: torch.Tensor) -> torch.Tensor:
if fused_gelu_active(self) and can_fuse_linear_gelu(self.fc_in, x):
x = fused_linear_gelu_tanh(x, self.fc_in.weight, self.fc_in.bias)
else:
x, _ = self.fc_in(x)
x = self.act(x)
x, _ = self.fc_out(x)
return x
def _hunyuan_pack_qkv(
img_q: torch.Tensor,
img_k: torch.Tensor,
img_v: torch.Tensor,
txt_q: torch.Tensor,
txt_k: torch.Tensor,
txt_v: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Apply image RoPE and pack image/text QKV in one bit-exact kernel."""
if torch.compiler.is_compiling():
return _hunyuan_pack_qkv_reference(
img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin
)
sig = (
img_q.dtype,
img_q.device,
img_q.shape[0],
img_q.shape[2],
img_q.shape[3],
tuple(img_q.stride()[2:]),
tuple(txt_q.stride()[2:]),
cos.dtype,
sin.dtype,
)
verified = sig in _HUNYUAN_QKV_PACK_SIGS
can_attempt = (
not _HUNYUAN_QKV_PACK.disabled
and img_q.is_cuda
and img_q.dtype == torch.bfloat16
and img_q.shape[-1] <= 128
and img_q.shape[-1] % 2 == 0
and all(x.stride(-1) == 1 for x in (img_q, img_k, img_v, txt_q, txt_k, txt_v))
and (verified or not torch.cuda.is_current_stream_capturing())
)
if not can_attempt:
return _hunyuan_pack_qkv_reference(
img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin
)
try:
out = hunyuan_qkv_rope_pack(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
except Exception as exc:
_HUNYUAN_QKV_PACK.on_exception(exc, logger=logger)
return _hunyuan_pack_qkv_reference(
img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin
)
if verified:
return out
return _HUNYUAN_QKV_PACK.accept_or_fallback(
out,
_hunyuan_pack_qkv_reference(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin),
sig=sig,
equal=tensors_equal,
logger=logger,
mismatch_msg=(
"HunyuanVideo fused QKV RoPE pack is not bit-exact on this platform"
),
)
def _hunyuan_pack_qkv_reference(
img_q: torch.Tensor,
img_k: torch.Tensor,
img_v: torch.Tensor,
txt_q: torch.Tensor,
txt_k: torch.Tensor,
txt_v: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return (
torch.cat(
(_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q),
dim=1,
),
torch.cat(
(_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k),
dim=1,
),
torch.cat((img_v, txt_v), dim=1),
)
class MixedRowParallelLinear(RowParallelLinear):
@@ -157,7 +302,7 @@ class MMDoubleStreamBlock(nn.Module):
quant_config=quant_config,
)
self.img_mlp = MLP(
self.img_mlp = HunyuanMLP(
hidden_size,
mlp_hidden_dim,
bias=True,
@@ -209,7 +354,7 @@ class MMDoubleStreamBlock(nn.Module):
quant_config=quant_config,
)
self.txt_mlp = MLP(
self.txt_mlp = HunyuanMLP(
hidden_size,
mlp_hidden_dim,
bias=True,
@@ -225,6 +370,7 @@ class MMDoubleStreamBlock(nn.Module):
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn",
)
mark_hunyuan_qknorm_site(self)
def forward(
self,
@@ -269,15 +415,10 @@ class MMDoubleStreamBlock(nn.Module):
img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2]
# Apply QK-Norm if needed
img_q = self.img_attn_q_norm(img_q.contiguous()).to(img_v)
img_k = self.img_attn_k_norm(img_k.contiguous()).to(img_v)
# Apply rotary embeddings
cos, sin = freqs_cis
img_q, img_k = (
_apply_rotary_emb(img_q, cos, sin, is_neox_style=False),
_apply_rotary_emb(img_k, cos, sin, is_neox_style=False),
img_q, img_k = _hunyuan_qknorm(
self, img_q, img_k, self.img_attn_q_norm, self.img_attn_k_norm
)
# Prepare text for attention using fused operation
txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
@@ -292,22 +433,26 @@ class MMDoubleStreamBlock(nn.Module):
txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2]
# Apply QK-Norm if needed
txt_q = self.txt_attn_q_norm(txt_q.contiguous()).to(txt_q.dtype)
txt_k = self.txt_attn_k_norm(txt_k.contiguous()).to(txt_k.dtype)
txt_q, txt_k = _hunyuan_qknorm(
self, txt_q, txt_k, self.txt_attn_q_norm, self.txt_attn_k_norm
)
cos, sin = freqs_cis
q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
# Run distributed attention
if txt_is_sharded:
attn = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
q,
k,
v,
seq_lens=seq_lens,
)
else:
attn = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
q,
k,
v,
num_replicated_suffix=text_seq_len,
)
img_attn, txt_attn = attn.split([image_seq_len, text_seq_len], dim=1)
@@ -423,6 +568,7 @@ class MMSingleStreamBlock(nn.Module):
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn",
)
mark_hunyuan_qknorm_site(self)
def forward(
self,
@@ -456,33 +602,28 @@ class MMSingleStreamBlock(nn.Module):
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
# Apply QK-Norm
q = self.q_norm(q.contiguous()).to(v.dtype)
k = self.k_norm(k.contiguous()).to(v.dtype)
q, k = _hunyuan_qknorm(self, q, k, self.q_norm, self.k_norm)
# Split into image and text parts
img_q, txt_q = q[:, :-txt_len], q[:, -txt_len:]
img_k, txt_k = k[:, :-txt_len], k[:, -txt_len:]
img_v, txt_v = v[:, :-txt_len], v[:, -txt_len:]
# Apply rotary embeddings to image parts
cos, sin = freqs_cis
img_q, img_k = (
_apply_rotary_emb(img_q, cos, sin, is_neox_style=False),
_apply_rotary_emb(img_k, cos, sin, is_neox_style=False),
)
q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
# Run distributed attention
if txt_is_sharded:
attn_output = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
q,
k,
v,
seq_lens=seq_lens,
)
else:
attn_output = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
q,
k,
v,
num_replicated_suffix=txt_len,
)
attn_output = attn_output.view(batch_size, seq_len, -1)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+9
View File
@@ -32,6 +32,10 @@ from sglang.kernels.ops.diffusion.fused_ln_modulate import (
mount_fused_ln_modulate,
unmount_fused_ln_modulate,
)
from sglang.kernels.ops.diffusion.hunyuan_qknorm import (
mount_hunyuan_qknorm,
unmount_hunyuan_qknorm,
)
from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
mount_ltx2_rms_norm_modulate,
unmount_ltx2_rms_norm_modulate,
@@ -171,6 +175,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
mount_fused_gate_rmsnorm,
unmount_fused_gate_rmsnorm,
),
(
"HunyuanVideo strided QK RMSNorm",
mount_hunyuan_qknorm,
unmount_hunyuan_qknorm,
),
)
test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py
+94
View File
@@ -0,0 +1,94 @@
"""HunyuanVideo eager QKV/RoPE and quality-gated QKNorm tests."""
import sys
from unittest.mock import patch
import pytest
import torch
import sglang.kernels.ops.diffusion.hunyuan_qknorm as hunyuan_qknorm
from sglang.kernels.ops.diffusion.hunyuan_qknorm import (
mark_hunyuan_qknorm_site,
mount_hunyuan_qknorm,
unmount_hunyuan_qknorm,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
_apply_rotary_emb,
)
from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
_hunyuan_pack_qkv,
_hunyuan_qknorm,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("img_tokens,txt_tokens", [(257, 31), (4096, 256)])
def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
torch.manual_seed(0)
shape_img = (1, img_tokens, 24, 128)
shape_txt = (1, txt_tokens, 24, 128)
img_q, img_k, img_v = (
torch.randn(shape_img, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
txt_q, txt_k, txt_v = (
torch.randn(shape_txt, device="cuda", dtype=torch.bfloat16) for _ in range(3)
)
cos = torch.randn(img_tokens, 64, device="cuda")
sin = torch.randn_like(cos)
q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
q_ref = torch.cat(
(_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q), dim=1
)
k_ref = torch.cat(
(_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k), dim=1
)
v_ref = torch.cat((img_v, txt_v), dim=1)
assert torch.equal(q, q_ref)
assert torch.equal(k, k_ref)
assert torch.equal(v, v_ref)
def test_hunyuan_quality_qknorm_matches_rmsnorm():
torch.manual_seed(1)
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
q_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
k_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
packed = torch.randn(1, 257, 3, 24, 128, device="cuda", dtype=torch.bfloat16)
q, k = packed[:, :, 0], packed[:, :, 1]
q_ref = q_norm(q.contiguous()).to(q)
k_ref = k_norm(k.contiguous()).to(k)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
assert mount_hunyuan_qknorm(site)
q_out, k_out = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
torch.testing.assert_close(q_out, q_ref, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(k_out, k_ref, atol=2e-2, rtol=2e-2)
unmount_hunyuan_qknorm(site)
q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
assert torch.equal(q_unmounted, q_ref)
assert torch.equal(k_unmounted, k_ref)
def test_hunyuan_quality_qknorm_stays_unmounted_without_cute_kernel():
site = torch.nn.Module()
mark_hunyuan_qknorm_site(site)
with patch.object(hunyuan_qknorm, "_get_qk_rmsnorm_cute", return_value=None):
assert not mount_hunyuan_qknorm(site)
assert not hunyuan_qknorm._FUSION.is_enabled(site)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))