@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 110 lines 4.06 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Test whether spec decoding handles the max model length properly."""
4
5 import pytest
6
7 from tests.utils import get_attn_backend_list_based_on_platform
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
12
13 _PROMPTS = [
14 "1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1",
15 "Repeat the following sentence 10 times: Consistency is key to mastering any skill.", # noqa: E501
16 "Who won the Turing Award in 2018, and for what contribution? Describe in detail.", # noqa: E501
17 ]
18
19
20 @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
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",
30 "prompt_lookup_max": 5,
31 "prompt_lookup_min": 3,
32 "num_speculative_tokens": num_speculative_tokens,
33 },
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(
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
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",
57 "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
58 "num_speculative_tokens": num_speculative_tokens,
59 "max_model_len": 80,
60 },
61 max_model_len=200,
62 attention_config={"backend": attn_backend},
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
72 sampling_params = SamplingParams(
73 max_tokens=200,
74 structured_outputs=StructuredOutputsParams(
75 regex="^" + "a b c d e " * 15 + "$"
76 ),
77 )
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])
94 def test_mtp_speculative_config_max_model_len(spec_max_model_len: int):
95 """Regression test for #41456: max_model_len in speculative config
96 should be respected for the draft model."""
97 model_config = ModelConfig(
98 model="XiaomiMiMo/MiMo-7B-Base",
99 runner="generate",
100 max_model_len=200,
101 trust_remote_code=True,
102 )
103 spec_config = SpeculativeConfig(
104 target_model_config=model_config,
105 target_parallel_config=ParallelConfig(),
106 method="mtp",
107 num_speculative_tokens=1,
108 max_model_len=spec_max_model_len,
109 )
110 assert spec_config.draft_model_config.max_model_len == spec_max_model_len