| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import dataclasses |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.config import CompilationMode |
| 8 | from vllm.platforms import current_platform |
| 9 | |
| 10 | from ...utils import compare_all_settings |
| 11 | |
| 12 | ATTN_BACKEND = "FLASH_ATTN" if not current_platform.is_rocm() else "ROCM_ATTN" |
| 13 | |
| 14 | |
| 15 | @dataclasses.dataclass |
| 16 | class TestSetting: |
| 17 | model: str |
| 18 | model_args: list[str] |
| 19 | pp_size: int |
| 20 | tp_size: int |
| 21 | attn_backend: str |
| 22 | method: str |
| 23 | |
| 24 | |
| 25 | # we cannot afford testing the full Cartesian product |
| 26 | # of all models and all modes |
| 27 | @pytest.mark.parametrize( |
| 28 | "test_setting", |
| 29 | [ |
| 30 | # MoE model |
| 31 | TestSetting( |
| 32 | model="ibm-granite/granite-3.0-1b-a400m-instruct", |
| 33 | model_args=["--max-model-len", "2048"], |
| 34 | pp_size=1, |
| 35 | tp_size=1, |
| 36 | attn_backend=ATTN_BACKEND, |
| 37 | method="generate", |
| 38 | ), |
| 39 | # embedding model |
| 40 | TestSetting( |
| 41 | model="BAAI/bge-multilingual-gemma2", |
| 42 | model_args=[ |
| 43 | "--runner", |
| 44 | "pooling", |
| 45 | "--dtype", |
| 46 | "bfloat16", |
| 47 | "--max-model-len", |
| 48 | "2048", |
| 49 | "--gpu-memory-utilization", |
| 50 | "0.98", |
| 51 | ], |
| 52 | pp_size=1, |
| 53 | tp_size=1, |
| 54 | attn_backend=ATTN_BACKEND, |
| 55 | method="encode", |
| 56 | ), |
| 57 | ], |
| 58 | ) |
| 59 | @pytest.mark.parametrize("use_v2_model_runner", [False, True]) |
| 60 | def test_compile_correctness( |
| 61 | test_setting: TestSetting, |
| 62 | use_v2_model_runner: bool, |
| 63 | ): |
| 64 | model = test_setting.model |
| 65 | model_args = test_setting.model_args |
| 66 | pp_size = test_setting.pp_size |
| 67 | tp_size = test_setting.tp_size |
| 68 | attn_backend = test_setting.attn_backend |
| 69 | method = test_setting.method |
| 70 | final_args = [ |
| 71 | *model_args, |
| 72 | "-pp", |
| 73 | str(pp_size), |
| 74 | "-tp", |
| 75 | str(tp_size), |
| 76 | "-cc.cudagraph_mode=none", |
| 77 | f"--attention-backend={attn_backend}", |
| 78 | ] |
| 79 | |
| 80 | # Pin every compared setting to the same model runner so that only the |
| 81 | # compilation mode differs within a comparison. |
| 82 | runner_env = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))} |
| 83 | |
| 84 | all_args: list[list[str]] = [] |
| 85 | all_envs: list[dict[str, str] | None] = [] |
| 86 | |
| 87 | # Test all compilation modes with inductor backend |
| 88 | for mode in [ |
| 89 | CompilationMode.NONE, |
| 90 | CompilationMode.STOCK_TORCH_COMPILE, |
| 91 | CompilationMode.DYNAMO_TRACE_ONCE, |
| 92 | CompilationMode.VLLM_COMPILE, |
| 93 | ]: |
| 94 | all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"]) |
| 95 | all_envs.append(dict(runner_env)) |
| 96 | # inductor will change the output, so we only compare if the output |
| 97 | # is close, not exactly the same. |
| 98 | compare_all_settings( |
| 99 | model, |
| 100 | all_args, |
| 101 | all_envs, |
| 102 | method=method if method != "generate" else "generate_close", |
| 103 | ) |
| 104 | |
| 105 | all_envs.clear() |
| 106 | all_args.clear() |
| 107 | |
| 108 | # Test all compilation modes with eager backend |
| 109 | for mode in [ |
| 110 | CompilationMode.NONE, |
| 111 | CompilationMode.STOCK_TORCH_COMPILE, |
| 112 | CompilationMode.DYNAMO_TRACE_ONCE, |
| 113 | CompilationMode.VLLM_COMPILE, |
| 114 | ]: |
| 115 | all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"]) |
| 116 | all_envs.append(dict(runner_env)) |
| 117 | compare_all_settings(model, all_args, all_envs, method=method) |