@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 112 lines 3.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Warm Qwen3-VL ViT kernels and M-RoPE for SupportsMRoPE models."""
4
5 from typing import TYPE_CHECKING
6
7 import torch
8
9 from vllm.logger import init_logger
10 from vllm.utils.math_utils import round_up
11
12 if TYPE_CHECKING:
13 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
14
15 logger = init_logger(__name__)
16
17 # Covers Triton's ==1 / %16==0 / other integer buckets.
18 _MROPE_TOKEN_COUNTS = (1, 2, 16)
19
20
21 def _warm_vision(model: torch.nn.Module, device: torch.device) -> None:
22 from vllm.model_executor.models.qwen3_vl import Qwen3_VisionTransformer
23
24 for visual in model.modules():
25 if not isinstance(visual, Qwen3_VisionTransformer):
26 continue
27 # CPU offloading keeps tower weights off the accelerator until the
28 # module forward hook copies them back, so warming here would launch
29 # Triton on host pointers.
30 if visual.pos_embed.weight.device.type != device.type:
31 continue
32 attention = visual.blocks[0].attn
33 num_heads = int(attention.num_attention_heads_per_partition)
34 head_size = int(attention.hidden_size_per_attention_head)
35 merge_size = int(visual.spatial_merge_size)
36 divisible = round_up(16, merge_size)
37 while divisible % 16:
38 divisible += merge_size
39 spatial_sizes = [divisible]
40 if merge_size % 16:
41 spatial_sizes.append(merge_size)
42 if merge_size == 1:
43 spatial_sizes.append(2)
44 grids = [(1, h, w) for h in spatial_sizes for w in spatial_sizes]
45
46 for grid in grids:
47 grid_thw = [list(grid)]
48 visual.fast_pos_embed_interpolate(grid_thw)
49 cos, sin = visual.rot_pos_emb(grid_thw)
50 qk = torch.empty(
51 (2, grid[0] * grid[1] * grid[2], num_heads, head_size),
52 dtype=cos.dtype,
53 device=cos.device,
54 )
55 attention.apply_rotary_emb(qk, cos, sin)
56
57 logger.info(
58 "Warmed position embedding and vision rotary kernels on grids=%s.",
59 grids,
60 )
61
62
63 def _warm_mrope(runner: "GPUModelRunner", model: torch.nn.Module) -> None:
64 from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
65 from vllm.model_executor.models.interfaces import supports_mrope
66
67 is_mrope_model = supports_mrope(model)
68 if not is_mrope_model:
69 return
70
71 rope = next(
72 (module for module in model.modules() if isinstance(module, MRotaryEmbedding)),
73 None,
74 )
75 if rope is None:
76 return
77
78 model_config = runner.model_config
79 parallel_config = runner.parallel_config
80 num_query_heads = int(model_config.get_num_attention_heads(parallel_config))
81 num_kv_heads = int(model_config.get_num_kv_heads(parallel_config))
82 for num_tokens in _MROPE_TOKEN_COUNTS:
83 positions = torch.arange(
84 num_tokens, dtype=torch.long, device=runner.device
85 ).expand(3, -1)
86 query = torch.empty(
87 (num_tokens, num_query_heads * rope.head_size),
88 dtype=runner.dtype,
89 device=runner.device,
90 )
91 key_tensor = torch.empty(
92 (num_tokens, num_kv_heads * rope.head_size),
93 dtype=runner.dtype,
94 device=runner.device,
95 )
96 rope(positions, query, key_tensor)
97
98 logger.info("Warmed M-RoPE Triton kernels.")
99
100
101 def _synchronize_device(device: torch.device) -> None:
102 if device.type == "cuda":
103 torch.accelerator.synchronize(device)
104
105
106 @torch.inference_mode()
107 def qwen_vl_triton_warmup(runner: "GPUModelRunner") -> None:
108 """Warm Qwen3-VL ViT kernels and M-RoPE for ``SupportsMRoPE`` models."""
109 model = runner.get_model()
110 _warm_vision(model, runner.device)
111 _warm_mrope(runner, model)
112 _synchronize_device(runner.device)