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

4
from abc import ABC, abstractmethod
5
from typing import Callable, List
6
7
8
9
10
11

from vllm.config import SchedulerConfig
from vllm.core.scheduler import Scheduler
from vllm.engine.output_processor.stop_checker import StopChecker
from vllm.sequence import Sequence, SequenceGroup, SequenceGroupOutput
from vllm.transformers_utils.detokenizer import Detokenizer
12
from vllm.transformers_utils.tokenizer import AnyTokenizer
13
from vllm.utils import Counter
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


class SequenceGroupOutputProcessor(ABC):
    """Interface for logic that processes new token ids in sequence groups,
    managing detokenization, stop checking, and freeing/forking sequences with
    the scheduler.

    This is highly coupled with the LLMEngine and should be seen as an extension
    of it. The logic is separated to simplify the LLMEngine class and allow
    separate implementations for single-step decoding (which supports beam
    search sequence forking) and multi-step decoding (which does not support
    beam search, but does support speculative decoding).
    """

    @staticmethod
    def create_output_processor(
        scheduler_config: SchedulerConfig,
        detokenizer: Detokenizer,
32
        scheduler: List[Scheduler],
33
        seq_counter: Counter,
34
        get_tokenizer_for_seq: Callable[[Sequence], AnyTokenizer],
35
36
37
38
        stop_checker: "StopChecker",
    ):
        """Create an output processor.

39
40
        Multi-step scheduling is no longer supported. Always return a
        single-step output processor.
41
        """
42
43
44
45
        from vllm.engine.output_processor.single_step import (
            SingleStepOutputProcessor)
        return SingleStepOutputProcessor(scheduler_config, detokenizer,
                                         scheduler, seq_counter, stop_checker)
46
47
48

    @abstractmethod
    def process_outputs(self, sequence_group: SequenceGroup,
49
                        outputs: List[SequenceGroupOutput],
50
                        is_async: bool) -> None:
51
52
53
54
55
        """Process new token ids for the sequence group. Handles logic such as
        detokenization, stop checking, and freeing/forking sequences in the
        scheduler.
        """
        pass
56
57
58
59
60
61

    @abstractmethod
    def process_prompt_logprob(self, seq_group: SequenceGroup,
                               outputs: List[SequenceGroupOutput]) -> None:
        """Update prompt logprobs received from outputs to seq_group."""
        pass