[Kernel] Fuse Q/K RMSNorm and RoPE (#5990)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/diffusion/layers/test_fused_qk_norm_rope.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+_HEAD_DIM = 128
+_ROTARY_DIM = 96
+_EPS = 1e-5
+
+
+def _reference(q, k, q_weight, k_weight, rope_table):
+    q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
+    k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
+    half = _ROTARY_DIM // 2
+    cos = rope_table[..., :half].unsqueeze(1)
+    sin = rope_table[..., half:].unsqueeze(1)
+
+    def apply(x):
+        first = x[..., :half]
+        second = x[..., half:_ROTARY_DIM]
+        return torch.cat(
+            (
+                first * cos - second * sin,
+                second * cos + first * sin,
+                x[..., _ROTARY_DIM:],
+            ),
+            dim=-1,
+        )
+
+    return apply(q), apply(k)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("seq_len", [1, 257, 1024])
+def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        fused_qk_norm_rope,
+    )
+
+    torch.manual_seed(17)
+    heads = 14
+    qkv = torch.randn(
+        seq_len,
+        heads * _HEAD_DIM * 3,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
+    k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
+        seq_len,
+        heads,
+        _HEAD_DIM,
+    )
+    q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
+    rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
+
+    expected_q, expected_k = _reference(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+    )
+    actual_q, actual_k = fused_qk_norm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+        _EPS,
+    )
+
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)

tests/e2e/accuracy/ltx/run_ltx_reference.py

@@ -260,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
         multimodal_output = frames.multimodal_output or {}
         audio = multimodal_output.get("audio")
         audio_sample_rate = multimodal_output.get("audio_sample_rate")
-        if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
-            frames = frames.request_output
-            multimodal_output = frames.multimodal_output or {}
-            audio = multimodal_output.get("audio", audio)
-            audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
-        if isinstance(frames, OmniRequestOutput):
-            if not frames.images:
-                raise ValueError("No video frames found in OmniRequestOutput")
-            frames = frames.images
+        if not frames.images:
+            raise ValueError("No video frames found in OmniRequestOutput")
+        frames = frames.images
 
     if isinstance(frames, list) and len(frames) == 1:
         frames = frames[0]

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1263,6 +1263,7 @@ def extract_minimax_h3_context(
         _BF16_DTYPE,
         _FORWARD_SUPPORTED_KWARGS,
         MINIMAX_H3_ADALN_MODALITY_NUM,
+        _build_rope_table,
         _modulate_scale_shift,
         _required_kwarg,
     )
