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

4
from collections.abc import Callable
5

Mor Zusman's avatar
Mor Zusman committed
6
7
import pytest

8
from tests.models.registry import HF_EXAMPLE_MODELS
9
from tests.utils import multi_gpu_test
10
from vllm import LLM
11
from vllm.engine.arg_utils import EngineArgs
12
from vllm.platforms import current_platform
13
from vllm.sampling_params import SamplingParams
14
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
15

16
from ...utils import check_logprobs_close, check_outputs_equal
17

18
19
20
# Mark all tests as hybrid
pytestmark = pytest.mark.hybrid_model

21
22
23
24
# 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

25
26
APC_MULTIPLY_BY = 300

27
28
29
SSM_MODELS = [
    "state-spaces/mamba-130m-hf",
    "tiiuae/falcon-mamba-tiny-dev",
30
31
    # mamba2-codestral in transformers is broken pending:
    # https://github.com/huggingface/transformers/pull/40861
32
    # "yujiepan/mamba2-codestral-v0.1-tiny-random",
33
]
34

35
36
HYBRID_MODELS = [
    "ai21labs/Jamba-tiny-dev",
37
    "pfnet/plamo-2-1b",
38
    "Zyphra/Zamba2-1.2B-instruct",
39
    "hmellor/tiny-random-BambaForCausalLM",
40
41
    "ibm-granite/granite-4.0-tiny-preview",
    "tiiuae/Falcon-H1-0.5B-Base",
42
    "LiquidAI/LFM2-1.2B",
43
    "tiny-random/qwen3-next-moe",
Chen Zhang's avatar
Chen Zhang committed
44
45
]

46
FULL_CUDA_GRAPH_MODELS = [
47
    "ai21labs/Jamba-tiny-dev",
48
    "pfnet/plamo-2-1b",
49
    "Zyphra/Zamba2-1.2B-instruct",
50
51
]

52
53
54
55
56
FP32_STATE_MODELS = [
    "state-spaces/mamba-130m-hf",
    "Zyphra/Zamba2-1.2B-instruct",
]

57
58
59
# Avoid OOM
MAX_NUM_SEQS = 4

Mor Zusman's avatar
Mor Zusman committed
60

61
62
63
@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
64
65
66
67
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
68
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
69
70
    model: str,
    max_tokens: int,
71
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
72
) -> None:
73
74
75
    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
76
        model_info.check_transformers_version(on_fail="skip")
77
    except ValueError:
78
        pass
79

80
    with hf_runner(model) as hf_model:
81
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
82
83
            example_prompts, max_tokens, num_logprobs
        )
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
84

85
86
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        vllm_outputs = vllm_model.generate_greedy_logprobs(
87
88
            example_prompts, max_tokens, num_logprobs
        )
Chen Zhang's avatar
Chen Zhang committed
89

90
91
92
93
94
95
    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
96
97


98
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
99
100
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
101
102
103
104
105
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
106
    num_logprobs: int,
107
) -> None:
108
109
110
111
112
113
114
    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

115
    for_loop_outputs = []
116
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
117
        for prompt in example_prompts:
118
119
120
            (single_output,) = vllm_model.generate_greedy_logprobs(
                [prompt], max_tokens, num_logprobs
            )
121
            for_loop_outputs.append(single_output)
122

123
        batched_outputs = vllm_model.generate_greedy_logprobs(
124
125
            example_prompts, max_tokens, num_logprobs
        )
126

127
    check_logprobs_close(
128
129
130
131
132
133
134
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


135
136
137
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
def test_chunked_prefill_with_parallel_sampling(
138
139
140
141
142
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
143
    """
144
145
    Tests chunked prefill in conjunction with n > 1.

146
147
148
149
150
151
152
    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)
    """
153
    sampling_params = SamplingParams(n=3, temperature=1, seed=0, max_tokens=max_tokens)
154
    with vllm_runner(
155
156
157
158
159
        model,
        enable_chunked_prefill=True,
        # forces prefill chunks with decoding
        max_num_batched_tokens=MAX_NUM_SEQS * 3,
        max_num_seqs=MAX_NUM_SEQS,
160
161
    ) as vllm_model:
        vllm_model.generate(example_prompts, sampling_params)
162
163


164
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
165
166
167
168
169
170
171
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
172
173
174
175
176
    """
    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.
    """
177
    vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config()
178
179
180
181
182
183
184
185
    cudagraph_dispatcher = CudagraphDispatcher(vllm_config)
    cudagraph_dispatcher.initialize_cudagraph_keys(
        vllm_config.compilation_config.cudagraph_mode
    )
    while (
        len(example_prompts)
        == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens
    ):
186
187
188
        example_prompts.append(example_prompts[0])

    try:
189
        with vllm_runner(model) as vllm_model:
190
191
192
193
194
            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. "
195
196
            "Could be related to mamba cache not padded correctly"
        )
197
198


199
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
200
201
202
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
203
    model: str,
204
) -> None:
205
206
207
208
209
210
    """
    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
211
    statelessness mechanism where it can clean up new incoming requests in
212
213
    a single step.
    """
