diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index a2bc66eaf..6cf5e5f6f 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -49,7 +49,7 @@ class FutureMap: chunked_prefill_size: int, context_len: int, device: torch.device, - spec_algo: Optional[SpeculativeAlgorithm] = None, + spec_algo: SpeculativeAlgorithm, ): # FIXME: the calculation of future_limit and future_buffer_len maybe too conservative self.future_ct = 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ef38e21de..14e9fd979 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -152,7 +152,6 @@ from sglang.srt.managers.mm_utils import ( unwrap_shm_features, ) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors -from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.prefill_delayer import ( PrefillDelayer, PrefillDelayerSinglePassExecutor, @@ -1343,12 +1342,11 @@ class Scheduler( self.future_map = None return - self.future_map = FutureMap( + self.future_map = self.spec_algorithm.create_future_map( self.max_running_requests, self.chunked_prefill_size, self.model_config.context_len, self.device, - self.spec_algorithm, ) self.batch_record_buf = [None] * 2 self.batch_record_ct = 0 diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 1c5a153a1..edf0dadcb 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -36,7 +36,6 @@ from sglang.srt.model_executor.forward_batch_info import ( PPProxyTensors, enable_num_token_non_padded, ) -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( log_info_on_rank0, require_attn_tp_gather, @@ -528,7 +527,7 @@ class CPUGraphRunner: not self.require_gathered_buffer ), "CPUGraphRunner does not support gathered buffer yet." assert ( - model_runner.spec_algorithm == SpeculativeAlgorithm.NONE + model_runner.spec_algorithm.is_none() ), "CPUGraphRunner does not support speculative inference yet." # TODO add compile support for encoder-decoder models assert ( diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 55118fa17..6e4e3ab0c 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -635,8 +635,10 @@ class CudaGraphRunner: self.num_tokens_per_bs = 1 if model_runner.spec_algorithm.is_speculative(): if self.model_runner.is_draft_worker: - # DFLASH draft workers reuse this runner for TARGET_VERIFY mode. - if not self.model_runner.spec_algorithm.is_dflash(): + # Draft workers can use TARGET_VERIFY mode. + if ( + not self.model_runner.spec_algorithm.supports_target_verify_for_draft() + ): raise RuntimeError("This should not happen") self.capture_forward_mode = ForwardMode.TARGET_VERIFY self.num_tokens_per_bs = self.speculative_num_draft_tokens diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 8a0565e32..0bbf357d5 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod from enum import Enum, IntEnum, auto from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union +import torch + from sglang.srt.speculative.spec_registry import ( CustomSpecAlgo, ServerArgsValidator, @@ -15,6 +17,7 @@ from sglang.srt.speculative.spec_registry import ( ) if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.schedule_batch import ModelWorkerBatch from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.server_args import ServerArgs @@ -109,6 +112,26 @@ class SpeculativeAlgorithm(Enum): def is_ngram(self) -> bool: return self == SpeculativeAlgorithm.NGRAM + def supports_target_verify_for_draft(self) -> bool: + return self.is_dflash() + + def create_future_map( + self, + max_running_requests: int, + chunked_prefill_size: int, + context_len: int, + device: torch.device, + ) -> FutureMap: + from sglang.srt.managers.overlap_utils import FutureMap + + return FutureMap( + max_running_requests, + chunked_prefill_size, + context_len, + device, + self, + ) + def supports_spec_v2(self) -> bool: return (self.is_eagle() and not self.is_frozen_kv_mtp()) or self.is_standalone()