test_model_runner.py 14.5 KB
Newer Older
1
2
from typing import List

youkaichao's avatar
youkaichao committed
3
import pytest
4
5
import torch

6
7
from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
                                             init_distributed_environment)
8
from vllm.engine.arg_utils import EngineArgs
9
from vllm.model_executor.sampling_metadata import SamplingMetadata
10
from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
11
from vllm.utils import get_open_port
youkaichao's avatar
youkaichao committed
12
from vllm.worker.model_runner import ModelRunner, _get_graph_batch_size
13
14


15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def _create_model_runner(model: str, *args, **kwargs) -> ModelRunner:
    engine_args = EngineArgs(model, *args, **kwargs)
    engine_config = engine_args.create_engine_config()
    model_runner = ModelRunner(
        model_config=engine_config.model_config,
        parallel_config=engine_config.parallel_config,
        scheduler_config=engine_config.scheduler_config,
        device_config=engine_config.device_config,
        cache_config=engine_config.cache_config,
        load_config=engine_config.load_config,
        lora_config=engine_config.lora_config,
        is_driver_worker=True,
    )
    return model_runner


youkaichao's avatar
youkaichao committed
31
32
@pytest.mark.parametrize("batch_size", list(range(1, 257)))
def test_prepare_prompt(batch_size):
33
34
35
36
37
38
    model_runner = _create_model_runner(
        "facebook/opt-125m",
        max_num_batched_tokens=100000,
        max_num_seqs=100000,
        enable_chunked_prefill=False,
    )
Woosuk Kwon's avatar
Woosuk Kwon committed
39

40
41
    seq_lens: List[int] = []
    seq_group_metadata_list: List[SequenceGroupMetadata] = []
42
    block_tables = {0: [1]}
43
44
    for i in range(batch_size):
        # make sure all tokens fit into one block
45
46
47
        seq_len = i % (model_runner.block_size - 1) + 1
        seq_lens.append(seq_len)
        seq_data = SequenceData(list(range(seq_len)))
