Allow custom speculative algorithm to support disaggregation (#26195)

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.
+        spec_info = self.spec_algorithm.build_disagg_draft_input(
-        if self.spec_algorithm.is_eagle():
+            self,
-            num_states = server_args.speculative_eagle_topk
+            server_args,
-            if server_args.enable_multi_layer_eagle:
+            last_tokens_tensor,
-                num_states *= server_args.speculative_num_steps
+            future_map,
-            topk_p = torch.stack(
+        )
-                [
+        if spec_info is not None:
-                    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)
             self.spec_info = spec_info
         else:
             # Non-spec: positive last token feeds decode directly. No FutureMap

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

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()
 

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] = {}