iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] fuse qk rmsnorm rope gate for qwen3.5 (#44176)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Jiangyun Zhu
2026-06-09 22:12:17 +08:00
committed by GitHub
parent dc10e467a9
commit 7a89b72564
4 changed files with 418 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+2 -1
View File
@@ -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
tests/kernels/test_fused_qk_norm_rope_gate.py
+159
View File
@@ -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)
vllm/model_executor/layers/fused_qk_norm_rope.py
+201
View File
@@ -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
vllm/model_executor/models/qwen3_next.py
+56 -11
View File
@@ -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)