214
    try:
215
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
216
217
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
218
219
220
221
        pytest.fail(
            "Hybrid inner state wasn't cleaned up properly between"
            "steps finished requests registered unnecessarily "
        )
222
223


224
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
225
226
227
def test_state_cleanup(
    vllm_runner,
    example_prompts,
228
    model: str,
Mor Zusman's avatar
Mor Zusman committed
229
) -> None:
230
    """
231
232
    This test is for verifying that the Hybrid state is cleaned up between
    steps.
233

234
    If it's not cleaned, an error would be expected.
235
    """
Mor Zusman's avatar
Mor Zusman committed
236
    try:
237
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
238
239
240
            for _ in range(10):
                vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
    except ValueError:
241
242
243
244
        pytest.fail(
            "Hybrid inner state wasn't cleaned up between states, "
            "could be related to finished_requests_ids"
        )
Mor Zusman's avatar
Mor Zusman committed
245
246


247
@multi_gpu_test(num_gpus=2)
248
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
249
@pytest.mark.parametrize("max_tokens", [64])
250
251
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
252
253
254
255
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
256
    num_logprobs: int,
257
) -> None:
258
259
260
    with vllm_runner(
        model, tensor_parallel_size=1, max_num_seqs=MAX_NUM_SEQS
    ) as vllm_model:
261
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
262
263
            example_prompts, max_tokens, num_logprobs
        )
264

265
266
267
    with vllm_runner(
        model, tensor_parallel_size=2, max_num_seqs=MAX_NUM_SEQS
    ) as vllm_model:
268
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
269
270
            example_prompts, max_tokens, num_logprobs
        )
271

272
    check_logprobs_close(
273
274
275
276
277
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )
278
279


280
@pytest.mark.parametrize("model", FULL_CUDA_GRAPH_MODELS)
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
@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:
300
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
301
302
            example_prompts, max_tokens, num_logprobs
        )
303

304
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
305
        vllm_outputs = vllm_model.generate_greedy_logprobs(
306
307
            example_prompts, max_tokens, num_logprobs
        )
308
309

    check_logprobs_close(
310
        outputs_0_lst=hf_outputs,
311
        outputs_1_lst=vllm_outputs,
312
        name_0="hf",
313
        name_1="vllm",
314
    )
315
316


317
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
318
319
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
320
321
322
@pytest.mark.parametrize(
    "cache_dtype_param", ["mamba_ssm_cache_dtype", "mamba_cache_dtype"]
)
323
def test_fp32_cache_state(
324
325
326
327
328
329
330
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    num_logprobs: int,
331
    cache_dtype_param: str,
332
333
334
335
336
337
338
339
340
) -> 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:
341
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
342
343
            example_prompts, max_tokens, num_logprobs
        )
344

345
346
347
    with vllm_runner(
        model, max_num_seqs=MAX_NUM_SEQS, **{cache_dtype_param: "float32"}
    ) as vllm_model:
348
        vllm_outputs = vllm_model.generate_greedy_logprobs(
349
350
            example_prompts, max_tokens, num_logprobs
        )
351

352
    check_logprobs_close(
353
        outputs_0_lst=hf_outputs,
354
        outputs_1_lst=vllm_outputs,
355
        name_0="hf",
356
        name_1="vllm",
357
    )
358
359
360


# Helper functions for the APC tests
361
362
363
364
365
def _get_vllm_runner_params(
    model: str,
    max_model_len: int,
    tensor_parallel_size: int = 1,
):
366
    return {
367
        "model_name": model,
368
        "enable_chunked_prefill": True,
369
370
371
372
        "enable_prefix_caching": False,
        "max_model_len": max_model_len,
        "tensor_parallel_size": tensor_parallel_size,
        "gpu_memory_utilization": 0.4,
373
374
375
    }


