Skip to content

vllm.spec_decode.interfaces

SpeculativeProposals dataclass

Datastructure used to represent proposal tokens from some proposer. It also tracks how many speculative tokens each sequence has.

Source code in vllm/spec_decode/interfaces.py
@dataclass
class SpeculativeProposals:
    """Datastructure used to represent proposal tokens from some proposer. It
    also tracks how many speculative tokens each sequence has.
    """

    # Speculative proposal tokens.
    proposal_token_ids: torch.Tensor

    # Probabilities of the proposal tokens according to the proposer.
    proposal_probs: torch.Tensor

    # The valid length of each proposal; can be zero.
    proposal_lens: torch.Tensor

    # A flag to mark that there's no available proposals
    no_proposals: bool = False

    def __repr__(self):
        return (f"SpeculativeProposals("
                f"proposal_token_ids={self.proposal_token_ids}, "
                f"proposal_probs={self.proposal_probs.shape}, "
                f"proposal_lens={self.proposal_lens})")

no_proposals class-attribute instance-attribute

no_proposals: bool = False

proposal_lens instance-attribute

proposal_lens: Tensor

proposal_probs instance-attribute

proposal_probs: Tensor

proposal_token_ids instance-attribute

proposal_token_ids: Tensor

__init__

__init__(
    proposal_token_ids: Tensor,
    proposal_probs: Tensor,
    proposal_lens: Tensor,
    no_proposals: bool = False,
) -> None

__repr__

__repr__()
Source code in vllm/spec_decode/interfaces.py
def __repr__(self):
    return (f"SpeculativeProposals("
            f"proposal_token_ids={self.proposal_token_ids}, "
            f"proposal_probs={self.proposal_probs.shape}, "
            f"proposal_lens={self.proposal_lens})")

SpeculativeProposer

Bases: ABC

Source code in vllm/spec_decode/interfaces.py
class SpeculativeProposer(ABC):

    @abstractmethod
    def get_spec_proposals(
        self,
        execute_model_req: ExecuteModelRequest,
        # If set, this contains all sequence IDs that were assigned
        # bonus tokens in their last forward pass.
        seq_ids_with_bonus_token_in_last_step: Set[int],
    ) -> SpeculativeProposals:
        raise NotImplementedError

get_spec_proposals abstractmethod

get_spec_proposals(
    execute_model_req: ExecuteModelRequest,
    seq_ids_with_bonus_token_in_last_step: Set[int],
) -> SpeculativeProposals
Source code in vllm/spec_decode/interfaces.py
@abstractmethod
def get_spec_proposals(
    self,
    execute_model_req: ExecuteModelRequest,
    # If set, this contains all sequence IDs that were assigned
    # bonus tokens in their last forward pass.
    seq_ids_with_bonus_token_in_last_step: Set[int],
) -> SpeculativeProposals:
    raise NotImplementedError

SpeculativeScorer

Bases: ABC

Source code in vllm/spec_decode/interfaces.py
class SpeculativeScorer(ABC):

    def __init__(self, scorer_worker: WorkerBase,
                 device: Union[torch.device, str], vocab_size: int):
        self._scorer_worker = scorer_worker
        if isinstance(device, torch.device):
            device = device.type
        self._device = device
        self._vocab_size = vocab_size

    @abstractmethod
    def score_proposals(
        self,
        execute_model_req: ExecuteModelRequest,
        proposals: SpeculativeProposals,
    ) -> SpeculativeScores:
        raise NotImplementedError

_device instance-attribute

_device = device

_scorer_worker instance-attribute

_scorer_worker = scorer_worker

_vocab_size instance-attribute

_vocab_size = vocab_size

__init__

__init__(
    scorer_worker: WorkerBase,
    device: Union[device, str],
    vocab_size: int,
)
Source code in vllm/spec_decode/interfaces.py
def __init__(self, scorer_worker: WorkerBase,
             device: Union[torch.device, str], vocab_size: int):
    self._scorer_worker = scorer_worker
    if isinstance(device, torch.device):
        device = device.type
    self._device = device
    self._vocab_size = vocab_size

score_proposals abstractmethod

score_proposals(
    execute_model_req: ExecuteModelRequest,
    proposals: SpeculativeProposals,
) -> SpeculativeScores
Source code in vllm/spec_decode/interfaces.py
@abstractmethod
def score_proposals(
    self,
    execute_model_req: ExecuteModelRequest,
    proposals: SpeculativeProposals,
) -> SpeculativeScores:
    raise NotImplementedError

SpeculativeScores dataclass

Datastructure used to represent the scores of speculative tokens according to the scoring model.

Source code in vllm/spec_decode/interfaces.py
@dataclass
class SpeculativeScores:
    """Datastructure used to represent the scores of speculative tokens
    according to the scoring model.
    """

    # Probabilities of the speculative tokens according to the scoring model.
    probs: torch.Tensor

    # Log-probabilities of the speculative tokens according to the scoring
    # model. These values can be used to generate Logprob objects that are
    # returned to the user.
    logprobs: torch.Tensor

    # Token ids sampled from the scoring model. Used for speculative bonus
    # tokens and also non-speculative normal decoding.
    token_ids: torch.Tensor

    # Optional last hidden states from the scoring model.
    hidden_states: Optional[torch.Tensor] = None

    # Scoring model may also return logprobs for prompt tokens
    # for each request, when chunked prefill is enabled.
    prompt_logprobs: Optional[List[PromptLogprobs]] = None

    def __repr__(self):
        return (f"SpeculativeScores("
                f"probs={self.probs.shape}, "
                f"token_ids={self.token_ids.shape})")

hidden_states class-attribute instance-attribute

hidden_states: Optional[Tensor] = None

logprobs instance-attribute

logprobs: Tensor

probs instance-attribute

probs: Tensor

prompt_logprobs class-attribute instance-attribute

prompt_logprobs: Optional[List[PromptLogprobs]] = None

token_ids instance-attribute

token_ids: Tensor

__init__

__init__(
    probs: Tensor,
    logprobs: Tensor,
    token_ids: Tensor,
    hidden_states: Optional[Tensor] = None,
    prompt_logprobs: Optional[List[PromptLogprobs]] = None,
) -> None

__repr__

__repr__()
Source code in vllm/spec_decode/interfaces.py
def __repr__(self):
    return (f"SpeculativeScores("
            f"probs={self.probs.shape}, "
            f"token_ids={self.token_ids.shape})")