@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 311 lines 9.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Warm up Qwen Triton kernels from the loaded model's compile keys."""
4
5 from dataclasses import dataclass
6 from typing import TYPE_CHECKING
7
8 import torch
9
10 from vllm.logger import init_logger
11
12 if TYPE_CHECKING:
13 from vllm.config import ModelConfig
14 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
15
16 logger = init_logger(__name__)
17
18 _QWEN_MODEL_TYPES = frozenset(
19 {
20 "qwen3_next",
21 "qwen3_5",
22 "qwen3_5_text",
23 "qwen3_5_moe",
24 "qwen3_5_moe_text",
25 "qwen4_exp",
26 "qwen4_exp_text",
27 }
28 )
29
30 # Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L.
31 _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16)
32
33
34 @dataclass(frozen=True)
35 class _QwenGDNWarmupConfig:
36 h: int
37 hv: int
38 k: int
39 v: int
40 conv_kernel_size: int
41 conv_state: torch.Tensor
42 conv_dtype: torch.dtype
43 norm_weight_dtype: torch.dtype
44 norm_before_gate: bool
45 norm_activation: str
46 a_log: torch.Tensor
47 dt_bias: torch.Tensor
48 state_stride_token: int
49 state_dtype: torch.dtype
50
51 @property
52 def conv_dim(self) -> int:
53 return 2 * self.h * self.k + self.hv * self.v
54
55
56 def _is_non_empty_tensor(value: object) -> bool:
57 return isinstance(value, torch.Tensor) and value.numel() > 0
58
59
60 def _is_qwen_gdn_layer(module: object) -> bool:
61 return all(
62 hasattr(module, attr)
63 for attr in (
64 "num_k_heads",
65 "num_v_heads",
66 "head_k_dim",
67 "head_v_dim",
68 "conv_kernel_size",
69 "tp_size",
70 "kv_cache",
71 "norm",
72 "A_log",
73 "dt_bias",
74 )
75 )
76
77
78 def _iter_qwen_gdn_layers(static_forward_context: object):
79 if not isinstance(static_forward_context, dict):
80 return
81
82 for module in static_forward_context.values():
83 if _is_qwen_gdn_layer(module):
84 yield module
85
86
87 def _split_qwen_gdn_cache(kv_cache: object) -> tuple[torch.Tensor, torch.Tensor] | None:
88 if isinstance(kv_cache, (list, tuple)) and len(kv_cache) >= 2:
89 conv_cache, ssm_state = kv_cache[:2]
90 if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state):
91 return conv_cache, ssm_state
92
93 if isinstance(kv_cache, torch.Tensor) and kv_cache.size(0) >= 2:
94 conv_cache = kv_cache[0]
95 ssm_state = kv_cache[1]
96 if _is_non_empty_tensor(conv_cache) and _is_non_empty_tensor(ssm_state):
97 return conv_cache, ssm_state
98 return None
99
100
101 def _qwen_gdn_warmup_config(
102 static_forward_context: object,
103 ) -> _QwenGDNWarmupConfig | None:
104 found_layer = False
105 for layer in _iter_qwen_gdn_layers(static_forward_context):
106 found_layer = True
107 cache_tensors = _split_qwen_gdn_cache(getattr(layer, "kv_cache", None))
108 if cache_tensors is None:
109 continue
110
111 conv_cache, ssm_state = cache_tensors
112 from vllm.model_executor.layers.mamba.mamba_utils import (
113 is_conv_state_dim_first,
114 )
115
116 conv_state = (
117 conv_cache if is_conv_state_dim_first() else conv_cache.transpose(-1, -2)
118 )
119 tp_size = int(layer.tp_size)
120 h = int(layer.num_k_heads) // tp_size
121 hv = int(layer.num_v_heads) // tp_size
122 norm = layer.norm
123
124 return _QwenGDNWarmupConfig(
125 h=h,
126 hv=hv,
127 k=int(layer.head_k_dim),
128 v=int(layer.head_v_dim),
129 conv_kernel_size=int(layer.conv_kernel_size),
130 conv_state=conv_state,
131 conv_dtype=conv_state.dtype,
132 norm_weight_dtype=norm.weight.dtype,
133 norm_before_gate=bool(norm.norm_before_gate),
134 norm_activation=str(norm.activation),
135 a_log=layer.A_log,
136 dt_bias=layer.dt_bias,
137 state_stride_token=int(ssm_state.stride(0)),
138 state_dtype=ssm_state.dtype,
139 )
140
141 if found_layer:
142 logger.info("Skipping Qwen GDN Triton warmup: no bound Qwen GDN cache found.")
143 else:
144 logger.info("Skipping Qwen GDN Triton warmup: no Qwen GDN layer found.")
145 return None
146
147
148 def _warm_gated_rms_norm_kernel(
149 device: torch.device,
150 config: _QwenGDNWarmupConfig,
151 max_num_tokens: int,
152 x_dtype: torch.dtype,
153 ) -> None:
154 from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
155 warmup_layer_norm_fwd,
156 )
157
158 warmup_layer_norm_fwd(
159 max_num_tokens=max_num_tokens,
160 rows_per_token=config.hv,
161 group_size=config.v,
162 x_dtype=x_dtype,
163 weight_dtype=config.norm_weight_dtype,
164 device=device,
165 norm_before_gate=config.norm_before_gate,
166 is_rms_norm=True,
167 activation=config.norm_activation,
168 )
169
170
171 def _warm_causal_conv1d_fwd_kernel(
172 device: torch.device, config: _QwenGDNWarmupConfig
173 ) -> None:
174 from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
175 causal_conv1d_fn,
176 )
177 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
178
179 x_storage = torch.empty(
180 (1, config.conv_dim), dtype=config.conv_dtype, device=device
181 )
182 x = x_storage.t()
183 weight = torch.empty(
184 (config.conv_dim, config.conv_kernel_size),
185 dtype=config.conv_dtype,
186 device=device,
187 )
188 cache_indices = torch.full((1,), NULL_BLOCK_ID, dtype=torch.int32, device=device)
189 has_initial_state = torch.empty(1, dtype=torch.bool, device=device)
190 query_start_loc = torch.tensor([0, 1], dtype=torch.int32, device=device)
191
192 causal_conv1d_fn(
193 x,
194 weight,
195 None,
196 config.conv_state,
197 query_start_loc,
198 cache_indices=cache_indices,
199 has_initial_state=has_initial_state,
200 activation="silu",
201 pad_slot_id=PAD_SLOT_ID,
202 null_block_id=NULL_BLOCK_ID,
203 metadata=None,
204 validate_data=False,
205 )
206
207
208 def _warm_fused_post_conv_kernel(
209 device: torch.device, config: _QwenGDNWarmupConfig
210 ) -> None:
211 from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( # noqa: E501
212 fused_post_conv_prep,
213 )
214
215 qkv_dim = 2 * config.h * config.k + config.hv * config.v
216 for length in _FLA_POST_CONV_WARMUP_LENGTHS:
217 conv_output = torch.empty(
218 (length, qkv_dim), dtype=config.conv_dtype, device=device
219 )
220 a = torch.empty((length, config.hv), dtype=config.conv_dtype, device=device)
221 b = torch.empty_like(a)
222
223 fused_post_conv_prep(
224 conv_output,
225 a,
226 b,
227 config.a_log,
228 config.dt_bias,
229 config.h,
230 config.k,
231 config.v,
232 apply_l2norm=True,
233 output_g_exp=False,
234 )
235
236
237 def _warm_fused_sigmoid_gating_delta_rule_update_kernel(
238 device: torch.device,
239 config: _QwenGDNWarmupConfig,
240 ) -> None:
241 from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import (
242 fused_sigmoid_gating_delta_rule_update,
243 )
244
245 q = torch.empty((1, 1, config.h, config.k), dtype=config.conv_dtype, device=device)
246 k = torch.empty_like(q)
247 v = torch.empty((1, 1, config.hv, config.v), dtype=config.conv_dtype, device=device)
248 a = torch.empty((1, 1, config.hv), dtype=config.conv_dtype, device=device)
249 b = torch.empty_like(a)
250 state = torch.empty(
251 (1, config.state_stride_token),
252 dtype=config.state_dtype,
253 device=device,
254 )
255 cu_seqlens = torch.tensor([0, 1], dtype=torch.int32, device=device)
256 ssm_state_indices = torch.empty((1, 1), dtype=torch.int32, device=device)
257 ssm_state_indices.zero_()
258
259 fused_sigmoid_gating_delta_rule_update(
260 A_log=config.a_log,
261 a=a,
262 b=b,
263 dt_bias=config.dt_bias,
264 q=q,
265 k=k,
266 v=v,
267 beta=1.0,
268 threshold=20.0,
269 initial_state=state,
270 inplace_final_state=True,
271 cu_seqlens=cu_seqlens,
272 ssm_state_indices=ssm_state_indices,
273 use_qk_l2norm_in_kernel=True,
274 is_kda=False,
275 )
276
277
278 def _synchronize_device(device: torch.device) -> None:
279 if device.type == "cuda":
280 torch.accelerator.synchronize(device)
281
282
283 @torch.inference_mode()
284 def qwen_triton_warmup(
285 runner: "GPUModelRunner",
286 model_config: "ModelConfig",
287 ) -> None:
288 """Warm Qwen GDN Triton kernels reported by the JIT monitor."""
289 model_type = getattr(model_config.hf_text_config, "model_type", "") or getattr(
290 model_config.hf_config, "model_type", ""
291 )
292 if model_type not in _QWEN_MODEL_TYPES:
293 return
294
295 device = runner.device
296 logger.info("Warming up Qwen GDN Triton kernels for model_type=%s.", model_type)
297
298 gdn_config = _qwen_gdn_warmup_config(
299 runner.compilation_config.static_forward_context
300 )
301 if gdn_config is None:
302 return
303
304 max_num_tokens = max(1, int(runner.max_num_tokens))
305 _warm_gated_rms_norm_kernel(device, gdn_config, max_num_tokens, model_config.dtype)
306 _warm_causal_conv1d_fwd_kernel(device, gdn_config)
307 _warm_fused_post_conv_kernel(device, gdn_config)
308 # Pooling only runs full prefills; the decode update kernel is unused.
309 if not runner.is_pooling_model:
310 _warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
311 _synchronize_device(device)