@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 42172ad18f
Fork 0

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

[FEAT] [Performance] Add triton mrope to replace the torch code path (#22375)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

TJian committed Aug 9, 2025 at 11:50 UTC 42172ad18fbc22c89d0063184e4570cc84186e16
3 files changed +766
benchmarks/kernels/benchmark_mrope.py new
+328
@@ -0,0 +1,328 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# This script benchmarks the mrope kernel (mainly for Qwen2VL and Qwen2.5VL models).
5 +# It generates test data, runs benchmarks, and saves results to a CSV file.
6 +#
7 +# The CSV file (named with current date/time) contains these columns:
8 +# model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position,
9 +# rope_theta, is_neox_style, rope_scaling, dtype, torch_mean, torch_median, torch_p99,
10 +# torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max,
11 +# speedup
12 +#
13 +# == Usage Examples ==
14 +#
15 +# Single model benchmark:
16 +# python3 benchmark_mrope.py --model-name Qwen/Qwen2-VL-7B-Instruct --tp-size 1 \
17 +# --warmup-iter 10 --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024
18 +#
19 +# All models benchmark:
20 +# python3 benchmark_mrope.py --model-name "" --tp-size 1 --warmup-iter 10 \
21 +# --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024
22 +#
23 +# All models with different TP sizes:
24 +# python3 benchmark_mrope.py --model-name "" --tp-size 1 2 4 8 --warmup-iter 10 \
25 +# --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024
26 +#
27 +# All models with different token counts:
28 +# python3 benchmark_mrope.py --model-name "" --tp-size 1 --warmup-iter 10 \
29 +# --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024 4096 16384
30 +import csv
31 +import os
32 +import time
33 +from datetime import datetime
34 +from typing import Any
35 +
36 +import numpy as np
37 +import torch
38 +
39 +from vllm.model_executor.layers.rotary_embedding import get_rope
40 +from vllm.platforms import current_platform
41 +from vllm.transformers_utils.config import get_config
42 +from vllm.utils import FlexibleArgumentParser
43 +
44 +device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
45 +
46 +
47 +def generate_test_data(
48 + num_tokens: int,
49 + num_q_heads: int,
50 + num_kv_heads: int,
51 + head_size: int,
52 + max_position_embeddings: int,
53 + dtype: torch.dtype,
54 + device: torch.device,
55 +):
56 + """Generate test data for given configuration."""
57 + # Create 2D positions (3, num_tokens) for multimodal case
58 + positions = torch.randint(
59 + 0, max_position_embeddings // 4, (3, num_tokens), device=device
60 + )
61 +
62 + # Create query and key tensors
63 + query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
64 + key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
65 +
66 + return positions, query, key
67 +
68 +
69 +def calculate_stats(times: list[float]) -> dict[str, float]:
70 + """Calculate statistics from a list of times."""
71 + times_array = np.array(times)
72 + return {
73 + "mean": np.mean(times_array),
74 + "median": np.median(times_array),
75 + "p99": np.percentile(times_array, 99),
76 + "min": np.min(times_array),
77 + "max": np.max(times_array),
78 + }
79 +
80 +
81 +def benchmark_mrope(
82 + model_name: str,
83 + num_tokens: int,
84 + head_dim: int,
85 + tp_size: int,
86 + num_heads: int,
87 + num_kv_heads: int,
88 + max_position: int = 8192,
89 + rope_theta: float = 10000,
90 + is_neox_style: bool = True,
91 + rope_scaling: dict[str, Any] = None,
92 + dtype: torch.dtype = torch.bfloat16,
93 + seed: int = 0,
94 + warmup_iter: int = 10,
95 + benchmark_iter: int = 100,
96 + csv_writer=None,
97 +):
98 + current_platform.seed_everything(seed)
99 + torch.set_default_device(device)
100 + # the parameters to compute the q k v size based on tp_size
101 + mrope_helper_class = get_rope(
102 + head_size=head_dim,
103 + rotary_dim=head_dim,
104 + max_position=max_position,
105 + base=rope_theta,
106 + is_neox_style=is_neox_style,
107 + rope_scaling=rope_scaling,
108 + dtype=dtype,
109 + ).to(device=device)
110 +
111 + print(80 * "=")
112 + print(
113 + f"Evaluating model: {model_name} "
114 + f"with tp_size: {tp_size} "
115 + f"and num_tokens: {num_tokens}, "
116 + f"dtype: {dtype}"
117 + )
118 +
119 + # create q k v input tensors
120 + # create rotary pos emb input tensors
121 + positions, query, key = generate_test_data(
122 + num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
123 + )
124 +
125 + # Warm up
126 + for _ in range(warmup_iter):
127 + mrope_helper_class.forward_native(
128 + positions,
129 + query.clone(),
130 + key.clone(),
131 + )
132 +
133 + mrope_helper_class.forward_cuda(
134 + positions,
135 + query.clone(),
136 + key.clone(),
137 + )
138 +
139 + torch.cuda.synchronize()
140 +
141 + # Time reference implementation
142 + torch_times = []
143 + for _ in range(benchmark_iter):
144 + query_clone = query.clone()
145 + key_clone = key.clone()
146 + torch.cuda.synchronize()
147 + start_time = time.time()
148 +
149 + mrope_helper_class.forward_native(
150 + positions,
151 + query_clone,
152 + key_clone,
153 + )
154 +
155 + torch.cuda.synchronize()
156 + torch_times.append(time.time() - start_time)
157 +
158 + # Time triton kernel implementation
159 + triton_times = []
160 + for _ in range(benchmark_iter):
161 + query_clone = query.clone()
162 + key_clone = key.clone()
163 + torch.cuda.synchronize()
164 + start_time = time.time()
165 + mrope_helper_class.forward_cuda(
166 + positions,
167 + query_clone,
168 + key_clone,
169 + )
170 + torch.cuda.synchronize()
171 + triton_times.append(time.time() - start_time)
172 +
173 + # Calculate statistics
174 + torch_stats = calculate_stats(torch_times)
175 + triton_stats = calculate_stats(triton_times)
176 + print(f"\nPerformance for config ({num_tokens}, {num_heads}, {num_kv_heads}):")
177 +
178 + print(
179 + f"Torch implementation: "
180 + f"mean={torch_stats['mean']:.8f}s, "
181 + f"median={torch_stats['median']:.8f}s, "
182 + f"p99={torch_stats['p99']:.8f}s"
183 + )
184 +
185 + print(
186 + f"Triton implementation: "
187 + f"mean={triton_stats['mean']:.8f}s, "
188 + f"median={triton_stats['median']:.8f}s, "
189 + f"p99={triton_stats['p99']:.8f}s"
190 + )
191 +
192 + print(
193 + f"Triton Speedup over Torch: {torch_stats['mean'] / triton_stats['mean']:.8f}x"
194 + )
195 +
196 + # Write to CSV
197 + if csv_writer:
198 + row = [
199 + model_name,
200 + tp_size,
201 + num_tokens,
202 + num_heads,
203 + num_kv_heads,
204 + head_dim,
205 + max_position,
206 + rope_theta,
207 + is_neox_style,
208 + str(rope_scaling),
209 + str(dtype).split(".")[-1],
210 + torch_stats["mean"],
211 + torch_stats["median"],
212 + torch_stats["p99"],
213 + torch_stats["min"],
214 + torch_stats["max"],
215 + triton_stats["mean"],
216 + triton_stats["median"],
217 + triton_stats["p99"],
218 + triton_stats["min"],
219 + triton_stats["max"],
220 + torch_stats["mean"] / triton_stats["mean"], # speedup
221 + ]
222 + csv_writer.writerow(row)
223 +
224 + return torch_stats, triton_stats
225 +
226 +
227 +if __name__ == "__main__":
228 + parser = FlexibleArgumentParser(
229 + description="Benchmark the rotary embedding kernels."
230 + )
231 + parser.add_argument("--model-name", type=str, default="")
232 + parser.add_argument("--tp-size", type=int, default=1)
233 + parser.add_argument("--warmup-iter", type=int, default=10)
234 + parser.add_argument("--benchmark-iter", type=int, default=100)
235 + parser.add_argument("--dtype", type=str, choices=["bfloat16"], default="bfloat16")
236 + parser.add_argument("--seed", type=int, default=0)
237 + parser.add_argument("--num-tokens", type=int, nargs="+", required=False)
238 + parser.add_argument("--trust-remote-code", action="store_true")
239 + parser.add_argument("--output-csv", type=str, default="mrope_benchmark_results.csv")
240 + args = parser.parse_args()
241 + print(args)
242 +
243 + # Create CSV file for results
244 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
245 + csv_filename = f"{os.path.splitext(args.output_csv)[0]}_{timestamp}.csv"
246 +
247 + with open(csv_filename, "w", newline="") as csvfile:
248 + csv_writer = csv.writer(csvfile)
249 + # Write header
250 + header = [
251 + "model_name",
252 + "tp_size",
253 + "num_tokens",
254 + "num_heads",
255 + "num_kv_heads",
256 + "head_dim",
257 + "max_position",
258 + "rope_theta",
259 + "is_neox_style",
260 + "rope_scaling",
261 + "dtype",
262 + "torch_mean",
263 + "torch_median",
264 + "torch_p99",
265 + "torch_min",
266 + "torch_max",
267 + "triton_mean",
268 + "triton_median",
269 + "triton_p99",
270 + "triton_min",
271 + "triton_max",
272 + "speedup",
273 + ]
274 + csv_writer.writerow(header)
275 +
276 + model_tp_dict = {}
277 + if args.model_name == "":
278 + model_tp_dict = {
279 + "Qwen/Qwen2-VL-2B-Instruct": [1],
280 + "Qwen/Qwen2-VL-7B-Instruct": [1],
281 + "Qwen/Qwen2-VL-72B-Instruct": [2, 4, 8],
282 + "Qwen/Qwen2.5-VL-3B-Instruct": [1, 2, 4, 8],
283 + "Qwen/Qwen2.5-VL-7B-Instruct": [1, 2, 4, 8],
284 + "Qwen/Qwen2.5-VL-72B-Instruct": [2, 4, 8],
285 + }
286 + else:
287 + model_tp_dict[args.model_name] = [args.tp_size]
288 +
289 + if args.num_tokens is None:
290 + num_tokens_list = [2**i for i in range(0, 18)]
291 + else:
292 + num_tokens_list = args.num_tokens
293 +
294 + for model_name, tp_list in model_tp_dict.items():
295 + config = get_config(model_name, trust_remote_code=args.trust_remote_code)
296 + for tp_size in tp_list:
297 + # get the model config
298 + total_num_kv_heads = config.num_key_value_heads
299 + total_num_heads = config.num_attention_heads
300 + num_heads = total_num_heads // tp_size
301 + num_kv_heads = max(1, total_num_kv_heads // tp_size)
302 + head_dim = config.hidden_size // total_num_heads
303 + q_size = num_heads * head_dim
304 + kv_size = num_kv_heads * head_dim
305 + is_neox_style = True
306 + rope_theta = config.rope_theta
307 + max_position = config.max_position_embeddings
308 +
309 + for num_tokens in num_tokens_list:
310 + benchmark_mrope(
311 + model_name=model_name,
312 + num_tokens=num_tokens,
313 + head_dim=head_dim,
314 + tp_size=tp_size,
315 + num_heads=num_heads,
316 + num_kv_heads=num_kv_heads,
317 + max_position=max_position,
318 + rope_theta=rope_theta,
319 + is_neox_style=is_neox_style,
320 + rope_scaling=config.rope_scaling,
321 + dtype=getattr(torch, args.dtype),
322 + seed=args.seed,
323 + warmup_iter=args.warmup_iter,
324 + benchmark_iter=args.benchmark_iter,
325 + csv_writer=csv_writer,
326 + )
327 +
328 + print(f"Benchmark results saved to {csv_filename}")
tests/kernels/test_mrope.py new
+207
@@ -0,0 +1,207 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +from transformers import AutoConfig
7 +
8 +from vllm.model_executor.layers.rotary_embedding import get_rope
9 +from vllm.platforms import current_platform
10 +
11 +device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12 +
13 +
14 +def generate_test_data(num_tokens: int, num_q_heads: int, num_kv_heads: int,
15 + head_size: int, max_position_embeddings: int,
16 + dtype: torch.dtype, device: torch.device):
17 + """Generate test data for given configuration."""
18 + # Create 2D positions (3, num_tokens) for multimodal case
19 + positions = torch.randint(0,
20 + max_position_embeddings // 4, (3, num_tokens),
21 + device=device)
22 +
23 + # Create query and key tensors
24 + query = torch.randn(num_tokens,
25 + num_q_heads * head_size,
26 + dtype=dtype,
27 + device=device)
28 + key = torch.randn(num_tokens,
29 + num_kv_heads * head_size,
30 + dtype=dtype,
31 + device=device)
32 +
33 + return positions, query, key
34 +
35 +
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 +}
47 +
48 +# https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317
49 +dtype_atol_rtol_list = [
50 + [torch.bfloat16, 1e-5, 1.6e-2],
51 +]
52 +
53 +num_tokens_list = [11, 8192]
54 +
55 +
56 +@pytest.mark.skipif(not current_platform.is_cuda_alike(),
57 + reason="Skipping CUDA/ROCm only tests.")
58 +@pytest.mark.parametrize("model_name, tp_size",
59 + unroll_model_tp_dict(model_tp_dict))
60 +@pytest.mark.parametrize("dtype, atol, rtol", dtype_atol_rtol_list)
61 +@pytest.mark.parametrize("num_tokens", num_tokens_list)
62 +def test_mrope(model_name, tp_size, dtype, atol, rtol, num_tokens):
63 +
64 + config = AutoConfig.from_pretrained(model_name)
65 +
66 + # get the model config
67 + total_num_kv_heads = config.num_key_value_heads
68 + total_num_heads = config.num_attention_heads
69 + num_heads = total_num_heads // tp_size
70 + num_kv_heads = max(1, total_num_kv_heads // tp_size)
71 + head_dim = config.hidden_size // total_num_heads
72 + is_neox_style = True
73 +
74 + rope_theta = config.rope_theta
75 + max_position = config.max_position_embeddings
76 +
77 + mrope_helper_class = get_rope(
78 + head_size=head_dim,
79 + rotary_dim=head_dim,
80 + max_position=max_position,
81 + base=rope_theta,
82 + is_neox_style=is_neox_style,
83 + rope_scaling=config.rope_scaling,
84 + dtype=dtype,
85 + ).to(device=device)
86 +
87 + # create q k v input tensors
88 + # create rotary pos emb input tensors
89 + positions, query, key = generate_test_data(num_tokens, num_heads,
90 + num_kv_heads, head_dim,
91 + max_position, dtype, device)
92 +
93 + query_native, key_native = mrope_helper_class.forward_native(
94 + positions,
95 + query.clone(),
96 + key.clone(),
97 + )
98 +
99 + query_cuda, key_cuda = mrope_helper_class.forward_cuda(
100 + positions,
101 + query.clone(),
102 + key.clone(),
103 + )
104 +
105 + torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol)
106 + torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol)
107 +
108 +
109 +@pytest.mark.skipif(not current_platform.is_cuda_alike(),
110 + reason="Skipping CUDA/ROCm only tests.")
111 +@pytest.mark.parametrize(
112 + "model_name, tp_size",
113 + unroll_model_tp_dict({"Qwen/Qwen2-VL-7B-Instruct": [1, 2]}))
114 +@pytest.mark.parametrize("dtype, atol, rtol", dtype_atol_rtol_list)
115 +@pytest.mark.parametrize("num_tokens", [4])
116 +def test_mrope_torch_compile_tracing(model_name, tp_size, dtype, atol, rtol,
117 + num_tokens):
118 + config = AutoConfig.from_pretrained(model_name)
119 +
120 + # get the model config
121 + total_num_kv_heads = config.num_key_value_heads
122 + total_num_heads = config.num_attention_heads
123 + num_heads = total_num_heads // tp_size
124 + num_kv_heads = max(1, total_num_kv_heads // tp_size)
125 + head_dim = config.hidden_size // total_num_heads
126 + is_neox_style = True
127 + rope_theta = config.rope_theta
128 + max_position = config.max_position_embeddings
129 +
130 + mrope_helper_class = get_rope(
131 + head_size=head_dim,
132 + rotary_dim=head_dim,
133 + max_position=max_position,
134 + base=rope_theta,
135 + is_neox_style=is_neox_style,
136 + rope_scaling=config.rope_scaling,
137 + dtype=dtype,
138 + ).to(device=device)
139 +
140 + # Generate test data
141 + positions, query, key = generate_test_data(num_tokens, num_heads,
142 + num_kv_heads, head_dim,
143 + max_position, dtype, device)
144 +
145 + # Create a wrapper that makes the in-place function appear functional
146 + def functional_forward_cuda(pos, q, k):
147 + """Wrapper that converts in-place operation to functional style
148 +
149 + CUDA Graph does not support in-place operations.
150 + This wrapper creates working copies of the
151 + input tensors and modifies them.
152 + """
153 + q_work = q.clone() # Create working copies
154 + k_work = k.clone()
155 + # Your in-place function modifies q_work and k_work
156 + mrope_helper_class.forward_cuda(pos, q_work, k_work)
157 + return q_work, k_work # Return the modified tensors
158 +
159 + # Get reference results
160 + query_native, key_native = mrope_helper_class.forward_native(
161 + positions,
162 + query.clone(),
163 + key.clone(),
164 + )
165 +
166 + try:
167 + compiled_forward_cuda = torch.compile(functional_forward_cuda,
168 + fullgraph=True,
169 + backend="inductor",
170 + mode="reduce-overhead",
171 + dynamic=False)
172 +
173 + # Run compiled version
174 + query_compiled_cuda, key_compiled_cuda = compiled_forward_cuda(
175 + positions,
176 + query,
177 + key,
178 + )
179 +
180 + # Run original version for comparison
181 + query_cuda = query.clone()
182 + key_cuda = key.clone()
183 + mrope_helper_class.forward_cuda(positions, query_cuda, key_cuda)
184 +
185 + # Verify results
186 + torch.testing.assert_close(query_compiled_cuda,
187 + query_cuda,
188 + atol=atol,
189 + rtol=rtol)
190 + torch.testing.assert_close(key_compiled_cuda,
191 + key_cuda,
192 + atol=atol,
193 + rtol=rtol)
194 + torch.testing.assert_close(query_compiled_cuda,
195 + query_native,
196 + atol=atol,
197 + rtol=rtol)
198 + torch.testing.assert_close(key_compiled_cuda,
199 + key_native,
200 + atol=atol,
201 + rtol=rtol)
202 +
203 + print("✓ forward_cuda successfully traced with torch.compile inductor")
204 +
205 + except Exception as e:
206 + pytest.fail(
207 + f"forward_cuda failed to trace with torch.compile inductor: {e}")
vllm/model_executor/layers/rotary_embedding/mrope.py
+231
@@ -8,10 +8,173 @@ import numpy as np
8 import torch
9 from transformers import PretrainedConfig
10
11 +from vllm.platforms import current_platform
12 +from vllm.triton_utils import tl, triton
13 +
14 from .base import RotaryEmbedding
15 from .common import apply_rotary_emb_dispatch
16
17
18 +@triton.jit
19 +def _triton_qwen2vl_mrope_forward(
20 + q_ptr,
21 + k_ptr,
22 + cos,
23 + sin,
24 + num_tokens,
25 + n_qh: tl.constexpr,
26 + n_kh: tl.constexpr,
27 + hd: tl.constexpr,
28 + pad_n_qh: tl.constexpr,
29 + pad_n_kh: tl.constexpr,
30 + pad_hd: tl.constexpr,
31 + mrope_section_t: tl.constexpr,
32 + mrope_section_h: tl.constexpr,
33 +):
34 + # Adapted from
35 + # https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py
36 + # This version supports flatten input tensors from vllm
37 + # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2)
38 + # instead of (3, bsz, seq_len, head_dim)
39 + pid = tl.program_id(0)
40 + # locate start address
41 + q_ptr = q_ptr + pid * (n_qh * hd)
42 + k_ptr = k_ptr + pid * (n_kh * hd)
43 +
44 + # ####################################################################
45 + # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position
46 + # m of this program instance
47 + # ####################################################################
48 + # Note: cos and sin now have shape (3, num_tokens, head_dim // 2)
49 +
50 + t_end = mrope_section_t
51 + h_end = t_end + mrope_section_h
52 +
53 + # 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
61 +
62 + # Updated offsets for half head_dim
63 + 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_hd)
67 +
68 + t_cos_row = tl.load(t_cos + cos_offsets, mask=t_mask, other=0)
69 + h_cos_row = tl.load(h_cos + cos_offsets, mask=h_mask, other=0)
70 + w_cos_row = tl.load(w_cos + cos_offsets, mask=w_mask, other=0)
71 + t_sin_row = tl.load(t_sin + cos_offsets, mask=t_mask, other=0)
72 + h_sin_row = tl.load(h_sin + cos_offsets, mask=h_mask, other=0)
73 + w_sin_row = tl.load(w_sin + cos_offsets, mask=w_mask, other=0)
74 +
75 + cos_row = t_cos_row + h_cos_row + w_cos_row
76 + sin_row = t_sin_row + h_sin_row + w_sin_row
77 +
78 + # ####################################################################
79 + # Load the left and right half of q and k for the current
80 + # program instance (i.e. for the current token) separately
81 + # ####################################################################
82 + # left half of the head
83 + first_half_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(
84 + 0, pad_hd // 2)[None, :]
85 + first_half_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(
86 + 0, pad_hd // 2)[None, :]
87 + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (tl.arange(
88 + 0, pad_hd // 2)[None, :] < hd // 2)
89 + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (tl.arange(
90 + 0, pad_hd // 2)[None, :] < hd // 2)
91 +
92 + q_tile_1 = tl.load(q_ptr + first_half_q_offsets,
93 + mask=first_q_mask,
94 + other=0).to(sin_row.dtype)
95 + k_tile_1 = tl.load(k_ptr + first_half_k_offsets,
96 + mask=first_k_mask,
97 + other=0).to(sin_row.dtype)
98 +
99 + # 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)
102 + second_q_mask = first_q_mask
103 + second_k_mask = first_k_mask
104 +
105 + q_tile_2 = tl.load(q_ptr + second_half_q_offsets,
106 + mask=second_q_mask,
107 + other=0).to(sin_row.dtype)
108 + k_tile_2 = tl.load(k_ptr + second_half_k_offsets,
109 + mask=second_k_mask,
110 + other=0).to(sin_row.dtype)
111 +
112 + # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin]
113 + # Since cos and sin are now half-size,
114 + # we use the same cos_row and sin_row for both halves
115 + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row
116 + tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask)
117 + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row
118 + tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask)
119 +
120 + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row
121 + tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask)
122 + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row
123 + tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask)
124 +
125 +
126 +def triton_mrope(
127 + q: torch.Tensor,
128 + k: torch.Tensor,
129 + cos: torch.Tensor,
130 + sin: torch.Tensor,
131 + mrope_section: list[int],
132 + head_size: int,
133 +) -> tuple[torch.Tensor, torch.Tensor]:
134 + """Qwen2VL mrope kernel.
135 +
136 + Args:
137 + query: [num_tokens, num_heads * head_size]
138 + key: [num_tokens, num_kv_heads * head_size]
139 + cos: [3, num_tokens, head_size //2 ]
140 + (T/H/W positions with multimodal inputs)
141 + sin: [3, num_tokens, head_size //2 ]
142 + (T/H/W positions with multimodal inputs)
143 + mrope_section: [t, h, w]
144 + head_size: int
145 + """
146 + n_row, n_q_head_head_dim = q.shape
147 + n_q_head = n_q_head_head_dim // head_size
148 + n_kv_head = k.shape[1] // head_size
149 + pad_hd = triton.next_power_of_2(head_size)
150 + pad_n_q_head = triton.next_power_of_2(n_q_head)
151 + pad_n_kv_head = triton.next_power_of_2(n_kv_head)
152 +
153 + # ensure tensors passed into the kernel are contiguous.
154 + # It will be no-op if they are already contiguous
155 + q = q.contiguous()
156 + k = k.contiguous()
157 + cos = cos.contiguous()
158 + sin = sin.contiguous()
159 +
160 + _triton_qwen2vl_mrope_forward[(n_row, )](
161 + q,
162 + k,
163 + cos,
164 + sin,
165 + n_row,
166 + n_q_head,
167 + n_kv_head,
168 + head_size,
169 + pad_n_q_head,
170 + pad_n_kv_head,
171 + pad_hd,
172 + mrope_section[0],
173 + mrope_section[1],
174 + )
175 + return q, k
176 +
177 +
178 class MRotaryEmbedding(RotaryEmbedding):
179 """Rotary Embedding with Multimodal Sections."""
180
@@ -36,11 +199,34 @@ class MRotaryEmbedding(RotaryEmbedding):
199 if self.mrope_section:
200 assert sum(self.mrope_section) == rotary_dim // 2
201
202 + self.use_triton = current_platform.is_cuda_alike()
203 +
204 def forward(
205 self,
206 positions: torch.Tensor,
207 query: torch.Tensor,
208 key: Optional[torch.Tensor] = None,
209 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
210 + """MRope forward.
211 +
212 + Args:
213 + positions:
214 + [num_tokens,] (text only) or
215 + [3, num_tokens] (T/H/W positions with multimodal inputs)
216 + query: [num_tokens, num_heads * head_size]
217 + key: [num_tokens, num_kv_heads * head_size]
218 + """
219 + if self.use_triton:
220 + return self.forward_cuda(positions, query, key)
221 + else:
222 + return self.forward_native(positions, query, key)
223 +
224 + def forward_native(
225 + self,
226 + positions: torch.Tensor,
227 + query: torch.Tensor,
228 + key: Optional[torch.Tensor] = None,
229 + offsets: Optional[torch.Tensor] = None,
230 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
231 """PyTorch-native implementation equivalent to forward().
232
@@ -88,6 +274,51 @@ class MRotaryEmbedding(RotaryEmbedding):
274 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
275 return query, key
276
277 + def forward_cuda(
278 + self,
279 + positions: torch.Tensor,
280 + query: torch.Tensor,
281 + key: Optional[torch.Tensor] = None,
282 + offsets: Optional[torch.Tensor] = None,
283 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
284 +
285 + assert positions.ndim == 1 or positions.ndim == 2
286 + assert key is not None
287 +
288 + num_tokens = positions.shape[-1]
289 + cos_sin = self.cos_sin_cache[positions]
290 + cos, sin = cos_sin.chunk(2, dim=-1)
291 + query_shape = query.shape
292 + key_shape = key.shape
293 + if positions.ndim == 2:
294 + assert self.mrope_section
295 +
296 + q, k = triton_mrope(
297 + query,
298 + key,
299 + cos,
300 + sin,
301 + self.mrope_section,
302 + self.head_size,
303 + )
304 +
305 + return q.reshape(query_shape), k.reshape(key_shape)
306 +
307 + query = query.view(num_tokens, -1, self.head_size)
308 + query_rot = query[..., :self.rotary_dim]
309 + query_pass = query[..., self.rotary_dim:]
310 + query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
311 + self.is_neox_style)
312 + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
313 +
314 + key = key.view(num_tokens, -1, self.head_size)
315 + key_rot = key[..., :self.rotary_dim]
316 + key_pass = key[..., self.rotary_dim:]
317 + key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
318 + self.is_neox_style)
319 + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
320 + return query, key
321 +
322 @classmethod
323 def get_input_positions(
324 cls,