test_hybrid.py 10.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

Mor Zusman's avatar
Mor Zusman committed
4
5
import pytest

6
from tests.models.registry import HF_EXAMPLE_MODELS
7
from tests.utils import multi_gpu_test
8
from vllm.engine.arg_utils import EngineArgs
9
from vllm.sampling_params import SamplingParams
10

11
from ...utils import check_logprobs_close
12

13
14
15
# Mark all tests as hybrid
pytestmark = pytest.mark.hybrid_model

16
17
18
19
20
21
22
# NOTE: The first model in each list is taken as the primary model,
# meaning that it will be used in all tests in this file
# The rest of the models will only be tested by test_models

SSM_MODELS = [
    "state-spaces/mamba-130m-hf",
    "tiiuae/falcon-mamba-tiny-dev",
23
24
25
    # mamba2-codestral in transformers is broken pending:
    # https://github.com/huggingface/transformers/pull/40861
    #"yujiepan/mamba2-codestral-v0.1-tiny-random",
26
]
27

28
29
HYBRID_MODELS = [
    "ai21labs/Jamba-tiny-dev",
30
    "pfnet/plamo-2-1b",
31
    "Zyphra/Zamba2-1.2B-instruct",
32
    "hmellor/tiny-random-BambaForCausalLM",
33
34
    "ibm-granite/granite-4.0-tiny-preview",
    "tiiuae/Falcon-H1-0.5B-Base",
35
    "LiquidAI/LFM2-1.2B",
36
    "tiny-random/qwen3-next-moe",
Chen Zhang's avatar
Chen Zhang committed
37
38
]

39
FULL_CUDA_GRAPH_MODELS = [
40
    "ai21labs/Jamba-tiny-dev",
41
    "pfnet/plamo-2-1b",
42
    "Zyphra/Zamba2-1.2B-instruct",
43
44
]

45
46
47
48
49
FP32_STATE_MODELS = [
    "state-spaces/mamba-130m-hf",
    "Zyphra/Zamba2-1.2B-instruct",
]

50
51
52
# Avoid OOM
MAX_NUM_SEQS = 4

Mor Zusman's avatar
Mor Zusman committed
53

54
55
56
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
Mor Zusman's avatar
Mor Zusman committed
57
58
59
60
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
61
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
62
63
    model: str,
    max_tokens: int,
64
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
65
) -> None:
66
67
68
69

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
70
        model_info.check_transformers_version(on_fail="skip")
71
    except ValueError:
72
        pass
73

74
    with hf_runner(model) as hf_model:
75
76
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
            example_prompts, max_tokens, num_logprobs)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
77

78
79
80
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        vllm_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
Chen Zhang's avatar
Chen Zhang committed
81

82
83
84
85
86
87
    check_logprobs_close(
        outputs_0_lst=hf_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="hf",
        name_1="vllm",
    )
Mor Zusman's avatar
Mor Zusman committed
88
89


90
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
91
92
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
93
94
95
96
97
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
98
    num_logprobs: int,
99
) -> None:
100
101
102
103
104
105
106
    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
        model_info.check_transformers_version(on_fail="skip")
    except ValueError:
        pass

107
    for_loop_outputs = []
108
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
109
        for prompt in example_prompts:
110
111
112
113
            single_output, = vllm_model.generate_greedy_logprobs([prompt],
                                                                 max_tokens,
                                                                 num_logprobs)
            for_loop_outputs.append(single_output)
114

115
116
        batched_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
117

