43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
import torch
|
|
from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
|
|
|
|
from vllm_ascend.utils import vllm_version_is
|
|
|
|
|
|
class AscendSuffixDecodingProposer(SuffixDecodingProposer):
|
|
def __init__(self, vllm_config, runner):
|
|
super().__init__(vllm_config)
|
|
self.runner = runner
|
|
|
|
def dummy_run(
|
|
self,
|
|
num_tokens,
|
|
with_prefill=None,
|
|
in_graph_capturing=None,
|
|
num_reqs=None,
|
|
num_tokens_across_dp=None,
|
|
aclgraph_runtime_mode=None,
|
|
batch_descriptor=None,
|
|
dummy_compute_logits=lambda hidden_states: None,
|
|
is_profile=False,
|
|
):
|
|
pass
|
|
|
|
def propose(
|
|
self,
|
|
sampled_token_ids: list[list[int]],
|
|
num_tokens_no_spec=None,
|
|
token_ids_cpu=None,
|
|
num_speculative_tokens: int = 0,
|
|
slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
|
|
):
|
|
if vllm_version_is("0.23.0"):
|
|
return super().propose(self.runner.input_batch, sampled_token_ids)
|
|
else:
|
|
return super().propose(
|
|
num_speculative_tokens,
|
|
self.runner.input_batch,
|
|
sampled_token_ids,
|
|
slot_mappings,
|
|
)
|