39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
from sglang.srt.environ import envs
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
|
|
class SchedulerRecvSkipper:
|
|
@staticmethod
|
|
def maybe_create(server_args: ServerArgs):
|
|
if server_args.scheduler_recv_interval <= 1:
|
|
return None
|
|
return SchedulerRecvSkipper(server_args)
|
|
|
|
def __init__(self, server_args: ServerArgs):
|
|
# Can be supported if needed, but may need e.g. `global_forward_mode`
|
|
assert not server_args.enable_dp_attention
|
|
self._counter = 0
|
|
self._threshold = server_args.scheduler_recv_interval
|
|
# All can be tuned if needed
|
|
self._default_weight = envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DEFAULT.get()
|
|
self._weight_of_forward_mode = {
|
|
ForwardMode.DECODE: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DECODE.get(),
|
|
ForwardMode.TARGET_VERIFY: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_TARGET_VERIFY.get(),
|
|
None: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE.get(),
|
|
}
|
|
|
|
def handle(self, last_forward_mode: ForwardMode):
|
|
should_recv = False
|
|
|
|
last_weight = self._weight_of_forward_mode.get(
|
|
last_forward_mode, self._default_weight
|
|
)
|
|
self._counter += last_weight
|
|
|
|
if self._counter >= self._threshold:
|
|
self._counter = 0
|
|
should_recv = True
|
|
|
|
return should_recv
|