A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Kernel][Performance] Add Triton kernel for Qwen3-VL interleaved MRoPE (#25055)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Sep 19, 2025 at 18:27 UTC
cea91a32f2364d19d5e708026e84ce21a450c53d
2 files changed
+88
-46
tests/kernels/core/test_mrope.py
+66
-32
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import NamedTuple
4
5
import pytest
6
import torch
7
+from packaging.version import Version
8
from transformers import AutoConfig
9
+from transformers import __version__ as TRANSFORMERS_VERSION
10
11
from vllm.model_executor.layers.rotary_embedding import get_rope
12
from vllm.platforms import current_platform
@@ -15,6 +18,7 @@ def generate_test_data(num_tokens: int, num_q_heads: int, num_kv_heads: int,
18
head_size: int, max_position_embeddings: int,
19
dtype: torch.dtype, device: torch.device):
20
"""Generate test data for given configuration."""
21
+ current_platform.seed_everything(42)
22
# Create 2D positions (3, num_tokens) for multimodal case
23
positions = torch.randint(0,
24
max_position_embeddings // 4, (3, num_tokens),
@@ -33,22 +37,37 @@ def generate_test_data(num_tokens: int, num_q_heads: int, num_kv_heads: int,
37
return positions, query, key
38
39
36
-def unroll_model_tp_dict(model_tp_dict):
37
- return [(model_name, tp_size)
38
- for model_name, tp_sizes in model_tp_dict.items()
39
- for tp_size in tp_sizes]
40
-
41
-
42
-model_tp_dict = {
43
- "Qwen/Qwen2-VL-7B-Instruct": [1, 2],
44
- "Qwen/Qwen2-VL-72B-Instruct": [1, 2],
45
- "Qwen/Qwen2.5-VL-72B-Instruct": [1, 2],
46
- "zai-org/GLM-4.1V-9B-Thinking": [1, 2],
47
-}
48
-
49
-# https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317
50
-dtype_atol_rtol_list = [
51
- [torch.bfloat16, 1e-2, 1.6e-2],
40
+class MRoPETestInfo(NamedTuple):
41
+ model_name: str
42
+ # https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317
43
+ atol: float = 1e-2
44
+ rtol: float = 1.6e-2
45
+ marks: list[pytest.MarkDecorator] = []
46
+
47
+
48
+TRANSFORMERS_BASE_VERSION = Version(TRANSFORMERS_VERSION).base_version
49
+
50
+MODELS_TO_TEST = [
51
+ MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"),
52
+ MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
53
+ MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
54
+ MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
55
+ MRoPETestInfo(
56
+ model_name="Qwen/Qwen3-VL-4B-Instruct",
57
+ marks=[
58
+ pytest.mark.skipif(
59
+ Version(TRANSFORMERS_BASE_VERSION) < Version("4.57.0"),
60
+ reason="Qwen3-VL only available after Transformers v4.57",
61
+ )
62
+ ]),
63
+ MRoPETestInfo(
64
+ model_name="Qwen/Qwen3-VL-30B-A3B-Instruct",
65
+ marks=[
66
+ pytest.mark.skipif(
67
+ Version(TRANSFORMERS_BASE_VERSION) < Version("4.57.0"),
68
+ reason="Qwen3-VL only available after Transformers v4.57",
69
+ )
70
+ ]),
71
]
72
73
num_tokens_list = [11, 8192]
@@ -56,20 +75,29 @@ num_tokens_list = [11, 8192]
75
76
@pytest.mark.skipif(not current_platform.is_cuda_alike(),
77
reason="Skipping CUDA/ROCm only tests.")
59
-@pytest.mark.parametrize("model_name, tp_size",
60
- unroll_model_tp_dict(model_tp_dict))
61
-@pytest.mark.parametrize("dtype, atol, rtol", dtype_atol_rtol_list)
78
+@pytest.mark.parametrize("model_info, model_name", [
79
+ pytest.param(test_config, test_config.model_name, marks=test_config.marks)
80
+ for test_config in MODELS_TO_TEST
81
+])
82
+@pytest.mark.parametrize("tp_size", [1, 2])
83
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
84
@pytest.mark.parametrize("num_tokens", num_tokens_list)
63
-def test_mrope(model_name, tp_size, dtype, atol, rtol, num_tokens):
85
+def test_mrope(model_name: str, model_info: MRoPETestInfo, tp_size: int,
86
+ dtype: torch.dtype, num_tokens: int):
87
+
88
+ atol = model_info.atol
89
+ rtol = model_info.rtol
90
91
config = AutoConfig.from_pretrained(model_name)
92
+ config = config.get_text_config()
93
94
# get the model config
95
total_num_kv_heads = config.num_key_value_heads
96
total_num_heads = config.num_attention_heads
97
num_heads = total_num_heads // tp_size
98
num_kv_heads = max(1, total_num_kv_heads // tp_size)
72
- head_dim = config.hidden_size // total_num_heads
99
+ head_dim = (config.head_dim if hasattr(config, "head_dim") else
100
+ config.hidden_size // total_num_heads)
101
is_neox_style = True
102
103
rope_theta = config.rope_theta
@@ -111,24 +139,30 @@ def test_mrope(model_name, tp_size, dtype, atol, rtol, num_tokens):
139
140
@pytest.mark.skipif(not current_platform.is_cuda_alike(),
141
reason="Skipping CUDA/ROCm only tests.")
114
-@pytest.mark.parametrize(
115
- "model_name, tp_size",
116
- unroll_model_tp_dict({
117
- "Qwen/Qwen2-VL-7B-Instruct": [1, 2],
118
- "zai-org/GLM-4.1V-9B-Thinking": [1, 2]
119
- }))
120
-@pytest.mark.parametrize("dtype, atol, rtol", dtype_atol_rtol_list)
121
-@pytest.mark.parametrize("num_tokens", [4])
122
-def test_mrope_torch_compile_tracing(model_name, tp_size, dtype, atol, rtol,
123
- num_tokens):
142
+@pytest.mark.parametrize("model_info, model_name", [
143
+ pytest.param(test_config, test_config.model_name, marks=test_config.marks)
144
+ for test_config in MODELS_TO_TEST
145
+])
146
+@pytest.mark.parametrize("tp_size", [1, 2])
147
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
148
+@pytest.mark.parametrize("num_tokens", num_tokens_list)
149
+def test_mrope_torch_compile_tracing(model_name: str,
150
+ model_info: MRoPETestInfo, tp_size: int,
151
+ dtype: torch.dtype, num_tokens: int):
152
+
153
+ atol = model_info.atol
154
+ rtol = model_info.rtol
155
+
156
config = AutoConfig.from_pretrained(model_name)
157
+ config = config.get_text_config()
158
159
# get the model config
160
total_num_kv_heads = config.num_key_value_heads
161
total_num_heads = config.num_attention_heads
162
num_heads = total_num_heads // tp_size
163
num_kv_heads = max(1, total_num_kv_heads // tp_size)
131
- head_dim = config.hidden_size // total_num_heads
164
+ head_dim = (config.head_dim if hasattr(config, "head_dim") else
165
+ config.hidden_size // total_num_heads)
166
is_neox_style = True
167
rope_theta = config.rope_theta
168
max_position = config.max_position_embeddings
vllm/model_executor/layers/rotary_embedding/mrope.py
+22
-14
@@ -15,7 +15,7 @@ from .common import apply_rotary_emb_dispatch
15
16
17
@triton.jit
18
-def _triton_qwen2vl_mrope_forward(
18
+def _triton_mrope_forward(
19
q_ptr,
20
k_ptr,
21
cos,
@@ -30,12 +30,14 @@ def _triton_qwen2vl_mrope_forward(
30
pad_hd: tl.constexpr,
31
mrope_section_t: tl.constexpr,
32
mrope_section_h: tl.constexpr,
33
+ mrope_section_w: tl.constexpr,
34
+ is_interleaved: tl.constexpr,
35
):
36
# Adapted from
37
# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py
38
# This version supports flatten input tensors from vllm
39
# and supports cos and sin cache with shape (3, num_tokens, head_dim // 2)
38
- # instead of (3, bsz, seq_len, head_dim)
40
+ # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary
41
pid = tl.program_id(0)
42
# locate start address
43
q_ptr = q_ptr + pid * (n_qh * hd)
@@ -47,9 +49,6 @@ def _triton_qwen2vl_mrope_forward(
49
# ####################################################################
50
# Note: cos and sin now have shape (3, num_tokens, head_dim // 2)
51
50
- t_end = mrope_section_t
51
- h_end = t_end + mrope_section_h
52
-
52
# Updated stride calculation for half head_dim
53
half_rd = rd // 2
54
t_cos = cos + pid * half_rd
@@ -61,9 +60,18 @@ def _triton_qwen2vl_mrope_forward(
60
61
# Updated offsets for half head_dim
62
cos_offsets = tl.arange(0, pad_hd // 2)
64
- t_mask = cos_offsets < t_end
65
- h_mask = (t_end <= cos_offsets) & (cos_offsets < h_end)
66
- w_mask = (h_end <= cos_offsets) & (cos_offsets < half_rd)
63
+ if is_interleaved:
64
+ h_mask = (((cos_offsets % 3) == 1) &
65
+ (cos_offsets <= 3 * mrope_section_h))
66
+ w_mask = (((cos_offsets % 3) == 2) &
67
+ (cos_offsets <= 3 * mrope_section_w))
68
+ t_mask = ~(h_mask | w_mask)
69
+ else:
70
+ t_end = mrope_section_t
71
+ h_end = t_end + mrope_section_h
72
+ t_mask = cos_offsets < mrope_section_t
73
+ h_mask = (t_end <= cos_offsets) & (cos_offsets < h_end)
74
+ w_mask = (h_end <= cos_offsets) & (cos_offsets < half_rd)
75
76
t_cos_row = tl.load(t_cos + cos_offsets, mask=t_mask, other=0)
77
h_cos_row = tl.load(h_cos + cos_offsets, mask=h_mask, other=0)
@@ -131,6 +139,7 @@ def triton_mrope(
139
mrope_section: list[int],
140
head_size: int,
141
rotary_dim: int,
142
+ mrope_interleaved: bool,
143
) -> tuple[torch.Tensor, torch.Tensor]:
144
"""Qwen2VL mrope kernel.
145
@@ -158,7 +167,7 @@ def triton_mrope(
167
cos = cos.contiguous()
168
sin = sin.contiguous()
169
161
- _triton_qwen2vl_mrope_forward[(n_row, )](
170
+ _triton_mrope_forward[(n_row, )](
171
q,
172
k,
173
cos,
@@ -173,6 +182,8 @@ def triton_mrope(
182
pad_hd,
183
mrope_section[0],
184
mrope_section[1],
185
+ mrope_section[2],
186
+ mrope_interleaved,
187
)
188
return q, k
189
@@ -201,7 +212,7 @@ class MRotaryEmbedding(RotaryEmbedding):
212
is_neox_style: bool,
213
dtype: torch.dtype,
214
mrope_section: Optional[list[int]] = None,
204
- mrope_interleaved: Optional[bool] = False,
215
+ mrope_interleaved: bool = False,
216
) -> None:
217
# In Qwen2.5-VL, the maximum index value is related to the duration of
218
# the input video. We enlarge max_position_embeddings to 4 times to get
@@ -282,10 +293,6 @@ class MRotaryEmbedding(RotaryEmbedding):
293
assert positions.ndim == 1 or positions.ndim == 2
294
assert key is not None
295
285
- if self.mrope_interleaved:
286
- # TODO: add triton implementation to support mrope-interleaved
287
- return self.forward_native(positions, query, key)
288
-
296
num_tokens = positions.shape[-1]
297
cos_sin = self.cos_sin_cache[positions]
298
cos, sin = cos_sin.chunk(2, dim=-1)
@@ -302,6 +309,7 @@ class MRotaryEmbedding(RotaryEmbedding):
309
self.mrope_section,
310
self.head_size,
311
self.rotary_dim,
312
+ self.mrope_interleaved,
313
)
314
315
return q.reshape(query_shape), k.reshape(key_shape)