From 7a89b72564a2c39d170b9478e947382216710ee2 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 9 Jun 2026 22:12:17 +0800 Subject: [PATCH] [Perf] fuse qk rmsnorm rope gate for qwen3.5 (#44176) Signed-off-by: zjy0516 --- .buildkite/test_areas/kernels.yaml | 3 +- tests/kernels/test_fused_qk_norm_rope_gate.py | 159 ++++++++++++++ .../layers/fused_qk_norm_rope.py | 201 ++++++++++++++++++ vllm/model_executor/models/qwen3_next.py | 67 +++++- 4 files changed, 418 insertions(+), 12 deletions(-) create mode 100644 tests/kernels/test_fused_qk_norm_rope_gate.py create mode 100644 vllm/model_executor/layers/fused_qk_norm_rope.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d52da1f0a4..d1a4ade2a7 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -21,8 +21,9 @@ steps: - csrc/ - tests/kernels/core - tests/kernels/test_concat_mla_q.py + - tests/kernels/test_fused_qk_norm_rope_gate.py commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py - label: Kernels MiniMax Reduce RMS Test (2 GPUs) key: kernels-minimax-reduce-rms-test-2-gpus diff --git a/tests/kernels/test_fused_qk_norm_rope_gate.py b/tests/kernels/test_fused_qk_norm_rope_gate.py new file mode 100644 index 0000000000..09ec90d488 --- /dev/null +++ b/tests/kernels/test_fused_qk_norm_rope_gate.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +# Qwen/Qwen3.6-27B config (huggingface.co/Qwen/Qwen3.6-27B), TP=1 shapes. +NUM_Q_HEADS = 24 +NUM_KV_HEADS = 4 +HEAD_DIM = 256 +PARTIAL_ROTARY_FACTOR = 0.25 +ROTARY_DIM = int(HEAD_DIM * PARTIAL_ROTARY_FACTOR) # 64 +RMS_NORM_EPS = 1e-6 +MAX_POSITION_EMBEDDINGS = 262144 +ROPE_THETA = 10000000.0 + +DTYPES = [torch.bfloat16] +SEEDS = [13] +NUM_TOKENS = [1, 4, 37] + + +def _ref_qk_rmsnorm_rope_gate( + q_gate: torch.Tensor, + k: torch.Tensor, + q_gamma: torch.Tensor, + k_gamma: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """PyTorch reference: split + RMSNorm + partial NeoX RoPE + gate extraction. + + Matches ``fused_qk_rmsnorm_rope_gate``'s contract: ``q_gamma`` / ``k_gamma`` + are the already-adjusted effective gammas (for GemmaRMSNorm the caller + has done ``weight + 1`` before passing them in). + """ + n_tokens = q_gate.shape[0] + half = rotary_dim // 2 + + # Per head the q projection is laid out as [q | gate]. + q_gate = q_gate.view(n_tokens, num_q_heads, 2 * head_dim) + q = q_gate[..., :head_dim] + gate = q_gate[..., head_dim:].reshape(n_tokens, num_q_heads * head_dim) + k = k.view(n_tokens, num_kv_heads, head_dim) + + def rms_norm(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + orig_dtype = x.dtype + x = x.float() + var = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(var + eps) + return (x * gamma.float()).to(orig_dtype) + + q = rms_norm(q, q_gamma) + k = rms_norm(k, k_gamma) + + # Partial NeoX RoPE on the first ``rotary_dim`` elements of each head; + # cos_sin_cache row is packed as [cos(half) | sin(half)]. + pos = positions.view(-1) + cos = cos_sin_cache[pos, :half].float()[:, None, :] + sin = cos_sin_cache[pos, half:rotary_dim].float()[:, None, :] + + def rope(x: torch.Tensor) -> torch.Tensor: + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x1 = x_rot[..., :half].float() + x2 = x_rot[..., half:].float() + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + rotated = torch.cat([o1, o2], dim=-1).to(x.dtype) + return torch.cat([rotated, x_pass], dim=-1) + + q = rope(q).reshape(n_tokens, num_q_heads * head_dim) + k = rope(k).reshape(n_tokens, num_kv_heads * head_dim) + return q, k, gate + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="fused_qk_rmsnorm_rope_gate Triton kernel requires CUDA/ROCm", +) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@torch.inference_mode() +def test_fused_qk_norm_rope_gate_matches_reference( + default_vllm_config, + dtype: torch.dtype, + seed: int, + num_tokens: int, +): + device = torch.device("cuda", torch.accelerator.current_device_index()) + torch.set_default_device(device) + set_random_seed(seed) + + q_gate = torch.randn( + num_tokens, NUM_Q_HEADS * 2 * HEAD_DIM, dtype=dtype, device=device + ) + k = torch.randn(num_tokens, NUM_KV_HEADS * HEAD_DIM, dtype=dtype, device=device) + # GemmaRMSNorm-style: the kernel takes the effective gamma (weight + 1). + q_gamma = ( + torch.empty(HEAD_DIM, dtype=dtype, device=device).normal_(mean=0.0, std=0.1) + + 1.0 + ) + k_gamma = ( + torch.empty(HEAD_DIM, dtype=dtype, device=device).normal_(mean=0.0, std=0.1) + + 1.0 + ) + + # fused_qk_rmsnorm_rope_gate only handles NeoX-style RoPE. + rope = RotaryEmbedding( + head_size=HEAD_DIM, + rotary_dim=ROTARY_DIM, + max_position_embeddings=MAX_POSITION_EMBEDDINGS, + base=ROPE_THETA, + is_neox_style=True, + dtype=dtype, + ).to(device) + positions = torch.arange(num_tokens, dtype=torch.long, device=device) + + q_ref, k_ref, gate_ref = _ref_qk_rmsnorm_rope_gate( + q_gate, + k, + q_gamma, + k_gamma, + rope.cos_sin_cache, + positions, + RMS_NORM_EPS, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + ROTARY_DIM, + ) + q_out, k_out, gate_out = fused_qk_rmsnorm_rope_gate( + q_gate, + k, + q_gamma, + k_gamma, + rope.cos_sin_cache, + positions, + RMS_NORM_EPS, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + ROTARY_DIM, + ) + + atol, rtol = 2e-3, 5e-3 + torch.testing.assert_close(q_out, q_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(k_out, k_ref, atol=atol, rtol=rtol) + # gate is a verbatim copy of the source slice — must match bit-exactly. + torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0) diff --git a/vllm/model_executor/layers/fused_qk_norm_rope.py b/vllm/model_executor/layers/fused_qk_norm_rope.py new file mode 100644 index 0000000000..21549dd783 --- /dev/null +++ b/vllm/model_executor/layers/fused_qk_norm_rope.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused QK-RMSNorm + (partial) RoPE + gate copy Triton kernel. + +Currently used by the Qwen3.5 attention path (``attn_output_gate`` with +NeoX-style partial RoPE). The unfused reference sequence is +``split -> GemmaRMSNorm -> RoPE -> gate chunk``; this collapses it into a +single Triton launch. See :func:`fused_qk_rmsnorm_rope_gate`. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_qk_rmsnorm_rope_gate_kernel( + q_gate_ptr, + k_ptr, + q_out_ptr, + k_out_ptr, + gate_out_ptr, + q_weight_ptr, + k_weight_ptr, + cos_sin_cache_ptr, + positions_ptr, + q_gate_stride_t, + k_stride_t, + q_out_stride_t, + k_out_stride_t, + gate_out_stride_t, + cache_stride_p, + num_q_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + head_dim: tl.constexpr, + rotary_dim: tl.constexpr, + half_rotary: tl.constexpr, + eps: tl.constexpr, + INPUT_DTYPE: tl.constexpr, + HEAD_BLOCK: tl.constexpr, + ROT_HALF_BLOCK: tl.constexpr, + HAS_PASS: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + is_k = head >= num_q_heads + local_head = tl.where(is_k, head - num_q_heads, head) + + if is_k: + in_base = k_ptr + token * k_stride_t + local_head * head_dim + w_ptr = k_weight_ptr + out_base = k_out_ptr + token * k_out_stride_t + local_head * head_dim + else: + in_base = q_gate_ptr + token * q_gate_stride_t + local_head * 2 * head_dim + w_ptr = q_weight_ptr + out_base = q_out_ptr + token * q_out_stride_t + local_head * head_dim + + # --- RMSNorm: variance over the full head_dim --- + head_offs = tl.arange(0, HEAD_BLOCK) + head_mask = head_offs < head_dim + x = tl.load(in_base + head_offs, mask=head_mask, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=0) / head_dim + inv_rms = tl.rsqrt(var + eps) + w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32) + # Round-trip through INPUT_DTYPE so the RoPE input matches the bf16-storage + # behavior of the unfused (qk_rmsnorm -> memory -> apply_rope) reference path. + x_norm = (x * inv_rms * w).to(INPUT_DTYPE).to(tl.float32) + + # --- Pass-through tail [rotary_dim, head_dim): RMSNorm-only, no rotation --- + # The rotary head [0, rotary_dim) will be overwritten by the RoPE store below. + if HAS_PASS: + pass_mask = head_mask & (head_offs >= rotary_dim) + tl.store(out_base + head_offs, x_norm, mask=pass_mask) + + # --- Partial RoPE on the first rotary_dim elements --- + # Triton lacks easy sub-vector slicing of x_norm, so we recompute the + # normalized rotary halves on a smaller block (next_pow2(half_rotary)). + # The extra ~rotary_dim element reload hits L1, so the cost is negligible. + rot_offs = tl.arange(0, ROT_HALF_BLOCK) + rot_mask = rot_offs < half_rotary + x_rot1 = tl.load(in_base + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + x_rot2 = tl.load(in_base + half_rotary + rot_offs, mask=rot_mask, other=0.0).to( + tl.float32 + ) + w_rot1 = tl.load(w_ptr + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + w_rot2 = tl.load(w_ptr + half_rotary + rot_offs, mask=rot_mask, other=0.0).to( + tl.float32 + ) + x_rot1 = (x_rot1 * inv_rms * w_rot1).to(INPUT_DTYPE).to(tl.float32) + x_rot2 = (x_rot2 * inv_rms * w_rot2).to(INPUT_DTYPE).to(tl.float32) + + # Always use int64 for position to avoid overflow in address computation. + pos = tl.load(positions_ptr + token).to(tl.int64) + cache_offset = pos * cache_stride_p + cos = tl.load( + cos_sin_cache_ptr + cache_offset + rot_offs, mask=rot_mask, other=0.0 + ).to(tl.float32) + sin = tl.load( + cos_sin_cache_ptr + cache_offset + half_rotary + rot_offs, + mask=rot_mask, + other=0.0, + ).to(tl.float32) + + o1 = x_rot1 * cos - x_rot2 * sin + o2 = x_rot2 * cos + x_rot1 * sin + tl.store(out_base + rot_offs, o1, mask=rot_mask) + tl.store(out_base + half_rotary + rot_offs, o2, mask=rot_mask) + + # --- Gate copy (q heads only, verbatim) --- + if not is_k: + gate_in_base = in_base + head_dim + gate_out_base = gate_out_ptr + token * gate_out_stride_t + local_head * head_dim + g = tl.load(gate_in_base + head_offs, mask=head_mask, other=0.0) + tl.store(gate_out_base + head_offs, g, mask=head_mask) + + +def fused_qk_rmsnorm_rope_gate( + q_gate: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused split + QK-RMSNorm + (partial) RoPE + gate copy for Qwen3.5 attn. + + Args: + q_gate: (n_tokens, num_q_heads * 2 * head_dim) -- per head: [q|gate] + k: (n_tokens, num_kv_heads * head_dim) + q_weight: (head_dim,) GemmaRMSNorm effective weight (already +1) + k_weight: (head_dim,) GemmaRMSNorm effective weight (already +1) + cos_sin_cache: (max_pos, rotary_dim) packed [cos|sin] + positions: (n_tokens,) int32 or int64 + eps: RMSNorm epsilon + num_q_heads: number of Q heads (after TP split) + num_kv_heads: number of KV heads (after TP split) + head_dim: per-head dimension + rotary_dim: rotary dimension; must be even and <= head_dim + + Returns: + (q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim). + ``gate_out`` is the raw (pre-sigmoid) gate. + """ + if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0: + raise ValueError( + f"rotary_dim must be a positive even integer <= head_dim, " + f"got rotary_dim={rotary_dim}, head_dim={head_dim}" + ) + + n_tokens = q_gate.shape[0] + q_out = torch.empty( + (n_tokens, num_q_heads * head_dim), dtype=q_gate.dtype, device=q_gate.device + ) + k_out = torch.empty( + (n_tokens, num_kv_heads * head_dim), dtype=k.dtype, device=k.device + ) + gate_out = torch.empty_like(q_out) + if n_tokens == 0: + return q_out, k_out, gate_out + + half_rotary = rotary_dim // 2 + head_block = triton.next_power_of_2(head_dim) + rot_half_block = triton.next_power_of_2(half_rotary) + num_warps = max(1, head_block // 64) + + grid = (n_tokens, num_q_heads + num_kv_heads) + _fused_qk_rmsnorm_rope_gate_kernel[grid]( + q_gate, + k, + q_out, + k_out, + gate_out, + q_weight, + k_weight, + cos_sin_cache, + positions, + q_gate.stride(0), + k.stride(0), + q_out.stride(0), + k_out.stride(0), + gate_out.stride(0), + cos_sin_cache.stride(0), + num_q_heads, + num_kv_heads, + head_dim, + rotary_dim, + half_rotary, + eps, + INPUT_DTYPE=tl.bfloat16 if q_gate.dtype == torch.bfloat16 else tl.float16, + HEAD_BLOCK=head_block, + ROT_HALF_BLOCK=rot_half_block, + HAS_PASS=rotary_dim < head_dim, + num_warps=num_warps, + num_stages=2, + ) + return q_out, k_out, gate_out diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 28e1846662..165a2d94cf 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoE, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate from vllm.model_executor.layers.layernorm import ( GemmaRMSNorm as Qwen3NextRMSNorm, ) @@ -58,6 +59,7 @@ from vllm.model_executor.model_loader.weight_utils import ( ) from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig @@ -283,13 +285,50 @@ class Qwen3NextAttention(nn.Module): self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps) - def forward( + # Fuse the gated split + QK-RMSNorm + (partial) NeoX RoPE + gate copy. + # TODO: support MRoPE + mm_config = model_config.multimodal_config if model_config else None + text_only = mm_config is None or mm_config.language_model_only + self.use_fused_qk_norm_rope_gate = ( + self.attn_output_gate + and getattr(self.rotary_emb, "is_neox_style", False) + and current_platform.is_cuda() + and text_only + ) + + def _project_qkv_gate( self, + qkv: torch.Tensor, positions: torch.Tensor, - output: torch.Tensor, - hidden_states: torch.Tensor, - ): - qkv, _ = self.qkv_proj(hidden_states) + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Return post-norm, post-RoPE (q, k, v) and the pre-sigmoid gate. + + Dispatches between the fused Triton kernel and the eager + split + QK-RMSNorm + RoPE path. ``gate`` is ``None`` when output + gating is disabled. + """ + if self.use_fused_qk_norm_rope_gate: + q_gate, k, v = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + # mRoPE passes positions as (3, n_tokens) for T/H/W. Fusion is only + # enabled text-only, where the three rows are identical, so taking + # the T row is exact. (1D positions pass through.) + pos = positions[0] if positions.ndim == 2 else positions + q, k, gate = fused_qk_rmsnorm_rope_gate( + q_gate, + k, + self.q_norm.weight.float() + 1.0, + self.k_norm.weight.float() + 1.0, + self.rotary_emb.cos_sin_cache, + pos, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.rotary_emb.rotary_dim, + ) + return q, k, v, gate if self.attn_output_gate: q_gate, k, v = qkv.split( @@ -302,6 +341,7 @@ class Qwen3NextAttention(nn.Module): gate = gate.reshape(*orig_shape, -1) else: q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + gate = None q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view( -1, self.num_heads * self.head_dim @@ -309,15 +349,20 @@ class Qwen3NextAttention(nn.Module): k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view( -1, self.num_kv_heads * self.head_dim ) - q, k = self.rotary_emb(positions, q, k) + return q, k, v, gate + def forward( + self, + positions: torch.Tensor, + output: torch.Tensor, + hidden_states: torch.Tensor, + ): + qkv, _ = self.qkv_proj(hidden_states) + q, k, v, gate = self._project_qkv_gate(qkv, positions) attn_output = self.attn(q, k, v) - - if self.attn_output_gate: - gate = torch.sigmoid(gate) - attn_output = attn_output * gate - + if gate is not None: + attn_output = attn_output * torch.sigmoid(gate) output[:], _ = self.o_proj(attn_output)