sampling_metadata.py 22.4 KB
Newer Older
1
from array import array
2
from dataclasses import dataclass
3
from typing import Dict, List, Optional, Tuple
4
5
6

import torch

7
from vllm.sampling_params import SamplingParams, SamplingType
8
9
from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData,
                           SequenceGroupMetadata)
10
from vllm.utils import (PyObjectCache, async_tensor_h2d,
11
                        is_pin_memory_available, make_tensor_with_pad)
12
13

_SAMPLING_EPS = 1e-5
14
15


16
17
@dataclass
class SequenceGroupToSample:
18
19
20
21
22
23
24
    # |---------- N-1 iteration --------|
    # |---------------- N iteration ---------------------|
    # |- tokenA -|......................|-- newTokens ---|
    # |---------- context_len ----------|
    # |-------------------- seq_len ----------------------|
    #                                   |-- query_len ---|

25
26
27
28
29
    # Sequence ids for the sequence group in a previous step.
    seq_ids: List[int]
    sampling_params: SamplingParams
    # seq_id -> sequence data.
    seq_data: Dict[int, SequenceData]
30
31
    # The length of the sequence (all tokens seen in the past + new token to
    # compute attention) of the sequence group. None if it is in a decode
32
    # stage.
33
34
35
36
37
    seq_len: Optional[int]
    # The length of new query tokens to compute in the current step. None if it
    # is in a decode stage. The length of query_len <= seq_len if chunked
    # prefill is enabled.
    query_len: Optional[int]
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
    # A random number generator for sampling.
    generator: Optional[torch.Generator]
    # True if the sequence group is in prefill stage. False if it is in a
    # decode stage.
    is_prompt: bool
    # Query token indices from logits. to compute prompt logprob. Empty if
    # prompt logprob is not required.
    prompt_logprob_indices: List[int]
    # Sample token indices from logits. Empty if sampling is not required.
    sample_indices: List[int]

    @property
    def do_sample(self):
        return len(self.sample_indices) > 0

    def __post_init__(self):
        if len(self.prompt_logprob_indices) > 0:
            assert self.sampling_params.prompt_logprobs is not None
        if self.is_prompt:
57
58
            assert self.seq_len is not None
            assert self.query_len is not None
59
60


61
62
63
64
65
66
67
68
69
70
def gen_seq_group_to_sample_builder(num_seqs: int):
    return lambda: SequenceGroupToSample(
        seq_ids=[0] * num_seqs,
        sampling_params=None,
        seq_data=None,  # type: ignore
        seq_len=0,
        query_len=0,
        generator=None,
        is_prompt=True,
        prompt_logprob_indices=[],
71
72
        sample_indices=[],
    )
73
74
75


class SamplingMetadataCache:
76
    """Used to cache SamplingMetadata objects between scheduler iterations"""
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

    def __init__(self):
        self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {}

    def get_cached_seq_group_to_sample(self, num_seqs):
        if num_seqs not in self._seq_group_to_sample_cache:
            self._seq_group_to_sample_cache[num_seqs] = PyObjectCache(
                gen_seq_group_to_sample_builder(num_seqs))

        obj = self._seq_group_to_sample_cache[num_seqs].get_object()
        return obj

    def reset(self):
        for cache in self._seq_group_to_sample_cache.values():
            cache.reset()


94
95
96
class SamplingMetadata:
    """Metadata for input sequences. Used in sampler.

97
98
99
100
101
102
103
104
105
106
    The usage is as follow;
    ```
    hidden_states = execute_model(...)
    logits = hidden_states[sampling_metadata.selected_token_indices]
    sample(logits)

    def sample(logits):
        # Use categorized_sample_indices for sampling....
    ```

107
    Args:
108
109
110
        seq_groups: List of batched sequence groups.
        selected_token_indices: (num_query_tokens_to_logprob). Indices to find
            logits from the initial model output hidden states.
111
        categorized_sample_indices: SamplingType -> token indices to sample.
112
113
114
115
116
117
118
119
120
            Each token indices is 2D tensor of (num_indices, num_indices) where
            the first item means the sample index within the returned logit
            (before pruning padding), and the second item means the sample
            index after pruning using selected_token_indices.
            For example, if the returned logit is [1, 2, 3], and we select
            [1, 2] for sampling, the pruned logit will be [2, 3]. In this case,
            The first tuple is [1, 2] (sampled index within original logit),
            and the second tuple is [0, 1] (sampled index within pruned logit).
        num_prompts: Number of prompt sequence groups in seq_groups.
121
        skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU
122
            serialization of token outputs.
123
        reuse_sampling_tensors: Indicates if we want to reuse sampling
124
125
            tensors that are part of the sampler forward pass. Currently,
            it is mainly used for multi-step decode.
126

127
128
129
130
    """

    def __init__(
        self,
131
        seq_groups: List[SequenceGroupToSample],
132
        selected_token_indices: torch.Tensor,
133
134
        categorized_sample_indices: Dict[SamplingType, torch.Tensor],
        num_prompts: int,
135
136
        skip_sampler_cpu_output: bool = False,
        reuse_sampling_tensors: bool = False,
137
138
139
140
    ) -> None:
        self.seq_groups = seq_groups
        self.selected_token_indices = selected_token_indices
        self.categorized_sample_indices = categorized_sample_indices
