diff --git a/python/sglang/kernels/ops/diffusion/hunyuan_qknorm.py b/python/sglang/kernels/ops/diffusion/hunyuan_qknorm.py new file mode 100644 index 000000000..ffd8716c2 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/hunyuan_qknorm.py @@ -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 diff --git a/python/sglang/kernels/ops/diffusion/triton/hunyuan_qkv_pack.py b/python/sglang/kernels/ops/diffusion/triton/hunyuan_qkv_pack.py new file mode 100644 index 000000000..0561dac95 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/triton/hunyuan_qkv_pack.py @@ -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)) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 5083bddb0..5f3bf4317 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 04711febc..1d5977532 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -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, + ), ) diff --git a/test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py b/test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py new file mode 100644 index 000000000..8f438beaf --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py @@ -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"]))