A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] fuse qk rmsnorm rope gate for qwen3.5 (#44176)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Jun 9, 2026 at 22:12 UTC
7a89b72564a2c39d170b9478e947382216710ee2
4 files changed
+418
-12
.buildkite/test_areas/kernels.yaml
+2
-1
@@ -21,8 +21,9 @@ steps:
21
- csrc/
22
- tests/kernels/core
23
- tests/kernels/test_concat_mla_q.py
24
+ - tests/kernels/test_fused_qk_norm_rope_gate.py
25
commands:
25
- - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
26
+ - 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
27
28
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
29
key: kernels-minimax-reduce-rms-test-2-gpus
tests/kernels/test_fused_qk_norm_rope_gate.py
new
+159
@@ -0,0 +1,159 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
8
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
9
+from vllm.platforms import current_platform
10
+from vllm.utils.torch_utils import set_random_seed
11
+
12
+# Qwen/Qwen3.6-27B config (huggingface.co/Qwen/Qwen3.6-27B), TP=1 shapes.
13
+NUM_Q_HEADS = 24
14
+NUM_KV_HEADS = 4
15
+HEAD_DIM = 256
16
+PARTIAL_ROTARY_FACTOR = 0.25
17
+ROTARY_DIM = int(HEAD_DIM * PARTIAL_ROTARY_FACTOR) # 64
18
+RMS_NORM_EPS = 1e-6
19
+MAX_POSITION_EMBEDDINGS = 262144
20
+ROPE_THETA = 10000000.0
21
+
22
+DTYPES = [torch.bfloat16]
23
+SEEDS = [13]
24
+NUM_TOKENS = [1, 4, 37]
25
+
26
+
27
+def _ref_qk_rmsnorm_rope_gate(
28
+ q_gate: torch.Tensor,
29
+ k: torch.Tensor,
30
+ q_gamma: torch.Tensor,
31
+ k_gamma: torch.Tensor,
32
+ cos_sin_cache: torch.Tensor,
33
+ positions: torch.Tensor,
34
+ eps: float,
35
+ num_q_heads: int,
36
+ num_kv_heads: int,
37
+ head_dim: int,
38
+ rotary_dim: int,
39
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
40
+ """PyTorch reference: split + RMSNorm + partial NeoX RoPE + gate extraction.
41
+
42
+ Matches ``fused_qk_rmsnorm_rope_gate``'s contract: ``q_gamma`` / ``k_gamma``
43
+ are the already-adjusted effective gammas (for GemmaRMSNorm the caller
44
+ has done ``weight + 1`` before passing them in).
45
+ """
46
+ n_tokens = q_gate.shape[0]
47
+ half = rotary_dim // 2
48
+
49
+ # Per head the q projection is laid out as [q | gate].
50
+ q_gate = q_gate.view(n_tokens, num_q_heads, 2 * head_dim)
51
+ q = q_gate[..., :head_dim]
52
+ gate = q_gate[..., head_dim:].reshape(n_tokens, num_q_heads * head_dim)
53
+ k = k.view(n_tokens, num_kv_heads, head_dim)
54
+
55
+ def rms_norm(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor:
56
+ orig_dtype = x.dtype
57
+ x = x.float()
58
+ var = x.pow(2).mean(dim=-1, keepdim=True)
59
+ x = x * torch.rsqrt(var + eps)
60
+ return (x * gamma.float()).to(orig_dtype)
61
+
62
+ q = rms_norm(q, q_gamma)
63
+ k = rms_norm(k, k_gamma)
64
+
65
+ # Partial NeoX RoPE on the first ``rotary_dim`` elements of each head;
66
+ # cos_sin_cache row is packed as [cos(half) | sin(half)].
67
+ pos = positions.view(-1)
68
+ cos = cos_sin_cache[pos, :half].float()[:, None, :]
69
+ sin = cos_sin_cache[pos, half:rotary_dim].float()[:, None, :]
70
+
71
+ def rope(x: torch.Tensor) -> torch.Tensor:
72
+ x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
73
+ x1 = x_rot[..., :half].float()
74
+ x2 = x_rot[..., half:].float()
75
+ o1 = x1 * cos - x2 * sin
76
+ o2 = x2 * cos + x1 * sin
77
+ rotated = torch.cat([o1, o2], dim=-1).to(x.dtype)
78
+ return torch.cat([rotated, x_pass], dim=-1)
79
+
80
+ q = rope(q).reshape(n_tokens, num_q_heads * head_dim)
81
+ k = rope(k).reshape(n_tokens, num_kv_heads * head_dim)
82
+ return q, k, gate
83
+
84
+
85
+@pytest.mark.skipif(
86
+ not current_platform.is_cuda_alike(),
87
+ reason="fused_qk_rmsnorm_rope_gate Triton kernel requires CUDA/ROCm",
88
+)
89
+@pytest.mark.parametrize("dtype", DTYPES)
90
+@pytest.mark.parametrize("seed", SEEDS)
91
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
92
+@torch.inference_mode()
93
+def test_fused_qk_norm_rope_gate_matches_reference(
94
+ default_vllm_config,
95
+ dtype: torch.dtype,
96
+ seed: int,
97
+ num_tokens: int,
98
+):
99
+ device = torch.device("cuda", torch.accelerator.current_device_index())
100
+ torch.set_default_device(device)
101
+ set_random_seed(seed)
102
+
103
+ q_gate = torch.randn(
104
+ num_tokens, NUM_Q_HEADS * 2 * HEAD_DIM, dtype=dtype, device=device
105
+ )
106
+ k = torch.randn(num_tokens, NUM_KV_HEADS * HEAD_DIM, dtype=dtype, device=device)
107
+ # GemmaRMSNorm-style: the kernel takes the effective gamma (weight + 1).
108
+ q_gamma = (
109
+ torch.empty(HEAD_DIM, dtype=dtype, device=device).normal_(mean=0.0, std=0.1)
110
+ + 1.0
111
+ )
112
+ k_gamma = (
113
+ torch.empty(HEAD_DIM, dtype=dtype, device=device).normal_(mean=0.0, std=0.1)
114
+ + 1.0
115
+ )
116
+
117
+ # fused_qk_rmsnorm_rope_gate only handles NeoX-style RoPE.
118
+ rope = RotaryEmbedding(
119
+ head_size=HEAD_DIM,
120
+ rotary_dim=ROTARY_DIM,
121
+ max_position_embeddings=MAX_POSITION_EMBEDDINGS,
122
+ base=ROPE_THETA,
123
+ is_neox_style=True,
124
+ dtype=dtype,
125
+ ).to(device)
126
+ positions = torch.arange(num_tokens, dtype=torch.long, device=device)
127
+
128
+ q_ref, k_ref, gate_ref = _ref_qk_rmsnorm_rope_gate(
129
+ q_gate,
130
+ k,
131
+ q_gamma,
132
+ k_gamma,
133
+ rope.cos_sin_cache,
134
+ positions,
135
+ RMS_NORM_EPS,
136
+ NUM_Q_HEADS,
137
+ NUM_KV_HEADS,
138
+ HEAD_DIM,
139
+ ROTARY_DIM,
140
+ )
141
+ q_out, k_out, gate_out = fused_qk_rmsnorm_rope_gate(
142
+ q_gate,
143
+ k,
144
+ q_gamma,
145
+ k_gamma,
146
+ rope.cos_sin_cache,
147
+ positions,
148
+ RMS_NORM_EPS,
149
+ NUM_Q_HEADS,
150
+ NUM_KV_HEADS,
151
+ HEAD_DIM,
152
+ ROTARY_DIM,
153
+ )
154
+
155
+ atol, rtol = 2e-3, 5e-3
156
+ torch.testing.assert_close(q_out, q_ref, atol=atol, rtol=rtol)
157
+ torch.testing.assert_close(k_out, k_ref, atol=atol, rtol=rtol)
158
+ # gate is a verbatim copy of the source slice — must match bit-exactly.
159
+ torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)
vllm/model_executor/layers/fused_qk_norm_rope.py
new
+201
@@ -0,0 +1,201 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused QK-RMSNorm + (partial) RoPE + gate copy Triton kernel.
4
+
5
+Currently used by the Qwen3.5 attention path (``attn_output_gate`` with
6
+NeoX-style partial RoPE). The unfused reference sequence is
7
+``split -> GemmaRMSNorm -> RoPE -> gate chunk``; this collapses it into a
8
+single Triton launch. See :func:`fused_qk_rmsnorm_rope_gate`.
9
+"""
10
+
11
+import torch
12
+
13
+from vllm.triton_utils import tl, triton
14
+
15
+
16
+@triton.jit
17
+def _fused_qk_rmsnorm_rope_gate_kernel(
18
+ q_gate_ptr,
19
+ k_ptr,
20
+ q_out_ptr,
21
+ k_out_ptr,
22
+ gate_out_ptr,
23
+ q_weight_ptr,
24
+ k_weight_ptr,
25
+ cos_sin_cache_ptr,
26
+ positions_ptr,
27
+ q_gate_stride_t,
28
+ k_stride_t,
29
+ q_out_stride_t,
30
+ k_out_stride_t,
31
+ gate_out_stride_t,
32
+ cache_stride_p,
33
+ num_q_heads: tl.constexpr,
34
+ num_kv_heads: tl.constexpr,
35
+ head_dim: tl.constexpr,
36
+ rotary_dim: tl.constexpr,
37
+ half_rotary: tl.constexpr,
38
+ eps: tl.constexpr,
39
+ INPUT_DTYPE: tl.constexpr,
40
+ HEAD_BLOCK: tl.constexpr,
41
+ ROT_HALF_BLOCK: tl.constexpr,
42
+ HAS_PASS: tl.constexpr,
43
+):
44
+ token = tl.program_id(0)
45
+ head = tl.program_id(1)
46
+ is_k = head >= num_q_heads
47
+ local_head = tl.where(is_k, head - num_q_heads, head)
48
+
49
+ if is_k:
50
+ in_base = k_ptr + token * k_stride_t + local_head * head_dim
51
+ w_ptr = k_weight_ptr
52
+ out_base = k_out_ptr + token * k_out_stride_t + local_head * head_dim
53
+ else:
54
+ in_base = q_gate_ptr + token * q_gate_stride_t + local_head * 2 * head_dim
55
+ w_ptr = q_weight_ptr
56
+ out_base = q_out_ptr + token * q_out_stride_t + local_head * head_dim
57
+
58
+ # --- RMSNorm: variance over the full head_dim ---
59
+ head_offs = tl.arange(0, HEAD_BLOCK)
60
+ head_mask = head_offs < head_dim
61
+ x = tl.load(in_base + head_offs, mask=head_mask, other=0.0).to(tl.float32)
62
+ var = tl.sum(x * x, axis=0) / head_dim
63
+ inv_rms = tl.rsqrt(var + eps)
64
+ w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32)
65
+ # Round-trip through INPUT_DTYPE so the RoPE input matches the bf16-storage
66
+ # behavior of the unfused (qk_rmsnorm -> memory -> apply_rope) reference path.
67
+ x_norm = (x * inv_rms * w).to(INPUT_DTYPE).to(tl.float32)
68
+
69
+ # --- Pass-through tail [rotary_dim, head_dim): RMSNorm-only, no rotation ---
70
+ # The rotary head [0, rotary_dim) will be overwritten by the RoPE store below.
71
+ if HAS_PASS:
72
+ pass_mask = head_mask & (head_offs >= rotary_dim)
73
+ tl.store(out_base + head_offs, x_norm, mask=pass_mask)
74
+
75
+ # --- Partial RoPE on the first rotary_dim elements ---
76
+ # Triton lacks easy sub-vector slicing of x_norm, so we recompute the
77
+ # normalized rotary halves on a smaller block (next_pow2(half_rotary)).
78
+ # The extra ~rotary_dim element reload hits L1, so the cost is negligible.
79
+ rot_offs = tl.arange(0, ROT_HALF_BLOCK)
80
+ rot_mask = rot_offs < half_rotary
81
+ x_rot1 = tl.load(in_base + rot_offs, mask=rot_mask, other=0.0).to(tl.float32)
82
+ x_rot2 = tl.load(in_base + half_rotary + rot_offs, mask=rot_mask, other=0.0).to(
83
+ tl.float32
84
+ )
85
+ w_rot1 = tl.load(w_ptr + rot_offs, mask=rot_mask, other=0.0).to(tl.float32)
86
+ w_rot2 = tl.load(w_ptr + half_rotary + rot_offs, mask=rot_mask, other=0.0).to(
87
+ tl.float32
88
+ )
89
+ x_rot1 = (x_rot1 * inv_rms * w_rot1).to(INPUT_DTYPE).to(tl.float32)
90
+ x_rot2 = (x_rot2 * inv_rms * w_rot2).to(INPUT_DTYPE).to(tl.float32)
91
+
92
+ # Always use int64 for position to avoid overflow in address computation.
93
+ pos = tl.load(positions_ptr + token).to(tl.int64)
94
+ cache_offset = pos * cache_stride_p
95
+ cos = tl.load(
96
+ cos_sin_cache_ptr + cache_offset + rot_offs, mask=rot_mask, other=0.0
97
+ ).to(tl.float32)
98
+ sin = tl.load(
99
+ cos_sin_cache_ptr + cache_offset + half_rotary + rot_offs,
100
+ mask=rot_mask,
101
+ other=0.0,
102
+ ).to(tl.float32)
103
+
104
+ o1 = x_rot1 * cos - x_rot2 * sin
105
+ o2 = x_rot2 * cos + x_rot1 * sin
106
+ tl.store(out_base + rot_offs, o1, mask=rot_mask)
107
+ tl.store(out_base + half_rotary + rot_offs, o2, mask=rot_mask)
108
+
109
+ # --- Gate copy (q heads only, verbatim) ---
110
+ if not is_k:
111
+ gate_in_base = in_base + head_dim
112
+ gate_out_base = gate_out_ptr + token * gate_out_stride_t + local_head * head_dim
113
+ g = tl.load(gate_in_base + head_offs, mask=head_mask, other=0.0)
114
+ tl.store(gate_out_base + head_offs, g, mask=head_mask)
115
+
116
+
117
+def fused_qk_rmsnorm_rope_gate(
118
+ q_gate: torch.Tensor,
119
+ k: torch.Tensor,
120
+ q_weight: torch.Tensor,
121
+ k_weight: torch.Tensor,
122
+ cos_sin_cache: torch.Tensor,
123
+ positions: torch.Tensor,
124
+ eps: float,
125
+ num_q_heads: int,
126
+ num_kv_heads: int,
127
+ head_dim: int,
128
+ rotary_dim: int,
129
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
130
+ """Fused split + QK-RMSNorm + (partial) RoPE + gate copy for Qwen3.5 attn.
131
+
132
+ Args:
133
+ q_gate: (n_tokens, num_q_heads * 2 * head_dim) -- per head: [q|gate]
134
+ k: (n_tokens, num_kv_heads * head_dim)
135
+ q_weight: (head_dim,) GemmaRMSNorm effective weight (already +1)
136
+ k_weight: (head_dim,) GemmaRMSNorm effective weight (already +1)
137
+ cos_sin_cache: (max_pos, rotary_dim) packed [cos|sin]
138
+ positions: (n_tokens,) int32 or int64
139
+ eps: RMSNorm epsilon
140
+ num_q_heads: number of Q heads (after TP split)
141
+ num_kv_heads: number of KV heads (after TP split)
142
+ head_dim: per-head dimension
143
+ rotary_dim: rotary dimension; must be even and <= head_dim
144
+
145
+ Returns:
146
+ (q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim).
147
+ ``gate_out`` is the raw (pre-sigmoid) gate.
148
+ """
149
+ if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
150
+ raise ValueError(
151
+ f"rotary_dim must be a positive even integer <= head_dim, "
152
+ f"got rotary_dim={rotary_dim}, head_dim={head_dim}"
153
+ )
154
+
155
+ n_tokens = q_gate.shape[0]
156
+ q_out = torch.empty(
157
+ (n_tokens, num_q_heads * head_dim), dtype=q_gate.dtype, device=q_gate.device
158
+ )
159
+ k_out = torch.empty(
160
+ (n_tokens, num_kv_heads * head_dim), dtype=k.dtype, device=k.device
161
+ )
162
+ gate_out = torch.empty_like(q_out)
163
+ if n_tokens == 0:
164
+ return q_out, k_out, gate_out
165
+
166
+ half_rotary = rotary_dim // 2
167
+ head_block = triton.next_power_of_2(head_dim)
168
+ rot_half_block = triton.next_power_of_2(half_rotary)
169
+ num_warps = max(1, head_block // 64)
170
+
171
+ grid = (n_tokens, num_q_heads + num_kv_heads)
172
+ _fused_qk_rmsnorm_rope_gate_kernel[grid](
173
+ q_gate,
174
+ k,
175
+ q_out,
176
+ k_out,
177
+ gate_out,
178
+ q_weight,
179
+ k_weight,
180
+ cos_sin_cache,
181
+ positions,
182
+ q_gate.stride(0),
183
+ k.stride(0),
184
+ q_out.stride(0),
185
+ k_out.stride(0),
186
+ gate_out.stride(0),
187
+ cos_sin_cache.stride(0),
188
+ num_q_heads,
189
+ num_kv_heads,
190
+ head_dim,
191
+ rotary_dim,
192
+ half_rotary,
193
+ eps,
194
+ INPUT_DTYPE=tl.bfloat16 if q_gate.dtype == torch.bfloat16 else tl.float16,
195
+ HEAD_BLOCK=head_block,
196
+ ROT_HALF_BLOCK=rot_half_block,
197
+ HAS_PASS=rotary_dim < head_dim,
198
+ num_warps=num_warps,
199
+ num_stages=2,
200
+ )
201
+ return q_out, k_out, gate_out
vllm/model_executor/models/qwen3_next.py
+56
-11
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe import (
28
FusedMoE,
29
fused_moe_make_expert_params_mapping,
30
)
31
+from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
32
from vllm.model_executor.layers.layernorm import (
33
GemmaRMSNorm as Qwen3NextRMSNorm,
34
)
@@ -58,6 +59,7 @@ from vllm.model_executor.model_loader.weight_utils import (
59
)
60
from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
61
from vllm.model_executor.models.utils import sequence_parallel_chunk
62
+from vllm.platforms import current_platform
63
from vllm.sequence import IntermediateTensors
64
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
65
@@ -283,13 +285,50 @@ class Qwen3NextAttention(nn.Module):
285
self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
286
self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
287
286
- def forward(
288
+ # Fuse the gated split + QK-RMSNorm + (partial) NeoX RoPE + gate copy.
289
+ # TODO: support MRoPE
290
+ mm_config = model_config.multimodal_config if model_config else None
291
+ text_only = mm_config is None or mm_config.language_model_only
292
+ self.use_fused_qk_norm_rope_gate = (
293
+ self.attn_output_gate
294
+ and getattr(self.rotary_emb, "is_neox_style", False)
295
+ and current_platform.is_cuda()
296
+ and text_only
297
+ )
298
+
299
+ def _project_qkv_gate(
300
self,
301
+ qkv: torch.Tensor,
302
positions: torch.Tensor,
289
- output: torch.Tensor,
290
- hidden_states: torch.Tensor,
291
- ):
292
- qkv, _ = self.qkv_proj(hidden_states)
303
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
304
+ """Return post-norm, post-RoPE (q, k, v) and the pre-sigmoid gate.
305
+
306
+ Dispatches between the fused Triton kernel and the eager
307
+ split + QK-RMSNorm + RoPE path. ``gate`` is ``None`` when output
308
+ gating is disabled.
309
+ """
310
+ if self.use_fused_qk_norm_rope_gate:
311
+ q_gate, k, v = qkv.split(
312
+ [self.q_size * 2, self.kv_size, self.kv_size], dim=-1
313
+ )
314
+ # mRoPE passes positions as (3, n_tokens) for T/H/W. Fusion is only
315
+ # enabled text-only, where the three rows are identical, so taking
316
+ # the T row is exact. (1D positions pass through.)
317
+ pos = positions[0] if positions.ndim == 2 else positions
318
+ q, k, gate = fused_qk_rmsnorm_rope_gate(
319
+ q_gate,
320
+ k,
321
+ self.q_norm.weight.float() + 1.0,
322
+ self.k_norm.weight.float() + 1.0,
323
+ self.rotary_emb.cos_sin_cache,
324
+ pos,
325
+ self.q_norm.variance_epsilon,
326
+ self.num_heads,
327
+ self.num_kv_heads,
328
+ self.head_dim,
329
+ self.rotary_emb.rotary_dim,
330
+ )
331
+ return q, k, v, gate
332
333
if self.attn_output_gate:
334
q_gate, k, v = qkv.split(
@@ -302,6 +341,7 @@ class Qwen3NextAttention(nn.Module):
341
gate = gate.reshape(*orig_shape, -1)
342
else:
343
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
344
+ gate = None
345
346
q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(
347
-1, self.num_heads * self.head_dim
@@ -309,15 +349,20 @@ class Qwen3NextAttention(nn.Module):
349
k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(
350
-1, self.num_kv_heads * self.head_dim
351
)
312
-
352
q, k = self.rotary_emb(positions, q, k)
353
+ return q, k, v, gate
354
355
+ def forward(
356
+ self,
357
+ positions: torch.Tensor,
358
+ output: torch.Tensor,
359
+ hidden_states: torch.Tensor,
360
+ ):
361
+ qkv, _ = self.qkv_proj(hidden_states)
362
+ q, k, v, gate = self._project_qkv_gate(qkv, positions)
363
attn_output = self.attn(q, k, v)
316
-
317
- if self.attn_output_gate:
318
- gate = torch.sigmoid(gate)
319
- attn_output = attn_output * gate
320
-
364
+ if gate is not None:
365
+ attn_output = attn_output * torch.sigmoid(gate)
366
output[:], _ = self.o_proj(attn_output)
367
368