118
    check_logprobs_close(
119
120
121
122
123
124
125
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


126
127
128
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
def test_chunked_prefill_with_parallel_sampling(
129
130
131
132
133
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
    """
    Tests chunked prefill in conjunction with n > 1. 
    
    In this case, prefill is populated with decoding tokens and
    we test that it doesn't fail.

    This test might fail if cache is not allocated correctly for n > 1
    decoding steps inside a chunked prefill forward pass
    (where we have both prefill and decode together)
    """
    sampling_params = SamplingParams(n=3,
                                     temperature=1,
                                     seed=0,
                                     max_tokens=max_tokens)
    with vllm_runner(
            model,
            enable_chunked_prefill=True,
            # forces prefill chunks with decoding
            max_num_batched_tokens=MAX_NUM_SEQS * 3,
            max_num_seqs=MAX_NUM_SEQS,
    ) as vllm_model:
        vllm_model.generate(example_prompts, sampling_params)
156
157


158
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
159
160
161
162
163
164
165
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
166
167
168
169
170
    """
    This test is for verifying that mamba cache is padded to CG captured
    batch size. If it's not, a torch RuntimeError will be raised because
    tensor dimensions aren't compatible.
    """
Shinichi Hemmi's avatar
Shinichi Hemmi committed
171
172
    vllm_config = EngineArgs(model=model,
                             trust_remote_code=True).create_engine_config()
173
    while len(example_prompts) == vllm_config.pad_for_cudagraph(
174
            len(example_prompts)):
175
176
177
        example_prompts.append(example_prompts[0])

    try:
178
        with vllm_runner(model) as vllm_model:
179
180
181
182
183
184
185
186
            vllm_model.generate_greedy(example_prompts, max_tokens)
    except RuntimeError:
        pytest.fail(
            "Couldn't run batch size which is not equal to a Cuda Graph "
            "captured batch size. "
            "Could be related to mamba cache not padded correctly")


187
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
188
189
190
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
191
    model: str,
192
) -> None:
193
194
195
196
197
198
    """
    This test is for verifying that the hybrid inner state management doesn't
    collapse in case where the number of incoming requests and
    finished_requests_ids is larger than the maximum mamba block capacity.

    This could generally happen due to the fact that hybrid does support
199
    statelessness mechanism where it can clean up new incoming requests in
200
201
    a single step.
    """
202
    try:
203
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
204
205
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
206
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
207
208
209
                    "steps finished requests registered unnecessarily ")


210
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
211
212
213
def test_state_cleanup(
    vllm_runner,
    example_prompts,
214
    model: str,
Mor Zusman's avatar
Mor Zusman committed
215
) -> None:
216
217
218
219
    """ 
    This test is for verifying that the Hybrid state is cleaned up between
    steps.
    
220
    If it's not cleaned, an error would be expected.
221
    """
Mor Zusman's avatar
Mor Zusman committed
222
    try:
223
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
224
225
226
            for _ in range(10):
                vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
227
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
228
229
230
                    "could be related to finished_requests_ids")


231
@multi_gpu_test(num_gpus=2)
232
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
233
@pytest.mark.parametrize("max_tokens", [64])
234
235
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
236
237
238
239
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
240
    num_logprobs: int,
241
) -> None:
242
    with vllm_runner(model, tensor_parallel_size=1,
243
                     max_num_seqs=MAX_NUM_SEQS) as vllm_model:
244
245
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
246

247
    with vllm_runner(model, tensor_parallel_size=2,
248
                     max_num_seqs=MAX_NUM_SEQS) as vllm_model:
249
250
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
251

252
    check_logprobs_close(
253
254
255
256
257
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )
258
259


260
@pytest.mark.parametrize("model", FULL_CUDA_GRAPH_MODELS)
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
def test_full_cuda_graph(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    num_logprobs: int,
) -> None:

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
        model_info.check_transformers_version(on_fail="skip")
    except ValueError:
        pass

    with hf_runner(model) as hf_model:
281
282
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
            example_prompts, max_tokens, num_logprobs)
283

284
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
285
        vllm_outputs = vllm_model.generate_greedy_logprobs(
286
            example_prompts, max_tokens, num_logprobs)
287
288

    check_logprobs_close(
289
        outputs_0_lst=hf_outputs,
290
        outputs_1_lst=vllm_outputs,
291
        name_0="hf",
292
        name_1="vllm",
293
    )
294
295


296
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
297
298
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
299
300
301
@pytest.mark.parametrize("cache_dtype_param",
                         ["mamba_ssm_cache_dtype", "mamba_cache_dtype"])
def test_fp32_cache_state(
302
303
304
305
306
307
308
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    num_logprobs: int,
309
    cache_dtype_param: str,
310
311
312
313
314
315
316
317
318
319
) -> None:

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
        model_info.check_transformers_version(on_fail="skip")
    except ValueError:
        pass

    with hf_runner(model) as hf_model:
320
321
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
            example_prompts, max_tokens, num_logprobs)
322

323
324
    with vllm_runner(model,
                     max_num_seqs=MAX_NUM_SEQS,
325
                     **{cache_dtype_param: "float32"}) as vllm_model:
326
        vllm_outputs = vllm_model.generate_greedy_logprobs(
327
328
            example_prompts, max_tokens, num_logprobs)

329
    check_logprobs_close(
330
        outputs_0_lst=hf_outputs,
331
        outputs_1_lst=vllm_outputs,
332
        name_0="hf",
333
        name_1="vllm",
334
    )