test_engine_args.py 2.17 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
5
6
7
8
import pytest

from vllm import envs
from vllm.config import VllmConfig
from vllm.engine.arg_utils import EngineArgs
from vllm.usage.usage_lib import UsageContext
9
from vllm.utils import FlexibleArgumentParser
10
11
12
13
14
15
16
17

if not envs.VLLM_USE_V1:
    pytest.skip(
        "Skipping V1 tests. Rerun with `VLLM_USE_V1=1` to test.",
        allow_module_level=True,
    )


18
19
20
def test_prefix_caching_from_cli():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    args = parser.parse_args([])
21
22
    vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
    assert (vllm_config.cache_config.enable_prefix_caching
23
24
25
26
            ), "V1 turns on prefix caching by default."

    # Turn it off possible with flag.
    args = parser.parse_args(["--no-enable-prefix-caching"])
27
28
    vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
    assert not vllm_config.cache_config.enable_prefix_caching
29
30
31

    # Turn it on with flag.
    args = parser.parse_args(["--enable-prefix-caching"])
32
33
    vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
    assert vllm_config.cache_config.enable_prefix_caching
34
35


36
37
38
39
40
def test_defaults_with_usage_context():
    engine_args = EngineArgs(model="facebook/opt-125m")
    vllm_config: VllmConfig = engine_args.create_engine_config(
        UsageContext.LLM_CLASS)

41
42
43
44
45
46
47
48
49
50
    from vllm.platforms import current_platform
    device_name = current_platform.get_device_name().lower()
    if "h100" in device_name or "h200" in device_name:
        # For H100 and H200, we use larger default values.
        default_llm_tokens = 16384
        default_server_tokens = 8192
    else:
        default_llm_tokens = 8192
        default_server_tokens = 2048

51
    assert vllm_config.scheduler_config.max_num_seqs == 1024
52
    assert vllm_config.scheduler_config.max_num_batched_tokens == default_llm_tokens  # noqa: E501
53
54
55
56
57

    engine_args = EngineArgs(model="facebook/opt-125m")
    vllm_config = engine_args.create_engine_config(
        UsageContext.OPENAI_API_SERVER)
    assert vllm_config.scheduler_config.max_num_seqs == 1024
58
    assert vllm_config.scheduler_config.max_num_batched_tokens == default_server_tokens  # noqa: E501