376
377
378
379
380
381
382
383
384
def _get_vLLM_output(
    vllm_runner,
    kwargs,
    prompts,
    max_tokens,
    num_logprobs,
    num_repetitions=1,
    vllm_model=None,
):
385
386
387
388
389
390
391
392
    outs = []
    if vllm_model is None:
        vllm_model = vllm_runner(**kwargs)
    for _ in range(num_repetitions):
        if num_logprobs < 0:
            vllm_output = vllm_model.generate_greedy(prompts, max_tokens)
        else:
            vllm_output = vllm_model.generate_greedy_logprobs(
393
394
                prompts, max_tokens, num_logprobs
            )
395
396
397
398
399
        outs.append(vllm_output)

    return outs, vllm_model


400
@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
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
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("n_repetitions", [2])
# If num_logprobs is set to -1, then the stringent version
# of the test is executed using `check_outputs_equal`
# instead of `check_logprobs_close`
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tensor_parallel_size", [1])
def test_apc_single_prompt(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    n_repetitions: int,
    num_logprobs: int,
    tensor_parallel_size: 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

426
427
428
    compare_operator: Callable = (
        check_logprobs_close if num_logprobs > 0 else check_outputs_equal  # type: ignore
    )
429
430

    # Sample prompts.
431
    generated_prompts = [APC_MULTIPLY_BY * example_prompts[0]]
432

433
    max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
434
    vllm_runner_kwargs = _get_vllm_runner_params(
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
        model, max_model_len, tensor_parallel_size=tensor_parallel_size
    )
    vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
    vllm_outputs_no_cache, _ = _get_vLLM_output(
        vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
    )

    vllm_runner_kwargs["enable_prefix_caching"] = True
    vllm_outputs_cache_rep, _ = _get_vLLM_output(
        vllm_runner,
        vllm_runner_kwargs,
        generated_prompts,
        max_tokens,
        num_logprobs,
        n_repetitions,
    )
451
452
453
454
455
456
457
458
459
460
461
462
463

    for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
        # In the first repetition, the caches are filled
        # In the second repetition, these caches are reused

        compare_operator(
            outputs_0_lst=vllm_outputs_no_cache[0],
            outputs_1_lst=vllm_outputs_cache_itn,
            name_0="vllm_no_cache",
            name_1=f"vllm_cache_it_{r_idx + 1}",
        )


464
@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("n_repetitions", [2])
# If num_logprobs is set to -1, then the stringent version
# of the test is executed using `check_outputs_equal`
# instead of `check_logprobs_close`
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tensor_parallel_size", [1])
def test_apc_single_prompt_block_align_alignment(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    n_repetitions: int,
    num_logprobs: int,
    tensor_parallel_size: 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

490
491
492
    compare_operator: Callable = (
        check_logprobs_close if num_logprobs > 0 else check_outputs_equal  # type: ignore
    )
493
494

    # Sample prompts. This custom prompt is used, as it causes the most issues
495
    generated_prompts = ["The president of the United States is " * APC_MULTIPLY_BY]
496

497
    max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
498
    vllm_runner_kwargs = _get_vllm_runner_params(
499
500
501
        model, max_model_len, tensor_parallel_size=tensor_parallel_size
    )
    vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
502

503
504
505
    vllm_outputs_no_cache, _ = _get_vLLM_output(
        vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
    )
506

507
    vllm_runner_kwargs["enable_prefix_caching"] = True
508
509
    with vllm_runner(**vllm_runner_kwargs) as vllm_model:
        # Retrieve the default mamba state block size
510
511
        vllm_config = vllm_model.llm.llm_engine.vllm_config
        mamba_block_size = vllm_config.cache_config.mamba_block_size
512
513
514
515
516
517
518

    # In case the hybrid model does not have the
    # "mamba_block_size" assume a fixed constant
    if mamba_block_size is None:
        mamba_block_size = 512

    mamba_block_size_multiplier = 10
519
520
521
522
523
524
525
526
527
528
529
530
    for offsets in [-3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3]:
        vllm_runner_kwargs["max_num_batched_tokens"] = (
            mamba_block_size_multiplier * mamba_block_size - offsets
        )
        vllm_outputs_cache_rep, _ = _get_vLLM_output(
            vllm_runner,
            vllm_runner_kwargs,
            generated_prompts,
            max_tokens,
            num_logprobs,
            n_repetitions,
        )
531
532
533
534
535
536
537
538
539
540
541
542
543
544

        # Check alignment of the output logits when using APC
        for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
            # In the first repetition, the caches are filled
            # In the second repetition, these caches are reused

            compare_operator(
                outputs_0_lst=vllm_outputs_no_cache[0],
                outputs_1_lst=vllm_outputs_cache_itn,
                name_0="vllm_no_cache",
                name_1=f"vllm_cache_it_{r_idx + 1}",
            )


545
@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("n_repetitions", [2])
# If num_logprobs is set to -1, then the stringent version
# of the test is executed using `check_outputs_equal`
# instead of `check_logprobs_close`
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tensor_parallel_size", [1])
def test_apc_multiple_prompts_all_cached_outputs(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    n_repetitions: int,
    num_logprobs: int,
    tensor_parallel_size: 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

571
572
573
    compare_operator: Callable = (
        check_logprobs_close if num_logprobs > 0 else check_outputs_equal  # type: ignore
    )
574
575

    # Sample prompts.
576
    generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
577

578
    max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
579
    vllm_runner_kwargs = _get_vllm_runner_params(
580
581
582
        model, max_model_len, tensor_parallel_size=tensor_parallel_size
    )
    vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
583
584
585
586
    # Reduce the effects of batch variance on ROCm since batch invariance is not
    # yet supported. See: https://github.com/vllm-project/vllm/issues/27433
    if current_platform.is_rocm():
        vllm_runner_kwargs["max_num_seqs"] = 4
587

588
589
590
    vllm_outputs_no_cache, _ = _get_vLLM_output(
        vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
    )
591

592
593
594
595
596
597
598
599
600
    vllm_runner_kwargs["enable_prefix_caching"] = True
    vllm_outputs_cache_rep, _ = _get_vLLM_output(
        vllm_runner,
        vllm_runner_kwargs,
        generated_prompts,
        max_tokens,
        num_logprobs,
        n_repetitions,
    )
601
602
603
604
605
606
607
608
609
610
611
612
613

    for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
        # In the first repetition, the caches are filled
        # In the second repetition, these caches are reused

        compare_operator(
            outputs_0_lst=vllm_outputs_no_cache[0],
            outputs_1_lst=vllm_outputs_cache_itn,
            name_0="vllm_no_cache",
            name_1=f"vllm_cache_it_{r_idx + 1}",
        )


614
@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("n_repetitions", [2])
# If num_logprobs is set to -1, then the stringent version
# of the test is executed using `check_outputs_equal`
# instead of `check_logprobs_close`
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tensor_parallel_size", [1])
def test_apc_multiple_prompts_block_align_alignment(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    n_repetitions: int,
    num_logprobs: int,
    tensor_parallel_size: 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

640
641
642
    compare_operator: Callable = (
        check_logprobs_close if num_logprobs > 0 else check_outputs_equal  # type: ignore
    )
643
644
645
646

    # Sample prompts. This custom prompt is used, as it causes the most issues
    prompt_text = "The president of the United States is "
    prompt_offsets = [0, 3, 7, 13, 17, 22, 25, 31]
647
648
649
    generated_prompts = [
        prompt_text[offset:] * APC_MULTIPLY_BY for offset in prompt_offsets
    ]
650
651
652
653
654
655
656
657
658
659
660
661

    max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
    vllm_runner_kwargs = _get_vllm_runner_params(
        model, max_model_len, tensor_parallel_size
    )
    vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"

    vllm_outputs_no_cache, _ = _get_vLLM_output(
        vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
    )

    vllm_runner_kwargs["enable_prefix_caching"] = True
662
663
    with vllm_runner(**vllm_runner_kwargs) as vllm_model:
        # Retrieve the default mamba state block size
664
665
        vllm_config = vllm_model.llm.llm_engine.vllm_config
        mamba_block_size = vllm_config.cache_config.mamba_block_size
666
667
668
669
670
671
672

    # In case the hybrid model does not have the
    # "mamba_block_size" assume a fixed constant
    if mamba_block_size is None:
        mamba_block_size = 512

    mamba_block_size_multiplier = 10
673
674
675
676
677
678
679
680
681
682
683
684
    for offsets in [-3, 3, mamba_block_size // 4 + 3, mamba_block_size // 2 - 3]:
        vllm_runner_kwargs["max_num_batched_tokens"] = (
            mamba_block_size_multiplier * mamba_block_size - offsets
        )
        vllm_outputs_cache_rep, _ = _get_vLLM_output(
            vllm_runner,
            vllm_runner_kwargs,
            generated_prompts,
            max_tokens,
            num_logprobs,
            n_repetitions,
        )
685
686
687
688
689
690
691
692
693
694
695
696
697
698

        # Check alignment of the output logits when using APC
        for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
            # In the first repetition, the caches are filled
            # In the second repetition, these caches are reused

            compare_operator(
                outputs_0_lst=vllm_outputs_no_cache[0],
                outputs_1_lst=vllm_outputs_cache_itn,
                name_0="vllm_no_cache",
                name_1=f"vllm_cache_it_{r_idx + 1}",
            )


699
@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("n_repetitions", [2])
# If num_logprobs is set to -1, then the stringent version
# of the test is executed using `check_outputs_equal`
# instead of `check_logprobs_close`
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("tensor_parallel_size", [1])
def test_apc_multiple_prompts_partial_cached_outputs(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    n_repetitions: int,
    num_logprobs: int,
    tensor_parallel_size: 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

725
726
727
    compare_operator: Callable = (
        check_logprobs_close if num_logprobs > 0 else check_outputs_equal  # type: ignore
    )
728
729

    # Sample prompts.
730
    generated_prompts = [APC_MULTIPLY_BY * prompt for prompt in example_prompts]
731

732
    max_model_len = max(len(prompt) + max_tokens for prompt in generated_prompts)
733
    vllm_runner_kwargs = _get_vllm_runner_params(
734
735
736
        model, max_model_len, tensor_parallel_size=tensor_parallel_size
    )
    vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
737

738
739
740
    vllm_outputs_no_cache, _ = _get_vLLM_output(
        vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
    )
741
742

    # Cache only part of all the prompts
743
    vllm_runner_kwargs["enable_prefix_caching"] = True
744
    vllm_outputs_partial_cache, vllm_model = _get_vLLM_output(
745
746
        vllm_runner, vllm_runner_kwargs, generated_prompts[:3], max_tokens, num_logprobs
    )
747
748
749
750
751
752
753
754

    compare_operator(
        outputs_0_lst=vllm_outputs_no_cache[0][:3],
        outputs_1_lst=vllm_outputs_partial_cache[0],
        name_0="vllm_no_cache",
        name_1="vllm_partial_cache",
    )

755
756
757
758
759
760
761
762
763
    vllm_outputs_cache_rep, _ = _get_vLLM_output(
        vllm_runner,
        vllm_runner_kwargs,
        generated_prompts,
        max_tokens,
        num_logprobs,
        n_repetitions,
        vllm_model=vllm_model,
    )
764
765
766
767
768
769
770
771
772
773
774

    for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
        # In the first repetition, the caches are filled
        # In the second repetition, these caches are reused

        compare_operator(
            outputs_0_lst=vllm_outputs_no_cache[0],
            outputs_1_lst=vllm_outputs_cache_itn,
            name_0="vllm_no_cache",
            name_1=f"vllm_cache_it_{r_idx + 1}",
        )
775
776


777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
# Test that outputs match whether prefix caching is enabled or not for mamba.
@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
def test_same_mamba_output_apc_on_vs_off(
    vllm_runner,
    model: str,
) -> None:
    num_logprobs = 5
    prompts = [
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
    ]
    max_tokens = 20
    max_model_len = max(len(p) for p in prompts) + max_tokens + 64

    base_kwargs = _get_vllm_runner_params(model, max_model_len)
    base_kwargs.update(
        enforce_eager=True, block_size=16, seed=42, gpu_memory_utilization=0.8
    )

    # No prefix caching
    kwargs_no_apc = {**base_kwargs, "enable_prefix_caching": False}
    with vllm_runner(**kwargs_no_apc) as vllm_model:
        outputs_no_apc, _ = _get_vLLM_output(
            vllm_runner,
            kwargs_no_apc,
            prompts,
            max_tokens,
            num_logprobs=num_logprobs,
            vllm_model=vllm_model,
        )
    # With prefix caching
    kwargs_with_apc = {
        **base_kwargs,
        "enable_prefix_caching": True,
        "mamba_block_size": 16,
    }
    with vllm_runner(**kwargs_with_apc) as vllm_model:
        outputs_with_apc, _ = _get_vLLM_output(
            vllm_runner,
            kwargs_with_apc,
            prompts,
            max_tokens,
            num_logprobs=num_logprobs,
            vllm_model=vllm_model,
        )

    check_logprobs_close(
        outputs_0_lst=outputs_no_apc[0],
        outputs_1_lst=outputs_with_apc[0],
        name_0="vllm_no_apc",
        name_1="vllm_with_apc",
    )


831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
# we have to use a real large model to get reasonable results
# the model can't be a hybrid model as we need block_size 16
@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
def test_apc_common_prefix_same_batch(
    model: str,
    monkeypatch,
) -> None:
    # Required to put the two requests in the same batch
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
    llm = LLM(
        model=model,
        enforce_eager=True,
        block_size=16,
        mamba_block_size=16,
        enable_prefix_caching=True,
        seed=42,
    )
    prompts = [
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
    ]
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
    outputs = llm.generate(prompts, sampling_params)
    for output in outputs:
        assert "two" in output.outputs[0].text