80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
from abc import ABC, abstractmethod
|
|
from enum import IntEnum, auto
|
|
from typing import List, Tuple
|
|
|
|
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
|
|
|
|
|
|
class SpeculativeAlgorithm(IntEnum):
|
|
NONE = auto()
|
|
EAGLE = auto()
|
|
EAGLE3 = auto()
|
|
STANDALONE = auto()
|
|
NGRAM = auto()
|
|
|
|
def is_none(self):
|
|
return self == SpeculativeAlgorithm.NONE
|
|
|
|
def is_eagle(self):
|
|
return self == SpeculativeAlgorithm.EAGLE or self == SpeculativeAlgorithm.EAGLE3
|
|
|
|
def is_eagle3(self):
|
|
return self == SpeculativeAlgorithm.EAGLE3
|
|
|
|
def is_standalone(self):
|
|
return self == SpeculativeAlgorithm.STANDALONE
|
|
|
|
def is_ngram(self):
|
|
return self == SpeculativeAlgorithm.NGRAM
|
|
|
|
@staticmethod
|
|
def from_string(name: str):
|
|
name_map = {
|
|
"EAGLE": SpeculativeAlgorithm.EAGLE,
|
|
"EAGLE3": SpeculativeAlgorithm.EAGLE3,
|
|
"STANDALONE": SpeculativeAlgorithm.STANDALONE,
|
|
"NGRAM": SpeculativeAlgorithm.NGRAM,
|
|
None: SpeculativeAlgorithm.NONE,
|
|
}
|
|
if name is not None:
|
|
name = name.upper()
|
|
return name_map[name]
|
|
|
|
|
|
class SpecInputType(IntEnum):
|
|
# NOTE: introduce this to distinguish the SpecInput types of multiple algorithms when asserting in attention backends.
|
|
# If all algorithms can share the same datastrucutre of draft_input and verify_input, consider simplify it
|
|
EAGLE_DRAFT = auto()
|
|
EAGLE_VERIFY = auto()
|
|
NGRAM_VERIFY = auto()
|
|
|
|
|
|
class SpecInput(ABC):
|
|
def __init__(self, spec_input_type: SpecInputType):
|
|
self.spec_input_type = spec_input_type
|
|
|
|
def is_draft_input(self) -> bool:
|
|
# FIXME: remove this function which is only used for assertion
|
|
# or use another variable name like `draft_input` to substitute `spec_info`
|
|
return self.spec_input_type == SpecInputType.EAGLE_DRAFT
|
|
|
|
def is_verify_input(self) -> bool:
|
|
return self.spec_input_type in {
|
|
SpecInputType.EAGLE_VERIFY,
|
|
SpecInputType.NGRAM_VERIFY,
|
|
}
|
|
|
|
@abstractmethod
|
|
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
|
|
pass
|
|
|
|
def get_spec_adjusted_global_num_tokens(
|
|
self, forward_batch: ModelWorkerBatch
|
|
) -> Tuple[List[int], List[int]]:
|
|
c1, c2 = self.get_spec_adjust_token_coefficient()
|
|
global_num_tokens = [x * c1 for x in forward_batch.global_num_tokens]
|
|
global_num_tokens_for_logprob = [
|
|
x * c2 for x in forward_batch.global_num_tokens_for_logprob
|
|
]
|
|
return global_num_tokens, global_num_tokens_for_logprob
|