Abstraction for spec worker and code cleanup (#11643)
This commit is contained in:
@@ -574,6 +574,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
|
||||
|
||||
@dataclass
|
||||
class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
|
||||
# Constant: alloc length per decode step
|
||||
ALLOC_LEN_PER_DECODE: ClassVar[int] = None
|
||||
|
||||
# The inputs for decode
|
||||
# shape: (b, topk)
|
||||
topk_p: torch.Tensor = None
|
||||
@@ -609,9 +612,6 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
|
||||
new_seq_lens: Optional[torch.Tensor] = None
|
||||
verify_done: Optional[torch.cuda.Event] = None
|
||||
|
||||
# FIXME(lsyin): remove this hack
|
||||
ALLOC_LEN_PER_DECODE: ClassVar[int] = None
|
||||
|
||||
def __post_init__(self):
|
||||
super().__init__(SpecInputType.EAGLE_DRAFT)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user