141
        self.num_prompts = num_prompts
142
143
        self.skip_sampler_cpu_output = skip_sampler_cpu_output
        self.reuse_sampling_tensors = reuse_sampling_tensors
144

145
146
147
    @staticmethod
    def prepare(
        seq_group_metadata_list: List[SequenceGroupMetadata],
148
        seq_lens: List[int],
149
        query_lens: List[int],
150
151
        device: str,
        pin_memory: bool,
152
        generators: Optional[Dict[str, torch.Generator]] = None,
153
        cache: Optional[SamplingMetadataCache] = None,
154
155
156
157
158
159
    ) -> "SamplingMetadata":
        (
            seq_groups,
            selected_token_indices,
            categorized_sample_indices,
            num_prompts,
160
        ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,
161
                                device, generators, cache)
162
163
164
165
166
167
        selected_token_indices = async_tensor_h2d(
            selected_token_indices,
            dtype=torch.long,
            target_device=device,
            pin_memory=pin_memory,
        )
168
        categorized_sample_indices = {
169
170
            t:
            async_tensor_h2d(
171
172
173
174
175
                seq_ids,
                dtype=torch.int,
                target_device=device,
                pin_memory=pin_memory,
            )
176
177
178
179
180
181
182
183
184
185
            for t, seq_ids in categorized_sample_indices.items()
        }

        sampling_metadata = SamplingMetadata(
            seq_groups=seq_groups,
            selected_token_indices=selected_token_indices,
            categorized_sample_indices=categorized_sample_indices,
            num_prompts=num_prompts,
        )
        return sampling_metadata
186
187
188
189
190
191

    def __repr__(self) -> str:
        return (
            "SamplingMetadata("
            f"seq_groups={self.seq_groups}, "
            f"selected_token_indices={self.selected_token_indices}, "
192
193
194
195
196
            f"categorized_sample_indices={self.categorized_sample_indices}), ")


def _prepare_seq_groups(
    seq_group_metadata_list: List[SequenceGroupMetadata],
197
    seq_lens: List[int],
198
    query_lens: List[int],
199
    device: str,
200
    generators: Optional[Dict[str, torch.Generator]] = None,
201
    cache: Optional[SamplingMetadataCache] = None,
202
203
204
205
206
207
) -> Tuple[
        List[SequenceGroupToSample],
        List[int],
        Dict[SamplingType, List[int]],
        int,
]:
208
209
210
211
    """Prepare sequence groups and indices for sampling.

    Args:
        seq_group_metadata_list: A list of sequence group to batch.
212
        seq_lens: A list of sequence lens per sequence group.
213
            Index of prompt len should match with seq_group_metadata_list.
214
        query_lens: A list of query lengths. Prompt lens include the length
215
            of entire prompt tokens, and it could be shorter.
216
        device: A device to use for random number generators,
217
            `SequenceGroupToSample.generator`.
218
219
        generators: A store of per-request random number generators used
            for seeded requests.
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237

    Returns:
        seq_groups: A list of sequence group to sample.
        selected_token_indices: See the definition from `SamplingMetadata`.
        categorized_sample_indices: See the definition from `SamplingMetadata`.
        num_prompts: Total number of prompts from `seq_group_metadata_list`.
    """
    # Batched sequence groups for the current model forward stsep.
    seq_groups: List[SequenceGroupToSample] = []
    # A list of token indices to sample/compute logprob. It is used to
    # prune the outcome logits from the model for the performance.
    selected_token_indices: List[int] = []
    # Used for selected_token_indices.
    model_output_idx = 0

    # Sampling type -> (
    # indices to sample/prompt logprob within pruned output logits,
    # indices to sample within pruned logits)