@@ -1318,7 +1319,7 @@ def extract_minimax_h3_context(
         raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
     device = x.device
 
-    rope_freqs = module.rope(img_position_ids).to(device)
+    rope_table = _build_rope_table(module.rope(img_position_ids).to(device))
 
     decoder_input, t_emb = module._embed(
         x=x,
@@ -1351,7 +1352,7 @@ def extract_minimax_h3_context(
     def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
         hidden, block_rope, block_combined = module.sp_prepare(
             decoder_input,
-            rope_freqs,
+            rope_table,
             combined_indices,
         )
         for block in module.blocks:
@@ -1359,7 +1360,7 @@ def extract_minimax_h3_context(
                 hidden,
                 t_emb=t_emb,
                 combined_indices=block_combined,
-                rope_freqs=block_rope,
+                rope_table=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,

vllm_omni/diffusion/layers/fused_qk_norm_rope.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Q/K RMSNorm followed by packed non-interleaved RoPE.
+
+The public contract is shared by diffusion attention implementations:
+
+* ``q`` and ``k`` are ``[tokens, heads, head_dim]``;
+* norm weights are one-dimensional ``[head_dim]`` tensors;
+* ``rope_table`` is ``[tokens, rotary_dim]`` and stores
+  ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``.
+
+The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K
+or rotary-product intermediates. Unsupported inputs use the eager reference.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch.library import Library
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
+
+_FUSED_HEAD_DIM = 128
+_FUSED_ROTARY_DIM = 96
+_HEADS_PER_PROGRAM = 8
+
+
+def _apply_rope_table(
+    x: torch.Tensor,
+    rope_table: torch.Tensor,
+    rotary_dim: int,
+) -> torch.Tensor:
+    half = rotary_dim // 2
+    cos = rope_table[..., :half].to(x.dtype).unsqueeze(1)
+    sin = rope_table[..., half:].to(x.dtype).unsqueeze(1)
+    first = x[..., :half]
+    second = x[..., half:rotary_dim]
+    return torch.cat(
+        (
+            first * cos - second * sin,
+            second * cos + first * sin,
+            x[..., rotary_dim:],
+        ),
+        dim=-1,
+    )
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _rms_norm_rope_kernel(
+        x_ptr,
+        weight_ptr,
+        rope_table_ptr,
+        out_ptr,
+        x_stride_t,
+        x_stride_h,
+        x_stride_d,
+        rope_stride_t,
+        out_stride_t,
+        out_stride_h,
+        out_stride_d,
+        num_heads: tl.constexpr,
+        head_dim: tl.constexpr,
+        rotary_half: tl.constexpr,
+        eps: tl.constexpr,
+        heads_per_program: tl.constexpr,
+    ):
+        token = tl.program_id(0)
+        head_group = tl.program_id(1)
+        heads = head_group * heads_per_program + tl.arange(0, heads_per_program)
+        dims = tl.arange(0, head_dim)
+        mask = heads[:, None] < num_heads
+        offsets = token * x_stride_t + heads[:, None] * x_stride_h + dims[None, :] * x_stride_d
+
+        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        weight = tl.load(weight_ptr + dims).to(tl.float32)
+        inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps)
+        normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16)
+
+        rotary_dim = rotary_half * 2
+        pair_dims = tl.where(
+            dims < rotary_half,
+            dims + rotary_half,
+            tl.where(dims < rotary_dim, dims - rotary_half, dims),
+        )
+        pair_offsets = token * x_stride_t + heads[:, None] * x_stride_h + pair_dims[None, :] * x_stride_d
+        pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32)
+        pair_weight = tl.load(weight_ptr + pair_dims).to(tl.float32)
+        pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16)
+
+        freq_dims = tl.where(
+            dims < rotary_half,
+            dims,
+            tl.where(dims < rotary_dim, dims - rotary_half, 0),
+        )
+        table_offsets = token * rope_stride_t + freq_dims
+        cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32)
+        sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32)
+        first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin
+        second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin
+        output = tl.where(
+            dims < rotary_dim,
+            tl.where(dims < rotary_half, first, second),
+            normalized.to(tl.float32),
+        )
+
+        out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d
+        tl.store(out_ptr + out_offsets, output, mask=mask)
+
+
+def _eager_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_norm = F.rms_norm(q, (head_dim,), q_weight, eps)
+    k_norm = F.rms_norm(k, (head_dim,), k_weight, eps)
+    return (
+        _apply_rope_table(q_norm, rope_table, rotary_dim),
+        _apply_rope_table(k_norm, rope_table, rotary_dim),
+    )
+
+
+def _fused_cuda_supported(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_dim: int,
+    rotary_dim: int,
+) -> bool:
+    return (
+        HAS_TRITON
+        and current_platform.is_cuda()
+        and q.is_cuda
+        and k.is_cuda
+        and q.dtype == torch.bfloat16
+        and k.dtype == torch.bfloat16
+        and head_dim == _FUSED_HEAD_DIM
+        and rotary_dim == _FUSED_ROTARY_DIM
+    )
+
+
+def _launch_fused_rms_norm_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    tokens, heads, head_dim = x.shape
+    rotary_half = rope_table.shape[-1] // 2
+    out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    if tokens == 0:
+        return out
+    grid = (tokens, triton.cdiv(heads, _HEADS_PER_PROGRAM))
+    _rms_norm_rope_kernel[grid](
+        x,
+        weight,
+        rope_table,
+        out,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        rope_table.stride(0),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        num_heads=heads,
+        head_dim=head_dim,
+        rotary_half=rotary_half,
+        eps=eps,
+        heads_per_program=_HEADS_PER_PROGRAM,
+        num_warps=8,
+    )
+    return out
+
+
+def _fused_qk_norm_rope_impl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+        return _eager_qk_norm_rope(
+            q,
+            k,
+            q_weight,
+            k_weight,
+            rope_table,
+            eps,
+            head_dim,
+            rotary_dim,
+        )
+    return (
+        _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps),
+        _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps),
+    )
+
+
+def _fused_qk_norm_rope_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim
+    return torch.empty_like(q), torch.empty_like(k)
+
+
+_OMNI_OP_LIB = Library("vllm_omni", "FRAGMENT")
+if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"):
+    direct_register_custom_op(
+        op_name="fused_qk_norm_rope",
+        op_func=_fused_qk_norm_rope_impl,
+        fake_impl=_fused_qk_norm_rope_fake,
+        mutates_args=[],
+        target_lib=_OMNI_OP_LIB,
+    )
+
+
+def fused_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    *,
+    head_dim: int | None = None,
+    rotary_dim: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Q/K RMSNorm and packed non-interleaved RoPE."""
+    if q.ndim != 3 or k.ndim != 3:
+        raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}")
+    if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]:
+        raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}")
+    if q.dtype != k.dtype or q.device != k.device:
+        raise ValueError("q and k must have the same dtype and device")
+    if q.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}")
+
+    head_dim = q.shape[-1] if head_dim is None else head_dim
+    rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim
+    if q.shape[-1] != head_dim:
+        raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}")
+    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2:
+        raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}")
+    if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,):
+        raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}")
+    if q_weight.device != q.device or k_weight.device != q.device:
+        raise ValueError("Q/K norm weights must be on the activation device")
+    if rope_table.device != q.device or rope_table.dtype != q.dtype:
+        raise ValueError("rope_table must have the same dtype and device as q/k")
+    if rope_table.shape != (q.shape[0], rotary_dim):
+        raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}")
+
+    q_weight = q_weight.contiguous()
+    k_weight = k_weight.contiguous()
+    rope_table = rope_table.contiguous()
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+        return _fused_qk_norm_rope_impl(
+            q,
+            k,
+            q_weight,
+            k_weight,
+            rope_table,
+            eps,
+            head_dim,
+            rotary_dim,
+        )
+    return torch.ops.vllm_omni.fused_qk_norm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+        eps,
+        head_dim,
+        rotary_dim,
+    )
+
+
+__all__ = ["fused_qk_norm_rope"]

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -33,6 +33,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
@@ -237,6 +238,15 @@ class MiniMaxH3Rope(nn.Module):
         return torch.cat((half, half), dim=-1)  # [S, 96]
 
 
+def _build_rope_table(freqs: torch.Tensor) -> torch.Tensor:
+    """Materialize H3's packed ``[cos(freqs[:48]), sin(freqs[:48])]`` table."""
+    half = freqs.shape[-1] // 2
+    return torch.cat(
+        (torch.cos(freqs[..., :half]), torch.sin(freqs[..., :half])),
+        dim=-1,
+    ).to(_BF16_DTYPE)
+
+
 class MiniMaxH3TimeEmbedder(nn.Module):
     def __init__(
         self,
@@ -453,7 +463,7 @@ class MiniMaxH3Attention(nn.Module):
         self,
         x: torch.Tensor,
         *,
-        rope_freqs: torch.Tensor | None,
+        rope_table: torch.Tensor | None,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int | None = None,
@@ -479,11 +489,18 @@ class MiniMaxH3Attention(nn.Module):
         q = q.view(total, self.num_heads, self.head_dim)
         k = k.view(total, self.num_kv_heads, self.head_dim)
         v = v.view(total, self.num_kv_heads, self.head_dim)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-        if rope_freqs is not None:
-            q = self._apply_rope(q, rope_freqs)
-            k = self._apply_rope(k, rope_freqs)
+        if rope_table is None:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        else:
+            q, k = fused_qk_norm_rope(
+                q,
+                k,
+                self.q_norm.weight,
+                self.k_norm.weight,
+                rope_table,
+                self.q_norm.variance_epsilon,
+            )
 
         # The packed layout uses a second document for alignment padding.
         # Local/Ulysses backends unpad it, while Ring keeps aligned rows for
@@ -628,7 +645,7 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
     ) -> torch.Tensor:
         x = x + self.attn(
             self.norm1(x),
-            rope_freqs=None,
+            rope_table=None,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
         )
