A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Bugfix] Fix hybrid model tests (#17182)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Apr 26, 2025 at 06:14 UTC
43faa0461afde316ffca5a2fa58d2e3c89acd623
3 files changed
+158
-534
tests/conftest.py
+4
-1
@@ -531,7 +531,10 @@ class HfRunner:
531
for _, hidden_state in enumerate(hidden_states):
532
last_hidden_states = hidden_state[-1][0]
533
logits = torch.matmul(
534
- last_hidden_states.to(output_embeddings.weight.device),
534
+ last_hidden_states.to(
535
+ device=output_embeddings.weight.device,
536
+ dtype=output_embeddings.weight.dtype,
537
+ ),
538
output_embeddings.weight.t(),
539
)
540
if getattr(output_embeddings, "bias", None) is not None:
tests/models/decoder_only/language/test_hybrid.py
+154
-196
@@ -6,71 +6,84 @@ from tests.utils import multi_gpu_test
6
from vllm.engine.arg_utils import EngineArgs
7
from vllm.sampling_params import SamplingParams
8
9
-from ...utils import check_outputs_equal
9
+from ...utils import check_logprobs_close, check_outputs_equal
10
+
11
+# NOTE: The first model in each list is taken as the primary model,
12
+# meaning that it will be used in all tests in this file
13
+# The rest of the models will only be tested by test_models
14
+
15
+SSM_MODELS = [
16
+ "state-spaces/mamba-130m-hf",
17
+ "tiiuae/falcon-mamba-tiny-dev",
18
+ # TODO: Compare to a Mamba2 model. The HF transformers implementation of
19
+ # Mamba2 is buggy for Codestral as it doesn't handle n_groups.
20
+ # See https://github.com/huggingface/transformers/pull/35943
21
+ # "mistralai/Mamba-Codestral-7B-v0.1",
22
+]
23
11
-# This test is for the hybrid models
12
-MODELS = [
13
- "ai21labs/Jamba-tiny-dev", "Zyphra/Zamba2-1.2B-instruct",
14
- "pfnet/plamo-2-1b"
24
+HYBRID_MODELS = [
25
+ "ai21labs/Jamba-tiny-dev",
26
+ # NOTE: Running Plamo2 in transformers implementation requires to install
27
+ # causal-conv1d package, which is not listed as a test dependency as it's
28
+ # not compatible with pip-compile.
29
+ "pfnet/plamo-2-1b",
30
+ "Zyphra/Zamba2-1.2B-instruct",
31
+ "ibm-ai-platform/Bamba-9B",
32
]
16
-# Bamba at Fp32 is too big for the CI (L4 GPU).
17
-# MODELS = ["ai21labs/Jamba-tiny-dev", "ibm-ai-platform/Bamba-9B"]
18
-# Note: Running Plamo2 in transformers implementation requires to install
19
-# causal-conv1d package, which is not listed as a test dependency as it's
20
-# not compatible with pip-compile.
33
+
34
+# Avoid OOM
35
+MAX_NUM_SEQS = 4
36
37
23
-@pytest.mark.parametrize("model", MODELS)
24
-@pytest.mark.parametrize("dtype", ["float"])
25
-@pytest.mark.parametrize("max_tokens", [96])
38
+@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
39
+@pytest.mark.parametrize("max_tokens", [64])
40
+@pytest.mark.parametrize("num_logprobs", [5])
41
def test_models(
42
hf_runner,
43
vllm_runner,
44
example_prompts,
45
model: str,
31
- dtype: str,
46
max_tokens: int,
47
+ num_logprobs: int,
48
) -> None:
34
- # numeric error produces different generation
35
- if "Bamba" in model:
36
- example_prompts.pop(3)
49
+ with hf_runner(model) as hf_model:
50
+ hf_outputs = hf_model.generate_greedy_logprobs_limit(
51
+ example_prompts, max_tokens, num_logprobs)
52
38
- with hf_runner(model, dtype=dtype) as hf_model:
39
- hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
40
-
41
- with vllm_runner(model, dtype=dtype) as vllm_model:
42
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
53
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
54
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
55
+ example_prompts, max_tokens, num_logprobs)
56
44
- for i in range(len(example_prompts)):
45
- hf_output_ids, hf_output_str = hf_outputs[i]
46
- vllm_output_ids, vllm_output_str = vllm_outputs[i]
47
- assert hf_output_str == vllm_output_str, (
48
- f"Test{i}:\nHF: {hf_output_str!r}\nvLLM: {vllm_output_str!r}")
49
- assert hf_output_ids == vllm_output_ids, (
50
- f"Test{i}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}")
57
+ check_logprobs_close(
58
+ outputs_0_lst=hf_outputs,
59
+ outputs_1_lst=vllm_outputs,
60
+ name_0="hf",
61
+ name_1="vllm",
62
+ )
63
64
53
-@pytest.mark.parametrize("model", MODELS)
54
-@pytest.mark.parametrize("dtype", ["float"])
55
-@pytest.mark.parametrize("max_tokens", [96])
65
+@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
66
+@pytest.mark.parametrize("max_tokens", [64])
67
+@pytest.mark.parametrize("num_logprobs", [5])
68
def test_batching(
69
vllm_runner,
70
example_prompts,
71
model: str,
60
- dtype: str,
72
max_tokens: int,
73
+ num_logprobs: int,
74
) -> None:
63
- # To pass the small model tests, we need full precision.
75
for_loop_outputs = []
65
- with vllm_runner(model, dtype=dtype) as vllm_model:
76
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
77
for prompt in example_prompts:
67
- for_loop_outputs.append(
68
- vllm_model.generate_greedy([prompt], max_tokens)[0])
78
+ single_output, = vllm_model.generate_greedy_logprobs([prompt],
79
+ max_tokens,
80
+ num_logprobs)
81
+ for_loop_outputs.append(single_output)
82
70
- batched_outputs = vllm_model.generate_greedy(example_prompts,
71
- max_tokens)
83
+ batched_outputs = vllm_model.generate_greedy_logprobs(
84
+ example_prompts, max_tokens, num_logprobs)
85
73
- check_outputs_equal(
86
+ check_logprobs_close(
87
outputs_0_lst=for_loop_outputs,
88
outputs_1_lst=batched_outputs,
89
name_0="for_loop_vllm",
@@ -78,72 +91,35 @@ def test_batching(
91
)
92
93
81
-@pytest.mark.parametrize("model", MODELS)
82
-@pytest.mark.parametrize("dtype", ["float16"])
83
-@pytest.mark.parametrize("max_tokens", [10])
84
-def test_mamba_prefill_chunking_with_parallel_sampling(
85
- hf_runner, vllm_runner, example_prompts, model: str, dtype: str,
86
- max_tokens: int) -> None:
87
- # Tests prefill chunking in conjunction with n>1, in this case,
88
- # prefill is populated with decoding tokens and we test that it
89
- # doesn't fail This test might fail if cache is not allocated
90
- # correctly for n > 1 decoding steps inside a
91
- # chunked prefill forward pass (where we have both prefills
92
- # and decoding together )
93
-
94
- if 'plamo-2' in model:
95
- dtype = "float" # use a different dtype for plamo
96
-
97
- sampling_params = SamplingParams(n=3,
98
- temperature=1,
99
- seed=0,
100
- max_tokens=max_tokens)
101
- with vllm_runner(
102
- model,
103
- dtype=dtype,
104
- enable_chunked_prefill=True,
105
- max_num_batched_tokens=30,
106
- max_num_seqs=10 # forces prefill chunks with decoding
107
- ) as vllm_model:
108
- vllm_model.generate(example_prompts, sampling_params)
109
-
110
-
111
-@pytest.mark.parametrize("model", MODELS)
112
-@pytest.mark.parametrize("dtype", ["bfloat16"])
113
-@pytest.mark.parametrize("max_tokens", [7])
114
-def test_mamba_prefill_chunking(hf_runner, vllm_runner, example_prompts,
115
- model: str, dtype: str,
116
- max_tokens: int) -> None:
117
- # numeric error during prefill chunking produces different generation
118
- # compared to w/o prefill chunking for those examples, removed them for now
119
- if "Jamba" in model:
120
- example_prompts.pop(7)
121
- example_prompts.pop(2)
122
- example_prompts.pop(1)
123
- elif "Bamba" in model:
124
- example_prompts.pop(6)
125
- example_prompts.pop(3)
126
- example_prompts.pop(2)
127
- dtype = "half" # use a different dtype for Bamba
128
-
129
- elif "Zamba2" in model:
130
- example_prompts.pop(7)
131
- dtype = "half"
132
- elif "plamo-2-1b" in model:
133
- example_prompts.pop(7)
134
-
135
- with hf_runner(model, dtype=dtype) as hf_model:
136
- non_chunked = hf_model.generate_greedy(example_prompts, max_tokens)
94
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
95
+@pytest.mark.parametrize("max_tokens", [32])
96
+@pytest.mark.parametrize("num_logprobs", [5])
97
+@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
98
+def test_chunked_prefill(
99
+ vllm_runner,
100
+ example_prompts,
101
+ model: str,
102
+ max_tokens: int,
103
+ num_logprobs: int,
104
+ chunked_prefill_token_size: int,
105
+) -> None:
106
+ max_num_seqs = chunked_prefill_token_size
107
+ max_num_batched_tokens = chunked_prefill_token_size
108
109
with vllm_runner(model,
139
- dtype=dtype,
110
enable_chunked_prefill=True,
141
- max_num_batched_tokens=5,
142
- max_num_seqs=2) as vllm_model:
143
- chunked = vllm_model.generate_greedy(example_prompts,
144
- max_tokens=max_tokens)
111
+ max_num_batched_tokens=max_num_batched_tokens,
112
+ max_num_seqs=max_num_seqs) as vllm_model:
113
+ chunked = vllm_model.generate_greedy_logprobs(example_prompts,
114
+ max_tokens, num_logprobs)
115
146
- check_outputs_equal(
116
+ with vllm_runner(model,
117
+ enable_chunked_prefill=False,
118
+ max_num_seqs=max_num_seqs) as vllm_model:
119
+ non_chunked = vllm_model.generate_greedy_logprobs(
120
+ example_prompts, max_tokens, num_logprobs)
121
+
122
+ check_logprobs_close(
123
outputs_0_lst=chunked,
124
outputs_1_lst=non_chunked,
125
name_0="chunked",
@@ -151,57 +127,51 @@ def test_mamba_prefill_chunking(hf_runner, vllm_runner, example_prompts,
127
)
128
129
154
-@pytest.mark.parametrize("model", MODELS)
155
-@pytest.mark.parametrize("dtype", ["float"])
156
-@pytest.mark.parametrize("max_tokens", [15])
157
-def test_parallel_sampling(
130
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
131
+@pytest.mark.parametrize("max_tokens", [10])
132
+def test_chunked_prefill_with_parallel_sampling(
133
vllm_runner,
134
example_prompts,
135
model: str,
161
- dtype: str,
136
max_tokens: int,
137
) -> None:
164
-
165
- with vllm_runner(model, dtype=dtype) as vllm_model:
166
- for_loop_outputs = []
167
- for _ in range(10):
168
- for_loop_outputs.append(
169
- # using example_prompts index 1 instead of 0 since with 0 the
170
- # logprobs get really close and the test doesn't pass
171
- vllm_model.generate_greedy([example_prompts[1]], max_tokens)
172
- [0])
173
- sampling_params = SamplingParams(n=10,
174
- temperature=0.001,
175
- seed=0,
176
- max_tokens=max_tokens)
177
- n_lt_1_outputs = vllm_model.generate([example_prompts[1]],
178
- sampling_params)
179
- token_ids, texts = n_lt_1_outputs[0]
180
- n_lt_1_outputs = [(token_id, text)
181
- for token_id, text in zip(token_ids, texts)]
182
-
183
- check_outputs_equal(
184
- outputs_0_lst=n_lt_1_outputs,
185
- outputs_1_lst=for_loop_outputs,
186
- name_0="vllm_n_lt_1_outputs",
187
- name_1="vllm",
188
- )
138
+ """
139
+ Tests chunked prefill in conjunction with n > 1.
140
+
141
+ In this case, prefill is populated with decoding tokens and
142
+ we test that it doesn't fail.
143
+
144
+ This test might fail if cache is not allocated correctly for n > 1
145
+ decoding steps inside a chunked prefill forward pass
146
+ (where we have both prefill and decode together)
147
+ """
148
+ sampling_params = SamplingParams(n=3,
149
+ temperature=1,
150
+ seed=0,
151
+ max_tokens=max_tokens)
152
+ with vllm_runner(
153
+ model,
154
+ enable_chunked_prefill=True,
155
+ # forces prefill chunks with decoding
156
+ max_num_batched_tokens=MAX_NUM_SEQS * 3,
157
+ max_num_seqs=MAX_NUM_SEQS,
158
+ ) as vllm_model:
159
+ vllm_model.generate(example_prompts, sampling_params)
160
161
191
-@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
192
-@pytest.mark.parametrize("model", MODELS)
193
-@pytest.mark.parametrize("dtype", ["bfloat16"])
162
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
163
@pytest.mark.parametrize("max_tokens", [20])
164
def test_mamba_cache_cg_padding(
165
vllm_runner,
166
example_prompts,
167
model: str,
199
- dtype: str,
168
max_tokens: int,
169
) -> None:
202
- # This test is for verifying that mamba cache is padded to CG captured
203
- # batch size. If it's not, a torch RuntimeError will be raised because
204
- # tensor dimensions aren't compatible
170
+ """
171
+ This test is for verifying that mamba cache is padded to CG captured
172
+ batch size. If it's not, a torch RuntimeError will be raised because
173
+ tensor dimensions aren't compatible.
174
+ """
175
vllm_config = EngineArgs(model=model,
176
trust_remote_code=True).create_engine_config()
177
while len(example_prompts) == vllm_config.pad_for_cudagraph(
@@ -209,7 +179,7 @@ def test_mamba_cache_cg_padding(
179
example_prompts.append(example_prompts[0])
180
181
try:
212
- with vllm_runner(model, dtype=dtype) as vllm_model:
182
+ with vllm_runner(model) as vllm_model:
183
vllm_model.generate_greedy(example_prompts, max_tokens)
184
except RuntimeError:
185
pytest.fail(
@@ -218,28 +188,24 @@ def test_mamba_cache_cg_padding(
188
"Could be related to mamba cache not padded correctly")
189
190
221
-@pytest.mark.parametrize("model", MODELS)
222
-@pytest.mark.parametrize("dtype", ["float"])
191
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
192
@pytest.mark.parametrize("max_tokens", [20])
193
def test_models_preemption_recompute(
225
- hf_runner,
194
vllm_runner,
195
example_prompts,
196
model: str,
229
- dtype: str,
197
max_tokens: int,
198
) -> None:
232
- # Tests that outputs are identical with and w/o preemtions (recompute)
233
- assert dtype == "float"
234
-
235
- with vllm_runner(model, dtype=dtype) as vllm_model:
236
- vllm_model.model.llm_engine.scheduler[
237
- 0].ENABLE_ARTIFICIAL_PREEMPT = True
199
+ """
200
+ Tests that outputs are identical with and w/o preemptions (recompute).
201
+ """
202
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
203
+ scheduler = vllm_model.model.llm_engine.scheduler[0]
204
+ scheduler.ENABLE_ARTIFICIAL_PREEMPT = True
205
preempt_vllm_outputs = vllm_model.generate_greedy(
206
example_prompts, max_tokens)
207
241
- vllm_model.model.llm_engine.scheduler[
242
- 0].ENABLE_ARTIFICIAL_PREEMPT = False
208
+ scheduler.ENABLE_ARTIFICIAL_PREEMPT = False
209
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
210
211
check_outputs_equal(
@@ -250,40 +216,43 @@ def test_models_preemption_recompute(
216
)
217
218
253
-@pytest.mark.parametrize("model", MODELS)
254
-@pytest.mark.parametrize("dtype", ["float"])
219
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
220
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
221
vllm_runner,
257
- model: str,
258
- dtype: str,
222
example_prompts,
223
+ model: str,
224
) -> None:
261
- # This test is for verifying that the hybrid inner state management doesn't
262
- # collapse in case where the number of incoming requests and
263
- # finished_requests_ids is larger than the maximum mamba block capacity.
264
- # This could generally happen due to the fact that hybrid does support
265
- # statelessness mechanism where it can cleanup new incoming requests in
266
- # a single step.
225
+ """
226
+ This test is for verifying that the hybrid inner state management doesn't
227
+ collapse in case where the number of incoming requests and
228
+ finished_requests_ids is larger than the maximum mamba block capacity.
229
+
230
+ This could generally happen due to the fact that hybrid does support
231
+ statelessness mechanism where it can cleanup new incoming requests in
232
+ a single step.
233
+ """
234
try:
268
- with vllm_runner(model, dtype=dtype, max_num_seqs=10) as vllm_model:
235
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
236
vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
237
except ValueError:
238
pytest.fail("Hybrid inner state wasn't cleaned up properly between"
239
"steps finished requests registered unnecessarily ")
240
241
275
-@pytest.mark.parametrize("model", MODELS)
276
-@pytest.mark.parametrize("dtype", ["float"])
242
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
243
def test_state_cleanup(
244
vllm_runner,
279
- model: str,
280
- dtype: str,
245
example_prompts,
246
+ model: str,
247
) -> None:
283
- # This test is for verifying that the Hybrid state is cleaned up between
284
- # steps, If its not cleaned, an error would be expected.
248
+ """
249
+ This test is for verifying that the Hybrid state is cleaned up between
250
+ steps.
251
+
252
+ If its not cleaned, an error would be expected.
253
+ """
254
try:
286
- with vllm_runner(model, dtype=dtype) as vllm_model:
255
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
256
for _ in range(10):
257
vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
258
except ValueError:
@@ -291,28 +260,14 @@ def test_state_cleanup(
260
"could be related to finished_requests_ids")
261
262
294
-@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
295
-@pytest.mark.parametrize("model", MODELS)
296
-@pytest.mark.parametrize("dtype", ["float"])
297
-def test_multistep(
263
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
264
+@pytest.mark.parametrize("max_tokens", [64])
265
+def test_multistep_correctness(
266
vllm_runner,
299
- model: str,
300
- dtype: str,
267
example_prompts,
268
+ model: str,
269
+ max_tokens: int,
270
) -> None:
303
- # This test is verifying that multistep works correctly
304
- #on mamba-like models
305
- with vllm_runner(model, num_scheduler_steps=8,
306
- max_num_seqs=2) as vllm_model:
307
- vllm_model.generate_greedy([example_prompts[0]] * 10, 1)
308
-
309
-
310
-@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
311
-@pytest.mark.parametrize("model", MODELS)
312
-@pytest.mark.parametrize("dtype", ["float"])
313
-@pytest.mark.parametrize("max_tokens", [64])
314
-def test_multistep_correctness(vllm_runner, model: str, dtype: str,
315
- max_tokens: int, example_prompts) -> None:
271
with vllm_runner(model, num_scheduler_steps=8,
272
max_num_seqs=2) as vllm_model:
273
vllm_outputs_multistep = vllm_model.generate_greedy(
@@ -332,18 +287,21 @@ def test_multistep_correctness(vllm_runner, model: str, dtype: str,
287
288
289
@multi_gpu_test(num_gpus=2)
335
-@pytest.mark.parametrize("model", MODELS)
336
-@pytest.mark.parametrize("dtype", ["float"])
290
+@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
291
@pytest.mark.parametrize("max_tokens", [64])
292
def test_hybrid_distributed_produces_identical_generation(
339
- vllm_runner, model: str, dtype: str, max_tokens: int,
340
- example_prompts) -> None:
341
-
342
- with vllm_runner(model, dtype=dtype, tensor_parallel_size=2) as vllm_model:
293
+ vllm_runner,
294
+ example_prompts,
295
+ model: str,
296
+ max_tokens: int,
297
+) -> None:
298
+ with vllm_runner(model, tensor_parallel_size=2,
299
+ max_num_seqs=2) as vllm_model:
300
vllm_outputs_tp_2 = vllm_model.generate_greedy(example_prompts,
301
max_tokens)
302
346
- with vllm_runner(model, dtype=dtype, tensor_parallel_size=1) as vllm_model:
303
+ with vllm_runner(model, tensor_parallel_size=1,
304
+ max_num_seqs=2) as vllm_model:
305
vllm_outputs_tp_1 = vllm_model.generate_greedy(example_prompts,
306
max_tokens)
307
tests/models/decoder_only/language/test_mamba.py
deleted
-337
@@ -1,337 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-"""Compare the outputs of HF and vLLM when using greedy sampling for Mamba.
3
-
4
-Run `pytest tests/models/test_mamba.py`.
5
-"""
6
-import pytest
7
-import torch
8
-from transformers import AutoModelForCausalLM, AutoTokenizer
9
-
10
-from vllm.engine.arg_utils import EngineArgs
11
-from vllm.sampling_params import SamplingParams
12
-
13
-from ...utils import check_outputs_equal
14
-
15
-MODELS = [
16
- "state-spaces/mamba-130m-hf",
17
- "tiiuae/falcon-mamba-tiny-dev",
18
- # TODO: Compare to a Mamba2 model. The HF transformers implementation of
19
- # Mamba2 is buggy for Codestral as it doesn't handle n_groups.
20
- # See https://github.com/huggingface/transformers/pull/35943
21
- # "mistralai/Mamba-Codestral-7B-v0.1",
22
-]
23
-
24
-
25
-# Use lower-level interfaces to create this greedy generator, as mamba will
26
-# choke on the model_kwarg 'attention_mask' if hf_model.generate_greedy is used.
27
-def generate_greedy(model_name, example_prompts, max_tokens):
28
- # Create a text generation pipeline
29
- tokenizer = AutoTokenizer.from_pretrained(model_name)
30
- model = AutoModelForCausalLM.from_pretrained(model_name)
31
-
32
- # Set the device (GPU if available, else CPU)
33
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
34
- model.to(device)
35
-
36
- # Generate texts from the prompts
37
- outputs = []
38
- for prompt in example_prompts:
39
- # Tokenize the input prompt with truncation
40
- inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
41
- input_ids = inputs["input_ids"].to(model.device)
42
-
43
- # Generate text using the model's generate method directly
44
- generated_ids = model.generate(input_ids,
45
- max_new_tokens=max_tokens,
46
- do_sample=False)
47
- generated_text = tokenizer.decode(generated_ids[0],
48
- skip_special_tokens=True)
49
-
50
- outputs.append((generated_ids[0].tolist(), generated_text))
51
-
52
- return outputs
53
-
54
-
55
-@pytest.mark.parametrize("model", MODELS)
56
-@pytest.mark.parametrize("dtype", ["float"])
57
-@pytest.mark.parametrize("max_tokens", [96])
58
-def test_models(
59
- vllm_runner,
60
- example_prompts,
61
- model: str,
62
- dtype: str,
63
- max_tokens: int,
64
-) -> None:
65
- hf_outputs = generate_greedy(model, example_prompts, max_tokens)
66
-
67
- # Set max_num_seqs to keep Codestral from going OOM at fp32
68
- with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:
69
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
70
-
71
- for i in range(len(example_prompts)):
72
- hf_output_ids, hf_output_str = hf_outputs[i]
73
- vllm_output_ids, vllm_output_str = vllm_outputs[i]
74
- assert hf_output_str == vllm_output_str, (
75
- f"Test{i}:\nHF: {hf_output_str!r}\nvLLM: {vllm_output_str!r}")
76
- assert hf_output_ids == vllm_output_ids, (
77
- f"Test{i}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}")
78
-
79
-
80
-@pytest.mark.parametrize("model", MODELS)
81
-@pytest.mark.parametrize("dtype", ["float"])
82
-@pytest.mark.parametrize("max_tokens", [96])
83
-def test_batching(
84
- vllm_runner,
85
- example_prompts,
86
- model: str,
87
- dtype: str,
88
- max_tokens: int,
89
-) -> None:
90
- # To pass the small model tests, we need full precision.
91
- for_loop_outputs = []
92
- with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:
93
- for prompt in example_prompts:
94
- for_loop_outputs.append(
95
- vllm_model.generate_greedy([prompt], max_tokens)[0])
96
-
97
- batched_outputs = vllm_model.generate_greedy(example_prompts,
98
- max_tokens)
99
-
100
- check_outputs_equal(
101
- outputs_0_lst=for_loop_outputs,
102
- outputs_1_lst=batched_outputs,
103
- name_0="for_loop_vllm",
104
- name_1="batched_vllm",
105
- )
106
-
107
-
108
-@pytest.mark.parametrize("model", MODELS)
109
-@pytest.mark.parametrize("dtype", ["float"])
110
-@pytest.mark.parametrize("max_tokens", [10])
111
-def test_chunked_prefill_with_parallel_sampling(vllm_runner, example_prompts,
112
- model: str, dtype: str,
113
- max_tokens: int) -> None:
114
- # Tests chunked prefill in conjunction with n>1. In this case, prefill is
115
- # populated with decoding tokens and we test that it doesn't fail.
116
- # This test might fail if cache is not allocated correctly for n > 1
117
- # decoding steps inside a chunked prefill forward pass (where we have both
118
- # prefill and decode together )
119
- sampling_params = SamplingParams(n=3,
120
- temperature=1,
121
- seed=0,
122
- max_tokens=max_tokens)
123
- with vllm_runner(
124
- model,
125
- dtype=dtype,
126
- enable_chunked_prefill=True,
127
- max_num_batched_tokens=30,
128
- max_num_seqs=10 # forces prefill chunks with decoding
129
- ) as vllm_model:
130
- vllm_model.generate(example_prompts, sampling_params)
131
-
132
-
133
-@pytest.mark.parametrize("model", MODELS)
134
-@pytest.mark.parametrize("dtype", ["float"])
135
-@pytest.mark.parametrize("max_tokens", [32])
136
-@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
137
-def test_chunked_prefill(vllm_runner, example_prompts, model: str, dtype: str,
138
- max_tokens: int,
139
- chunked_prefill_token_size: int) -> None:
140
- """
141
- Checks exact match decode between huggingface model and vllm runner with
142
- chunked prefill.
143
- """
144
- max_num_seqs = chunked_prefill_token_size
145
- max_num_batched_tokens = chunked_prefill_token_size
146
-
147
- non_chunked = generate_greedy(model, example_prompts, max_tokens)
148
-
149
- with vllm_runner(model,
150
- dtype=dtype,
151
- enable_chunked_prefill=True,
152
- max_num_batched_tokens=max_num_batched_tokens,
153
- max_num_seqs=max_num_seqs) as vllm_model:
154
- chunked = vllm_model.generate_greedy(example_prompts,
155
- max_tokens=max_tokens)
156
-
157
- check_outputs_equal(
158
- outputs_0_lst=chunked,
159
- outputs_1_lst=non_chunked,
160
- name_0="chunked",
161
- name_1="non_chunked",
162
- )
163
-
164
-
165
-@pytest.mark.parametrize("model", MODELS)
166
-@pytest.mark.parametrize("dtype", ["float"])
167
-@pytest.mark.parametrize("max_tokens", [15])
168
-def test_parallel_sampling(
169
- vllm_runner,
170
- example_prompts,
171
- model: str,
172
- dtype: str,
173
- max_tokens: int,
174
-) -> None:
175
-
176
- # Numerical differences produce slightly different output for these
177
- if 'state-spaces' in model:
178
- example_prompts.pop(0)
179
- example_prompts.pop(0)
180
- example_prompts.pop(0)
181
-
182
- with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:
183
- for_loop_outputs = []
184
- for _ in range(10):
185
- for_loop_outputs.append(
186
- vllm_model.generate_greedy(example_prompts, max_tokens)[0])
187
- sampling_params = SamplingParams(n=10,
188
- temperature=0.001,
189
- seed=0,
190
- max_tokens=max_tokens)
191
- n_lt_1_outputs = vllm_model.generate(example_prompts, sampling_params)
192
- token_ids, texts = n_lt_1_outputs[0]
193
- n_lt_1_outputs = [(token_id, text)
194
- for token_id, text in zip(token_ids, texts)]
195
-
196
- check_outputs_equal(
197
- outputs_0_lst=n_lt_1_outputs,
198
- outputs_1_lst=for_loop_outputs,
199
- name_0="vllm_n_lt_1_outputs",
200
- name_1="vllm",
201
- )
202
-
203
-
204
-@pytest.mark.parametrize("model", MODELS)
205
-@pytest.mark.parametrize("dtype", ["bfloat16"])
206
-@pytest.mark.parametrize("max_tokens", [20])
207
-def test_mamba_cache_cg_padding(
208
- vllm_runner,
209
- example_prompts,
210
- model: str,
211
- dtype: str,
212
- max_tokens: int,
213
-) -> None:
214
- # This test is for verifying that mamba cache is padded to CG captured
215
- # batch size. If it's not, a torch RuntimeError will be raised because
216
- # tensor dimensions aren't compatible
217
- vllm_config = EngineArgs(model=model).create_engine_config()
218
- while len(example_prompts) == vllm_config.pad_for_cudagraph(
219
- len(example_prompts)):
220
- example_prompts.append(example_prompts[0])
221
-
222
- try:
223
- with vllm_runner(model, dtype=dtype) as vllm_model:
224
- vllm_model.generate_greedy(example_prompts, max_tokens)
225
- except RuntimeError:
226
- pytest.fail(
227
- "Couldn't run batch size which is not equal to a Cuda Graph "
228
- "captured batch size. "
229
- "Could be related to mamba cache not padded correctly")
230
-
231
-
232
-@pytest.mark.parametrize("model", MODELS)
233
-@pytest.mark.parametrize("dtype", ["float"])
234
-@pytest.mark.parametrize("max_tokens", [20])
235
-def test_models_preemption_recompute(
236
- vllm_runner,
237
- example_prompts,
238
- model: str,
239
- dtype: str,
240
- max_tokens: int,
241
-) -> None:
242
- # Tests that outputs are identical with and w/o preemtions (recompute)
243
- assert dtype == "float"
244
-
245
- with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:
246
- vllm_model.model.llm_engine.scheduler[
247
- 0].ENABLE_ARTIFICIAL_PREEMPT = True
248
- preempt_vllm_outputs = vllm_model.generate_greedy(
249
- example_prompts, max_tokens)
250
-
251
- vllm_model.model.llm_engine.scheduler[
252
- 0].ENABLE_ARTIFICIAL_PREEMPT = False
253
- vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
254
-
255
- check_outputs_equal(
256
- outputs_0_lst=preempt_vllm_outputs,
257
- outputs_1_lst=vllm_outputs,
258
- name_0="vllm_preepmtions",
259
- name_1="vllm",
260
- )
261
-
262
-
263
-@pytest.mark.parametrize("model", MODELS)
264
-@pytest.mark.parametrize("dtype", ["float"])
265
-def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
266
- vllm_runner,
267
- model: str,
268
- dtype: str,
269
- example_prompts,
270
-) -> None:
271
- # This test is for verifying that the Mamba inner state management doesn't
272
- # collapse in case where the number of incoming requests and
273
- # finished_requests_ids is larger than the maximum Mamba block capacity.
274
- # This could generally happen due to the fact that Mamba does support
275
- # statelessness mechanism where it can cleanup new incoming requests in
276
- # a single step.
277
- try:
278
- with vllm_runner(model, dtype=dtype, max_num_seqs=10) as vllm_model:
279
- vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
280
- except ValueError:
281
- pytest.fail("Mamba inner state wasn't cleaned up properly between"
282
- "steps finished requests registered unnecessarily ")
283
-
284
-
285
-@pytest.mark.parametrize("model", MODELS)
286
-@pytest.mark.parametrize("dtype", ["float"])
287
-def test_state_cleanup(
288
- vllm_runner,
289
- model: str,
290
- dtype: str,
291
- example_prompts,
292
-) -> None:
293
- # This test is for verifying that the Mamba state is cleaned up between
294
- # steps, If its not cleaned, an error would be expected.
295
- try:
296
- with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:
297
- for _ in range(10):
298
- vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
299
- except ValueError:
300
- pytest.fail("Mamba inner state wasn't cleaned up between states, "
301
- "could be related to finished_requests_ids")
302
-
303
-
304
-@pytest.mark.parametrize("model", MODELS)
305
-@pytest.mark.parametrize("dtype", ["float"])
306
-def test_multistep(
307
- vllm_runner,
308
- model: str,
309
- dtype: str,
310
- example_prompts,
311
-) -> None:
312
- with vllm_runner(model, num_scheduler_steps=8,
313
- max_num_seqs=2) as vllm_model:
314
- vllm_model.generate_greedy([example_prompts[0]] * 10, 1)
315
-
316
-
317
-@pytest.mark.parametrize("model", MODELS)
318
-@pytest.mark.parametrize("dtype", ["float"])
319
-@pytest.mark.parametrize("max_tokens", [64])
320
-def test_multistep_correctness(vllm_runner, model: str, dtype: str,
321
- max_tokens: int, example_prompts) -> None:
322
- with vllm_runner(model, num_scheduler_steps=8,
323
- max_num_seqs=2) as vllm_model:
324
- vllm_outputs_multistep = vllm_model.generate_greedy(
325
- example_prompts, max_tokens)
326
-
327
- with vllm_runner(model, num_scheduler_steps=1,
328
- max_num_seqs=2) as vllm_model:
329
- vllm_outputs_single_step = vllm_model.generate_greedy(
330
- example_prompts, max_tokens)
331
-
332
- check_outputs_equal(
333
- outputs_0_lst=vllm_outputs_multistep,
334
- outputs_1_lst=vllm_outputs_single_step,
335
- name_0="vllm_outputs_multistep",
336
- name_1="vllm_outputs_single_step",
337
- )