[ROCm][Bugfix] Register MRV2 sampler JIT warmups (#58092)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 23, 2026 at 03:17 UTC
96e203b08ed4b935051151989f550c9b6c6eba6e
4 files changed
+114
.buildkite/test-amd.yaml
+24
@@ -5252,3 +5252,27 @@ steps:
5252
- amdgfx950nightly
5253
- amdmi355
5254
optional: true
5255
+
5256
+- label: ":amd: (MI355) MRV2 Sampler JIT Warmup"
5257
+ timeout_in_minutes: 10
5258
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
5259
+ dind: false
5260
+ agent_pool: mi355_1
5261
+ num_gpus: 1
5262
+ optional: true
5263
+ working_dir: "/vllm-workspace/"
5264
+ source_file_dependencies:
5265
+ - tests/jit_monitor/
5266
+ - tests/models/utils.py
5267
+ - tests/models/registry.py
5268
+ - tests/utils.py
5269
+ - vllm/v1/worker/gpu/
5270
+ - vllm/v1/worker/gpu_worker.py
5271
+ - vllm/model_executor/warmup/
5272
+ - vllm/v1/sample/ops/topk_topp_sampler.py
5273
+ - vllm/utils/jit_monitor.py
5274
+ - vllm/config/observability.py
5275
+ - vllm/platforms/rocm.py
5276
+ - vllm/envs.py
5277
+ commands:
5278
+ - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
.buildkite/test_areas/jit_monitor.yaml
+21
@@ -20,3 +20,24 @@ steps:
20
# instead of running until the build timeout.
21
- export PYTHONFAULTHANDLER=1
22
- pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread
23
+ mirror:
24
+ amd:
25
+ # Focused MRV2 regression; the broader model/backend matrix is pending.
26
+ label: ':amd: (MI355) MRV2 Sampler JIT Warmup'
27
+ dind: false
28
+ device: mi355_1
29
+ num_devices: 1
30
+ timeout_in_minutes: 10
31
+ working_dir: /vllm-workspace/
32
+ depends_on:
33
+ - image-build-amd
34
+ source_file_dependencies:
35
+ - tests/jit_monitor/
36
+ - tests/models/utils.py
37
+ - tests/utils.py
38
+ - vllm/v1/worker/gpu/
39
+ - vllm/v1/sample/ops/topk_topp_sampler.py
40
+ - vllm/platforms/rocm.py
41
+ - vllm/envs.py
42
+ commands:
43
+ - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
tests/jit_monitor/test_no_runtime_jit_rocm.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Keep the MRV2 sampler warmup registry populated on ROCm."""
4
+
5
+import pytest
6
+
7
+from vllm import LLM
8
+from vllm.platforms import current_platform
9
+
10
+from ..models.utils import dummy_hf_overrides
11
+from ..utils import create_new_process_for_each_test
12
+from .test_no_runtime_jit import _run_shape_battery
13
+
14
+pytestmark = pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm")
15
+
16
+
17
+def _worker_monitor_state(worker):
18
+ from vllm.utils import jit_monitor
19
+
20
+ return (
21
+ jit_monitor.is_active(),
22
+ worker.observability_config.jit_monitor_mode,
23
+ )
24
+
25
+
26
+@create_new_process_for_each_test("spawn")
27
+def _run_rocm_shape_battery() -> None:
28
+ llm = LLM(
29
+ "Qwen/Qwen3-0.6B",
30
+ max_model_len=2048,
31
+ max_num_seqs=8,
32
+ gpu_memory_utilization=0.03,
33
+ kv_cache_memory_bytes=256 * 1024 * 1024,
34
+ load_format="dummy",
35
+ hf_overrides=dummy_hf_overrides,
36
+ enforce_eager=False,
37
+ jit_monitor_mode="error",
38
+ )
39
+ try:
40
+ states = llm.collective_rpc(_worker_monitor_state, timeout=30)
41
+ assert states and all(state == (True, "error") for state in states)
42
+ _run_shape_battery(llm)
43
+ finally:
44
+ llm.llm_engine.engine_core.shutdown(timeout=30)
45
+
46
+
47
+def test_v2_sampler_warmup_rocm(monkeypatch, tmp_path):
48
+ """Warmup must cover the existing prefill, decode, and sampler battery.
49
+
50
+ A fresh cache and child process prevent previous model executions from
51
+ hiding a missing warmup specialization.
52
+ """
53
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
54
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "0")
55
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
56
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
57
+ monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton"))
58
+ monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm"))
59
+ monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor"))
60
+ _run_rocm_shape_battery()
vllm/v1/worker/gpu/model_runner.py
+9
@@ -442,6 +442,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
442
self.vllm_config, self.req_states, self.is_pooling_model, config_processors
443
)
444
if self.is_last_pp_rank and not self.is_pooling_model:
445
+ from vllm.v1.sample.ops.topk_topp_sampler import (
446
+ register_top_k_top_p_warmups,
447
+ )
448
+
449
+ # V2 bypasses TopKTopPSampler, which registers native warmups.
450
+ # CUDA also needs these for its FlashInfer fallback paths.
451
+ with self.jit_warmup_registry.activate():
452
+ register_top_k_top_p_warmups()
453
+
454
sampler_kwargs: dict[str, Any] = {
455
"vllm_config": self.vllm_config,
456
"max_num_reqs": self.max_num_reqs,