| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import gc |
| 5 | from collections.abc import Generator |
| 6 | from functools import partial |
| 7 | from unittest.mock import patch |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from vllm import LLM |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.mem_constants import GiB_bytes |
| 16 | from vllm.v1.attention.backends.utils import resolve_kv_cache_layout |
| 17 | from vllm.v1.core.kv_cache_utils import ( |
| 18 | generate_scheduler_kv_cache_config, |
| 19 | get_kv_cache_configs, |
| 20 | ) |
| 21 | from vllm.v1.engine.core import EngineCore as V1EngineCore |
| 22 | |
| 23 | from ..utils import ( |
| 24 | create_new_process_for_each_test, |
| 25 | requires_spawn_multiprocessing, |
| 26 | wait_for_gpu_memory_to_clear, |
| 27 | wait_for_memory_to_settle, |
| 28 | ) |
| 29 | from .registry import ( |
| 30 | _TRANSFORMERS_BACKEND_MODELS, |
| 31 | AUTO_EXAMPLE_MODELS, |
| 32 | HF_EXAMPLE_MODELS, |
| 33 | HfExampleModels, |
| 34 | ) |
| 35 | from .utils import dummy_hf_overrides |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | # This minimal list of model architectures is smaller than the total list of |
| 40 | # supported models. The intention is that in the "typical" regression testing |
| 41 | # scenario, we only test initializing these models. This subset was chosen |
| 42 | # to include representative examples of model varieties/workloads (conditional |
| 43 | # generation, sequence classification, causal LM, ranking, chat, reward model, |
| 44 | # multimodal, geospatial, voice, embedding, MTP) |
| 45 | MINIMAL_MODEL_ARCH_LIST = [ |
| 46 | "LlavaForConditionalGeneration", |
| 47 | "Llama4ForConditionalGeneration", |
| 48 | "BertForSequenceClassification", |
| 49 | "Gemma3nForCausalLM", |
| 50 | "JinaVLForRanking", |
| 51 | "InternVLChatModel", |
| 52 | "InternLM2ForRewardModel", |
| 53 | "TransformersMultiModalForCausalLM", |
| 54 | "Terratorch", |
| 55 | "UltravoxModel", |
| 56 | "DeepSeekMTPModel", |
| 57 | "XLMRobertaModel", |
| 58 | ] |
| 59 | |
| 60 | # This list is the complement of the minimal list above. The intention is that |
| 61 | # this list of models is only tested in a "special case" i.e. most PRs should |
| 62 | # not test these models |
| 63 | OTHER_MODEL_ARCH_LIST = set(HF_EXAMPLE_MODELS.get_supported_archs()) - set( |
| 64 | MINIMAL_MODEL_ARCH_LIST |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | @pytest.fixture(autouse=True) |
| 69 | def cleanup_gpu_memory_between_tests() -> Generator[None, None, None]: |
| 70 | """Reclaim GPU memory left over by each test's forked subprocess. |
| 71 | |
| 72 | These tests are known to leave behind GPU memory (see the comment at |
| 73 | ``gpu_memory_utilization=0.80`` below). Each parametrized case builds an |
| 74 | LLM in a forked subprocess whose VRAM is only reclaimed asynchronously |
| 75 | when the subprocess exits, so the leftovers can OOM a later test in this |
| 76 | module (observed with the heavier multimodal archs, e.g. |
| 77 | InternVLChatModel). Mirrors ``clean_gpu_memory_between_tests`` in the |
| 78 | top-level conftest, but unconditionally for this module. |
| 79 | """ |
| 80 | yield |
| 81 | gc.collect() |
| 82 | if not torch.cuda.is_available(): |
| 83 | return |
| 84 | torch.accelerator.empty_cache() |
| 85 | if current_platform.is_rocm() or current_platform.is_xpu(): |
| 86 | wait_for_memory_to_settle() |
| 87 | return |
| 88 | try: |
| 89 | wait_for_gpu_memory_to_clear( |
| 90 | devices=list(range(torch.accelerator.device_count())), |
| 91 | threshold_ratio=0.1, |
| 92 | ) |
| 93 | except Exception as e: |
| 94 | # NVML is not usable on MIG slices (raises NVMLError_NoPermission); |
| 95 | # cleanup is best-effort, so never fail a test over it. |
| 96 | logger.info("Failed to clean GPU memory: %s", e) |
| 97 | |
| 98 | |
| 99 | @create_new_process_for_each_test() |
| 100 | def can_initialize( |
| 101 | model_arch: str, monkeypatch: pytest.MonkeyPatch, EXAMPLE_MODELS: HfExampleModels |
| 102 | ): |
| 103 | """The reason for using create_new_process_for_each_test is to avoid |
| 104 | the WARNING: |
| 105 | "We must use the 'spawn' multiprocessing start method. Overriding |
| 106 | VLLM_WORKER_MULTIPROC_METHOD to 'spawn'." |
| 107 | The spawn process causes the _initialize_kv_caches_v1 function below to |
| 108 | become ineffective. |
| 109 | """ |
| 110 | model_info = EXAMPLE_MODELS.get_hf_info(model_arch) |
| 111 | model_info.check_available_online(on_fail="skip") |
| 112 | model_info.check_transformers_version( |
| 113 | on_fail="skip", |
| 114 | check_max_version=False, |
| 115 | check_version_reason="vllm", |
| 116 | ) |
| 117 | |
| 118 | hf_overrides_fn = partial( |
| 119 | dummy_hf_overrides, |
| 120 | model_arch=model_arch, |
| 121 | exist_overrides=model_info.hf_overrides, |
| 122 | use_original_num_layers=getattr(model_info, "use_original_num_layers", False), |
| 123 | ) |
| 124 | |
| 125 | # Avoid calling model.forward() |
| 126 | def _initialize_kv_caches_v1(self, vllm_config): |
| 127 | kv_cache_specs = self.model_executor.get_kv_cache_specs() |
| 128 | layout = resolve_kv_cache_layout( |
| 129 | vllm_config, |
| 130 | self.model_executor.get_supported_kv_cache_layouts(), |
| 131 | [spec for worker_specs in kv_cache_specs for spec in worker_specs.values()], |
| 132 | ) |
| 133 | self.model_executor.set_kv_cache_layout(layout.name) |
| 134 | kv_cache_configs = get_kv_cache_configs( |
| 135 | vllm_config, |
| 136 | kv_cache_specs, |
| 137 | [10 * GiB_bytes], |
| 138 | ) |
| 139 | scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) |
| 140 | vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks |
| 141 | kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups |
| 142 | if kv_cache_groups: |
| 143 | vllm_config.cache_config.block_size = min( |
| 144 | g.kv_cache_spec.block_size for g in kv_cache_groups |
| 145 | ) |
| 146 | |
| 147 | vllm_config.validate_block_size() |
| 148 | return scheduler_kv_cache_config |
| 149 | |
| 150 | if model_arch == "MoonshotKimiaForCausalLM": |
| 151 | pytest.skip( |
| 152 | "Kimi-Audio requires SpeechToTextConfig " |
| 153 | "which is not configured in test environment" |
| 154 | ) |
| 155 | |
| 156 | if model_arch == "Terratorch": |
| 157 | import importlib.util |
| 158 | |
| 159 | if importlib.util.find_spec("terratorch") is None: |
| 160 | pytest.skip( |
| 161 | "terratorch is not installed; " |
| 162 | "temporarily skipped while PyPI has `lightning` quarantined " |
| 163 | "(see #41376)" |
| 164 | ) |
| 165 | |
| 166 | if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]: |
| 167 | from vllm.platforms import current_platform |
| 168 | |
| 169 | capability = current_platform.get_device_capability() |
| 170 | if capability and capability.major < 9: |
| 171 | pytest.skip( |
| 172 | f"DeepseekV32 requires Hopper (9.0+) or Blackwell (10.0+) " |
| 173 | f"for FLASHMLA_SPARSE backend. Current device has compute " |
| 174 | f"capability {capability.major}.{capability.minor}" |
| 175 | ) |
| 176 | |
| 177 | if model_arch == "DeepseekV4ForConditionalGeneration": |
| 178 | from vllm.platforms import current_platform |
| 179 | |
| 180 | if not (current_platform.is_cuda() or current_platform.is_rocm()): |
| 181 | pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm") |
| 182 | |
| 183 | with ( |
| 184 | patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1), |
| 185 | monkeypatch.context() as m, |
| 186 | ): |
| 187 | if requires_spawn_multiprocessing() or model_arch == "Glm5NextForCausalLM": |
| 188 | # The EngineCore subprocess re-imports the class and does not |
| 189 | # inherit the KV-cache patch above, so it OOMs. Run in-process |
| 190 | # so the patch applies. |
| 191 | m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 192 | |
| 193 | # FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU |
| 194 | # has cc==8.9 which hasn't supported FA3 yet. Remove this hack when |
| 195 | # L4 supports FA3. |
| 196 | # Step1ForCausalLM requires TRITON_ATTN for use_alibi_sqrt support. |
| 197 | attention_config = ( |
| 198 | {"backend": "TRITON_ATTN"} |
| 199 | if model_arch in ("GptOssForCausalLM", "Step1ForCausalLM") |
| 200 | else None |
| 201 | ) |
| 202 | if model_arch == "WhisperForConditionalGeneration": |
| 203 | m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 204 | |
| 205 | kwargs = {} |
| 206 | if not model_info.enable_prefix_caching: |
| 207 | kwargs["enable_prefix_caching"] = False |
| 208 | |
| 209 | LLM( |
| 210 | model_info.default, |
| 211 | tokenizer=model_info.tokenizer, |
| 212 | tokenizer_mode=model_info.tokenizer_mode, |
| 213 | revision=model_info.revision, |
| 214 | enforce_eager=model_info.enforce_eager, |
| 215 | skip_tokenizer_init=model_info.require_embed_inputs, |
| 216 | enable_prompt_embeds=model_info.require_embed_inputs, |
| 217 | enable_mm_embeds=model_info.require_embed_inputs, |
| 218 | dtype=model_info.dtype, |
| 219 | speculative_config={ |
| 220 | "model": model_info.speculative_model, |
| 221 | "method": model_info.speculative_method, |
| 222 | "num_speculative_tokens": 1, |
| 223 | } |
| 224 | if model_info.speculative_model |
| 225 | else None, |
| 226 | trust_remote_code=model_info.trust_remote_code, |
| 227 | max_model_len=model_info.max_model_len, |
| 228 | max_num_batched_tokens=model_info.max_num_batched_tokens, |
| 229 | # these tests seem to produce leftover memory |
| 230 | gpu_memory_utilization=0.80, |
| 231 | load_format="dummy", |
| 232 | model_impl="transformers" |
| 233 | if model_arch in _TRANSFORMERS_BACKEND_MODELS |
| 234 | else "vllm", |
| 235 | hf_overrides=hf_overrides_fn, |
| 236 | max_num_seqs=model_info.max_num_seqs, |
| 237 | attention_config=attention_config, |
| 238 | **kwargs, |
| 239 | ) |
| 240 | |
| 241 | |
| 242 | @pytest.mark.parametrize("model_arch", MINIMAL_MODEL_ARCH_LIST) |
| 243 | def test_can_initialize_small_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch): |
| 244 | """Test initializing small subset of supported models.""" |
| 245 | can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS) |
| 246 | |
| 247 | |
| 248 | @pytest.mark.parametrize("model_arch", OTHER_MODEL_ARCH_LIST) |
| 249 | def test_can_initialize_large_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch): |
| 250 | """Test initializing large subset of supported models. |
| 251 | |
| 252 | This test covers the complement of the tests covered in the "small subset" |
| 253 | test. |
| 254 | """ |
| 255 | if model_arch in ("HYV4ForCausalLM", "HYV4MTPModel"): |
| 256 | from vllm.platforms import current_platform |
| 257 | |
| 258 | if current_platform.is_rocm(): |
| 259 | pytest.skip("HY V4 ROCm initialization requires #54405") |
| 260 | |
| 261 | can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS) |
| 262 | |
| 263 | |
| 264 | @pytest.mark.parametrize("model_arch", AUTO_EXAMPLE_MODELS.get_supported_archs()) |
| 265 | def test_implicit_converted_models(model_arch: str, monkeypatch: pytest.MonkeyPatch): |
| 266 | can_initialize(model_arch, monkeypatch, AUTO_EXAMPLE_MODELS) |