| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | |
| 9 | from ....utils import large_gpu_mark |
| 10 | from ...registry import HF_EXAMPLE_MODELS |
| 11 | from ...utils import check_logprobs_close |
| 12 | |
| 13 | # Models that require embedding scaling for prompt_embeds test |
| 14 | EMBED_SCALING_MODELS = { |
| 15 | "openbmb/MiniCPM4.1-8B", |
| 16 | } |
| 17 | |
| 18 | # This list contains the model that are using AITER kernel. |
| 19 | # Skip model that are not using AITER tests. |
| 20 | # When more AITER kernels are added, this list will not be |
| 21 | # needed as all the models will be calling AITER kernels |
| 22 | # in parts of the operators |
| 23 | AITER_MODEL_LIST = [ |
| 24 | "meta-llama/Llama-3.2-1B-Instruct", |
| 25 | "openbmb/MiniCPM3-4B", |
| 26 | "Qwen/Qwen2.5-0.5B-Instruct", |
| 27 | "axolotl-ai-co/tiny-mixtral-30m", |
| 28 | "Qwen/Qwen3-8B", |
| 29 | ] |
| 30 | |
| 31 | |
| 32 | # @maybe_test_rocm_aiter |
| 33 | @pytest.mark.parametrize( |
| 34 | "model", |
| 35 | [ |
| 36 | pytest.param( |
| 37 | "bigscience/bloom-560m", # bloom - testing alibi slopes |
| 38 | marks=[ |
| 39 | pytest.mark.core_model, |
| 40 | pytest.mark.slow_test, |
| 41 | pytest.mark.cpu_model, |
| 42 | ], |
| 43 | ), |
| 44 | pytest.param( |
| 45 | "openai-community/gpt2", # gpt2 |
| 46 | marks=[pytest.mark.core_model], |
| 47 | ), |
| 48 | pytest.param("Milos/slovak-gpt-j-405M"), # gptj |
| 49 | pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode |
| 50 | pytest.param("EleutherAI/pythia-70m"), # gpt_neox |
| 51 | pytest.param( |
| 52 | "google/gemma-1.1-2b-it", # gemma |
| 53 | marks=[ |
| 54 | pytest.mark.core_model, |
| 55 | pytest.mark.cpu_model, |
| 56 | pytest.mark.slow_test, |
| 57 | ], |
| 58 | ), |
| 59 | pytest.param( |
| 60 | "google/gemma-2-2b-it", # test hybrid attention |
| 61 | marks=[pytest.mark.cpu_model], |
| 62 | ), |
| 63 | pytest.param( |
| 64 | "zai-org/chatglm3-6b", # chatglm (text-only) |
| 65 | ), |
| 66 | pytest.param( |
| 67 | "meta-llama/Llama-3.2-1B-Instruct", # llama |
| 68 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 69 | ), |
| 70 | pytest.param( |
| 71 | "openbmb/MiniCPM4.1-8B", # minicpm |
| 72 | marks=[pytest.mark.core_model, large_gpu_mark(min_gb=48)], |
| 73 | ), |
| 74 | pytest.param( |
| 75 | "facebook/opt-125m", # opt |
| 76 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 77 | ), |
| 78 | pytest.param( |
| 79 | "microsoft/phi-2", # phi |
| 80 | marks=[pytest.mark.core_model, pytest.mark.slow_test], |
| 81 | ), |
| 82 | pytest.param( |
| 83 | "Qwen/Qwen2.5-0.5B-Instruct", # qwen2 |
| 84 | marks=[ |
| 85 | pytest.mark.core_model, |
| 86 | pytest.mark.cpu_model, |
| 87 | pytest.mark.slow_test, |
| 88 | ], |
| 89 | ), |
| 90 | pytest.param( |
| 91 | "Qwen/Qwen3-8B", # qwen (text-only) |
| 92 | ), |
| 93 | pytest.param("stabilityai/stablelm-3b-4e1t"), # stablelm |
| 94 | pytest.param("bigcode/starcoder2-3b"), # starcoder2 |
| 95 | pytest.param( |
| 96 | "axolotl-ai-co/tiny-mixtral-30m", # mixtral |
| 97 | marks=[pytest.mark.core_model], |
| 98 | ), |
| 99 | pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus |
| 100 | pytest.param( |
| 101 | "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", # hyperclovax |
| 102 | marks=[large_gpu_mark(min_gb=32)], |
| 103 | ), |
| 104 | ], |
| 105 | ) |
| 106 | @pytest.mark.parametrize("max_tokens", [32]) |
| 107 | @pytest.mark.parametrize("num_logprobs", [5]) |
| 108 | @pytest.mark.parametrize( |
| 109 | "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 110 | ) |
| 111 | @pytest.mark.parametrize("use_prompt_embeds", [True, False]) |
| 112 | def test_models( |
| 113 | hf_runner, |
| 114 | vllm_runner, |
| 115 | example_prompts, |
| 116 | model: str, |
| 117 | max_tokens: int, |
| 118 | num_logprobs: int, |
| 119 | use_rocm_aiter: bool, |
| 120 | use_prompt_embeds: bool, |
| 121 | monkeypatch, |
| 122 | ) -> None: |
| 123 | model_info = HF_EXAMPLE_MODELS.find_hf_info(model) |
| 124 | model_info.check_available_online(on_fail="skip") |
| 125 | model_info.check_transformers_version(on_fail="skip") |
| 126 | |
| 127 | if use_rocm_aiter and (model in AITER_MODEL_LIST): |
| 128 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 129 | elif use_rocm_aiter and model not in AITER_MODEL_LIST: |
| 130 | # Skip model that are not using AITER tests. |
| 131 | # When more AITER kernels are added, this list will not be |
| 132 | # needed as all the models will be calling AITER kernels |
| 133 | # in parts of the operators |
| 134 | pytest.skip(f"Skipping '{model}' model test with AITER kernel.") |
| 135 | |
| 136 | if model == "bigcode/starcoder2-3b": |
| 137 | # Replace example.txt's Test1 (an NL prompt) with a code prompt: |
| 138 | # starcoder2-3b is a code model, so NL prompts give near-uniform |
| 139 | # digit logits where HF<->vLLM bf16 drift can reorder top-K. |
| 140 | example_prompts = list(example_prompts) |
| 141 | example_prompts[1] = ( |
| 142 | "def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - " |
| 143 | ) |
| 144 | |
| 145 | with hf_runner( |
| 146 | model, |
| 147 | revision=model_info.revision, |
| 148 | trust_remote_code=model_info.trust_remote_code, |
| 149 | ) as hf_model: |
| 150 | hf_outputs = hf_model.generate_greedy_logprobs_limit( |
| 151 | example_prompts, max_tokens, num_logprobs |
| 152 | ) |
| 153 | |
| 154 | prompt_embeds: list[torch.Tensor] | None = [] if use_prompt_embeds else None |
| 155 | |
| 156 | for prompt in example_prompts: |
| 157 | token_ids = hf_model.tokenizer(prompt, return_tensors="pt").input_ids.to( |
| 158 | hf_model.model.device |
| 159 | ) |
| 160 | if prompt_embeds is not None: |
| 161 | # Retained prompt values must not keep reference weights alive. |
| 162 | with torch.no_grad(): |
| 163 | embed = hf_model.model.get_input_embeddings()(token_ids) |
| 164 | |
| 165 | # MiniCPM models apply scale_emb to embeddings internally. |
| 166 | # vLLM expects pre-scaled embeddings when using inputs_embeds. |
| 167 | if model in EMBED_SCALING_MODELS: |
| 168 | config = hf_model.model.config |
| 169 | embed = embed * config.scale_emb |
| 170 | |
| 171 | prompt_embeds.append(embed.squeeze(0)) |
| 172 | |
| 173 | vllm_kwargs = {} |
| 174 | if ( |
| 175 | model == "bigscience/bloom-560m" |
| 176 | and current_platform.is_device_capability_family(90) |
| 177 | ): |
| 178 | # On SM90, the metadata builder otherwise selects FA3 AOT scheduling |
| 179 | # before Bloom's ALiBi layers fall back to FA2. Pinning FA2 keeps the |
| 180 | # builder and layer consistent and preserves the L4 test path. |
| 181 | vllm_kwargs["attention_config"] = {"flash_attn_version": 2} |
| 182 | |
| 183 | with vllm_runner( |
| 184 | model, |
| 185 | tokenizer_name=model_info.tokenizer or model, |
| 186 | tokenizer_mode=model_info.tokenizer_mode, |
| 187 | revision=model_info.revision, |
| 188 | trust_remote_code=model_info.trust_remote_code, |
| 189 | # Remove the effects of batch variance on ROCm since batch invariance |
| 190 | # is not yet supported. |
| 191 | # See: https://github.com/vllm-project/vllm/issues/27433 |
| 192 | max_num_seqs=1 if current_platform.is_rocm() else 2, |
| 193 | enable_prompt_embeds=use_prompt_embeds, |
| 194 | compilation_config={"cudagraph_capture_sizes": [1, 2]}, |
| 195 | **vllm_kwargs, |
| 196 | ) as vllm_model: |
| 197 | vllm_outputs = vllm_model.generate_greedy_logprobs( |
| 198 | example_prompts, max_tokens, num_logprobs |
| 199 | ) |
| 200 | if prompt_embeds is not None: |
| 201 | vllm_outputs_from_embeds = vllm_model.generate_greedy_logprobs( |
| 202 | prompt_embeds, max_tokens, num_logprobs |
| 203 | ) |
| 204 | |
| 205 | check_logprobs_close( |
| 206 | outputs_0_lst=hf_outputs, |
| 207 | outputs_1_lst=vllm_outputs, |
| 208 | name_0="hf", |
| 209 | name_1="vllm", |
| 210 | ) |
| 211 | if prompt_embeds is not None: |
| 212 | check_logprobs_close( |
| 213 | outputs_0_lst=vllm_outputs, |
| 214 | outputs_1_lst=vllm_outputs_from_embeds, |
| 215 | name_0="vllm", |
| 216 | name_1="vllm_from_embeds", |
| 217 | ) |
| 218 | |
| 219 | if use_rocm_aiter: |
| 220 | # this is to ensure that vllm engine |
| 221 | # has deallocated the memory before running the next |
| 222 | # unit tests. On ROCm, when using AITER |
| 223 | # the memory might not be deallocated completely |
| 224 | # before running the next test case |
| 225 | torch.accelerator.synchronize() |