238
    categorized_sample_indices: Dict[SamplingType, List[int]] = {
239
240
241
242
243
244
245
246
247
248
        t: []
        for t in SamplingType
    }
    # Index of logits to compute logprob. Logits include both prompt logprob
    # and sample logprob indices.
    logit_idx = 0
    # Total number of prompts from given sequence groups.
    num_prompts = 0

    for i, seq_group_metadata in enumerate(seq_group_metadata_list):
249
250
251
252
253
254
255
256
257
258
259
        seq_ids = seq_group_metadata.seq_data.keys()

        if cache is not None:
            sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids))

            for j, seq_id in enumerate(seq_ids):
                sample_obj.seq_ids[j] = seq_id

            sample_obj.prompt_logprob_indices.clear()
            sample_obj.sample_indices.clear()

260
261
262
263
        sampling_params = seq_group_metadata.sampling_params
        is_prompt = seq_group_metadata.is_prompt
        generator: Optional[torch.Generator] = None
        # If the current seq group is in decode stage, it is None.
264
265
        seq_len: Optional[int] = None
        query_len: Optional[int] = None
266
267
268
269
        prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices
                                             if cache is not None else [])
        sample_indices: List[int] = (sample_obj.sample_indices
                                     if cache is not None else [])
270
271
272
273
        do_sample = seq_group_metadata.do_sample

        if seq_group_metadata.is_prompt:
            if sampling_params.seed is not None:
274
275
276
277
                generator = torch.Generator(device=device).manual_seed(
                    sampling_params.seed)
                if generators is not None:
                    generators[seq_group_metadata.request_id] = generator
278
279
280
281

            num_prompts += 1
            num_prefill_sample = len(seq_ids)
            assert num_prefill_sample == 1
282
283
            assert query_lens is not None and seq_lens is not None
            query_len, seq_len = query_lens[i], seq_lens[i]
284
285
            # If we need sampling, exclude num_prefill_sample tokens from
            # prompt logprob.
286
287
            prompt_logprob_len = (query_len - num_prefill_sample
                                  if do_sample else query_len)
288
289
290
291
            sample_len = num_prefill_sample if do_sample else 0
        else:
            # Decode
            prompt_logprob_len = 0
292
293
            query_len = query_lens[i] if query_lens is not None and len(
                query_lens) > 0 else 1
294
            sample_len = len(seq_ids) * query_len if do_sample else 0
295

296
297
298
            if sampling_params.seed is not None and generators is not None:
                generator = generators.get(seq_group_metadata.request_id)

299
300
301
302
303
304
305
306
307
        # Update indices to select from the model output.
        """
        This blocks computes selected_token_indices which is used in the
        following way.

        hidden_states = model(...)
        logits = hidden_states[selected_token_indices]
        """

308
        if sampling_params.prompt_logprobs is not None:
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
            selected_token_indices.extend(
                range(model_output_idx, model_output_idx + prompt_logprob_len))
        model_output_idx += prompt_logprob_len
        if do_sample:
            selected_token_indices.extend(
                range(model_output_idx, model_output_idx + sample_len))
        model_output_idx += sample_len

        # We now find indices for logprob computation and sampling.
        """
        This block computes categorized_sample_indices which is used in the
        following way.

        hidden_states = model(...)
        logits = hidden_states[selected_token_indices]
        def sample(logits):
           # Use categorized_sample_indices for sampling.
           # prompt_logprob_indices to find prompt logprob indices.
           # sample_indices to find sample indices.
        """

        if sampling_params.prompt_logprobs is not None:
            prompt_logprob_indices.extend(
                range(logit_idx, logit_idx + prompt_logprob_len))
            logit_idx += prompt_logprob_len
        if do_sample:
            sample_indices.extend(range(logit_idx, logit_idx + sample_len))
            categorized_sample_indices[sampling_params.sampling_type].extend(
337
                list(range(logit_idx, logit_idx + sample_len)))
338
339
            logit_idx += sample_len