@@ -707,7 +724,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         *,
         t_emb: torch.Tensor,
         combined_indices: torch.Tensor,
-        rope_freqs: torch.Tensor,
+        rope_table: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
@@ -735,7 +752,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE)
         h = self.attn(
             h,
-            rope_freqs=rope_freqs,
+            rope_table=rope_table,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
             packed_total=packed_total,
@@ -817,10 +834,10 @@ class MiniMaxH3SPPrepare(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        rope_freqs: torch.Tensor,
+        rope_table: torch.Tensor,
         combined_indices: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        return hidden_states, rope_freqs, combined_indices
+        return hidden_states, rope_table, combined_indices
 
 
 class MiniMaxH3SPGather(nn.Module):
@@ -1153,8 +1170,8 @@ class MiniMaxH3DiTModel(nn.Module):
         if inverse_indices.shape[0] != seq_len:
             raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
         device = x.device
-        # Compute RoPE frequencies over the full packed sequence.
-        rope_freqs = self.rope(img_position_ids).to(device)
+        # Compute the packed RoPE table once for all transformer blocks.
+        rope_table = _build_rope_table(self.rope(img_position_ids).to(device))
 
         decoder_input, t_emb = self._embed(
             x=x,
@@ -1175,7 +1192,7 @@ class MiniMaxH3DiTModel(nn.Module):
 
         hidden = decoder_input
         cu_seqlens = cu_seqlens.to(device)
-        block_rope = rope_freqs
+        block_rope = rope_table
         block_combined = combined_indices
 
         hidden, block_rope, block_combined = self.sp_prepare(
@@ -1188,7 +1205,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 hidden,
                 t_emb=t_emb,
                 combined_indices=block_combined,
-                rope_freqs=block_rope,
+                rope_table=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,