[Kernel][Perf] Enable fused QK-norm + partial MRoPE + gate for Qwen3.6 (#52676)
Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
kiroxu committed
Aug 25, 2026 at 09:28 UTC
07ef21bc69842665cf5056371ea15a6d014f3b05
3 files changed
+192
-138
tests/kernels/test_fused_qk_norm_rope_gate.py
+77
-118
@@ -5,155 +5,114 @@ 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
8
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
9
+from vllm.model_executor.layers.rotary_embedding import (
10
+ MRotaryEmbedding,
11
+ RotaryEmbedding,
12
+)
13
from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import set_random_seed
15
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
16
+# Qwen3.6 TP=1 attention geometry.
17
HEAD_DIM = 256
16
-PARTIAL_ROTARY_FACTOR = 0.25
17
-ROTARY_DIM = int(HEAD_DIM * PARTIAL_ROTARY_FACTOR) # 64
18
+ROTARY_DIM = 64
19
RMS_NORM_EPS = 1e-6
20
MAX_POSITION_EMBEDDINGS = 262144
21
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
22
+DTYPE = torch.bfloat16
23
+SEED = 13
24
+MROPE_SECTION = (11, 11, 10)
25
+ROPE_CASES = [
26
+ pytest.param(24, 4, None, id="rope"),
27
+ pytest.param(16, 2, MROPE_SECTION, id="interleaved-mrope"),
28
+]
29
30
31
@pytest.mark.skipif(
32
not current_platform.is_cuda_alike(),
33
reason="fused_qk_rmsnorm_rope_gate Triton kernel requires CUDA/ROCm",
34
)
89
-@pytest.mark.parametrize("dtype", DTYPES)
90
-@pytest.mark.parametrize("seed", SEEDS)
91
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
35
+@pytest.mark.parametrize("num_q_heads,num_kv_heads,mrope_section", ROPE_CASES)
36
+@pytest.mark.parametrize("num_tokens", [1, 4, 37])
37
@torch.inference_mode()
38
def test_fused_qk_norm_rope_gate_matches_reference(
39
default_vllm_config,
95
- dtype: torch.dtype,
96
- seed: int,
40
num_tokens: int,
98
-):
41
+ num_q_heads: int,
42
+ num_kv_heads: int,
43
+ mrope_section: tuple[int, int, int] | None,
44
+) -> None:
45
device = torch.device("cuda", torch.accelerator.current_device_index())
46
torch.set_default_device(device)
101
- set_random_seed(seed)
47
+ set_random_seed(SEED)
48
49
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
50
+ num_tokens, num_q_heads * 2 * HEAD_DIM, dtype=DTYPE, device=device
51
)
52
+ k = torch.randn(num_tokens, num_kv_heads * HEAD_DIM, dtype=DTYPE, device=device)
53
+
54
+ q_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
55
+ k_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
56
+ q_norm.weight.normal_(std=0.1)
57
+ k_norm.weight.normal_(std=0.1)
58
+
59
+ q_gate_heads = q_gate.view(num_tokens, num_q_heads, 2 * HEAD_DIM)
60
+ q = q_gate_heads[..., :HEAD_DIM]
61
+ gate_ref = q_gate_heads[..., HEAD_DIM:].reshape(num_tokens, num_q_heads * HEAD_DIM)
62
+ q_ref = q_norm.forward_native(q)
63
+ k_ref = k_norm.forward_native(k.view(num_tokens, num_kv_heads, HEAD_DIM))
64
+ assert isinstance(q_ref, torch.Tensor)
65
+ assert isinstance(k_ref, torch.Tensor)
66
+ q_ref = q_ref.reshape(num_tokens, num_q_heads * HEAD_DIM)
67
+ k_ref = k_ref.reshape(num_tokens, num_kv_heads * HEAD_DIM)
68
+
69
+ if mrope_section is None:
70
+ rope = RotaryEmbedding(
71
+ HEAD_DIM,
72
+ ROTARY_DIM,
73
+ MAX_POSITION_EMBEDDINGS,
74
+ ROPE_THETA,
75
+ True,
76
+ DTYPE,
77
+ ).to(device)
78
+ positions = torch.arange(num_tokens, dtype=torch.long, device=device)
79
+ else:
80
+ rope = MRotaryEmbedding(
81
+ HEAD_DIM,
82
+ ROTARY_DIM,
83
+ MAX_POSITION_EMBEDDINGS,
84
+ ROPE_THETA,
85
+ True,
86
+ DTYPE,
87
+ mrope_section=list(mrope_section),
88
+ mrope_interleaved=True,
89
+ ).to(device)
90
+ positions = torch.arange(3 * num_tokens, dtype=torch.long, device=device).view(
91
+ 3, num_tokens
92
+ )
93
+ assert torch.unique(positions[:, 0]).numel() == 3
94
+
95
+ q_ref, k_ref = rope.forward_native(positions, q_ref, k_ref)
96
+ assert k_ref is not None
97
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
- )
98
q_out, k_out, gate_out = fused_qk_rmsnorm_rope_gate(
99
q_gate,
100
k,
144
- q_gamma,
145
- k_gamma,
101
+ q_norm.weight,
102
+ k_norm.weight,
103
rope.cos_sin_cache,
104
positions,
105
RMS_NORM_EPS,
149
- NUM_Q_HEADS,
150
- NUM_KV_HEADS,
106
+ num_q_heads,
107
+ num_kv_heads,
108
HEAD_DIM,
109
ROTARY_DIM,
110
+ mrope_section=mrope_section,
111
+ norm_beta=1.0,
112
)
113
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.
114
+ # The built-in reference performs RoPE in BF16, while the fused kernel
115
+ # promotes the BF16-normalized values to FP32 for RoPE before storing BF16.
116
+ torch.testing.assert_close(q_out, q_ref, atol=1e-2, rtol=1e-2)
117
+ torch.testing.assert_close(k_out, k_ref, atol=1e-2, rtol=1e-2)
118
torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)
vllm/model_executor/layers/fused_qk_norm_rope.py
+87
-10
@@ -30,16 +30,22 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
30
k_out_stride_t,
31
gate_out_stride_t,
32
cache_stride_p,
33
+ positions_stride_m,
34
+ positions_stride_t,
35
num_q_heads: tl.constexpr,
36
num_kv_heads: tl.constexpr,
37
head_dim: tl.constexpr,
38
rotary_dim: tl.constexpr,
39
half_rotary: tl.constexpr,
40
eps: tl.constexpr,
41
+ norm_beta: tl.constexpr,
42
INPUT_DTYPE: tl.constexpr,
43
HEAD_BLOCK: tl.constexpr,
44
ROT_HALF_BLOCK: tl.constexpr,
45
HAS_PASS: tl.constexpr,
46
+ HAS_MROPE: tl.constexpr,
47
+ MROPE_SECTION_H: tl.constexpr,
48
+ MROPE_SECTION_W: tl.constexpr,
49
):
50
token = tl.program_id(0)
51
head = tl.program_id(1)
@@ -61,7 +67,7 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
67
x = tl.load(in_base + head_offs, mask=head_mask, other=0.0).to(tl.float32)
68
var = tl.sum(x * x, axis=0) / head_dim
69
inv_rms = tl.rsqrt(var + eps)
64
- w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32)
70
+ w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32) + norm_beta
71
# Round-trip through INPUT_DTYPE so the RoPE input matches the bf16-storage
72
# behavior of the unfused (qk_rmsnorm -> memory -> apply_rope) reference path.
73
x_norm = (x * inv_rms * w).to(INPUT_DTYPE).to(tl.float32)
@@ -82,15 +88,30 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
88
x_rot2 = tl.load(in_base + half_rotary + rot_offs, mask=rot_mask, other=0.0).to(
89
tl.float32
90
)
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
91
+ w_rot1 = (
92
+ tl.load(w_ptr + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + norm_beta
93
+ )
94
+ w_rot2 = (
95
+ tl.load(w_ptr + half_rotary + rot_offs, mask=rot_mask, other=0.0).to(tl.float32)
96
+ + norm_beta
97
)
98
x_rot1 = (x_rot1 * inv_rms * w_rot1).to(INPUT_DTYPE).to(tl.float32)
99
x_rot2 = (x_rot2 * inv_rms * w_rot2).to(INPUT_DTYPE).to(tl.float32)
100
92
- # Always use int64 for position to avoid overflow in address computation.
93
- pos = tl.load(positions_ptr + token).to(tl.int64)
101
+ # Always use int64 for positions to avoid overflow in address computation.
102
+ pos_t = tl.load(positions_ptr + token * positions_stride_t).to(tl.int64)
103
+ if HAS_MROPE:
104
+ pos_h = tl.load(
105
+ positions_ptr + positions_stride_m + token * positions_stride_t
106
+ ).to(tl.int64)
107
+ pos_w = tl.load(
108
+ positions_ptr + 2 * positions_stride_m + token * positions_stride_t
109
+ ).to(tl.int64)
110
+ is_h = (rot_offs % 3 == 1) & (rot_offs < 3 * MROPE_SECTION_H)
111
+ is_w = (rot_offs % 3 == 2) & (rot_offs < 3 * MROPE_SECTION_W)
112
+ pos = tl.where(is_h, pos_h, tl.where(is_w, pos_w, pos_t))
113
+ else:
114
+ pos = pos_t
115
cache_offset = pos * cache_stride_p
116
cos = tl.load(
117
cos_sin_cache_ptr + cache_offset + rot_offs, mask=rot_mask, other=0.0
@@ -126,21 +147,25 @@ def fused_qk_rmsnorm_rope_gate(
147
num_kv_heads: int,
148
head_dim: int,
149
rotary_dim: int,
150
+ mrope_section: list[int] | tuple[int, int, int] | None = None,
151
+ norm_beta: float = 0.0,
152
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
130
- """Fused split + QK-RMSNorm + (partial) RoPE + gate copy for Qwen3.5 attn.
153
+ """Fused split + QK-RMSNorm + (partial) RoPE + gate copy for Qwen attn.
154
155
Args:
156
q_gate: (n_tokens, num_q_heads * 2 * head_dim) -- per head: [q|gate]
157
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)
158
+ q_weight: (head_dim,) RMSNorm weight
159
+ k_weight: (head_dim,) RMSNorm weight
160
cos_sin_cache: (max_pos, rotary_dim) packed [cos|sin]
138
- positions: (n_tokens,) int32 or int64
161
+ positions: (n_tokens,) or (3, n_tokens) int32 or int64
162
eps: RMSNorm epsilon
163
num_q_heads: number of Q heads (after TP split)
164
num_kv_heads: number of KV heads (after TP split)
165
head_dim: per-head dimension
166
rotary_dim: rotary dimension; must be even and <= head_dim
167
+ mrope_section: interleaved T/H/W frequency counts for 2D positions
168
+ norm_beta: scalar added to the RMSNorm weight
169
170
Returns:
171
(q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim).
@@ -151,6 +176,52 @@ def fused_qk_rmsnorm_rope_gate(
176
f"rotary_dim must be a positive even integer <= head_dim, "
177
f"got rotary_dim={rotary_dim}, head_dim={head_dim}"
178
)
179
+ if q_gate.dtype not in (torch.float16, torch.bfloat16) or k.dtype != q_gate.dtype:
180
+ raise ValueError(
181
+ "q_gate and k must have the same FP16 or BF16 dtype, "
182
+ f"got {q_gate.dtype} and {k.dtype}"
183
+ )
184
+ for name, tensor in (
185
+ ("q_gate", q_gate),
186
+ ("k", k),
187
+ ("q_weight", q_weight),
188
+ ("k_weight", k_weight),
189
+ ("cos_sin_cache", cos_sin_cache),
190
+ ):
191
+ if tensor.stride(-1) != 1:
192
+ raise ValueError(f"{name} must be contiguous in its last dimension")
193
+
194
+ if positions.ndim not in (1, 2):
195
+ raise ValueError(f"positions must be 1D or 2D, got shape={positions.shape}")
196
+ if positions.shape[-1] != q_gate.shape[0]:
197
+ raise ValueError(
198
+ "positions token dimension must match q_gate, "
199
+ f"got {positions.shape[-1]} and {q_gate.shape[0]}"
200
+ )
201
+
202
+ has_mrope = positions.ndim == 2
203
+ if has_mrope:
204
+ if positions.shape[0] != 3:
205
+ raise ValueError(
206
+ f"MRoPE positions must have shape (3, n_tokens), got {positions.shape}"
207
+ )
208
+ if mrope_section is None or len(mrope_section) != 3:
209
+ raise ValueError("mrope_section must contain the T/H/W frequency counts")
210
+ if sum(mrope_section) != rotary_dim // 2:
211
+ raise ValueError(
212
+ "mrope_section must sum to rotary_dim // 2, "
213
+ f"got {mrope_section} and rotary_dim={rotary_dim}"
214
+ )
215
+ mrope_section_h = mrope_section[1]
216
+ mrope_section_w = mrope_section[2]
217
+ positions_stride_m, positions_stride_t = positions.stride()
218
+ else:
219
+ if mrope_section is not None:
220
+ raise ValueError("mrope_section requires 2D MRoPE positions")
221
+ mrope_section_h = 0
222
+ mrope_section_w = 0
223
+ positions_stride_m = 0
224
+ positions_stride_t = positions.stride(0)
225
226
n_tokens = q_gate.shape[0]
227
q_out = torch.empty(
@@ -185,16 +256,22 @@ def fused_qk_rmsnorm_rope_gate(
256
k_out.stride(0),
257
gate_out.stride(0),
258
cos_sin_cache.stride(0),
259
+ positions_stride_m,
260
+ positions_stride_t,
261
num_q_heads,
262
num_kv_heads,
263
head_dim,
264
rotary_dim,
265
half_rotary,
266
eps,
267
+ norm_beta=norm_beta,
268
INPUT_DTYPE=tl.bfloat16 if q_gate.dtype == torch.bfloat16 else tl.float16,
269
HEAD_BLOCK=head_block,
270
ROT_HALF_BLOCK=rot_half_block,
271
HAS_PASS=rotary_dim < head_dim,
272
+ HAS_MROPE=has_mrope,
273
+ MROPE_SECTION_H=mrope_section_h,
274
+ MROPE_SECTION_W=mrope_section_w,
275
num_warps=num_warps,
276
num_stages=2,
277
)
vllm/model_executor/models/qwen3_next.py
+28
-10
@@ -43,7 +43,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
43
MambaStateShapeCalculator,
44
)
45
from vllm.model_executor.layers.quantization import QuantizationConfig
46
-from vllm.model_executor.layers.rotary_embedding import get_rope
46
+from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
47
from vllm.model_executor.layers.vocab_parallel_embedding import (
48
ParallelLMHead,
49
VocabParallelEmbedding,
@@ -300,14 +300,26 @@ class Qwen3NextAttention(nn.Module):
300
self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
301
302
# Fuse the gated split + QK-RMSNorm + (partial) NeoX RoPE + gate copy.
303
- # TODO: support MRoPE
303
mm_config = model_config.multimodal_config if model_config else None
304
text_only = mm_config is None or mm_config.language_model_only
305
+ mrope_section = getattr(self.rotary_emb, "mrope_section", None)
306
+ supports_mrope = bool(
307
+ type(self.rotary_emb) is MRotaryEmbedding
308
+ and mrope_section
309
+ and len(mrope_section) == 3
310
+ and sum(mrope_section) == self.rotary_emb.rotary_dim // 2
311
+ and getattr(self.rotary_emb, "mrope_interleaved", False)
312
+ )
313
+ supports_dtype = getattr(self.rotary_emb, "dtype", None) in (
314
+ torch.float16,
315
+ torch.bfloat16,
316
+ )
317
self.use_fused_qk_norm_rope_gate = (
318
self.attn_output_gate
319
and getattr(self.rotary_emb, "is_neox_style", False)
320
and current_platform.is_cuda()
310
- and text_only
321
+ and supports_dtype
322
+ and (text_only or supports_mrope)
323
)
324
325
def _project_qkv_gate(
@@ -325,22 +337,28 @@ class Qwen3NextAttention(nn.Module):
337
q_gate, k, v = qkv.split(
338
[self.q_size * 2, self.kv_size, self.kv_size], dim=-1
339
)
328
- # mRoPE passes positions as (3, n_tokens) for T/H/W. Fusion is only
329
- # enabled text-only, where the three rows are identical, so taking
330
- # the T row is exact. (1D positions pass through.)
331
- pos = positions[0] if positions.ndim == 2 else positions
340
+ if positions.ndim == 2 and not getattr(
341
+ self.rotary_emb, "mrope_section", None
342
+ ):
343
+ positions = positions[0]
344
q, k, gate = fused_qk_rmsnorm_rope_gate(
345
q_gate,
346
k,
335
- self.q_norm.weight.float() + 1.0,
336
- self.k_norm.weight.float() + 1.0,
347
+ self.q_norm.weight,
348
+ self.k_norm.weight,
349
self.rotary_emb.cos_sin_cache,
338
- pos,
350
+ positions,
351
self.q_norm.variance_epsilon,
352
self.num_heads,
353
self.num_kv_heads,
354
self.head_dim,
355
self.rotary_emb.rotary_dim,
356
+ norm_beta=1.0,
357
+ mrope_section=(
358
+ getattr(self.rotary_emb, "mrope_section", None)
359
+ if positions.ndim == 2
360
+ else None
361
+ ),
362
)
363
return q, k, v, gate
364