File size: 886 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | from __future__ import annotations
from diffulex.config import Config
from diffulex.engine.scheduler import AutoScheduler
from diffulex.engine.request import DllmReq
from diffulex.strategy_template.multi_block.engine.scheduler import MultiBlockSchedulerTemplate
@AutoScheduler.register("d2f", is_default=True)
class D2fScheduler(MultiBlockSchedulerTemplate):
def __init__(self, config: Config):
super().__init__(config)
self.init_multi_block()
def add(self, req: DllmReq) -> None:
self.add_multi_block(req)
def schedule(self) -> tuple[list[DllmReq], bool]:
return self.schedule_multi_block()
def preempt(self, req: DllmReq) -> None:
self.preempt_multi_block(req)
def postprocess(
self,
reqs: list[DllmReq],
sample_output,
) -> None:
self.postprocess_multi_block(reqs, sample_output)
|