118 lines
5.1 KiB
Python
118 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
from typing import Set
|
||
from vllm.sequence import ExecuteModelRequest
|
||
from vllm.spec_decode.interfaces import (SpeculativeProposals,
|
||
SpeculativeProposer)
|
||
import torch
|
||
|
||
from .var import *
|
||
|
||
class Top1Proposer(SpeculativeProposer):
|
||
|
||
def get_spec_proposals(
|
||
self,
|
||
execute_model_req: ExecuteModelRequest,
|
||
seq_ids_with_bonus_token_in_last_step: Set[int],
|
||
) -> SpeculativeProposals:
|
||
"""Get speculative proposals given the input batch.
|
||
|
||
Sequences which would exceed the max model length are skipped during
|
||
speculation.
|
||
"""
|
||
proposal_len = execute_model_req.num_lookahead_slots
|
||
seq_group_metadata_list = execute_model_req.seq_group_metadata_list
|
||
|
||
# Split speculative- and non-speculative- sequences.
|
||
(
|
||
proposal_lens,
|
||
nonzero_proposal_len_seqs,
|
||
nonzero_proposal_len_indices,
|
||
) = self._split_by_proposal_len(seq_group_metadata_list, proposal_len)
|
||
|
||
if nonzero_proposal_len_seqs:
|
||
# Speculate tokens using the draft worker for the speculative
|
||
# sequences.
|
||
# If sampler_transposed is true, then maybe_sampler_output's
|
||
# token_ids is like [batch] format in proposal_len size list,
|
||
# while if it is false, the format would be [proposal_len]
|
||
# in batch size list
|
||
hidden_states = execute_model_req.previous_hidden_states
|
||
if hidden_states is not None:
|
||
hidden_states.prune(nonzero_proposal_len_seqs)
|
||
nonzero_execute_model_req = ExecuteModelRequest(
|
||
seq_group_metadata_list=nonzero_proposal_len_seqs,
|
||
num_lookahead_slots=proposal_len,
|
||
previous_hidden_states=hidden_states,
|
||
)
|
||
|
||
#思路，将sampler out和merge output合成1个OP
|
||
#remove no proposal_seqs 该流程没必要，直接全是草稿输出，不存在混合
|
||
#功能1： 筛选 indices
|
||
#功能2： 返回 tensor
|
||
#功能3： 对tensor进行transpose
|
||
if USE_FUSED_MTP_SAMPLER:
|
||
sampler_outputs, token_indices, transposed = self._worker.sampler_output(
|
||
execute_model_req=nonzero_execute_model_req,
|
||
sample_len=proposal_len,
|
||
seq_ids_with_bonus_token_in_last_step=\
|
||
seq_ids_with_bonus_token_in_last_step,
|
||
)
|
||
|
||
outputs = sampler_outputs[0].outputs
|
||
token_probs = sampler_outputs[0].sampled_token_probs
|
||
token_ids = sampler_outputs[0].sampled_token_ids
|
||
|
||
bs = len(seq_group_metadata_list)
|
||
proposal_lens = torch.ones((bs), dtype=torch.int, device=self._worker.device)
|
||
|
||
proposal_tokens = token_ids[token_indices]
|
||
proposal_probs = token_probs[token_indices]
|
||
|
||
s0, s1 = proposal_probs.shape
|
||
proposal_probs = proposal_probs.view(s0, 1, s1)
|
||
|
||
# 筛选indices，构建新的SamplerOut
|
||
proposals = SpeculativeProposals(proposal_token_ids=proposal_tokens,
|
||
proposal_probs=proposal_probs,
|
||
proposal_lens=proposal_lens,
|
||
no_proposals=outputs
|
||
is None)
|
||
return proposals
|
||
|
||
maybe_sampler_output, transposed = self._worker.sampler_output(
|
||
execute_model_req=nonzero_execute_model_req,
|
||
sample_len=proposal_len,
|
||
seq_ids_with_bonus_token_in_last_step=\
|
||
seq_ids_with_bonus_token_in_last_step,
|
||
)
|
||
(
|
||
proposal_lens,
|
||
maybe_sampler_output,
|
||
nonzero_proposal_len_indices,
|
||
) = self._remove_no_proposal_seqs(proposal_lens,
|
||
maybe_sampler_output,
|
||
nonzero_proposal_len_indices,
|
||
transposed)
|
||
else:
|
||
# If no sequences can be speculated, set sampler output to None.
|
||
maybe_sampler_output = None
|
||
transposed = False
|
||
|
||
# Combine speculative- and non-speculative sequences into the same
|
||
# representation.
|
||
proposal_tokens, proposal_probs, proposal_lens = self._merge_outputs(
|
||
batch_size=len(seq_group_metadata_list),
|
||
proposal_len=proposal_len,
|
||
maybe_sampler_output=maybe_sampler_output,
|
||
proposal_lens=proposal_lens,
|
||
nonzero_proposal_len_indices=nonzero_proposal_len_indices,
|
||
sampler_transposed=transposed,
|
||
)
|
||
|
||
proposals = SpeculativeProposals(proposal_token_ids=proposal_tokens,
|
||
proposal_probs=proposal_probs,
|
||
proposal_lens=proposal_lens,
|
||
no_proposals=maybe_sampler_output
|
||
is None)
|
||
return proposals |