| 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 | |
| 8 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 9 | from vllm.model_executor.layers.rotary_embedding.mrope import ( |
| 10 | apply_interleaved_rope, |
| 11 | triton_mrope, |
| 12 | ) |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.transformers_utils.config import get_config |
| 15 | from vllm.utils.torch_utils import set_random_seed |
| 16 | |
| 17 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 18 | |
| 19 | |
| 20 | def generate_test_data( |
| 21 | num_tokens: int, |
| 22 | num_q_heads: int, |
| 23 | num_kv_heads: int, |
| 24 | head_size: int, |
| 25 | max_position_embeddings: int, |
| 26 | dtype: torch.dtype, |
| 27 | device: torch.device, |
| 28 | ): |
| 29 | """Generate test data for given configuration.""" |
| 30 | set_random_seed(42) |
| 31 | # Create 2D positions (3, num_tokens) for multimodal case |
| 32 | positions = torch.randint( |
| 33 | 0, max_position_embeddings // 4, (3, num_tokens), device=device |
| 34 | ) |
| 35 | |
| 36 | # Create query and key tensors |
| 37 | query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device) |
| 38 | key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device) |
| 39 | |
| 40 | return positions, query, key |
| 41 | |
| 42 | |
| 43 | class MRoPETestInfo(NamedTuple): |
| 44 | model_name: str |
| 45 | is_neox_style: bool = True |
| 46 | # https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317 |
| 47 | atol: float = 1e-2 |
| 48 | rtol: float = 1.6e-2 |
| 49 | marks: list[pytest.MarkDecorator] = [] |
| 50 | |
| 51 | |
| 52 | MODELS_TO_TEST = [ |
| 53 | MRoPETestInfo( |
| 54 | model_name="zai-org/GLM-4.1V-9B-Thinking", |
| 55 | is_neox_style=False, |
| 56 | ), |
| 57 | MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"), |
| 58 | MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"), |
| 59 | MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"), |
| 60 | MRoPETestInfo(model_name="Qwen/Qwen3-VL-4B-Instruct"), |
| 61 | MRoPETestInfo(model_name="Qwen/Qwen3-VL-30B-A3B-Instruct"), |
| 62 | ] |
| 63 | |
| 64 | num_tokens_list = [11, 8192] |
| 65 | |
| 66 | |
| 67 | def test_apply_interleaved_rope(): |
| 68 | mrope_section = [3, 1, 1] |
| 69 | x = torch.tensor( |
| 70 | [ |
| 71 | [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], |
| 72 | [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], |
| 73 | [[20, 21, 22, 23, 24], [25, 26, 27, 28, 29]], |
| 74 | ] |
| 75 | ) |
| 76 | |
| 77 | result = apply_interleaved_rope(x, mrope_section) |
| 78 | |
| 79 | expected = torch.tensor([[0, 11, 22, 3, 4], [5, 16, 27, 8, 9]]) |
| 80 | torch.testing.assert_close(result, expected, rtol=0, atol=0) |
| 81 | |
| 82 | |
| 83 | @pytest.mark.skipif( |
| 84 | not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only test." |
| 85 | ) |
| 86 | def test_apply_interleaved_rope_torch_compile(): |
| 87 | mrope_section = [24, 20, 20] |
| 88 | num_tokens = 8192 |
| 89 | rotary_dim = sum(mrope_section) * 2 |
| 90 | cache = torch.randn( |
| 91 | 3, |
| 92 | num_tokens, |
| 93 | rotary_dim, |
| 94 | device=device, |
| 95 | dtype=torch.bfloat16, |
| 96 | ) |
| 97 | x = cache[..., : rotary_dim // 2] |
| 98 | |
| 99 | expected = apply_interleaved_rope(x, mrope_section) |
| 100 | compiled_fn = torch.compile( |
| 101 | apply_interleaved_rope, |
| 102 | backend="inductor", |
| 103 | fullgraph=True, |
| 104 | ) |
| 105 | |
| 106 | result = compiled_fn(x, mrope_section) |
| 107 | |
| 108 | torch.testing.assert_close(result, expected, rtol=0, atol=0) |
| 109 | |
| 110 | |
| 111 | @pytest.mark.skipif( |
| 112 | not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests." |
| 113 | ) |
| 114 | @pytest.mark.parametrize( |
| 115 | "model_info, model_name", |
| 116 | [ |
| 117 | pytest.param(test_config, test_config.model_name, marks=test_config.marks) |
| 118 | for test_config in MODELS_TO_TEST |
| 119 | ], |
| 120 | ) |
| 121 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 122 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 123 | @pytest.mark.parametrize("num_tokens", num_tokens_list) |
| 124 | def test_mrope( |
| 125 | default_vllm_config, |
| 126 | model_name: str, |
| 127 | model_info: MRoPETestInfo, |
| 128 | tp_size: int, |
| 129 | dtype: torch.dtype, |
| 130 | num_tokens: int, |
| 131 | ): |
| 132 | atol = model_info.atol |
| 133 | rtol = model_info.rtol |
| 134 | |
| 135 | config = get_config(model_name, False).get_text_config() |
| 136 | |
| 137 | # get the model config |
| 138 | total_num_kv_heads = config.num_key_value_heads |
| 139 | total_num_heads = config.num_attention_heads |
| 140 | num_heads = total_num_heads // tp_size |
| 141 | num_kv_heads = max(1, total_num_kv_heads // tp_size) |
| 142 | head_dim = ( |
| 143 | config.head_dim |
| 144 | if hasattr(config, "head_dim") |
| 145 | else config.hidden_size // total_num_heads |
| 146 | ) |
| 147 | is_neox_style = model_info.is_neox_style |
| 148 | |
| 149 | max_position = config.max_position_embeddings |
| 150 | |
| 151 | mrope_helper_class = get_rope( |
| 152 | head_size=head_dim, |
| 153 | max_position=max_position, |
| 154 | is_neox_style=is_neox_style, |
| 155 | rope_parameters=config.rope_parameters, |
| 156 | dtype=dtype, |
| 157 | ).to(device=device) |
| 158 | |
| 159 | # create q k v input tensors |
| 160 | # create rotary pos emb input tensors |
| 161 | positions, query, key = generate_test_data( |
| 162 | num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device |
| 163 | ) |
| 164 | |
| 165 | query_native, key_native = mrope_helper_class.forward_native( |
| 166 | positions, |
| 167 | query.clone(), |
| 168 | key.clone(), |
| 169 | ) |
| 170 | |
| 171 | query_cuda, key_cuda = mrope_helper_class.forward_cuda( |
| 172 | positions, |
| 173 | query.clone(), |
| 174 | key.clone(), |
| 175 | ) |
| 176 | |
| 177 | torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol) |
| 178 | torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol) |
| 179 | |
| 180 | # Strided q/k views of a packed buffer (the MiniMax M3 ViT usage): |
| 181 | # bitwise-identical to the contiguous run, neighboring columns untouched. |
| 182 | pad = 16 |
| 183 | packed = torch.zeros( |
| 184 | num_tokens, |
| 185 | query.shape[1] + key.shape[1] + pad, |
| 186 | dtype=dtype, |
| 187 | device=device, |
| 188 | ) |
| 189 | packed[:, : query.shape[1]] = query |
| 190 | packed[:, query.shape[1] : query.shape[1] + key.shape[1]] = key |
| 191 | sentinel = packed[:, -pad:].clone() |
| 192 | q_strided = packed[:, : query.shape[1]] |
| 193 | k_strided = packed[:, query.shape[1] : query.shape[1] + key.shape[1]] |
| 194 | mrope_helper_class.forward_cuda(positions, q_strided, k_strided) |
| 195 | torch.testing.assert_close(q_strided, query_cuda, atol=0, rtol=0) |
| 196 | torch.testing.assert_close(k_strided, key_cuda, atol=0, rtol=0) |
| 197 | torch.testing.assert_close(packed[:, -pad:], sentinel, atol=0, rtol=0) |
| 198 | |
| 199 | |
| 200 | @pytest.mark.skipif( |
| 201 | not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests." |
| 202 | ) |
| 203 | def test_mrope_row_offset_int64(): |
| 204 | # MiniMax M3 ViT worst case: a max-size video item packs ~768k patches, |
| 205 | # and num_tokens * packed-qkv row stride (3840) exceeds 2^31 at 559,241 |
| 206 | # rows. Row-base math must be int64 or those rows wrap to negative |
| 207 | # offsets (illegal memory access at startup profiling). |
| 208 | set_random_seed(0) |
| 209 | n_qh, n_kh, hd, rd = 16, 16, 80, 78 |
| 210 | mrope_section = [13, 13, 13] |
| 211 | row_stride = (n_qh + n_kh + n_kh) * hd # packed qkv row |
| 212 | num_tokens = 2**31 // row_stride + 8 |
| 213 | free, _ = current_platform.mem_get_info() |
| 214 | if free < 2 * num_tokens * row_stride * 2: |
| 215 | pytest.skip("not enough free device memory") |
| 216 | |
| 217 | packed = torch.empty(num_tokens, row_stride, dtype=torch.bfloat16, device=device) |
| 218 | q_fill = torch.randn(n_qh * hd, dtype=torch.bfloat16, device=device) |
| 219 | k_fill = torch.randn(n_kh * hd, dtype=torch.bfloat16, device=device) |
| 220 | packed[:, : n_qh * hd] = q_fill |
| 221 | packed[:, n_qh * hd : (n_qh + n_kh) * hd] = k_fill |
| 222 | half_rd = rd // 2 |
| 223 | cos = torch.randn(3, num_tokens, half_rd, dtype=torch.float32, device=device) |
| 224 | sin = torch.randn(3, num_tokens, half_rd, dtype=torch.float32, device=device) |
| 225 | q = packed[:, : n_qh * hd] |
| 226 | k = packed[:, n_qh * hd : (n_qh + n_kh) * hd] |
| 227 | triton_mrope(q, k, cos, sin, mrope_section, hd, rd, False, True) |
| 228 | |
| 229 | t_end = mrope_section[0] |
| 230 | h_end = t_end + mrope_section[1] |
| 231 | for r in (0, num_tokens - 8, num_tokens - 1): |
| 232 | cos_row = torch.cat( |
| 233 | [cos[0, r, :t_end], cos[1, r, t_end:h_end], cos[2, r, h_end:]] |
| 234 | ) |
| 235 | sin_row = torch.cat( |
| 236 | [sin[0, r, :t_end], sin[1, r, t_end:h_end], sin[2, r, h_end:]] |
| 237 | ) |
| 238 | for fill, out in ((q_fill, q[r]), (k_fill, k[r])): |
| 239 | row = fill.view(-1, hd).float() |
| 240 | r1, r2 = row[:, :half_rd], row[:, half_rd:rd] |
| 241 | exp = torch.cat( |
| 242 | [r1 * cos_row - r2 * sin_row, r2 * cos_row + r1 * sin_row, row[:, rd:]], |
| 243 | dim=1, |
| 244 | ) |
| 245 | torch.testing.assert_close( |
| 246 | out, exp.view(-1).to(torch.bfloat16), atol=1e-2, rtol=1.6e-2 |
| 247 | ) |
| 248 | |
| 249 | |
| 250 | @pytest.mark.skipif( |
| 251 | not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests." |
| 252 | ) |
| 253 | @pytest.mark.parametrize( |
| 254 | "model_info, model_name", |
| 255 | [ |
| 256 | pytest.param(test_config, test_config.model_name, marks=test_config.marks) |
| 257 | for test_config in MODELS_TO_TEST |
| 258 | ], |
| 259 | ) |
| 260 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 261 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 262 | @pytest.mark.parametrize("num_tokens", num_tokens_list) |
| 263 | def test_mrope_torch_compile_tracing( |
| 264 | default_vllm_config, |
| 265 | model_name: str, |
| 266 | model_info: MRoPETestInfo, |
| 267 | tp_size: int, |
| 268 | dtype: torch.dtype, |
| 269 | num_tokens: int, |
| 270 | ): |
| 271 | atol = model_info.atol |
| 272 | rtol = model_info.rtol |
| 273 | |
| 274 | config = get_config(model_name, False).get_text_config() |
| 275 | |
| 276 | # get the model config |
| 277 | total_num_kv_heads = config.num_key_value_heads |
| 278 | total_num_heads = config.num_attention_heads |
| 279 | num_heads = total_num_heads // tp_size |
| 280 | num_kv_heads = max(1, total_num_kv_heads // tp_size) |
| 281 | head_dim = ( |
| 282 | config.head_dim |
| 283 | if hasattr(config, "head_dim") |
| 284 | else config.hidden_size // total_num_heads |
| 285 | ) |
| 286 | is_neox_style = model_info.is_neox_style |
| 287 | max_position = config.max_position_embeddings |
| 288 | |
| 289 | mrope_helper_class = get_rope( |
| 290 | head_size=head_dim, |
| 291 | max_position=max_position, |
| 292 | is_neox_style=is_neox_style, |
| 293 | rope_parameters=config.rope_parameters, |
| 294 | dtype=dtype, |
| 295 | ).to(device=device) |
| 296 | |
| 297 | # Generate test data |
| 298 | positions, query, key = generate_test_data( |
| 299 | num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device |
| 300 | ) |
| 301 | |
| 302 | # Create a wrapper that makes the in-place function appear functional |
| 303 | def functional_forward_cuda(pos, q, k): |
| 304 | """Wrapper that converts in-place operation to functional style. |
| 305 | |
| 306 | CUDA Graph does not support in-place operations. |
| 307 | This wrapper creates working copies of the |
| 308 | input tensors and modifies them. |
| 309 | """ |
| 310 | q_work = q.clone() # Create working copies |
| 311 | k_work = k.clone() |
| 312 | # Your in-place function modifies q_work and k_work |
| 313 | mrope_helper_class.forward_cuda(pos, q_work, k_work) |
| 314 | return q_work, k_work # Return the modified tensors |
| 315 | |
| 316 | # Get reference results |
| 317 | query_native, key_native = mrope_helper_class.forward_native( |
| 318 | positions, |
| 319 | query.clone(), |
| 320 | key.clone(), |
| 321 | ) |
| 322 | |
| 323 | try: |
| 324 | compiled_forward_cuda = torch.compile( |
| 325 | functional_forward_cuda, |
| 326 | fullgraph=True, |
| 327 | backend="inductor", |
| 328 | mode="reduce-overhead", |
| 329 | dynamic=False, |
| 330 | ) |
| 331 | |
| 332 | # Run compiled version |
| 333 | query_compiled_cuda, key_compiled_cuda = compiled_forward_cuda( |
| 334 | positions, |
| 335 | query, |
| 336 | key, |
| 337 | ) |
| 338 | |
| 339 | # Run original version for comparison |
| 340 | query_cuda = query.clone() |
| 341 | key_cuda = key.clone() |
| 342 | mrope_helper_class.forward_cuda(positions, query_cuda, key_cuda) |
| 343 | |
| 344 | # Verify results |
| 345 | torch.testing.assert_close( |
| 346 | query_compiled_cuda, query_cuda, atol=atol, rtol=rtol |
| 347 | ) |
| 348 | torch.testing.assert_close(key_compiled_cuda, key_cuda, atol=atol, rtol=rtol) |
| 349 | torch.testing.assert_close( |
| 350 | query_compiled_cuda, query_native, atol=atol, rtol=rtol |
| 351 | ) |
| 352 | torch.testing.assert_close(key_compiled_cuda, key_native, atol=atol, rtol=rtol) |
| 353 | |
| 354 | print("✓ forward_cuda successfully traced with torch.compile inductor") |
| 355 | |
| 356 | except Exception as e: |
| 357 | pytest.fail(f"forward_cuda failed to trace with torch.compile inductor: {e}") |