test_integration.py 5.68 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
5
6
"""Tests which cover integration of the speculative decoding framework with
other features, e.g. cuda graphs.
"""

import pytest
7
import os
8

9
from .conftest import run_equality_correctness_test
10
from ...utils import models_path_prefix
11

12
MAIN_MODEL = os.path.join(models_path_prefix, "JackFram/llama-68m")
13
14
15
16
17
18
19
20


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{

        # Verify equality when cuda graphs allowed.
        "enforce_eager": False,
21
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
22
23
24
25
26
27
    }])
@pytest.mark.parametrize(
    "per_test_common_llm_kwargs",
    [
        {
            # Identical models.
28
            "speculative_config": {
zhuwenwen's avatar
zhuwenwen committed
29
                "model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
30
31
                "num_speculative_tokens": 5,
            },
32
33
34
35
36
37
38
        },
    ])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize("test_llm_kwargs", [{}])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize("output_len", [32])
@pytest.mark.parametrize("seed", [1])
39
40
41
42
def test_spec_decode_cuda_graph(vllm_runner, common_llm_kwargs,
                                per_test_common_llm_kwargs,
                                baseline_llm_kwargs, test_llm_kwargs,
                                batch_size: int, output_len: int, seed: int):
43
44
    """Verify spec decode equality when cuda graphs are enabled.
    """
45
46
47
48
49
50
51
52
53
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  max_output_len=output_len,
                                  seed=seed,
                                  temperature=0.0)
54
55
56
57
58


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
59
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-160m"),
60
61
62
63

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
zhuwenwen's avatar
zhuwenwen committed
64

65
@pytest.mark.parametrize("per_test_common_llm_kwargs", [])
66
67
68
69
70
@pytest.mark.parametrize(
    "test_llm_kwargs",
    [
        # Explicitly specify draft model quantization
        {
71
            "speculative_config": {
zhuwenwen's avatar
zhuwenwen committed
72
                "model": os.path.join(models_path_prefix, "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit"),
73
74
75
                "num_speculative_tokens": 5,
                "quantization": "gptq",
            },
76
77
78
        },
        # Explicitly specify GPTQ-based draft model to use marlin quantization
        {
79
            "speculative_config": {
zhuwenwen's avatar
zhuwenwen committed
80
                "model": os.path.join(models_path_prefix, "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit"),
81
82
83
                "num_speculative_tokens": 5,
                "quantization": "marlin",
            },
84
85
86
        },
        # Not explicitly specify draft model quantization
        {
87
            "speculative_config": {
zhuwenwen's avatar
zhuwenwen committed
88
                "model": os.path.join(models_path_prefix, "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit"),
89
90
91
                "num_speculative_tokens": 5,
                "quantization": None,
            },
92
93
94
95
96
        },
    ])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize("batch_size", [2])
@pytest.mark.parametrize("seed", [1])
97
98
99
100
101
def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
                                               per_test_common_llm_kwargs,
                                               baseline_llm_kwargs,
                                               test_llm_kwargs,
                                               batch_size: int, seed: int):
102
103
    """Verify spec decode works well with draft model quantization configs.
    """
104
105
106
107
108
109
110
111
112
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  max_output_len=32,
                                  seed=seed,
                                  temperature=0.0)
113
114
115
116
117
118
119
120
121
122
123
124


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
        "model_name": MAIN_MODEL,

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
125
126
@pytest.mark.parametrize("test_llm_kwargs", [{
    "speculative_config": {
zhuwenwen's avatar
zhuwenwen committed
127
        "model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
128
129
130
131
        "num_speculative_tokens": 3,
        "disable_mqa_scorer": True,
    },
}])
132
133
134
135
136
137
138
139
140
141
142
@pytest.mark.parametrize("batch_size", [1, 5])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
                    baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
                    output_len: int, seed: int):
143
    """Verify that speculative decoding generates the same output
144
145
146
147
148
149
150
151
152
153
    with batch expansion scorer and mqa scorer.
    """
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  max_output_len=output_len,
                                  seed=seed,
zhuwenwen's avatar
zhuwenwen committed
154
                                  temperature=0.0)