[CI][Bugfix] Limit MRV2 sampler JIT warmup registration to ROCm (#58465)
Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Kevin H. Luu committed
Sep 23, 2026 at 16:11 UTC
21ee743f474a35a056e8d5ade2c9aa744e726e6e
1 file changed
+9
-7
vllm/v1/worker/gpu/model_runner.py
+9
-7
@@ -55,6 +55,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
55
from vllm.multimodal.encoder_budget import (
56
MultiModalBudget,
57
)
58
+from vllm.platforms import current_platform
59
from vllm.sequence import IntermediateTensors
60
from vllm.tasks import SupportedTask
61
from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
@@ -439,14 +440,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
440
self.vllm_config, self.req_states, self.is_pooling_model, config_processors
441
)
442
if self.is_last_pp_rank and not self.is_pooling_model:
442
- from vllm.v1.sample.ops.topk_topp_sampler import (
443
- register_top_k_top_p_warmups,
444
- )
445
-
443
# V2 bypasses TopKTopPSampler, which registers native warmups.
447
- # CUDA also needs these for its FlashInfer fallback paths.
448
- with self.jit_warmup_registry.activate():
449
- register_top_k_top_p_warmups()
444
+ # ROCm only: on CUDA these warmups add ~2 min to every engine start.
445
+ if current_platform.is_rocm():
446
+ from vllm.v1.sample.ops.topk_topp_sampler import (
447
+ register_top_k_top_p_warmups,
448
+ )
449
+
450
+ with self.jit_warmup_registry.activate():
451
+ register_top_k_top_p_warmups()
452
453
sampler_kwargs: dict[str, Any] = {
454
"vllm_config": self.vllm_config,