[CI Bugfix] Lazily import Qwen warmup dependencies (#47539)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Roberto L. Castro committed
Jul 3, 2026 at 17:10 UTC
d7192cfccf0f4fa677fc108c6d9a7828d32a8eea
1 file changed
+21
-13
vllm/model_executor/warmup/qwen_triton_warmup.py
+21
-13
@@ -8,19 +8,6 @@ from typing import TYPE_CHECKING
8
import torch
9
10
from vllm.logger import init_logger
11
-from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
12
- fused_post_conv_prep,
13
-)
14
-from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import (
15
- fused_sigmoid_gating_delta_rule_update,
16
-)
17
-from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
18
-from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
19
- causal_conv1d_fn,
20
-)
21
-from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
22
-from vllm.v1.worker.block_table import BlockTable
23
-from vllm.v1.worker.utils import _zero_kv_blocks_kernel
11
12
if TYPE_CHECKING:
13
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
@@ -128,6 +115,10 @@ def _qwen_gdn_warmup_config(
115
continue
116
117
conv_cache, ssm_state = cache_tensors
118
+ from vllm.model_executor.layers.mamba.mamba_utils import (
119
+ is_conv_state_dim_first,
120
+ )
121
+
122
conv_state = (
123
conv_cache if is_conv_state_dim_first() else conv_cache.transpose(-1, -2)
124
)
@@ -191,6 +182,8 @@ def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
182
def _warm_zero_kv_blocks_kernel(
183
device: torch.device, config: _ZeroKvWarmupConfig
184
) -> None:
185
+ from vllm.v1.worker.utils import _zero_kv_blocks_kernel
186
+
187
max_n_blocks = max(_ZERO_KV_N_BLOCKS)
188
scratch = torch.empty(
189
max_n_blocks * config.page_size_el,
@@ -219,6 +212,8 @@ def _warm_zero_kv_blocks_kernel(
212
213
214
def _warm_compute_slot_mapping_kernel(device: torch.device) -> None:
215
+ from vllm.v1.worker.block_table import BlockTable
216
+
217
# num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny.
218
num_tokens = 1
219
query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
@@ -244,6 +239,11 @@ def _warm_compute_slot_mapping_kernel(device: torch.device) -> None:
239
def _warm_causal_conv1d_fwd_kernel(
240
device: torch.device, config: _QwenGDNWarmupConfig
241
) -> None:
242
+ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
243
+ causal_conv1d_fn,
244
+ )
245
+ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
246
+
247
x_storage = torch.empty(
248
(1, config.conv_dim), dtype=config.conv_dtype, device=device
249
)
@@ -276,6 +276,10 @@ def _warm_causal_conv1d_fwd_kernel(
276
def _warm_fused_post_conv_kernel(
277
device: torch.device, config: _QwenGDNWarmupConfig
278
) -> None:
279
+ from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
280
+ fused_post_conv_prep,
281
+ )
282
+
283
qkv_dim = 2 * config.h * config.k + config.hv * config.v
284
for length in _FLA_POST_CONV_WARMUP_LENGTHS:
285
conv_output = torch.empty(
@@ -302,6 +306,10 @@ def _warm_fused_sigmoid_gating_delta_rule_update_kernel(
306
device: torch.device,
307
config: _QwenGDNWarmupConfig,
308
) -> None:
309
+ from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import (
310
+ fused_sigmoid_gating_delta_rule_update,
311
+ )
312
+
313
q = torch.empty((1, 1, config.h, config.k), dtype=config.conv_dtype, device=device)
314
k = torch.empty_like(q)
315
v = torch.empty((1, 1, config.hv, config.v), dtype=config.conv_dtype, device=device)