From 4028680b7edb0882f1e50442ad1a58d92fcada5e Mon Sep 17 00:00:00 2001 From: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Date: Tue, 30 Jun 2026 16:32:47 +0800 Subject: [PATCH] [Perf]: Fused RMSNorm+3D Rope kernel in Sensenova-U1 (#4669) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- .../test_sensenova_triton_kernel.py | 232 +++++++++++++++++ .../models/sensenova_u1/fused_rmsnorm_rope.py | 244 ++++++++++++++++++ .../sensenova_u1/sensenova_u1_transformer.py | 35 ++- 3 files changed, 507 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py create mode 100644 vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py b/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py new file mode 100644 index 000000000..d966e5238 --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Correctness tests for Sensenova-U1 Triton kernels. + +Verifies that fused-RMSNorm+3D Rope Triton kernels +produce results numerically equivalent to the PyTorch reference. +All computations follow the kernel's internal float32 promotion rules. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + +SEED = 42 +B = 1 +H_Q = 32 +H_K = 8 +HEAD_DIM = 128 +T_DIM = 64 +HW_DIM = 64 +HW_ROPE_DIM = 32 +QKV_DIM = (H_Q + H_K + H_K) * HEAD_DIM +EPS = 1e-6 +DEVICE = torch.device("cuda:0") + + +@dataclass(frozen=True) +class KernelInput: + seq_len: int + q: torch.Tensor + k: torch.Tensor + q_norm_weight: torch.Tensor + k_norm_weight: torch.Tensor + q_norm_hw_weight: torch.Tensor + k_norm_hw_weight: torch.Tensor + cos_t: torch.Tensor + sin_t: torch.Tensor + cos_h: torch.Tensor + sin_h: torch.Tensor + cos_w: torch.Tensor + sin_w: torch.Tensor + + def args(self) -> tuple[torch.Tensor, ...]: + return ( + self.q, + self.k, + self.q_norm_weight, + self.k_norm_weight, + self.q_norm_hw_weight, + self.k_norm_hw_weight, + self.cos_t, + self.sin_t, + self.cos_h, + self.sin_h, + self.cos_w, + self.sin_w, + ) + + +try: + from vllm_omni.diffusion.models.sensenova_u1.fused_rmsnorm_rope import triton_qk_norm_rope + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + +triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform") + + +# --------------------------------------------------------------------------- +# PyTorch reference implementations (mirror the kernel's internal arithmetic) +# --------------------------------------------------------------------------- + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = EPS) -> torch.Tensor: + input_dtype = x.dtype + xf = x.float() + var = xf.pow(2).mean(dim=-1, keepdim=True) + xf = xf * torch.rsqrt(var + eps) + return xf.to(input_dtype) * weight + + +def assert_close_with_error_stats( + actual: torch.Tensor, + expected: torch.Tensor, + *, + name: str, + atol: float, + rtol: float, +) -> None: + try: + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + except AssertionError as exc: + abs_diff = (actual.float() - expected.float()).abs().flatten() + tolerance = atol + rtol * expected.float().abs().flatten() + mismatch = abs_diff > tolerance + mismatch_count = int(mismatch.sum().item()) + total = abs_diff.numel() + p99 = torch.quantile(abs_diff, 0.99).item() if total else 0.0 + stats = ( + f"{name} error stats: " + f"max={abs_diff.max().item() if total else 0.0:.6g}, " + f"p99={p99:.6g}, " + f"mean={abs_diff.mean().item() if total else 0.0:.6g}, " + f"mismatch={mismatch_count}/{total} ({mismatch_count / total:.4%}), " + f"atol={atol}, rtol={rtol}" + ) + raise AssertionError(f"{stats}\n{exc}") from exc + + +def reference_kernel( + q: torch.Tensor, + k: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_hw_weight: torch.Tensor, + k_norm_hw_weight: torch.Tensor, + cos_t: torch.Tensor, + sin_t: torch.Tensor, + cos_h: torch.Tensor, + sin_h: torch.Tensor, + cos_w: torch.Tensor, + sin_w: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference matching the model's current math.""" + q_t, q_hw = q.chunk(2, dim=-1) + k_t, k_hw = k.chunk(2, dim=-1) + + q_t = rmsnorm(q_t, q_norm_weight).transpose(1, 2) + k_t = rmsnorm(k_t, k_norm_weight).transpose(1, 2) + q_hw = rmsnorm(q_hw, q_norm_hw_weight).transpose(1, 2) + k_hw = rmsnorm(k_hw, k_norm_hw_weight).transpose(1, 2) + + q_h, q_w = q_hw.chunk(2, dim=-1) + k_h, k_w = k_hw.chunk(2, dim=-1) + + q_t, k_t = apply_rope(q_t, k_t, cos_t, sin_t) + q_h, k_h = apply_rope(q_h, k_h, cos_h, sin_h) + q_w, k_w = apply_rope(q_w, k_w, cos_w, sin_w) + + query = torch.cat([q_t, q_h, q_w], dim=-1) + key = torch.cat([k_t, k_h, k_w], dim=-1) + return query, key + + +def make_input(seq_len: int, dtype: torch.dtype, device: torch.device, seed: int) -> KernelInput: + gen = torch.Generator(device=device) + gen.manual_seed(seed + seq_len) + + # Match real model layout: q/k/v are views split from fused qkv. + qkv = torch.randn(B, seq_len, QKV_DIM, device=device, dtype=dtype, generator=gen) + q, k, _v = qkv.split([H_Q * HEAD_DIM, H_K * HEAD_DIM, H_K * HEAD_DIM], dim=-1) + q = q.view(B, seq_len, H_Q, HEAD_DIM) + k = k.view(B, seq_len, H_K, HEAD_DIM) + + q_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen) + k_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen) + q_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen) + k_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen) + + # Match Qwen3RotaryEmbedding: emb = torch.cat((freqs, freqs), dim=-1). + def make_rope_pair(dim: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen) + sin_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen) + return torch.cat((cos_half, cos_half), dim=-1), torch.cat((sin_half, sin_half), dim=-1) + + cos_t, sin_t = make_rope_pair(T_DIM) + cos_h, sin_h = make_rope_pair(HW_ROPE_DIM) + cos_w, sin_w = make_rope_pair(HW_ROPE_DIM) + + return KernelInput( + seq_len=seq_len, + q=q, + k=k, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + q_norm_hw_weight=q_norm_hw_weight, + k_norm_hw_weight=k_norm_hw_weight, + cos_t=cos_t, + sin_t=sin_t, + cos_h=cos_h, + sin_h=sin_h, + cos_w=cos_w, + sin_w=sin_w, + ) + + +# --------------------------------------------------------------------------- +# Fused RMSNorm + 3D Rope +# --------------------------------------------------------------------------- + + +@triton_available +@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 9, 16, 32, 64, 128, 256, 260, 271, 512, 1024, 2048, 4096]) +@pytest.mark.parametrize( + ("dtype", "atol", "rtol"), + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 3e-2, 3e-2), + ], +) +def test_fused_qk_norm_rope_matches_reference(seq_len: int, dtype: torch.dtype, atol: float, rtol: float): + data = make_input(seq_len, dtype=dtype, device=DEVICE, seed=SEED) + call_args = data.args() + ref_q, ref_k = reference_kernel(*call_args) + out_q, out_k = triton_qk_norm_rope(*call_args, EPS) + + assert_close_with_error_stats(out_q, ref_q, name="query", atol=atol, rtol=rtol) + assert_close_with_error_stats(out_k, ref_k, name="key", atol=atol, rtol=rtol) diff --git a/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py b/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py new file mode 100644 index 000000000..5b81b6ce4 --- /dev/null +++ b/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def qk_norm_rope_kernel( + q, + k, + query, + key, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + q_stride_b: tl.constexpr, + q_stride_s: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + k_stride_b: tl.constexpr, + k_stride_s: tl.constexpr, + k_stride_h: tl.constexpr, + k_stride_d: tl.constexpr, + query_stride_b: tl.constexpr, + query_stride_h: tl.constexpr, + query_stride_s: tl.constexpr, + query_stride_d: tl.constexpr, + key_stride_b: tl.constexpr, + key_stride_h: tl.constexpr, + key_stride_s: tl.constexpr, + key_stride_d: tl.constexpr, + cos_t_stride_b: tl.constexpr, + cos_t_stride_s: tl.constexpr, + cos_t_stride_d: tl.constexpr, + cos_h_stride_b: tl.constexpr, + cos_h_stride_s: tl.constexpr, + cos_h_stride_d: tl.constexpr, + cos_w_stride_b: tl.constexpr, + cos_w_stride_s: tl.constexpr, + cos_w_stride_d: tl.constexpr, + head_q: tl.constexpr, + head_dim: tl.constexpr, + eps: tl.constexpr = 1e-6, +): + token_idx = tl.program_id(0) + head_pid = tl.program_id(1) + batch_idx = tl.program_id(2) + is_q = head_pid < head_q + head_idx = tl.where(is_q, head_pid, head_pid - head_q) + + offs = tl.arange(0, head_dim) + t_mask = offs < head_dim // 2 + hw_mask = ~t_mask + + in_base = tl.where( + is_q, + q + batch_idx * q_stride_b + token_idx * q_stride_s + head_idx * q_stride_h, + k + batch_idx * k_stride_b + token_idx * k_stride_s + head_idx * k_stride_h, + ) + vals = tl.load(in_base + offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32) + + t_vals = tl.where(t_mask, vals, 0.0) + hw_vals = tl.where(hw_mask, vals, 0.0) + rms_t = tl.rsqrt(tl.sum(t_vals * t_vals, axis=0) / (head_dim // 2) + eps) + rms_hw = tl.rsqrt(tl.sum(hw_vals * hw_vals, axis=0) / (head_dim // 2) + eps) + rms = tl.where(t_mask, rms_t, rms_hw) + + t_weight = tl.load(tl.where(is_q, q_norm_weight, k_norm_weight) + offs, mask=t_mask, other=0.0) + hw_weight = tl.load( + tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (offs - head_dim // 2), + mask=hw_mask, + other=0.0, + ) + weight = tl.where(t_mask, t_weight, hw_weight) + normed = (vals * rms).to(query.dtype.element_ty) * weight + + pair_offs = tl.where( + offs < head_dim // 4, + offs + head_dim // 4, + tl.where( + offs < head_dim // 2, + offs - head_dim // 4, + tl.where( + offs < head_dim // 8 * 5, + offs + head_dim // 8, + tl.where( + offs < head_dim // 4 * 3, + offs - head_dim // 8, + tl.where(offs < head_dim // 8 * 7, offs + head_dim // 8, offs - head_dim // 8), + ), + ), + ), + ) + pair_vals = tl.load(in_base + pair_offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32) + pair_rms = tl.where(pair_offs < head_dim // 2, rms_t, rms_hw) + pair_t_mask = pair_offs < head_dim // 2 + pair_hw_mask = ~pair_t_mask + pair_t_weight = tl.load( + tl.where(is_q, q_norm_weight, k_norm_weight) + pair_offs, + mask=pair_t_mask, + other=0.0, + ) + pair_hw_weight = tl.load( + tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (pair_offs - head_dim // 2), + mask=pair_hw_mask, + other=0.0, + ) + pair_weight = tl.where(pair_offs < head_dim // 2, pair_t_weight, pair_hw_weight) + pair_normed = (pair_vals * pair_rms).to(query.dtype.element_ty) * pair_weight + + cos_t_vals = tl.load( + cos_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d, + mask=t_mask, + other=0.0, + ) + sin_t_vals = tl.load( + sin_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d, + mask=t_mask, + other=0.0, + ) + hw_local = offs - head_dim // 2 + h_local = hw_local + w_local = hw_local - head_dim // 4 + h_mask = (offs >= head_dim // 2) & (offs < head_dim // 4 * 3) + w_mask = offs >= head_dim // 4 * 3 + cos_h_vals = tl.load( + cos_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d, + mask=h_mask, + other=0.0, + ) + sin_h_vals = tl.load( + sin_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d, + mask=h_mask, + other=0.0, + ) + cos_w_vals = tl.load( + cos_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d, + mask=w_mask, + other=0.0, + ) + sin_w_vals = tl.load( + sin_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d, + mask=w_mask, + other=0.0, + ) + cos_vals = cos_t_vals + cos_h_vals + cos_w_vals + sin_vals = sin_t_vals + sin_h_vals + sin_w_vals + + sign = tl.where( + (offs < head_dim // 4) + | ((offs >= head_dim // 2) & (offs < head_dim // 8 * 5)) + | ((offs >= head_dim // 4 * 3) & (offs < head_dim // 8 * 7)), + -1.0, + 1.0, + ) + out = normed * cos_vals + sign * pair_normed * sin_vals + + out_base = tl.where( + is_q, + query + batch_idx * query_stride_b + head_idx * query_stride_h + token_idx * query_stride_s, + key + batch_idx * key_stride_b + head_idx * key_stride_h + token_idx * key_stride_s, + ) + tl.store(out_base + offs * tl.where(is_q, query_stride_d, key_stride_d), out) + + +def triton_qk_norm_rope( + q, + k, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + eps, +): + batch_size, seq_len, head_q, head_dim = q.shape + head_k = k.shape[2] + + query = torch.empty((batch_size, head_q, seq_len, head_dim), device=q.device, dtype=q.dtype) + key = torch.empty((batch_size, head_k, seq_len, head_dim), device=k.device, dtype=k.dtype) + grid = ( + seq_len, + head_q + head_k, + batch_size, + ) + + qk_norm_rope_kernel[grid]( + q, + k, + query, + key, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + query.stride(0), + query.stride(1), + query.stride(2), + query.stride(3), + key.stride(0), + key.stride(1), + key.stride(2), + key.stride(3), + cos_t.stride(0), + cos_t.stride(1), + cos_t.stride(2), + cos_h.stride(0), + cos_h.stride(1), + cos_h.stride(2), + cos_w.stride(0), + cos_w.stride(1), + cos_w.stride(2), + head_q=head_q, + head_dim=head_dim, + eps=eps, + ) + return query, key diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py index 53e1a47c2..843358a70 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py +++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py @@ -18,6 +18,7 @@ import torch import torch.nn as nn from cache_dit import ForwardPattern from transformers.cache_utils import DynamicCache +from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -33,6 +34,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter +logger = init_logger(__name__) + @dataclass class SenseNovaU1ModelOutput: @@ -373,6 +376,34 @@ class SenseNovaU1Attention(nn.Module): k = k.view(*input_shape, self.num_kv_heads, self.head_dim) v = v.view(*input_shape, self.num_kv_heads, self.head_dim) + cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)) + cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)) + cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)) + + value_states = v.transpose(1, 2) # [B, H, S, D] + try: + from .fused_rmsnorm_rope import triton_qk_norm_rope + except ImportError: + triton_qk_norm_rope = None + if triton_qk_norm_rope is not None: + query_states, key_states = triton_qk_norm_rope( + q, + k, + q_norm.weight, + k_norm.weight, + q_norm_hw.weight, + k_norm_hw.weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + self.config.rms_norm_eps, + ) + return query_states, key_states, value_states + logger.debug("Triton fused SenseNova qk norm rope unavailable; using PyTorch fallback") + # Split head_dim into t and hw halves q_t, q_hw = q.chunk(2, dim=-1) k_t, k_hw = k.chunk(2, dim=-1) @@ -388,17 +419,13 @@ class SenseNovaU1Attention(nn.Module): k_h, k_w = k_hw.chunk(2, dim=-1) # RoPE for each dimension - cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)) q_t, k_t = apply_rotary_pos_emb(q_t, k_t, cos_t, sin_t) - cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)) q_h, k_h = apply_rotary_pos_emb(q_h, k_h, cos_h, sin_h) - cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)) q_w, k_w = apply_rotary_pos_emb(q_w, k_w, cos_w, sin_w) # Reassemble: [B, H, S, head_dim] query_states = torch.cat([q_t, q_h, q_w], dim=-1) key_states = torch.cat([k_t, k_h, k_w], dim=-1) - value_states = v.transpose(1, 2) # [B, H, S, D] return query_states, key_states, value_states def forward_und(self, hidden_states, indexes, attention_mask, past_key_values=None, **kwargs):