test_hybrid.py 14.5 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
12
from ...utils import check_logprobs_close, check_outputs_equal

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
    "yujiepan/mamba2-codestral-v0.1-tiny-random",
24
]
25

26
27
HYBRID_MODELS = [
    "ai21labs/Jamba-tiny-dev",
28
29
    # skipping until vLLM implementation issues are resolved
    # "pfnet/plamo-2-1b",
30
    "Zyphra/Zamba2-1.2B-instruct",
31
    "hmellor/tiny-random-BambaForCausalLM",
32
33
34
35
36
37
38
39
40
    "ibm-granite/granite-4.0-tiny-preview",
    "tiiuae/Falcon-H1-0.5B-Base",
]

HF_UNSUPPORTED_MODELS = [
    # The HF transformers implementation of
    # Mamba2 is buggy for Codestral as it doesn't handle n_groups, so the test
    # doesn't compare vLLM output with HF output.
    # See https://github.com/huggingface/transformers/pull/35943
41
42
43
44
    "yujiepan/mamba2-codestral-v0.1-tiny-random",
    # transformers 4.55 is still producing garbage for this model
    # TODO(tdoublep): follow-up on transformers side
    "ibm-granite/granite-4.0-tiny-preview"
Shinichi Hemmi's avatar
Shinichi Hemmi committed
45
]
46

Chen Zhang's avatar
Chen Zhang committed
47
V1_SUPPORTED_MODELS = [
48
49
    "state-spaces/mamba-130m-hf",
    "ai21labs/Jamba-tiny-dev",
50
    "yujiepan/mamba2-codestral-v0.1-tiny-random",
51
    "Zyphra/Zamba2-1.2B-instruct",
52
    "hmellor/tiny-random-BambaForCausalLM",
53
54
    "ibm-granite/granite-4.0-tiny-preview",
    "tiiuae/Falcon-H1-0.5B-Base",
Chen Zhang's avatar
Chen Zhang committed
55
56
]

57
58
# Avoid OOM
MAX_NUM_SEQS = 4
Mor Zusman's avatar
Mor Zusman committed
59

60
61
62
63
64
65
66
# Once we add support for FCG in Mamba1, this list will be removed and tests
# all test cases will use enforce_eager=False
ENFORCE_EAGER_MODELS_V1 = [
    "state-spaces/mamba-130m-hf",
    "ai21labs/Jamba-tiny-dev",
]

Mor Zusman's avatar
Mor Zusman committed
67

68
69
70
@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
71
72
73
74
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
75
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
76
77
    model: str,
    max_tokens: int,
78
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
79
) -> None:
80
81
82
83

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
84
85
        hf_version_check = model_info.check_transformers_version(
            on_fail="return")
86
    except ValueError:
87
88
89
90
        hf_version_check = None

    if hf_version_check is not None:
        print(f"Skipping transformers comparison because: {hf_version_check}")
91

92
    with hf_runner(model) as hf_model:
93
        if model not in HF_UNSUPPORTED_MODELS and hf_version_check is None:
Chen Zhang's avatar
Chen Zhang committed
94
95
96
97
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
98

99
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Chen Zhang's avatar
Chen Zhang committed
100
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
101
            example_prompts, max_tokens, num_logprobs)
102

Chen Zhang's avatar
Chen Zhang committed
103
    if model in V1_SUPPORTED_MODELS:
104
        enforce_eager = False
Chen Zhang's avatar
Chen Zhang committed
105
106
        with monkeypatch.context() as m:
            m.setenv("VLLM_USE_V1", "1")
107
108
109
            if model in HYBRID_MODELS:
                # required due to reorder_batch behaviour
                m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
110
111
112
113

            if model in ENFORCE_EAGER_MODELS_V1:
                enforce_eager = True

Chen Zhang's avatar
Chen Zhang committed
114
115
            with vllm_runner(model,
                             max_num_seqs=MAX_NUM_SEQS,
116
                             enforce_eager=enforce_eager,
117
                             enable_prefix_caching=False) as vllm_model:
Chen Zhang's avatar
Chen Zhang committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
                vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                    example_prompts, max_tokens, num_logprobs)
    else:
        vllm_v1_outputs = None

    if hf_outputs is not None:
        check_logprobs_close(
            outputs_0_lst=hf_outputs,
            outputs_1_lst=vllm_v0_outputs,
            name_0="hf",
            name_1="vllm-v0",
        )

    if model in V1_SUPPORTED_MODELS:
        ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
        check_logprobs_close(
            outputs_0_lst=ref_outputs,
            outputs_1_lst=vllm_v1_outputs,
            name_0="hf" if hf_outputs is not None else "vllm-v0",
            name_1="vllm-v1",
        )
Mor Zusman's avatar
Mor Zusman committed
139
140


141
142
143
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
144
145
146
147
148
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
149
    num_logprobs: int,
150
) -> None:
151
152
153
154
155
156
157
158

    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

159
    for_loop_outputs = []
160
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
161
        for prompt in example_prompts:
162
163
164
165
            single_output, = vllm_model.generate_greedy_logprobs([prompt],
                                                                 max_tokens,
                                                                 num_logprobs)
            for_loop_outputs.append(single_output)
166

167
168
        batched_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
169

170
    check_logprobs_close(
171
172
173
174
175
176
177
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


178
179
180
181
182
183
184
185
186
187
188
189
190
191
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [32])
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
def test_chunked_prefill(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
    num_logprobs: int,
    chunked_prefill_token_size: int,
) -> None:
    max_num_seqs = chunked_prefill_token_size
    max_num_batched_tokens = chunked_prefill_token_size
192
193
194

    with vllm_runner(model,
                     enable_chunked_prefill=True,
195
196
197
198
                     max_num_batched_tokens=max_num_batched_tokens,
                     max_num_seqs=max_num_seqs) as vllm_model:
        chunked = vllm_model.generate_greedy_logprobs(example_prompts,
                                                      max_tokens, num_logprobs)
199

200
201
202
203
204
205
206
    with vllm_runner(model,
                     enable_chunked_prefill=False,
                     max_num_seqs=max_num_seqs) as vllm_model:
        non_chunked = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    check_logprobs_close(
207
208
209
210
211
212
213
        outputs_0_lst=chunked,
        outputs_1_lst=non_chunked,
        name_0="chunked",
        name_1="non_chunked",
    )


214
215
216
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
def test_chunked_prefill_with_parallel_sampling(
217
218
219
220
221
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
    """
    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)
244
245


246
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
247
248
249
250
251
252
253
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
254
255
256
257
258
    """
    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
259
260
    vllm_config = EngineArgs(model=model,
                             trust_remote_code=True).create_engine_config()
261
    while len(example_prompts) == vllm_config.pad_for_cudagraph(
262
            len(example_prompts)):
263
264
265
        example_prompts.append(example_prompts[0])

    try:
266
        with vllm_runner(model) as vllm_model:
267
268
269
270
271
272
273
274
            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")


275
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
276
277
278
279
280
281
282
@pytest.mark.parametrize("max_tokens", [20])
def test_models_preemption_recompute(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
283
284
285
286
    """
    Tests that outputs are identical with and w/o preemptions (recompute).
    """
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
287
        scheduler = vllm_model.llm.llm_engine.scheduler[0]
288
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = True
289
290
291
        preempt_vllm_outputs = vllm_model.generate_greedy(
            example_prompts, max_tokens)

292
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = False
293
294
295
296
297
298
299
300
301
302
        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)

    check_outputs_equal(
        outputs_0_lst=preempt_vllm_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="vllm_preepmtions",
        name_1="vllm",
    )


303
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
304
305
306
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
307
    model: str,
308
) -> None:
309
310
311
312
313
314
315
316
317
    """
    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
    statelessness mechanism where it can cleanup new incoming requests in
    a single step.
    """
318
    try:
319
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
320
321
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
322
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
323
324
325
                    "steps finished requests registered unnecessarily ")


326
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
327
328
329
def test_state_cleanup(
    vllm_runner,
    example_prompts,
330
    model: str,
Mor Zusman's avatar
Mor Zusman committed
331
) -> None:
332
333
334
335
336
337
    """ 
    This test is for verifying that the Hybrid state is cleaned up between
    steps.
    
    If its not cleaned, an error would be expected.
    """
Mor Zusman's avatar
Mor Zusman committed
338
    try:
339
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
340
341
342
            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
343
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
344
345
346
                    "could be related to finished_requests_ids")


347
@multi_gpu_test(num_gpus=2)
348
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
349
@pytest.mark.parametrize("max_tokens", [64])
350
351
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
352
353
354
355
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
356
    num_logprobs: int,
357
) -> None:
358
    with vllm_runner(model, tensor_parallel_size=1,
359
                     max_num_seqs=2) as vllm_model:
360
361
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
362

363
    with vllm_runner(model, tensor_parallel_size=2,
364
                     max_num_seqs=2) as vllm_model:
365
366
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
367

368
    check_logprobs_close(
369
370
371
372
373
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433


@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
@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:
        if model not in HF_UNSUPPORTED_MODELS:
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None

    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "1")
        if model in HYBRID_MODELS:
            # required due to reorder_batch behaviour
            m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
        with vllm_runner(model,
                         max_num_seqs=MAX_NUM_SEQS,
                         compilation_config={'full_cuda_graph': True},
                         enable_prefix_caching=False) as vllm_model:
            vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                example_prompts, max_tokens, num_logprobs)

    if hf_outputs is not None:
        check_logprobs_close(
            outputs_0_lst=hf_outputs,
            outputs_1_lst=vllm_v0_outputs,
            name_0="hf",
            name_1="vllm-v0",
        )

    ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
    check_logprobs_close(
        outputs_0_lst=ref_outputs,
        outputs_1_lst=vllm_v1_outputs,
        name_0="hf" if hf_outputs is not None else "vllm-v0",
        name_1="vllm-v1",
    )