340
341
342
343
344
345
346
347
348
349
        if cache is not None:
            sample_obj.sampling_params = sampling_params
            sample_obj.seq_data = seq_group_metadata.seq_data
            sample_obj.seq_len = seq_len
            sample_obj.query_len = query_len
            sample_obj.generator = generator
            sample_obj.is_prompt = is_prompt
        else:
            sample_obj = SequenceGroupToSample(
                seq_ids=list(seq_ids),
350
351
                sampling_params=sampling_params,
                seq_data=seq_group_metadata.seq_data,
352
353
                seq_len=seq_len,
                query_len=query_len,
354
355
356
                generator=generator,
                is_prompt=is_prompt,
                prompt_logprob_indices=list(prompt_logprob_indices),
357
358
                sample_indices=list(sample_indices),
            )
359
360
361
362
363
364

        seq_groups.append(sample_obj)

    if cache is not None:
        cache.reset()

365
366
    return (seq_groups, selected_token_indices, categorized_sample_indices,
            num_prompts)
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384


@dataclass
class SamplingTensors:
    """Tensors for sampling."""

    temperatures: torch.Tensor
    top_ps: torch.Tensor
    top_ks: torch.Tensor
    min_ps: torch.Tensor
    presence_penalties: torch.Tensor
    frequency_penalties: torch.Tensor
    repetition_penalties: torch.Tensor
    prompt_tokens: torch.Tensor
    output_tokens: torch.Tensor

    @classmethod
    def from_sampling_metadata(
385
386
387
388
389
390
        cls,
        sampling_metadata: "SamplingMetadata",
        vocab_size: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> Tuple["SamplingTensors", bool, bool, bool]:
391
392
        prompt_tokens: List[array] = []
        output_tokens: List[array] = []
393
394
395
396
397
398
399
400
401
402
        top_ks: List[int] = []
        temperatures: List[float] = []
        top_ps: List[float] = []
        min_ps: List[float] = []
        presence_penalties: List[float] = []
        frequency_penalties: List[float] = []
        repetition_penalties: List[float] = []
        do_penalties = False
        do_top_p_top_k = False
        do_min_p = False
403

404
        assert sampling_metadata.seq_groups is not None
405
406
407
        for seq_group in sampling_metadata.seq_groups:
            seq_ids = seq_group.seq_ids
            sampling_params = seq_group.sampling_params
408
409
410
411
412
413
            temperature = sampling_params.temperature
            p = sampling_params.presence_penalty
            f = sampling_params.frequency_penalty
            r = sampling_params.repetition_penalty
            top_p = sampling_params.top_p
            min_p = sampling_params.min_p
414

415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
            # k should not be greater than the vocab size.
            top_k = min(sampling_params.top_k, vocab_size)
            top_k = vocab_size if top_k == -1 else top_k
            if temperature < _SAMPLING_EPS:
                # NOTE: Zero temperature means deterministic sampling
                # (i.e., greedy sampling or beam search).
                # Set the temperature to 1 to avoid division by zero.
                temperature = 1.0
            if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
                                       or top_k != vocab_size):
                do_top_p_top_k = True
            if not do_min_p and min_p > _SAMPLING_EPS:
                do_min_p = True
            if not do_penalties and (abs(p) >= _SAMPLING_EPS
                                     or abs(f) >= _SAMPLING_EPS
                                     or abs(r - 1.0) >= _SAMPLING_EPS):
                do_penalties = True
432

433
            is_prompt = seq_group.is_prompt
434
            if is_prompt and sampling_params.prompt_logprobs is not None:
435
436
                # For tokens in the prompt that we only need to get
                # their logprobs
437
438
                query_len = seq_group.query_len
                assert query_len is not None
439
440
441
442
443
444
445
446
447
448
449
                prefill_len = len(seq_group.prompt_logprob_indices)
                temperatures += [temperature] * prefill_len
                top_ps += [top_p] * prefill_len
                top_ks += [top_k] * prefill_len
                min_ps += [min_p] * prefill_len
                presence_penalties += [0] * prefill_len
                frequency_penalties += [0] * prefill_len
                repetition_penalties += [1] * prefill_len

            if seq_group.do_sample:
                sample_lens = len(seq_group.sample_indices)