48
49
50
51
52
53
54
55
56
        seq_group_metadata = SequenceGroupMetadata(
            request_id=f"test_{i}",
            is_prompt=True,
            seq_data={0: seq_data},
            sampling_params=SamplingParams(temperature=0),
            block_tables=block_tables,
        )
        assert seq_group_metadata.token_chunk_size == seq_data.get_len()
        seq_group_metadata_list.append(seq_group_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
57

58
59
    expected_selected_token_indices = []
    selected_token_start_idx = 0
60
    for seq_len in seq_lens:
61
        expected_selected_token_indices.append(selected_token_start_idx +
62
63
                                               seq_len - 1)
        selected_token_start_idx += seq_len
64
65
66
67
68
69
    model_input = model_runner._prepare_model_input(seq_group_metadata_list)
    input_tokens = model_input.input_tokens
    input_positions = model_input.input_positions
    attn_metadata = model_input.attn_metadata
    return_seq_lens = model_input.seq_lens
    slot_mapping = model_input.slot_mapping
70
    assert return_seq_lens == seq_lens
71
    assert len(slot_mapping) == len(input_tokens)
72
73
74

    # Verify input metadata is correct for prompts.
    device = model_runner.device
75
76
    assert attn_metadata.num_prefills > 0
    assert attn_metadata.num_decode_tokens == 0
77
78
79
80
    assert torch.allclose(
        attn_metadata.seq_lens_tensor,
        torch.tensor(seq_lens, device=device, dtype=torch.int))
    assert attn_metadata.seq_lens == seq_lens
81
82
    assert attn_metadata.max_prefill_seq_len == max(seq_lens)
    assert attn_metadata.max_decode_seq_len == 0
83
84
85
86

    # Test subquery start locs.
    start_idx = 0
    start_loc = [start_idx]
87
88
    for seq_len in seq_lens:
        start_idx += seq_len
89
90
        start_loc.append(start_idx)
    assert torch.allclose(
91
        attn_metadata.query_start_loc,
92
93
94
        torch.tensor(start_loc, dtype=torch.int32, device=device))

    # Test seq start locs. Note that for normal prefill it is
95
    # equivalent to query_start_loc.
96
97
    start_idx = 0
    seq_start_loc = [start_idx]
98
99
    for seq_len in seq_lens:
        start_idx += seq_len
100
101
102
        seq_start_loc.append(start_idx)

    assert torch.allclose(
103
        attn_metadata.seq_start_loc,
104
105
        torch.tensor(start_loc, dtype=torch.int32, device=device))
    assert torch.allclose(
106
107
        attn_metadata.context_lens_tensor,
        torch.zeros(attn_metadata.context_lens_tensor.shape[0],
108
109
110
111
112
113
                    dtype=torch.int,
                    device=device))

    expected = torch.tensor([[] for _ in range(len(seq_group_metadata_list))],
                            dtype=torch.int32,
                            device=model_runner.device)
114
    assert torch.allclose(attn_metadata.block_tables, expected)
115
    # Cuda graph should not be used for prerill.
116
    assert attn_metadata.use_cuda_graph is False
117

118
119
    assert len(input_tokens) == sum(seq_lens)
    assert len(input_positions) == sum(seq_lens)
120
121
    torch.testing.assert_close(input_tokens, input_positions)

122
123
    sampling_metadata = SamplingMetadata.prepare(
        seq_group_metadata_list,
124
125
        seq_lens,
        query_lens=seq_lens,
126
127
        device=model_runner.device,
        pin_memory=model_runner.pin_memory)
128
129
    assert len(input_tokens) == sum(seq_lens)
    assert len(input_positions) == sum(seq_lens)
130
131
132
133
134
    actual = sampling_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    torch.testing.assert_close(actual, expected)
135
    torch.allclose(input_tokens, input_positions)
136
137
138
139
140
141
142
143

    actual = sampling_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    torch.testing.assert_close(actual, expected)


youkaichao's avatar
youkaichao committed
144
145
@pytest.mark.parametrize("batch_size", list(range(1, 257)))
def test_prepare_decode_cuda_graph(batch_size):
146
    model_runner = _create_model_runner(
147
148
149
150
        "facebook/opt-125m",
        seed=0,
        dtype="float16",
        enforce_eager=False,
151
152
153
        max_num_batched_tokens=100000,
        max_num_seqs=100000,
        enable_chunked_prefill=False,
154
155
    )

156
157
    context_lens: List[int] = []
    seq_group_metadata_list: List[SequenceGroupMetadata] = []
158
    # Assume each seq group finishes prefill.
159
160
    for i in range(batch_size):
        # make sure all tokens fit into one block
161
162
        context_len = i % (model_runner.block_size - 1) + 1
        context_lens.append(context_len)
163
        seq_data = SequenceData(list(range(context_len)))
164
165
166
        seq_data.update_num_computed_tokens(context_len)
        # Append one token ID since prefill is finished.
        seq_data.append_token_id(1, 0)
167
168
169
170
171
172
173
174
175
        seq_group_metadata = SequenceGroupMetadata(
            request_id=f"test_{i}",
            is_prompt=False,
            seq_data={0: seq_data},
            sampling_params=SamplingParams(temperature=0),
            block_tables={0: [1]},
        )
        assert seq_group_metadata.token_chunk_size == 1
        seq_group_metadata_list.append(seq_group_metadata)
176

177
178
179
180
    model_input = model_runner._prepare_model_input(seq_group_metadata_list)
    input_tokens, input_positions, attn_metadata, slot_mapping = (
        model_input.input_tokens, model_input.input_positions,
        model_input.attn_metadata, model_input.slot_mapping)
181
    assert len(slot_mapping) == len(input_tokens)
182

youkaichao's avatar
youkaichao committed
183
    expected_bs = _get_graph_batch_size(len(seq_group_metadata_list))
184
185
    # Verify input metadata is correct for prompts.
    device = model_runner.device
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    assert attn_metadata.num_prefills == 0
    assert attn_metadata.num_prefill_tokens == 0
    seq_lens = [context_len + 1 for context_len in context_lens]
    # seq_lens are padded to expected_bs
    for _ in range(expected_bs - len(seq_lens)):
        seq_lens.append(1)
    assert attn_metadata.seq_lens == seq_lens
    start_idx = 0
    start_loc = [start_idx]
    for _ in context_lens:
        # decode has only 1 token for query.
        start_idx += 1
        start_loc.append(start_idx)
    assert torch.allclose(
        attn_metadata.query_start_loc,
        torch.tensor(start_loc, dtype=torch.int32, device=device))

    start_idx = 0
    seq_start_loc = [start_idx]
    for seq_len in seq_lens:
        start_idx += seq_len
        seq_start_loc.append(start_idx)
    assert torch.allclose(
        attn_metadata.seq_start_loc,
        torch.tensor(seq_start_loc, dtype=torch.int32, device=device))

    assert torch.allclose(
        attn_metadata.context_lens_tensor,
        torch.tensor(context_lens, dtype=torch.int, device=device))
    assert attn_metadata.max_decode_seq_len == max(seq_lens)
216
    assert torch.allclose(
217
218
        attn_metadata.seq_lens_tensor[:len(seq_lens)],
        torch.tensor(seq_lens, dtype=torch.int, device=device))
219
220
221

    # block table's first index corresponds to each batch, meaning in
    # decoding it is each token.
222
    assert attn_metadata.block_tables.shape[0] == len(input_tokens)
223
224
    # Block table's second dim correspondsd to each token's block number.
    # It is padded up to
225
    assert attn_metadata.block_tables.shape[1] == (
226
        model_runner.get_max_block_per_batch())
227
    assert attn_metadata.use_cuda_graph is True
228

229
230
    assert len(input_tokens) == expected_bs
    assert len(input_positions) == expected_bs
231
    torch.allclose(input_tokens, input_positions)
Woosuk Kwon's avatar
Woosuk Kwon committed
232

233
234
235
    # Verify Sampling
    expected_selected_token_indices = []
    selected_token_start_idx = 0
236
    for _ in context_lens:
237
238
        expected_selected_token_indices.append(selected_token_start_idx)
        selected_token_start_idx += 1
239
240
    sampling_metadata = SamplingMetadata.prepare(
        seq_group_metadata_list,
241
        seq_lens,
242
243
        # query lens is all 1 for decode.
        query_lens=[1 for _ in range(len(context_lens))],
244
245
        device=model_runner.device,
        pin_memory=model_runner.pin_memory)
Woosuk Kwon's avatar
Woosuk Kwon committed
246
    actual = sampling_metadata.selected_token_indices
247
248
249
250
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    torch.testing.assert_close(actual, expected)
251
252
253
254


def test_empty_seq_group():
    """Verify prepare prompt and decode returns empty output."""
255
    model_runner = _create_model_runner(
256
257
258
259
260
        "facebook/opt-125m",
        seed=0,
        dtype="float16",
        enforce_eager=False,
    )
261
    seq_group_metadata_list: List[SequenceGroupMetadata] = []
262
263
264
265
266
267
268
    model_input = model_runner._prepare_model_input(seq_group_metadata_list)
    input_tokens, input_positions, attn_metadata, slot_mapping = (
        model_input.input_tokens,
        model_input.input_positions,
        model_input.attn_metadata,
        model_input.slot_mapping,
    )
269
270
271
272
273
    assert len(input_tokens) == 0
    assert len(input_positions) == 0
    assert attn_metadata is None
    assert len(slot_mapping) == 0

274
275
276
277
278
279
280
281
282
    model_input = model_runner._prepare_model_input(seq_group_metadata_list)
    (input_tokens, input_positions, attn_metadata, slot_mapping,
     return_seq_lens) = (
         model_input.input_tokens,
         model_input.input_positions,
         model_input.attn_metadata,
         model_input.slot_mapping,
         model_input.seq_lens,
     )
283
284
285
286
    assert len(input_tokens) == 0
    assert len(input_positions) == 0
    assert attn_metadata is None
    assert len(slot_mapping) == 0
287
    assert len(return_seq_lens) == 0
288
289


290
291
292
293
294
295
296
@pytest.fixture
def distributed_init():
    init_distributed_environment(
        world_size=1,
        rank=0,
        distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
        local_rank=0)
297
    ensure_model_parallel_initialized(1, 1)
298
299


300
301
302
@pytest.mark.parametrize("batch_size", list(range(2, 128)))
@pytest.mark.parametrize("enforce_eager", [True, False])
def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
303
    model_runner = _create_model_runner(
304
305
306
307
        "facebook/opt-125m",
        seed=0,
        dtype="float16",
        enforce_eager=enforce_eager,
308
309
310
        max_num_batched_tokens=100000,
        max_num_seqs=100000,
        enable_chunked_prefill=True,
311
312
313
    )

    # Add prefill requests.
314
315
316
317
    seq_lens: List[int] = []
    seq_group_metadata_list: List[SequenceGroupMetadata] = []
    prefill_metadata_list: List[SequenceGroupMetadata] = []
    decode_metadata_list: List[SequenceGroupMetadata] = []
318
319
320
321
322
    block_tables = {0: [1]}
    prefill_batch_size = batch_size // 2
    decode_batch_size = batch_size - prefill_batch_size
    for i in range(prefill_batch_size):
        # make sure all tokens fit into one block
323
324
325
        seq_len = i % (model_runner.block_size - 1) + 1
        seq_lens.append(seq_len)
        seq_data = SequenceData(list(range(seq_len)))
326
327
328
329
330
331
332
333
334
335
336
337
338
339
        seq_group_metadata = SequenceGroupMetadata(
            request_id=f"test_{i}",
            is_prompt=True,
            seq_data={0: seq_data},
            sampling_params=SamplingParams(temperature=0),
            block_tables=block_tables,
        )
        assert seq_group_metadata.token_chunk_size == seq_data.get_len()
        seq_group_metadata_list.append(seq_group_metadata)
        prefill_metadata_list.append(seq_group_metadata)

    # Add decode requests
    for i in range(prefill_batch_size, batch_size):
        # make sure all tokens fit into one block
340
341
        context_len = i % (model_runner.block_size - 1) + 1
        prompt_toks = list(range(context_len))
342
        seq_data = SequenceData(prompt_toks)
343
344
        seq_data.append_token_id(1, 0)
        seq_data.update_num_computed_tokens(context_len)
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
        seq_group_metadata = SequenceGroupMetadata(
            request_id=f"test_{i}",
            is_prompt=False,
            seq_data={0: seq_data},
            sampling_params=SamplingParams(temperature=0),
            block_tables={0: [1]},
        )
        assert seq_group_metadata.token_chunk_size == 1
        seq_group_metadata_list.append(seq_group_metadata)
        decode_metadata_list.append(seq_group_metadata)

    (input_tokens, input_positions, attn_metadata, _, _, _,
     _) = model_runner.prepare_input_tensors(seq_group_metadata_list)

    prefill_meta_actual = attn_metadata.prefill_metadata
    decode_meta_actual = attn_metadata.decode_metadata

    assert len(attn_metadata.slot_mapping) == len(input_tokens)
    assert len(input_positions) == len(input_tokens)
    assert attn_metadata.num_prefills == prefill_batch_size
365
    assert attn_metadata.num_decode_tokens == decode_batch_size
366
    assert attn_metadata.num_prefill_tokens == sum(seq_lens)
367
368
369

    # Verify attn metadata is consistent. We don't need to test individual
    # values here because they are tested above.
370
371
    attn_metadata = model_runner._prepare_model_input(
        seq_group_metadata_list).attn_metadata
372

373
    for attr_expected, attr_actual in zip(vars(attn_metadata.prefill_metadata),
374
375
                                          vars(prefill_meta_actual)):
        assert attr_expected[1] == attr_actual[1]
376
    for attr_expected, attr_actual in zip(vars(attn_metadata.decode_metadata),
377
378
                                          vars(decode_meta_actual)):
        assert attr_expected[1] == attr_actual[1]