[Session] Work with PD and Fix empty continuations (#39038)

Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
Co-authored-by: Manik Singhal <3400497+Manikvsin@users.noreply.github.com>

python/sglang/srt/managers/io_struct.py

@@ -454,8 +454,18 @@ class GenerateReqInput:
             self.input_embeds = None
         elif self.input_ids is not None:
             if len(self.input_ids) == 0:
-                raise ValueError("input_ids cannot be empty.")
+                # Session history may supply the entire prompt. The scheduler
-            if isinstance(self.input_ids[0], int):
+                # rejects requests that are still empty after reconstruction.
+                session_id = (
+                    self.session_params.get("id")
+                    if isinstance(self.session_params, dict)
+                    else None
+                )
+                if not session_id:
+                    raise ValueError("input_ids cannot be empty.")
+                self.is_single = True
+                self.batch_size = 1
+            elif isinstance(self.input_ids[0], int):
                 self.is_single = True
                 self.batch_size = 1
             else:

python/sglang/srt/managers/scheduler.py

@@ -2727,6 +2727,8 @@ class Scheduler(
         session_id = (
             recv_req.session_params.id if recv_req.session_params is not None else None
         )
+        if recv_req.bootstrap_port is None:
+            recv_req.bootstrap_port = get_disagg().disaggregation_bootstrap_port
         # Radix-native sessions use only the top-level session_id.
         radix_native_session = (
             recv_req.session_id is not None and self.enable_session_radix_cache
@@ -2739,10 +2741,6 @@ class Scheduler(
                 seq_length = len(recv_req.input_embeds)
                 recv_req.input_ids = array("q", [1]) * seq_length
 
-            if recv_req.bootstrap_port is None:
-                # Use default bootstrap port
-                recv_req.bootstrap_port = get_disagg().disaggregation_bootstrap_port
-
             is_beam = BeamCoordinator.request_beam_width(recv_req) > 1
             req = Req(
                 recv_req.rid,
@@ -2833,6 +2831,7 @@ class Scheduler(
                 self.tokenizer,
                 self.model_config.vocab_size,
                 eos_token_ids=self.model_config.hf_eos_token_id,
+                disagg_mode=self.disaggregation_mode,
             )
             if self.enable_session_radix_cache:
                 req.session_generation = self.tree_cache.ensure_session_generation(

python/sglang/srt/session/session_controller.py

@@ -18,6 +18,7 @@ import uuid
 from array import array
 from typing import TYPE_CHECKING, Dict, Optional
 
+from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.managers.io_struct import (
     CloseSessionReqInput,
     OpenSessionReqInput,
@@ -206,6 +207,7 @@ class Session:
         tokenizer,
         vocab_size: int,
         eos_token_ids=None,
+        disagg_mode: Optional[DisaggregationMode] = None,
     ):
         assert req.session_params is not None
         session_params = req.session_params
@@ -288,6 +290,12 @@ class Session:
             input_ids = req.input_ids
             input_ids_unpadded = req.input_ids
 
+        if not abort and len(input_ids) == 0:
+            abort = True
+            abort_message = (
+                "A session request must contain input tokens after restoring history."
+            )
+
         new_req = Req(
             rid=req.rid,
             origin_input_text=None,
@@ -308,6 +316,12 @@ class Session:
             return_hidden_states=req.return_hidden_states,
             return_routed_experts=req.return_routed_experts,
             routed_experts_start_len=req.routed_experts_start_len,
+            bootstrap_host=req.bootstrap_host,
+            bootstrap_port=req.bootstrap_port,
+            bootstrap_room=req.bootstrap_room,
+            disagg_mode=disagg_mode,
+            routed_dp_rank=req.routed_dp_rank,
+            disagg_prefill_dp_rank=req.disagg_prefill_dp_rank,
             priority=req.priority,
             routing_key=req.routing_key,
             extra_key=req.extra_key,

test/registered/unit/mem_cache/test_session_token_share_unit.py

@@ -43,6 +43,11 @@ def _recv(rid, input_ids, max_new_tokens=8):
         return_hidden_states=False,
         return_routed_experts=False,
         routed_experts_start_len=0,
+        bootstrap_host=None,
+        bootstrap_port=None,
+        bootstrap_room=None,
+        routed_dp_rank=None,
+        disagg_prefill_dp_rank=None,
         priority=None,
         routing_key=None,
         extra_key=None,