450
451
452
453
454
455
456
457
                assert sample_lens >= len(seq_ids)
                temperatures += [temperature] * sample_lens
                top_ps += [top_p] * sample_lens
                top_ks += [top_k] * sample_lens
                min_ps += [min_p] * sample_lens
                presence_penalties += [p] * sample_lens
                frequency_penalties += [f] * sample_lens
                repetition_penalties += [r] * sample_lens
458

459
460
461
        if do_penalties:
            for seq_group in sampling_metadata.seq_groups:
                seq_ids = seq_group.seq_ids
462
                sampling_params = seq_group.sampling_params
463
464
465
                if (seq_group.is_prompt
                        and sampling_params.prompt_logprobs is not None):
                    prefill_len = len(seq_group.prompt_logprob_indices)
466
                    prompt_tokens.extend(
467
468
                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
                        for _ in range(prefill_len))
469
                    output_tokens.extend(
470
471
                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
                        for _ in range(prefill_len))
472
473
474
                if seq_group.do_sample:
                    for seq_id in seq_ids:
                        seq_data = seq_group.seq_data[seq_id]
475
476
                        prompt_tokens.append(seq_data.prompt_token_ids_array)
                        output_tokens.append(seq_data.output_token_ids_array)
477

478
        sampling_tensors = SamplingTensors.from_lists(
479
480
481
482
483
484
485
486
487
488
489
490
491
            temperatures,
            top_ps,
            top_ks,
            min_ps,
            presence_penalties,
            frequency_penalties,
            repetition_penalties,
            prompt_tokens,
            output_tokens,
            vocab_size,
            device,
            dtype,
        )
492
493
494
        return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)

    @classmethod
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
    def from_lists(
        cls,
        temperatures: List[float],
        top_ps: List[float],
        top_ks: List[int],
        min_ps: List[float],
        presence_penalties: List[float],
        frequency_penalties: List[float],
        repetition_penalties: List[float],
        prompt_tokens: List[array],
        output_tokens: List[array],
        vocab_size: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> "SamplingTensors":
510
511
        # Note that the performance will be very bad without
        # pinned memory.
512
        pin_memory = is_pin_memory_available()
513
514
515
516

        do_penalties = prompt_tokens or output_tokens

        if do_penalties:
517
518
            prompt_t = make_tensor_with_pad(
                prompt_tokens,
519
                vocab_size,
520
521
522
523
524
525
                device="cpu",
                dtype=torch.int64,
                pin_memory=pin_memory,
            )
            output_t = make_tensor_with_pad(
                output_tokens,
526
                vocab_size,
527
528
529
530
531
532
533
534
                device="cpu",
                dtype=torch.int64,
                pin_memory=pin_memory,
            )
        else:
            empty_tensor = torch.empty(0, device=device, dtype=torch.long)
            prompt_t = empty_tensor
            output_t = empty_tensor
535
536
537
538
539
540
541
542
543
544
545
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
571
572
573
574
575
576
577
578
579

        temperatures_t = torch.tensor(
            temperatures,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        top_ps_t = torch.tensor(
            top_ps,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        min_ps_t = torch.tensor(
            min_ps,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        presence_penalties_t = torch.tensor(
            presence_penalties,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        frequency_penalties_t = torch.tensor(
            frequency_penalties,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        repetition_penalties_t = torch.tensor(
            repetition_penalties,
            device="cpu",
            dtype=dtype,
            pin_memory=pin_memory,
        )
        top_ks_t = torch.tensor(
            top_ks,
            device="cpu",
            dtype=torch.int,
            pin_memory=pin_memory,
        )
        # Because the memory is pinned, we can do non-blocking
        # transfer to device.
580

581
582
583
584
585
586
587
588
589
590
591
        return cls(
            temperatures=temperatures_t.to(device=device, non_blocking=True),
            top_ps=top_ps_t.to(device=device, non_blocking=True),
            top_ks=top_ks_t.to(device=device, non_blocking=True),
            min_ps=min_ps_t.to(device=device, non_blocking=True),
            presence_penalties=presence_penalties_t.to(device=device,
                                                       non_blocking=True),
            frequency_penalties=frequency_penalties_t.to(device=device,
                                                         non_blocking=True),
            repetition_penalties=repetition_penalties_t.to(device=device,
                                                           non_blocking=True),
592
593
            prompt_tokens=prompt_t.to(device=device, non_blocking=True),
            output_tokens=output_t.to(device=device, non_blocking=True),
594
        )