[ROCm] Fix and optimize GPT-J-style MRoPE (#49906)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 28, 2026 at 11:50 UTC
05a08148631ab2d852c1984ddb3abebad8039a29
2 files changed
+100
-59
tests/kernels/core/test_mrope.py
+7
-3
@@ -38,6 +38,7 @@ def generate_test_data(
38
39
class MRoPETestInfo(NamedTuple):
40
model_name: str
41
+ is_neox_style: bool = True
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,7 +46,10 @@ class MRoPETestInfo(NamedTuple):
46
47
48
MODELS_TO_TEST = [
48
- MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"),
49
+ MRoPETestInfo(
50
+ model_name="zai-org/GLM-4.1V-9B-Thinking",
51
+ is_neox_style=False,
52
+ ),
53
MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
54
MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
55
MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
@@ -92,7 +96,7 @@ def test_mrope(
96
if hasattr(config, "head_dim")
97
else config.hidden_size // total_num_heads
98
)
95
- is_neox_style = True
99
+ is_neox_style = model_info.is_neox_style
100
101
max_position = config.max_position_embeddings
102
@@ -162,7 +166,7 @@ def test_mrope_torch_compile_tracing(
166
if hasattr(config, "head_dim")
167
else config.hidden_size // total_num_heads
168
)
165
- is_neox_style = True
169
+ is_neox_style = model_info.is_neox_style
170
max_position = config.max_position_embeddings
171
172
mrope_helper_class = get_rope(
vllm/model_executor/layers/rotary_embedding/mrope.py
+93
-56
@@ -5,6 +5,7 @@
5
import numpy as np
6
import torch
7
8
+from vllm.platforms import current_platform
9
from vllm.triton_utils import tl, triton
10
11
from .base import RotaryEmbeddingBase
@@ -24,16 +25,17 @@ def _triton_mrope_forward(
25
rd: tl.constexpr,
26
pad_n_qh: tl.constexpr,
27
pad_n_kh: tl.constexpr,
27
- pad_hd: tl.constexpr,
28
+ pad_rd: tl.constexpr,
29
mrope_section_t: tl.constexpr,
30
mrope_section_h: tl.constexpr,
31
mrope_section_w: tl.constexpr,
32
is_interleaved: tl.constexpr,
33
+ is_neox_style: tl.constexpr,
34
):
35
# Adapted from
36
# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py
37
# This version supports flatten input tensors from vllm
36
- # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2)
38
+ # and supports cos and sin cache with shape (3, num_tokens, rotary_dim // 2)
39
# instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary
40
pid = tl.program_id(0)
41
# locate start address
@@ -44,9 +46,9 @@ def _triton_mrope_forward(
46
# get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position
47
# m of this program instance
48
# ####################################################################
47
- # Note: cos and sin now have shape (3, num_tokens, head_dim // 2)
49
+ # Note: cos and sin now have shape (3, num_tokens, rotary_dim // 2)
50
49
- # Updated stride calculation for half head_dim
51
+ # Updated stride calculation for half rotary_dim
52
half_rd = rd // 2
53
t_cos = cos + pid * half_rd
54
h_cos = t_cos + num_tokens * half_rd
@@ -55,12 +57,17 @@ def _triton_mrope_forward(
57
h_sin = t_sin + num_tokens * half_rd
58
w_sin = h_sin + num_tokens * half_rd
59
58
- # Updated offsets for half head_dim
59
- cos_offsets = tl.arange(0, pad_hd // 2)
60
+ # Updated offsets for half rotary_dim
61
+ cos_offsets = tl.arange(0, pad_rd // 2)
62
if is_interleaved:
61
- h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h)
62
- w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w)
63
- t_mask = ~(h_mask | w_mask)
63
+ valid_mask = cos_offsets < half_rd
64
+ h_mask = (
65
+ valid_mask & ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h)
66
+ )
67
+ w_mask = (
68
+ valid_mask & ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w)
69
+ )
70
+ t_mask = valid_mask & ~(h_mask | w_mask)
71
else:
72
t_end = mrope_section_t
73
h_end = t_end + mrope_section_h
@@ -79,55 +86,74 @@ def _triton_mrope_forward(
86
sin_row = t_sin_row + h_sin_row + w_sin_row
87
88
# ####################################################################
82
- # Load the left and right half of q and k for the current
83
- # program instance (i.e. for the current token) separately
89
+ # Load the two values in each rotary pair for the current token.
90
+ # NeoX pairs the first and second halves, while GPT-J pairs
91
+ # adjacent values.
92
# ####################################################################
85
- # left half of the head
86
- first_half_q_offsets = (
87
- tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :]
88
- )
89
- first_half_k_offsets = (
90
- tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :]
91
- )
92
- first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (
93
- tl.arange(0, pad_hd // 2)[None, :] < rd // 2
94
- )
95
- first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (
96
- tl.arange(0, pad_hd // 2)[None, :] < rd // 2
97
- )
93
+ if is_neox_style:
94
+ rotary_offsets = tl.arange(0, pad_rd // 2)
95
+ first_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :]
96
+ first_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :]
97
+ first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (
98
+ rotary_offsets[None, :] < rd // 2
99
+ )
100
+ first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (
101
+ rotary_offsets[None, :] < rd // 2
102
+ )
103
99
- q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to(
100
- sin_row.dtype
101
- )
102
- k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to(
103
- sin_row.dtype
104
- )
104
+ q_tile_1 = tl.load(q_ptr + first_q_offsets, mask=first_q_mask, other=0).to(
105
+ sin_row.dtype
106
+ )
107
+ k_tile_1 = tl.load(k_ptr + first_k_offsets, mask=first_k_mask, other=0).to(
108
+ sin_row.dtype
109
+ )
110
106
- # right half of the head
107
- second_half_q_offsets = first_half_q_offsets + (rd // 2)
108
- second_half_k_offsets = first_half_k_offsets + (rd // 2)
109
- second_q_mask = first_q_mask
110
- second_k_mask = first_k_mask
111
+ second_q_offsets = first_q_offsets + (rd // 2)
112
+ second_k_offsets = first_k_offsets + (rd // 2)
113
+ q_tile_2 = tl.load(q_ptr + second_q_offsets, mask=first_q_mask, other=0).to(
114
+ sin_row.dtype
115
+ )
116
+ k_tile_2 = tl.load(k_ptr + second_k_offsets, mask=first_k_mask, other=0).to(
117
+ sin_row.dtype
118
+ )
119
112
- q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to(
113
- sin_row.dtype
114
- )
115
- k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to(
116
- sin_row.dtype
117
- )
120
+ new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row
121
+ tl.store(q_ptr + first_q_offsets, new_q_tile_1, mask=first_q_mask)
122
+ new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row
123
+ tl.store(q_ptr + second_q_offsets, new_q_tile_2, mask=first_q_mask)
124
+
125
+ new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row
126
+ tl.store(k_ptr + first_k_offsets, new_k_tile_1, mask=first_k_mask)
127
+ new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row
128
+ tl.store(k_ptr + second_k_offsets, new_k_tile_2, mask=first_k_mask)
129
+ else:
130
+ # Load and store adjacent rotary pairs contiguously. Using stride-two
131
+ # even/odd offsets makes Triton emit scalar 16-bit memory operations on
132
+ # AMD, while split/interleave only rearranges values in registers.
133
+ rotary_offsets = tl.arange(0, pad_rd)
134
+ q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :]
135
+ k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :]
136
+ q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (
137
+ rotary_offsets[None, :] < rd
138
+ )
139
+ k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (
140
+ rotary_offsets[None, :] < rd
141
+ )
142
+
143
+ q_tile = tl.load(q_ptr + q_offsets, mask=q_mask, other=0).to(sin_row.dtype)
144
+ k_tile = tl.load(k_ptr + k_offsets, mask=k_mask, other=0).to(sin_row.dtype)
145
+ q_tile_1, q_tile_2 = tl.split(tl.reshape(q_tile, (pad_n_qh, pad_rd // 2, 2)))
146
+ k_tile_1, k_tile_2 = tl.split(tl.reshape(k_tile, (pad_n_kh, pad_rd // 2, 2)))
147
119
- # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin]
120
- # Since cos and sin are now half-size,
121
- # we use the same cos_row and sin_row for both halves
122
- new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row
123
- tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask)
124
- new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row
125
- tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask)
148
+ new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row
149
+ new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row
150
+ new_q_tile = tl.interleave(new_q_tile_1, new_q_tile_2)
151
+ tl.store(q_ptr + q_offsets, new_q_tile, mask=q_mask)
152
127
- new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row
128
- tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask)
129
- new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row
130
- tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask)
153
+ new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row
154
+ new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row
155
+ new_k_tile = tl.interleave(new_k_tile_1, new_k_tile_2)
156
+ tl.store(k_ptr + k_offsets, new_k_tile, mask=k_mask)
157
158
159
def triton_mrope(
@@ -139,23 +165,26 @@ def triton_mrope(
165
head_size: int,
166
rotary_dim: int,
167
mrope_interleaved: bool,
168
+ is_neox_style: bool,
169
) -> tuple[torch.Tensor, torch.Tensor]:
170
"""Qwen2VL mrope kernel.
171
172
Args:
173
q: [num_tokens, num_heads * head_size]
174
k: [num_tokens, num_kv_heads * head_size]
148
- cos: [3, num_tokens, head_size //2 ]
175
+ cos: [3, num_tokens, rotary_dim // 2]
176
(T/H/W positions with multimodal inputs)
150
- sin: [3, num_tokens, head_size //2 ]
177
+ sin: [3, num_tokens, rotary_dim // 2]
178
(T/H/W positions with multimodal inputs)
179
mrope_section: [t, h, w]
180
head_size: int
181
+ is_neox_style: Whether rotary pairs use split-half (NeoX) or
182
+ adjacent (GPT-J) layout.
183
"""
184
n_row, n_q_head_head_dim = q.shape
185
n_q_head = n_q_head_head_dim // head_size
186
n_kv_head = k.shape[1] // head_size
158
- pad_hd = triton.next_power_of_2(head_size)
187
+ pad_rd = triton.next_power_of_2(rotary_dim)
188
pad_n_q_head = triton.next_power_of_2(n_q_head)
189
pad_n_kv_head = triton.next_power_of_2(n_kv_head)
190
@@ -166,6 +195,11 @@ def triton_mrope(
195
cos = cos.contiguous()
196
sin = sin.contiguous()
197
198
+ # Small adjacent-pair tiles perform best with one wave per program on
199
+ # ROCm. Keep the existing launch shape for larger rotary dimensions,
200
+ # NeoX, and other backends.
201
+ use_single_wave = current_platform.is_rocm() and not is_neox_style and pad_rd <= 64
202
+ num_warps = 1 if use_single_wave else 4
203
_triton_mrope_forward[(n_row,)](
204
q,
205
k,
@@ -178,11 +212,13 @@ def triton_mrope(
212
rotary_dim,
213
pad_n_q_head,
214
pad_n_kv_head,
181
- pad_hd,
215
+ pad_rd,
216
mrope_section[0],
217
mrope_section[1],
218
mrope_section[2],
219
mrope_interleaved,
220
+ is_neox_style,
221
+ num_warps=num_warps,
222
)
223
return q, k
224
@@ -349,6 +385,7 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
385
self.head_size,
386
self.rotary_dim,
387
self.mrope_interleaved,
388
+ self.is_neox_style,
389
)
390
391
return q.reshape(query_shape), k.reshape(key_shape)