minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Allow custom speculative algorithm to support disaggregation (#26195)

Browse Source
This commit is contained in:
jasonjk-park
2026-05-27 09:54:53 -07:00
committed by GitHub
parent a95b4e2e09
commit d6e1692410
4 changed files with 105 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+7 -45
View File
@@ -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,
spec_info = self.spec_algorithm.build_disagg_draft_input(
self,
server_args,
last_tokens_tensor,
future_map,
)
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)
if spec_info is not None:
self.spec_info = spec_info
else:
# Non-spec: positive last token feeds decode directly. No FutureMap
python/sglang/srt/speculative/eagle_disaggregation.py
+66
View File
@@ -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
python/sglang/srt/speculative/spec_info.py
+17
View File
@@ -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()
python/sglang/srt/speculative/spec_registry.py
+14
View File
@@ -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] = {}