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
[Bugfix][Kernel] Support partial rotary embedding for MRoPE triton kernel (#22593)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Aug 11, 2025 at 00:00 UTC
b76753f0b58a070f626549115d1414ec421e7e49
2 files changed
+30
-18
tests/kernels/core/test_mrope.py
renamed
+14
-6
@@ -42,12 +42,13 @@ def unroll_model_tp_dict(model_tp_dict):
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]
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 = [
50
- [torch.bfloat16, 1e-5, 1.6e-2],
51
+ [torch.bfloat16, 1e-2, 1.6e-2],
52
]
53
54
num_tokens_list = [11, 8192]
@@ -73,10 +74,12 @@ def test_mrope(model_name, tp_size, dtype, atol, rtol, num_tokens):
74
75
rope_theta = config.rope_theta
76
max_position = config.max_position_embeddings
77
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
78
+ rotary_dim = int(head_dim * partial_rotary_factor)
79
80
mrope_helper_class = get_rope(
81
head_size=head_dim,
79
- rotary_dim=head_dim,
82
+ rotary_dim=rotary_dim,
83
max_position=max_position,
84
base=rope_theta,
85
is_neox_style=is_neox_style,
@@ -110,7 +113,10 @@ def test_mrope(model_name, tp_size, dtype, atol, rtol, num_tokens):
113
reason="Skipping CUDA/ROCm only tests.")
114
@pytest.mark.parametrize(
115
"model_name, tp_size",
113
- unroll_model_tp_dict({"Qwen/Qwen2-VL-7B-Instruct": [1, 2]}))
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,
@@ -126,10 +132,12 @@ def test_mrope_torch_compile_tracing(model_name, tp_size, dtype, atol, rtol,
132
is_neox_style = True
133
rope_theta = config.rope_theta
134
max_position = config.max_position_embeddings
135
+ partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
136
+ rotary_dim = int(head_dim * partial_rotary_factor)
137
138
mrope_helper_class = get_rope(
139
head_size=head_dim,
132
- rotary_dim=head_dim,
140
+ rotary_dim=rotary_dim,
141
max_position=max_position,
142
base=rope_theta,
143
is_neox_style=is_neox_style,
@@ -145,7 +153,7 @@ def test_mrope_torch_compile_tracing(model_name, tp_size, dtype, atol, rtol,
153
# Create a wrapper that makes the in-place function appear functional
154
def functional_forward_cuda(pos, q, k):
155
"""Wrapper that converts in-place operation to functional style
148
-
156
+
157
CUDA Graph does not support in-place operations.
158
This wrapper creates working copies of the
159
input tensors and modifies them.
vllm/model_executor/layers/rotary_embedding/mrope.py
+16
-12
@@ -25,6 +25,7 @@ def _triton_qwen2vl_mrope_forward(
25
n_qh: tl.constexpr,
26
n_kh: tl.constexpr,
27
hd: tl.constexpr,
28
+ rd: tl.constexpr,
29
pad_n_qh: tl.constexpr,
30
pad_n_kh: tl.constexpr,
31
pad_hd: tl.constexpr,
@@ -51,19 +52,19 @@ def _triton_qwen2vl_mrope_forward(
52
h_end = t_end + mrope_section_h
53
54
# Updated stride calculation for half head_dim
54
- half_hd = hd // 2
55
- t_cos = cos + pid * half_hd
56
- h_cos = t_cos + num_tokens * half_hd
57
- w_cos = h_cos + num_tokens * half_hd
58
- t_sin = sin + pid * half_hd
59
- h_sin = t_sin + num_tokens * half_hd
60
- w_sin = h_sin + num_tokens * half_hd
55
+ half_rd = rd // 2
56
+ t_cos = cos + pid * half_rd
57
+ h_cos = t_cos + num_tokens * half_rd
58
+ w_cos = h_cos + num_tokens * half_rd
59
+ t_sin = sin + pid * half_rd
60
+ h_sin = t_sin + num_tokens * half_rd
61
+ w_sin = h_sin + num_tokens * half_rd
62
63
# Updated offsets for half head_dim
64
cos_offsets = tl.arange(0, pad_hd // 2)
65
t_mask = cos_offsets < t_end
66
h_mask = (t_end <= cos_offsets) & (cos_offsets < h_end)
66
- w_mask = (h_end <= cos_offsets) & (cos_offsets < half_hd)
67
+ w_mask = (h_end <= cos_offsets) & (cos_offsets < half_rd)
68
69
t_cos_row = tl.load(t_cos + cos_offsets, mask=t_mask, other=0)
70
h_cos_row = tl.load(h_cos + cos_offsets, mask=h_mask, other=0)
@@ -85,9 +86,9 @@ def _triton_qwen2vl_mrope_forward(
86
first_half_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(
87
0, pad_hd // 2)[None, :]
88
first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (tl.arange(
88
- 0, pad_hd // 2)[None, :] < hd // 2)
89
+ 0, pad_hd // 2)[None, :] < rd // 2)
90
first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (tl.arange(
90
- 0, pad_hd // 2)[None, :] < hd // 2)
91
+ 0, pad_hd // 2)[None, :] < rd // 2)
92
93
q_tile_1 = tl.load(q_ptr + first_half_q_offsets,
94
mask=first_q_mask,
@@ -97,8 +98,8 @@ def _triton_qwen2vl_mrope_forward(
98
other=0).to(sin_row.dtype)
99
100
# right half of the head
100
- second_half_q_offsets = first_half_q_offsets + (hd // 2)
101
- second_half_k_offsets = first_half_k_offsets + (hd // 2)
101
+ second_half_q_offsets = first_half_q_offsets + (rd // 2)
102
+ second_half_k_offsets = first_half_k_offsets + (rd // 2)
103
second_q_mask = first_q_mask
104
second_k_mask = first_k_mask
105
@@ -130,6 +131,7 @@ def triton_mrope(
131
sin: torch.Tensor,
132
mrope_section: list[int],
133
head_size: int,
134
+ rotary_dim: int,
135
) -> tuple[torch.Tensor, torch.Tensor]:
136
"""Qwen2VL mrope kernel.
137
@@ -166,6 +168,7 @@ def triton_mrope(
168
n_q_head,
169
n_kv_head,
170
head_size,
171
+ rotary_dim,
172
pad_n_q_head,
173
pad_n_kv_head,
174
pad_hd,
@@ -300,6 +303,7 @@ class MRotaryEmbedding(RotaryEmbedding):
303
sin,
304
self.mrope_section,
305
self.head_size,
306
+ self.rotary_dim,
307
)
308
309
return q.reshape(query_shape), k.reshape(key_shape)