Abstraction for spec worker and code cleanup (#11643)

This commit is contained in:
Liangsheng Yin
2025-10-17 23:31:36 +08:00
committed by GitHub
parent 3e4c7da2f5
commit cde5a6e30f
14 changed files with 707 additions and 461 deletions
+3 -3
View File
@@ -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)