A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 15:00 UTC
[Test] Fix test_dynamic_shapes_compilation for torch 2.12 (#40743)
Signed-off-by: Angela Yi <angelayi@meta.com>
Angela Yi committed
Apr 27, 2026 at 17:51 UTC
03aeed802f374c0319ad9eca34fae8e7e784769a
1 file changed
+26
-19
tests/compile/test_dynamic_shapes_compilation.py
+26
-19
@@ -8,6 +8,7 @@ from contextlib import contextmanager
8
import pytest
9
import torch
10
11
+from tests.models.utils import check_logprobs_close
12
from vllm import LLM, SamplingParams
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
@@ -17,7 +18,6 @@ from vllm.config.compilation import (
18
DynamicShapesType,
19
)
20
from vllm.forward_context import set_forward_context
20
-from vllm.tokenizers import get_tokenizer
21
from vllm.utils.torch_utils import is_torch_equal_or_newer
22
23
@@ -84,30 +84,37 @@ def test_dynamic_shapes_compilation(
84
max_model_len=1024,
85
)
86
87
- output = model.generate(prompt)
88
- result = output[0].outputs[0].text
89
- # Example of setting the sampling parameters
90
- tokenizer = get_tokenizer(model_name)
91
- yes_tokens = tokenizer.encode("yes", add_special_tokens=False)
92
- no_tokens = tokenizer.encode("no", add_special_tokens=False)
93
- allowed_ids = list(set(yes_tokens + no_tokens))
94
- sampling_params = SamplingParams(
95
- max_tokens=1, temperature=0, allowed_token_ids=allowed_ids
96
- )
87
+ sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10)
88
+ test_prompts = [prompt, "The capital of France is"]
89
98
- output = model.generate(
99
- "answer with yes or no is " + result + " rubbish for prompt " + prompt + "?",
100
- sampling_params=sampling_params,
101
- )
102
- result = output[0].outputs[0].text
103
- assert result == "yes"
90
+ compiled_outputs = []
91
+ for p in test_prompts:
92
+ output = model.generate(p, sampling_params)[0].outputs[0]
93
+ assert len(output.text.strip()) > 0, "Compiled model produced empty output"
94
+ compiled_outputs.append((output.token_ids, output.text, output.logprobs))
95
105
- # Clean up GPU memory
96
del model
97
gc.collect()
98
torch.accelerator.empty_cache()
99
torch.accelerator.synchronize()
110
- print("GPU memory cleared")
100
+
101
+ eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024)
102
+ eager_outputs = []
103
+ for p in test_prompts:
104
+ output = eager_model.generate(p, sampling_params)[0].outputs[0]
105
+ assert len(output.text.strip()) > 0, "Eager model produced empty output"
106
+ eager_outputs.append((output.token_ids, output.text, output.logprobs))
107
+ del eager_model
108
+ gc.collect()
109
+ torch.accelerator.empty_cache()
110
+ torch.accelerator.synchronize()
111
+
112
+ check_logprobs_close(
113
+ outputs_0_lst=eager_outputs,
114
+ outputs_1_lst=compiled_outputs,
115
+ name_0="eager",
116
+ name_1="compiled",
117
+ )
118
119
120
@pytest.mark.parametrize("use_aot_compile", ["0", "1"])