From 596c16a550aa134faf7f3dcfa0f8adf513ccd9ce Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:46:59 +0800 Subject: [PATCH] [Kernel] Fuse Q/K RMSNorm and RoPE (#5990) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Signed-off-by: lishunyang12 Signed-off-by: Mengjie Zhao Co-authored-by: lishunyang12 Co-authored-by: Claude Opus 5 --- .../layers/test_fused_qk_norm_rope.py | 82 +++++ tests/e2e/accuracy/ltx/run_ltx_reference.py | 12 +- .../diffusion/cache/teacache/extractors.py | 7 +- .../diffusion/layers/fused_qk_norm_rope.py | 299 ++++++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 47 ++- 5 files changed, 420 insertions(+), 27 deletions(-) create mode 100644 tests/diffusion/layers/test_fused_qk_norm_rope.py create mode 100644 vllm_omni/diffusion/layers/fused_qk_norm_rope.py diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py new file mode 100644 index 000000000..40888556a --- /dev/null +++ b/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) diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py index 3ab794662..ccabd4596 100644 --- a/tests/e2e/accuracy/ltx/run_ltx_reference.py +++ b/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] diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 5f1ca1d22..8689c2c54 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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, diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py new file mode 100644 index 000000000..64462d6c2 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 5e180aa23..659b64bc4 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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,