| 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() |