[CI] Explicitly tear down speculative decode runners (#49910)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 27, 2026 at 00:32 UTC
5f89a03dcb52702a62644e15b93f766765d06b28
1 file changed
+43
-33
tests/v1/spec_decode/test_max_len.py
+43
-33
@@ -5,7 +5,7 @@
5
import pytest
6
7
from tests.utils import get_attn_backend_list_based_on_platform
8
-from vllm import LLM, SamplingParams
8
+from vllm import SamplingParams
9
from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
10
from vllm.platforms import current_platform
11
from vllm.sampling_params import StructuredOutputsParams
@@ -18,10 +18,12 @@ _PROMPTS = [
18
19
20
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
21
-def test_ngram_max_len(num_speculative_tokens: int):
22
- llm = LLM(
23
- model="facebook/opt-125m",
21
+def test_ngram_max_len(num_speculative_tokens: int, vllm_runner):
22
+ with vllm_runner(
23
+ "facebook/opt-125m",
24
+ trust_remote_code=False,
25
max_model_len=100,
26
+ enable_chunked_prefill=None,
27
enforce_eager=True, # For faster initialization.
28
speculative_config={
29
"method": "ngram",
@@ -29,21 +31,26 @@ def test_ngram_max_len(num_speculative_tokens: int):
31
"prompt_lookup_min": 3,
32
"num_speculative_tokens": num_speculative_tokens,
33
},
32
- )
33
- sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
34
- llm.generate(_PROMPTS, sampling_params)
34
+ ) as runner:
35
+ sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
36
+ runner.llm.generate(_PROMPTS, sampling_params)
37
38
39
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
40
@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
41
def test_eagle_max_len(
40
- monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str
42
+ monkeypatch: pytest.MonkeyPatch,
43
+ num_speculative_tokens: int,
44
+ attn_backend: str,
45
+ vllm_runner,
46
):
47
if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
48
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
49
45
- llm = LLM(
46
- model="meta-llama/Meta-Llama-3-8B-Instruct",
50
+ with vllm_runner(
51
+ "meta-llama/Meta-Llama-3-8B-Instruct",
52
+ trust_remote_code=False,
53
+ enable_chunked_prefill=None,
54
enforce_eager=True, # For faster initialization.
55
speculative_config={
56
"method": "eagle",
@@ -53,31 +60,34 @@ def test_eagle_max_len(
60
},
61
max_model_len=200,
62
attention_config={"backend": attn_backend},
56
- )
57
- sampling_params = SamplingParams(max_tokens=200, ignore_eos=True)
58
- outputs = llm.generate(_PROMPTS, sampling_params)
59
- for o in outputs:
60
- assert o.outputs[0].finish_reason == "length", (
61
- "This test is only meaningful if the output is truncated due to max length"
62
- )
63
+ ) as runner:
64
+ sampling_params = SamplingParams(max_tokens=200, ignore_eos=True)
65
+ outputs = runner.llm.generate(_PROMPTS, sampling_params)
66
+ for o in outputs:
67
+ assert o.outputs[0].finish_reason == "length", (
68
+ "This test is only meaningful if the output is truncated "
69
+ "due to max length"
70
+ )
71
64
- sampling_params = SamplingParams(
65
- max_tokens=200,
66
- structured_outputs=StructuredOutputsParams(regex="^" + "a b c d e " * 15 + "$"),
67
- )
68
- output = llm.generate(_PROMPTS, sampling_params)
69
- for o in output:
70
- assert o.prompt_token_ids is not None
71
- assert (
72
- len(o.prompt_token_ids)
73
- < 80
74
- < len(o.prompt_token_ids) + len(o.outputs[0].token_ids)
75
- <= 200
76
- ), (
77
- "This test is only meaningful if the output "
78
- "is longer than the eagle max length"
72
+ sampling_params = SamplingParams(
73
+ max_tokens=200,
74
+ structured_outputs=StructuredOutputsParams(
75
+ regex="^" + "a b c d e " * 15 + "$"
76
+ ),
77
)
80
- assert o.outputs[0].text == "a b c d e " * 15
78
+ output = runner.llm.generate(_PROMPTS, sampling_params)
79
+ for o in output:
80
+ assert o.prompt_token_ids is not None
81
+ assert (
82
+ len(o.prompt_token_ids)
83
+ < 80
84
+ < len(o.prompt_token_ids) + len(o.outputs[0].token_ids)
85
+ <= 200
86
+ ), (
87
+ "This test is only meaningful if the output "
88
+ "is longer than the eagle max length"
89
+ )
90
+ assert o.outputs[0].text == "a b c d e " * 15
91
92
93
@pytest.mark.parametrize("spec_max_model_len", [80, 150])