| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.warmup.qwen_vl_triton_warmup import ( |
| 9 | _warm_mrope, |
| 10 | _warm_vision, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | def test_vision_warmup_calls_only_position_and_rotary_paths() -> None: |
| 15 | from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer |
| 16 | |
| 17 | calls: list[tuple[str, object]] = [] |
| 18 | |
| 19 | class FakeAttention: |
| 20 | num_attention_heads_per_partition = 4 |
| 21 | hidden_size_per_attention_head = 8 |
| 22 | |
| 23 | def apply_rotary_emb(self, qk, cos, sin): |
| 24 | calls.append(("rotary", qk.shape)) |
| 25 | |
| 26 | class FakeVisual(Qwen3_VisionTransformer): |
| 27 | spatial_merge_size = 2 |
| 28 | |
| 29 | def __init__(self) -> None: |
| 30 | torch.nn.Module.__init__(self) |
| 31 | self.blocks = [SimpleNamespace(attn=FakeAttention())] |
| 32 | self.pos_embed = torch.nn.Embedding(4, 4, device="cpu") |
| 33 | |
| 34 | def fast_pos_embed_interpolate(self, grid_thw): |
| 35 | calls.append(("position", grid_thw[0])) |
| 36 | |
| 37 | def rot_pos_emb(self, grid_thw): |
| 38 | _, h, w = grid_thw[0] |
| 39 | shape = (h * w, 4) |
| 40 | return torch.empty(shape), torch.empty(shape) |
| 41 | |
| 42 | def forward(self, *args, **kwargs): |
| 43 | raise AssertionError("vision warmup must not run the full tower") |
| 44 | |
| 45 | model = torch.nn.Module() |
| 46 | model.visual = FakeVisual() |
| 47 | _warm_vision(model, torch.device("cpu")) |
| 48 | |
| 49 | grids = [(1, 16, 16), (1, 16, 2), (1, 2, 16), (1, 2, 2)] |
| 50 | assert [value for name, value in calls if name == "position"] == [ |
| 51 | list(grid) for grid in grids |
| 52 | ] |
| 53 | assert [value for name, value in calls if name == "rotary"] == [ |
| 54 | torch.Size((2, h * w, 4, 8)) for _, h, w in grids |
| 55 | ] |
| 56 | |
| 57 | # CPU offloading leaves tower weights off the accelerator, so the Triton |
| 58 | # paths must not be launched on host pointers. |
| 59 | calls.clear() |
| 60 | _warm_vision(model, torch.device("cuda")) |
| 61 | assert calls == [] |
| 62 | |
| 63 | |
| 64 | def test_mrope_warmup_launches_triton_shapes() -> None: |
| 65 | from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding |
| 66 | |
| 67 | launched: list[tuple[torch.Size, torch.Size]] = [] |
| 68 | |
| 69 | class FakeRope(MRotaryEmbedding): |
| 70 | def __init__(self) -> None: |
| 71 | torch.nn.Module.__init__(self) |
| 72 | self.head_size = 8 |
| 73 | self.rotary_dim = 8 |
| 74 | self.mrope_section = [2, 3, 3] |
| 75 | self.mrope_interleaved = False |
| 76 | self.is_neox_style = True |
| 77 | |
| 78 | def forward(self, positions, query, key): |
| 79 | launched.append((positions.shape, query.shape)) |
| 80 | return query, key |
| 81 | |
| 82 | class FakeMropeModel(torch.nn.Module): |
| 83 | supports_mrope = True |
| 84 | |
| 85 | def __init__(self) -> None: |
| 86 | super().__init__() |
| 87 | self.rotary_emb = FakeRope() |
| 88 | |
| 89 | def get_mrope_input_positions(self, input_tokens, mm_features): |
| 90 | raise AssertionError("warmup must not build live M-RoPE positions") |
| 91 | |
| 92 | runner = SimpleNamespace( |
| 93 | model_config=SimpleNamespace( |
| 94 | get_num_attention_heads=lambda parallel_config: 4, |
| 95 | get_num_kv_heads=lambda parallel_config: 2, |
| 96 | ), |
| 97 | parallel_config=object(), |
| 98 | dtype=torch.bfloat16, |
| 99 | device=torch.device("cpu"), |
| 100 | ) |
| 101 | _warm_mrope(runner, FakeMropeModel()) |
| 102 | assert [shape for shape, _ in launched] == [ |
| 103 | torch.Size((3, 1)), |
| 104 | torch.Size((3, 2)), |
| 105 | torch.Size((3, 16)), |
| 106 | ] |
| 107 | |
| 108 | |
| 109 | def test_mrope_warmup_skips_models_without_supports_mrope() -> None: |
| 110 | def fail(*_args, **_kwargs): |
| 111 | raise AssertionError("M-RoPE warmup must not run without SupportsMRoPE") |
| 112 | |
| 113 | runner = SimpleNamespace( |
| 114 | model_config=SimpleNamespace(get_num_attention_heads=fail), |
| 115 | get_model=fail, |
| 116 | ) |
| 117 | _warm_mrope(runner, torch.nn.Module()) |