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

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

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))

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, img_k = _hunyuan_qknorm(
-        img_q = self.img_attn_q_norm(img_q.contiguous()).to(img_v)
+            self, img_q, img_k, self.img_attn_q_norm, self.img_attn_k_norm
-        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),
         )
+
         # 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_q, txt_k = _hunyuan_qknorm(
-        txt_k = self.txt_attn_k_norm(txt_k.contiguous()).to(txt_k.dtype)
+            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),
+                q,
-                torch.cat((img_k, txt_k), dim=1),
+                k,
-                torch.cat((img_v, txt_v), dim=1),
+                v,
                 seq_lens=seq_lens,
             )
         else:
             attn = self.attn(
-                torch.cat((img_q, txt_q), dim=1),
+                q,
-                torch.cat((img_k, txt_k), dim=1),
+                k,
-                torch.cat((img_v, txt_v), dim=1),
+                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)
+        q, k = _hunyuan_qknorm(self, q, k, self.q_norm, self.k_norm)
-        k = self.k_norm(k.contiguous()).to(v.dtype)
 
         # 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 = (
+        q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
-            _apply_rotary_emb(img_q, cos, sin, is_neox_style=False),
-            _apply_rotary_emb(img_k, cos, sin, is_neox_style=False),
-        )
 
         # Run distributed attention
         if txt_is_sharded:
             attn_output = self.attn(
-                torch.cat((img_q, txt_q), dim=1),
+                q,
-                torch.cat((img_k, txt_k), dim=1),
+                k,
-                torch.cat((img_v, txt_v), dim=1),
+                v,
                 seq_lens=seq_lens,
             )
         else:
             attn_output = self.attn(
-                torch.cat((img_q, txt_q), dim=1),
+                q,
-                torch.cat((img_k, txt_k), dim=1),
+                k,
-                torch.cat((img_v, txt_v), dim=1),
+                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

@@ -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

@@ -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"]))