test_hybrid.py 16.3 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
    "ibm-granite/granite-4.0-tiny-preview",
    "tiiuae/Falcon-H1-0.5B-Base",
34
    "LiquidAI/LFM2-1.2B",
35
36
37
38
39
40
41
]

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
42
43
44
45
    "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
46
]
47

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

59
FULL_CUDA_GRAPH_MODELS = [
60
    "ai21labs/Jamba-tiny-dev",
61
    "Zyphra/Zamba2-1.2B-instruct",
62
63
]

64
65
66
67
V0_UNSUPPORTED_MODELS = [
    "LiquidAI/LFM2-1.2B",
]

68
69
70
71
72
FP32_STATE_MODELS = [
    "state-spaces/mamba-130m-hf",
    "Zyphra/Zamba2-1.2B-instruct",
]

73
74
75
# Avoid OOM
MAX_NUM_SEQS = 4

Mor Zusman's avatar
Mor Zusman committed
76

77
78
79
@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
80
81
82
83
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
84
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
85
86
    model: str,
    max_tokens: int,
87
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
88
) -> None:
89
90
91
92

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
93
94
        hf_version_check = model_info.check_transformers_version(
            on_fail="return")
95
    except ValueError:
96
97
98
99
        hf_version_check = None

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

101
    with hf_runner(model) as hf_model:
102
        if model not in HF_UNSUPPORTED_MODELS and hf_version_check is None:
Chen Zhang's avatar
Chen Zhang committed
103
104
105
106
            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
107

108
109
110
111
112
113
114
115
    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "0")
        if model not in V0_UNSUPPORTED_MODELS:
            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)
        else:
            vllm_v0_outputs = None
116

Chen Zhang's avatar
Chen Zhang committed
117
    if model in V1_SUPPORTED_MODELS:
118
119
120
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
            vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                example_prompts, max_tokens, num_logprobs)
Chen Zhang's avatar
Chen Zhang committed
121
122
123
    else:
        vllm_v1_outputs = None

124
    if hf_outputs is not None and vllm_v0_outputs is not None:
Chen Zhang's avatar
Chen Zhang committed
125
126
127
128
129
130
131
132
133
        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
134
        assert ref_outputs is not None
Chen Zhang's avatar
Chen Zhang committed
135
136
137
138
139
140
        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
141
142


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

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

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

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


179
180
181
182
183
184
185
186
187
188
189
@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,
190
    monkeypatch,
191
192
193
) -> None:
    max_num_seqs = chunked_prefill_token_size
    max_num_batched_tokens = chunked_prefill_token_size
194

195
196
197
198
199
200
201
202
    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "0")
        with vllm_runner(model,
                         enable_chunked_prefill=True,
                         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)
203

204
205
206
207
208
        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)
209

210
211
212
213
214
215
        check_logprobs_close(
            outputs_0_lst=chunked,
            outputs_1_lst=non_chunked,
            name_0="chunked",
            name_1="non_chunked",
        )
216
217


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


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

    try:
270
        with vllm_runner(model) as vllm_model:
271
272
273
274
275
276
277
278
            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")


279
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
280
281
282
283
284
285
@pytest.mark.parametrize("max_tokens", [20])
def test_models_preemption_recompute(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
286
    monkeypatch,
287
) -> None:
288
289
290
    """
    Tests that outputs are identical with and w/o preemptions (recompute).
    """
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "0")
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
            scheduler = vllm_model.llm.llm_engine.scheduler[0]
            scheduler.ENABLE_ARTIFICIAL_PREEMPT = True
            preempt_vllm_outputs = vllm_model.generate_greedy(
                example_prompts, max_tokens)

            scheduler.ENABLE_ARTIFICIAL_PREEMPT = False
            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",
        )
309
310


311
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
312
313
314
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
315
    model: str,
316
) -> None:
317
318
319
320
321
322
323
324
325
    """
    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.
    """
326
    try:
327
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
328
329
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
330
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
331
332
333
                    "steps finished requests registered unnecessarily ")


334
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
335
336
337
def test_state_cleanup(
    vllm_runner,
    example_prompts,
338
    model: str,
Mor Zusman's avatar
Mor Zusman committed
339
) -> None:
340
341
342
343
344
345
    """ 
    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
346
    try:
347
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
348
349
350
            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
351
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
352
353
354
                    "could be related to finished_requests_ids")


355
@multi_gpu_test(num_gpus=2)
356
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
357
@pytest.mark.parametrize("max_tokens", [64])
358
359
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
360
361
362
363
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
364
    num_logprobs: int,
365
) -> None:
366
    with vllm_runner(model, tensor_parallel_size=1,
367
                     max_num_seqs=2) as vllm_model:
368
369
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
370

371
    with vllm_runner(model, tensor_parallel_size=2,
372
                     max_num_seqs=2) as vllm_model:
373
374
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
375

376
    check_logprobs_close(
377
378
379
380
381
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )
382
383


384
@pytest.mark.parametrize("model", FULL_CUDA_GRAPH_MODELS)
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
@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 monkeypatch.context() as m:
412
413
414
415
416
417
418
419
420
421
422
        m.setenv("VLLM_USE_V1", "0")
        if model not in V0_UNSUPPORTED_MODELS:
            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)
        else:
            vllm_v0_outputs = None

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

424
    if hf_outputs is not None and vllm_v0_outputs is not None:
425
426
427
428
429
430
431
432
        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
433
    assert ref_outputs is not None
434
435
436
437
438
439
    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",
    )
440
441


442
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
def test_fp32_state(
    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 monkeypatch.context() as m:
470
        m.setenv("VLLM_USE_V1", "0")
471
472
        with vllm_runner(model,
                         max_num_seqs=MAX_NUM_SEQS,
473
474
                         mamba_ssm_cache_dtype="float32") as vllm_model:
            vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
475
476
                example_prompts, max_tokens, num_logprobs)

477
478
479
480
481
482
    with vllm_runner(model,
                     max_num_seqs=MAX_NUM_SEQS,
                     mamba_ssm_cache_dtype="float32") as vllm_model:
        vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
    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",
    )