| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import tempfile |
| 5 | from contextlib import contextmanager |
| 6 | from dataclasses import dataclass |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.models.utils import check_logprobs_close |
| 12 | from vllm import SamplingParams |
| 13 | from vllm.compilation.decorators import support_torch_compile |
| 14 | from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config |
| 15 | from vllm.config.compilation import ( |
| 16 | CompilationMode, |
| 17 | DynamicShapesConfig, |
| 18 | DynamicShapesType, |
| 19 | ) |
| 20 | from vllm.forward_context import set_forward_context |
| 21 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 22 | |
| 23 | |
| 24 | def get_test_models(): |
| 25 | """Get list of models to test based on PyTorch version.""" |
| 26 | models = [ |
| 27 | "Qwen/Qwen3-0.6B", |
| 28 | "openai-community/gpt2", |
| 29 | ] |
| 30 | return models |
| 31 | |
| 32 | |
| 33 | @dataclass(frozen=True) |
| 34 | class DynamicShapesTestCase: |
| 35 | model_name: str |
| 36 | shapes_type: DynamicShapesType |
| 37 | use_aot_compile: bool |
| 38 | use_bytecode_hook: bool |
| 39 | |
| 40 | def __str__(self) -> str: |
| 41 | model_id = self.model_name.rsplit("/", 1)[-1] |
| 42 | return ( |
| 43 | f"{model_id}-{self.shapes_type.value}-" |
| 44 | f"aot{int(self.use_aot_compile)}-hook{int(self.use_bytecode_hook)}" |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def get_dynamic_shapes_test_cases(): |
| 49 | """Pairwise-cover compilation options and smoke-test model classes.""" |
| 50 | reference_model = "Qwen/Qwen3-0.6B" |
| 51 | cases = [ |
| 52 | DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, False, True), |
| 53 | DynamicShapesTestCase(reference_model, DynamicShapesType.BACKED, True, False), |
| 54 | DynamicShapesTestCase( |
| 55 | reference_model, DynamicShapesType.UNBACKED, False, False |
| 56 | ), |
| 57 | DynamicShapesTestCase(reference_model, DynamicShapesType.UNBACKED, True, True), |
| 58 | DynamicShapesTestCase( |
| 59 | reference_model, |
| 60 | DynamicShapesType.BACKED_SIZE_OBLIVIOUS, |
| 61 | False, |
| 62 | True, |
| 63 | ), |
| 64 | DynamicShapesTestCase( |
| 65 | reference_model, |
| 66 | DynamicShapesType.BACKED_SIZE_OBLIVIOUS, |
| 67 | True, |
| 68 | False, |
| 69 | ), |
| 70 | ] |
| 71 | cases.extend( |
| 72 | DynamicShapesTestCase(model_name, DynamicShapesType.BACKED, False, True) |
| 73 | for model_name in get_test_models() |
| 74 | if model_name != reference_model |
| 75 | ) |
| 76 | return cases |
| 77 | |
| 78 | |
| 79 | TEST_PROMPTS = ["Hello, my name is", "The capital of France is"] |
| 80 | |
| 81 | |
| 82 | def generate_outputs(vllm_model): |
| 83 | sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) |
| 84 | outputs = [] |
| 85 | for prompt in TEST_PROMPTS: |
| 86 | output = vllm_model.llm.generate(prompt, sampling_params)[0].outputs[0] |
| 87 | assert len(output.text.strip()) > 0, "Model produced empty output" |
| 88 | outputs.append((output.token_ids, output.text, output.logprobs)) |
| 89 | return outputs |
| 90 | |
| 91 | |
| 92 | @pytest.fixture(scope="module") |
| 93 | def get_eager_outputs(vllm_runner): |
| 94 | cache = {} |
| 95 | |
| 96 | def get(model_name): |
| 97 | if model_name not in cache: |
| 98 | with vllm_runner( |
| 99 | model_name, |
| 100 | enforce_eager=True, |
| 101 | max_model_len=1024, |
| 102 | enable_chunked_prefill=None, |
| 103 | ) as vllm_model: |
| 104 | cache[model_name] = generate_outputs(vllm_model) |
| 105 | return cache[model_name] |
| 106 | |
| 107 | return get |
| 108 | |
| 109 | |
| 110 | @pytest.mark.parametrize("test_case", get_dynamic_shapes_test_cases(), ids=str) |
| 111 | @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") |
| 112 | def test_dynamic_shapes_compilation( |
| 113 | monkeypatch, |
| 114 | vllm_runner, |
| 115 | get_eager_outputs, |
| 116 | test_case, |
| 117 | ): |
| 118 | """Test representative dynamic-shapes configurations end to end.""" |
| 119 | if ( |
| 120 | test_case.shapes_type == DynamicShapesType.UNBACKED |
| 121 | and not is_torch_equal_or_newer("2.11.0") |
| 122 | ): |
| 123 | # NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only |
| 124 | # landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On |
| 125 | # older torch there's no way to express it, so unbacked shapes go |
| 126 | # data-dependent and compilation blows up -- nothing to test. |
| 127 | pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11") |
| 128 | |
| 129 | monkeypatch.setenv( |
| 130 | "VLLM_USE_AOT_COMPILE", "1" if test_case.use_aot_compile else "0" |
| 131 | ) |
| 132 | monkeypatch.setenv( |
| 133 | "VLLM_USE_BYTECODE_HOOK", "1" if test_case.use_bytecode_hook else "0" |
| 134 | ) |
| 135 | |
| 136 | print(f"Testing {test_case.shapes_type.name} dynamic shapes...") |
| 137 | |
| 138 | # The compiled engine shuts down before an uncached eager baseline starts. |
| 139 | with vllm_runner( |
| 140 | test_case.model_name, |
| 141 | compilation_config={ |
| 142 | "mode": CompilationMode.VLLM_COMPILE, |
| 143 | "dynamic_shapes_config": { |
| 144 | "type": test_case.shapes_type.value, |
| 145 | }, |
| 146 | }, |
| 147 | max_model_len=1024, |
| 148 | enable_chunked_prefill=None, |
| 149 | ) as vllm_model: |
| 150 | compiled_outputs = generate_outputs(vllm_model) |
| 151 | |
| 152 | check_logprobs_close( |
| 153 | outputs_0_lst=get_eager_outputs(test_case.model_name), |
| 154 | outputs_1_lst=compiled_outputs, |
| 155 | name_0="eager", |
| 156 | name_1="compiled", |
| 157 | ) |
| 158 | |
| 159 | |
| 160 | @pytest.mark.parametrize("use_aot_compile", [False, True]) |
| 161 | @pytest.mark.parametrize( |
| 162 | "dynamic_shapes_type", |
| 163 | [ |
| 164 | DynamicShapesType.BACKED, |
| 165 | DynamicShapesType.BACKED_SIZE_OBLIVIOUS, |
| 166 | ], |
| 167 | ) |
| 168 | @pytest.mark.parametrize("evaluate_guards", [False, True]) |
| 169 | def test_model_specialization_with_evaluate_guards( |
| 170 | monkeypatch, use_aot_compile, dynamic_shapes_type, evaluate_guards |
| 171 | ): |
| 172 | """Test that evaluate_guards correctly detects shape specialization |
| 173 | violations. |
| 174 | """ |
| 175 | if ( |
| 176 | use_aot_compile |
| 177 | and dynamic_shapes_type == DynamicShapesType.BACKED |
| 178 | and evaluate_guards |
| 179 | ): |
| 180 | pytest.skip("evaluate_guards for backed does not work with aot_compile=1") |
| 181 | |
| 182 | @support_torch_compile |
| 183 | class ModelWithSizeCheck(torch.nn.Module): |
| 184 | def __init__(self, **kwargs): |
| 185 | super().__init__() |
| 186 | |
| 187 | def forward(self, x: torch.Tensor): |
| 188 | # This will cause specialization - torch.compile will guard on |
| 189 | # sx.shape[0] |
| 190 | if x.shape[0] >= 10: |
| 191 | return x * 10 |
| 192 | else: |
| 193 | return x * 10 |
| 194 | |
| 195 | @support_torch_compile |
| 196 | class ModelWithOneSizeCheck(torch.nn.Module): |
| 197 | def __init__(self, **kwargs): |
| 198 | super().__init__() |
| 199 | |
| 200 | def forward(self, x: torch.Tensor): |
| 201 | # This will cause 0/1 specializations. |
| 202 | if x.shape[0] == 0: |
| 203 | return x * 10 |
| 204 | if x.shape[0] == 1: |
| 205 | return x * 10 |
| 206 | else: |
| 207 | return x * 10 |
| 208 | |
| 209 | @contextmanager |
| 210 | def use_vllm_config(vllm_config: VllmConfig): |
| 211 | with set_forward_context({}, vllm_config), set_current_vllm_config(vllm_config): |
| 212 | yield |
| 213 | |
| 214 | monkeypatch.setenv("TOKENIZERS_PARALLELISM", "true") |
| 215 | monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1" if use_aot_compile else "0") |
| 216 | monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "0") |
| 217 | |
| 218 | vllm_config = VllmConfig( |
| 219 | compilation_config=CompilationConfig( |
| 220 | mode=CompilationMode.VLLM_COMPILE, |
| 221 | dynamic_shapes_config=DynamicShapesConfig( |
| 222 | type=dynamic_shapes_type, |
| 223 | evaluate_guards=evaluate_guards, |
| 224 | ), |
| 225 | ) |
| 226 | ) |
| 227 | |
| 228 | def test(model_class, input1, input2, is_01_specialization=False): |
| 229 | with ( |
| 230 | torch.no_grad(), |
| 231 | use_vllm_config(vllm_config), |
| 232 | tempfile.TemporaryDirectory() as tmpdirname, |
| 233 | ): |
| 234 | monkeypatch.setenv("VLLM_CACHE_ROOT", tmpdirname) |
| 235 | |
| 236 | model = model_class(vllm_config=vllm_config).cuda() |
| 237 | |
| 238 | model(input1) |
| 239 | |
| 240 | if evaluate_guards and ( |
| 241 | not ( |
| 242 | is_01_specialization |
| 243 | and dynamic_shapes_type == DynamicShapesType.BACKED |
| 244 | ) |
| 245 | ): |
| 246 | # This should fail because guards were added. |
| 247 | with pytest.raises(RuntimeError) as excinfo: |
| 248 | model(input2) |
| 249 | |
| 250 | # Expected failure - guard was violated |
| 251 | error_msg = str(excinfo.value) |
| 252 | assert ( |
| 253 | "GuardManager check failed" in error_msg |
| 254 | or "Detected recompile when torch.compile stance" in error_msg |
| 255 | ), error_msg |
| 256 | |
| 257 | else: |
| 258 | model(input2) |
| 259 | |
| 260 | test(ModelWithSizeCheck, torch.randn(20, 10).cuda(), torch.randn(5, 10).cuda()) |
| 261 | test(ModelWithSizeCheck, torch.randn(5, 10).cuda(), torch.randn(20, 10).cuda()) |
| 262 | test( |
| 263 | ModelWithOneSizeCheck, |
| 264 | torch.randn(20, 10).cuda(), |
| 265 | torch.randn(1, 10).cuda(), |
| 266 | is_01_specialization=True, |
| 267 | ) |
| 268 | |
| 269 | |
| 270 | @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") |
| 271 | def test_piecewise_backend_empty_sym_shape_indices(vllm_runner): |
| 272 | """Test that PiecewiseBackend handles empty sym_shape_indices correctly. |
| 273 | |
| 274 | When all inputs have static shapes (no torch.SymInt), sym_shape_indices |
| 275 | will be empty. The fix in PiecewiseBackend.__call__ handles this case |
| 276 | by using the first compiled range_entry. |
| 277 | """ |
| 278 | # Use small max_model_len and max_num_batched_tokens to encourage |
| 279 | # static shape compilation with empty sym_shape_indices |
| 280 | with vllm_runner( |
| 281 | "Qwen/Qwen3-0.6B", |
| 282 | max_model_len=512, |
| 283 | max_num_batched_tokens=1, |
| 284 | enable_chunked_prefill=None, |
| 285 | compilation_config={ |
| 286 | "mode": CompilationMode.VLLM_COMPILE, |
| 287 | "dynamic_shapes_config": { |
| 288 | "type": DynamicShapesType.BACKED.value, |
| 289 | }, |
| 290 | }, |
| 291 | ) as vllm_model: |
| 292 | sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) |
| 293 | |
| 294 | # Generate with static shape inputs |
| 295 | output = vllm_model.llm.generate( |
| 296 | "Hello, my name is", sampling_params=sampling_params |
| 297 | ) |
| 298 | result = output[0].outputs[0].text |
| 299 | assert len(result) > 0, "Should generate non-empty output" |
| 300 | |
| 301 | # Generate again to verify compilation works with empty sym_shape_indices |
| 302 | output = vllm_model.llm.generate( |
| 303 | "The capital of France is", sampling_params=sampling_params |
| 304 | ) |
| 305 | result = output[0].outputs[0].text |
| 306 | assert len(result) > 0, "Should generate non-empty output on second run" |