35 lines
698 B
Python
35 lines
698 B
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.tp_worker import TpModelWorker
|
|
|
|
|
|
class BaseDraftWorker(ABC):
|
|
@abstractmethod
|
|
def draft():
|
|
pass
|
|
|
|
@abstractmethod
|
|
def draft_extend():
|
|
pass
|
|
|
|
|
|
class BaseSpecWorker(ABC):
|
|
@property
|
|
@abstractmethod
|
|
def target_worker(self) -> TpModelWorker:
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def draft_worker(self) -> BaseDraftWorker:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def clear_cache_pool(self):
|
|
# TODO: move this abstract method to BaseTpWorker and call through self.model_runner
|
|
pass
|