diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index ecb54d9c4..86b9aeb3b 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, List import torch from sglang.srt.mem_cache.common import maybe_cache_unfinished_req -from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode +from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo logger = logging.getLogger(__name__) @@ -136,51 +136,13 @@ class ScheduleBatchDisaggregationDecodeMixin: last_tokens, dtype=torch.int64, device=self.device ) - # Simulate the eagle run. - if self.spec_algorithm.is_eagle(): - num_states = server_args.speculative_eagle_topk - if server_args.enable_multi_layer_eagle: - num_states *= server_args.speculative_num_steps - topk_p = torch.stack( - [ - torch.as_tensor( - req.output_topk_p[:num_states], - device=self.device, - dtype=torch.float32, - ) - for req in self.reqs - ], - dim=0, - ) - topk_index = torch.stack( - [ - torch.as_tensor( - req.output_topk_index[:num_states], - device=self.device, - dtype=torch.int64, - ) - for req in self.reqs - ], - dim=0, - ) - - hidden_states_list = [req.hidden_states_tensor for req in self.reqs] - hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device) - - # local import to avoid circular import - from sglang.srt.speculative.eagle_info import EagleDraftInput - - spec_info = EagleDraftInput( - topk_p=topk_p, - topk_index=topk_index, - hidden_states=hidden_states, - bonus_tokens=last_tokens_tensor, - ) - spec_info.capture_hidden_mode = CaptureHiddenMode.LAST - if self.enable_overlap: - spec_info.future_indices = self.req_pool_indices - future_map.publish(spec_info.future_indices, self.seq_lens) - future_map.stash(spec_info.future_indices, spec_info) + spec_info = self.spec_algorithm.build_disagg_draft_input( + self, + server_args, + last_tokens_tensor, + future_map, + ) + if spec_info is not None: self.spec_info = spec_info else: # Non-spec: positive last token feeds decode directly. No FutureMap diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py new file mode 100644 index 000000000..dc9493f28 --- /dev/null +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.speculative.eagle_info import EagleDraftInput + +if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.server_args import ServerArgs + + +def build_eagle_disagg_draft_input( + batch: ScheduleBatch, + server_args: ServerArgs, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, +) -> EagleDraftInput: + num_states = server_args.speculative_eagle_topk + if server_args.enable_multi_layer_eagle: + num_states *= server_args.speculative_num_steps + + topk_p = torch.stack( + [ + torch.as_tensor( + req.output_topk_p[:num_states], + device=batch.device, + dtype=torch.float32, + ) + for req in batch.reqs + ], + dim=0, + ) + topk_index = torch.stack( + [ + torch.as_tensor( + req.output_topk_index[:num_states], + device=batch.device, + dtype=torch.int64, + ) + for req in batch.reqs + ], + dim=0, + ) + + hidden_states = torch.stack( + [req.hidden_states_tensor for req in batch.reqs], dim=0 + ).to(batch.device) + + spec_info = EagleDraftInput( + topk_p=topk_p, + topk_index=topk_index, + hidden_states=hidden_states, + bonus_tokens=last_tokens_tensor, + ) + spec_info.capture_hidden_mode = CaptureHiddenMode.LAST + + if batch.enable_overlap: + spec_info.future_indices = batch.req_pool_indices + future_map.publish(spec_info.future_indices, batch.seq_lens) + future_map.stash(spec_info.future_indices, spec_info) + + return spec_info diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 75b9af39f..65daccd13 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -127,6 +127,23 @@ class SpeculativeAlgorithm(Enum): return FutureMap(device, self, req_to_token_pool) + def build_disagg_draft_input( + self, + batch: ScheduleBatch, + server_args: ServerArgs, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, + ) -> Optional[SpecInput]: + if self.is_eagle(): + from sglang.srt.speculative.eagle_disaggregation import ( + build_eagle_disagg_draft_input, + ) + + return build_eagle_disagg_draft_input( + batch, server_args, last_tokens_tensor, future_map + ) + return None + def supports_spec_v2(self) -> bool: return (self.is_eagle() and not self.is_frozen_kv_mtp()) or self.is_standalone() diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index f15d82acb..e438cbbc0 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -6,8 +6,13 @@ from __future__ import annotations from typing import TYPE_CHECKING, Callable, Dict, Optional, Type +import torch + if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap + from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.spec_info import SpecInput WorkerFactory = Callable[["ServerArgs"], Type] ServerArgsValidator = Callable[["ServerArgs"], None] @@ -87,6 +92,15 @@ class CustomSpecAlgo: # Here, we expose this interface to allow the other use cases. return num_draft_tokens + def build_disagg_draft_input( + self, + batch: ScheduleBatch, + server_args: ServerArgs, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, + ) -> Optional[SpecInput]: + return None + _REGISTRY: Dict[str, CustomSpecAlgo] = {}