main
py 29 lines 938 Bytes
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Mooncake P/D accuracy with same-host HIP transport on ROCm."""
4
5 import pytest
6
7 from tests.utils import spawn_new_process_for_each_test
8 from tests.v1.kv_connector.rocm_pd_accuracy_utils import (
9 QWEN_MODEL,
10 run_rocm_pd_accuracy,
11 )
12 from vllm.platforms import current_platform
13
14 pytestmark = [
15 pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm P/D equivalents"),
16 pytest.mark.distributed(num_gpus=4),
17 ]
18
19
20 @pytest.mark.parametrize(
21 "model,prefill_tp,decode_tp",
22 [
23 pytest.param(QWEN_MODEL, 1, 1, id="mooncake-hip-qwen-p1-d1"),
24 pytest.param(QWEN_MODEL, 2, 2, id="mooncake-hip-qwen-p2-d2"),
25 ],
26 )
27 @spawn_new_process_for_each_test
28 def test_rocm_pd_accuracy(model, prefill_tp, decode_tp, monkeypatch):
29 run_rocm_pd_accuracy("mooncake", model, prefill_tp, decode_tp, monkeypatch)