[PD] Introduce runtime role switching between prefill and decode (#28403)

Signed-off-by: huanglong <huanglong@linux.alibaba.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Co-authored-by: huanglong <huanglong@linux.alibaba.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Huang Long <121648372+LLLL114@users.noreply.github.com>

python/sglang/srt/arg_groups/fields/disagg.py

@@ -114,6 +114,10 @@ class Disagg(msgspec.Struct):
         int,
         "The interval to poll requests in decode server. Can be set to >1 to reduce the overhead of this.",
     ] = 1
+    enable_pd_role_switch: A[
+        bool,
+        "Allow runtime prefill<->decode role switch via /pd_role_switch (PD mode).",
+    ] = False
     optimistic_prefill_attempts: A[
         int, "Number of optimistic prefill forward passes that skip the bootstrap wait."
     ] = 0

python/sglang/srt/arg_groups/pd_disaggregation_hook.py

@@ -147,6 +147,39 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None:
                 f"got '{cfg.disaggregation_transfer_backend}'."
             )
 
+        # Reject features whose role-specific state is not rebuilt on a flip.
+        if cfg.enable_pd_role_switch:
+            view = resolved_view(server_args)
+            unsupported = []
+            if view.enable_dp_attention:
+                unsupported.append("DP attention (--enable-dp-attention)")
+            if view.ep_size > 1:
+                unsupported.append(f"expert parallelism (--ep-size {view.ep_size})")
+            if view.moe_a2a_backend != "none":
+                unsupported.append(
+                    f"MoE all-to-all (--moe-a2a-backend {view.moe_a2a_backend})"
+                )
+            if view.pp_size > 1:
+                unsupported.append(f"pipeline parallelism (--pp-size {view.pp_size})")
+            if view.dp_size > 1:
+                unsupported.append(f"data parallelism (--dp-size {view.dp_size})")
+            if view.dcp_size > 1:
+                unsupported.append(
+                    f"decode context parallelism (--dcp-size {view.dcp_size})"
+                )
+            if view.speculative_algorithm is not None:
+                unsupported.append(
+                    "speculative decoding "
+                    f"(--speculative-algorithm {view.speculative_algorithm})"
+                )
+            if unsupported:
+                raise ValueError(
+                    "--enable-pd-role-switch does not rebuild role-specific "
+                    "state for the following features: "
+                    + ", ".join(unsupported)
+                    + ". Remove these options or drop --enable-pd-role-switch."
+                )
+
 
 def _alias_bootstrap_port_to_api_port(server_args: ServerArgs) -> None:
     """Rust-server prefill serves the KV bootstrap registry on the api listener

python/sglang/srt/disaggregation/base/conn.py

@@ -127,6 +127,15 @@ class BaseKVManager(ABC):
         """Register prefill server info to the bootstrap server."""
         ...
 
+    # Opt-in per backend: set True and implement teardown() to support runtime PD
+    # role switch (release transfer resources; the scheduler owns the KV pool).
+    supports_role_switch: bool = False
+
+    def teardown(self) -> None:
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support PD role switch teardown"
+        )
+
 
 class BaseKVSender(ABC):
     @abstractmethod

python/sglang/srt/disaggregation/common/conn.py

@@ -287,6 +287,9 @@ class CommonKVManager(BaseKVManager):
             self.max_failures = max(
                 envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
             )
+            # Event used to signal the heartbeat checker thread to exit
+            # during teardown (e.g. runtime P<->D role switch).
+            self._heartbeat_shutdown = threading.Event()
             # If a timeout happens on the decode side, it means decode instances
             # fail to receive the KV Cache transfer done signal after bootstrapping.
             # These timeout requests should be aborted to release the tree cache.
@@ -707,6 +710,25 @@ class CommonKVManager(BaseKVManager):
             return
         self._kv_replica_factor = info.required_dst_info_num
 
+    def _make_worker_recv(self, socket, timeout_ms: int = 500):
+        """Build the blocking multipart recv used by a worker thread.
+
+        Plain blocking recv unless role switching is enabled: teardown flips a
+        stop flag that a blocked recv can never observe, so in that mode poll
+        with a timeout and return None when it expires. Deployments without
+        --enable-pd-role-switch keep the original blocking recv and pay nothing.
+        """
+        if not self.server_args.enable_pd_role_switch:
+            return socket.recv_multipart
+
+        poller = zmq.Poller()
+        poller.register(socket, zmq.POLLIN)
+
+        def recv():
+            return socket.recv_multipart() if poller.poll(timeout_ms) else None
+
+        return recv
+
     def _ensure_prefill_recompute_executor(
         self,
     ) -> concurrent.futures.ThreadPoolExecutor:
@@ -1308,12 +1330,18 @@ class CommonKVManager(BaseKVManager):
 
         return src_kv_ptrs, sliced_dst
 
-    def _start_heartbeat_checker_thread(self):
-        """Start the heartbeat checker thread for Decode worker."""
+    def _start_heartbeat_checker_thread(self) -> threading.Thread:
+        """Start the heartbeat checker thread for Decode worker.
+
+        Returns the thread object so callers can track/join it during teardown.
+        """
 
         def heartbeat_checker():
-            while True:
-                time.sleep(self.heartbeat_interval)
+            while not self._heartbeat_shutdown.is_set():
+                # Use Event.wait() instead of time.sleep() so teardown can
+                # wake this thread immediately by setting the event.
+                if self._heartbeat_shutdown.wait(self.heartbeat_interval):
+                    break
                 with self.connection_lock:
                     addresses = list(self.prefill_info_table.keys())
 
@@ -1352,7 +1380,13 @@ class CommonKVManager(BaseKVManager):
                             if bootstrap_addr in self.session_pool:
                                 del self.session_pool[bootstrap_addr]
 
-        threading.Thread(target=heartbeat_checker, daemon=True).start()
+        t = threading.Thread(
+            target=heartbeat_checker,
+            name="HeartbeatChecker",
+            daemon=True,
+        )
+        t.start()
+        return t
 
     def _on_heartbeat_success(self, bootstrap_addr: str):
         """Hook called on successful heartbeat. Override for backend-specific cleanup."""
@@ -1800,6 +1834,31 @@ class CommonKVReceiver(BaseKVReceiver):
                 sock.close()
             logger.debug(f"Disconnected stale ZMQ PUSH socket (receiver): {endpoint}")
 
+    @classmethod
+    def close_all_sockets(cls):
+        """Close all cached PUSH sockets on role switch, keeping ``_ctx`` reusable."""
+        with cls._global_lock:
+            entries = list(cls._socket_cache.items())
+            locks = cls._socket_locks.copy()
+            cls._socket_cache.clear()
+            cls._socket_locks.clear()
+
+        # Close outside _global_lock: _connect drops it before the per-endpoint lock.
+        for endpoint, sock in entries:
+            lock = locks.get(endpoint)
+            try:
+                if lock:
+                    with lock:
+                        sock.close(linger=0)
+                else:
+                    sock.close(linger=0)
+            except Exception:
+                logger.exception(
+                    f"Failed to close ZMQ PUSH socket (receiver): {endpoint}"
+                )
+        if entries:
+            logger.debug(f"Closed {len(entries)} receiver ZMQ PUSH socket(s)")
+
     @classmethod
     def _connect_to_bootstrap_server(cls, bootstrap_info: dict):
         ip_address = bootstrap_info["rank_ip"]

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -208,6 +208,8 @@ class KVArgsRegisterInfo:
 
 class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
     AUX_DATA_HEADER = b"AUX_DATA"
+    # Implements teardown() below, so runtime PD role switching is supported.
+    supports_role_switch = True
 
     def __init__(
         self,
@@ -224,6 +226,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
             envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
         )
         self.enable_trace = get_observability().enable_trace
+        # Set by teardown() to make worker threads exit (P<->D role switch).
+        self._stopped = False
+        self._worker_threads: List[threading.Thread] = []
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
             self.session_failures = defaultdict(int)
             self.failed_sessions = set()
@@ -262,7 +267,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
             for i, (queue, executor) in enumerate(
                 zip(self.transfer_queues, self.executors)
             ):
-                threading.Thread(
+                # Track the thread so teardown() can join it: otherwise every
+                # P->D->P flip that re-enters PREFILL leaks threads
+                # (each parked forever in FastQueue.get()).
+                t = threading.Thread(
                     target=self.transfer_worker,
                     args=(
                         queue,
@@ -275,7 +283,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                         i,
                     ),
                     daemon=True,
-                ).start()
+                )
+                t.start()
+                self._worker_threads.append(t)
             self.enable_failed_session_probe = (
                 envs.SGLANG_ENABLE_FAILED_SESSION_PROBE.get()
             )
@@ -284,11 +294,13 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                     envs.SGLANG_FAILED_SESSION_PROBE_INTERVAL_S.get()
                 )
                 self._failed_session_probe_shutdown = threading.Event()
-                threading.Thread(
+                t = threading.Thread(
                     target=self._failed_session_probe_loop,
                     name="MooncakeFailedSessionProbe",
                     daemon=True,
-                ).start()
+                )
+                t.start()
+                self._worker_threads.append(t)
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self._staging_ctx = DecodeStagingContext() if self.enable_staging else None
             if self.enable_staging:
@@ -338,6 +350,94 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
             with self.connection_lock:
                 self.connection_pool.clear()
 
+    def teardown(self) -> None:
+        """Stop worker threads and release transport resources so this
+        KVManager can be discarded during a P<->D role switch.
+
+        The KV cache pool memory is owned by the scheduler and is NOT freed
+        here; only mooncake-side registrations / sockets are released.
+        """
+        self._stopped = True
+
+        # Stop the failed-session probe loop (PREFILL role) if running.
+        probe_shutdown = getattr(self, "_failed_session_probe_shutdown", None)
+        if probe_shutdown is not None:
+            probe_shutdown.set()
+
+        # Stop the heartbeat checker thread (DECODE role) if running.
+        heartbeat_shutdown = getattr(self, "_heartbeat_shutdown", None)
+        if heartbeat_shutdown is not None:
+            heartbeat_shutdown.set()
+
+        # Transfer workers (PREFILL role) park in FastQueue.get(), which has no
+        # timeout; push a None sentinel per shard to wake and stop them so the
+        # join below returns instead of leaking the thread.
+        for queue in getattr(self, "transfer_queues", []):
+            try:
+                queue.put(None)
+            except Exception:
+                logger.exception(
+                    "Failed to signal mooncake transfer worker on teardown"
+                )
+
+        # Shutdown thread pool executors (PREFILL role).
+        for executor in getattr(self, "executors", []):
+            try:
+                executor.shutdown(wait=False, cancel_futures=True)
+            except TypeError:
+                # Python < 3.9 does not support cancel_futures
+                executor.shutdown(wait=False)
+            except Exception:
+                logger.exception("Failed to shutdown executor on teardown")
+        self.executors = []
+
+        # Join workers before touching their sockets: ZMQ sockets aren't
+        # thread-safe, so don't close server_socket while a worker may poll it.
+        for t in self._worker_threads:
+            t.join(timeout=3.0)
+        self._worker_threads = []
+
+        # Drop the queues so their buffered tasks/senders are released too.
+        self.transfer_queues = []
+
+        # Close cached PUSH sockets (used by _connect for status sync).
+        with self._socket_lock:
+            for sock in self._socket_cache.values():
+                try:
+                    sock.close(linger=0)
+                except Exception:
+                    pass
+            for monitor in self._monitor_cache.values():
+                try:
+                    monitor.close()
+                except Exception:
+                    pass
+            self._socket_cache.clear()
+            self._monitor_cache.clear()
+
+        try:
+            self.server_socket.close(linger=0)
+        except Exception:
+            logger.exception("Failed to close mooncake server_socket during teardown")
+
+        # destroy() force-closes every socket in the context; plain term()
+        # would block waiting on them.
+        try:
+            self._zmq_ctx.destroy(linger=0)
+        except Exception:
+            logger.exception("Failed to destroy mooncake zmq context during teardown")
+
+        # Deregister memory from the transfer engine.
+        try:
+            self.deregister_buffer_to_engine()
+        except Exception:
+            logger.exception("Failed to deregister buffers during teardown")
+
+        logger.info(
+            "MooncakeKVManager torn down (was role=%s)",
+            self.disaggregation_mode.value,
+        )
+
     # ------------------------------------------------------------------
     # Staging buffer methods (all delegate to staging_handler.py)
     # ------------------------------------------------------------------
@@ -1862,6 +1962,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         while True:
             try:
                 kv_chunk: TransferKVChunk = queue.get()
+                # teardown() pushes a None sentinel to unblock get() and stop
+                # the worker: FastQueue.get() blocks indefinitely, so checking
+                # _stopped alone can never wake a parked worker during a role switch.
+                if kv_chunk is None:
+                    break
                 if self.enable_trace:
                     kv_chunk.trace_ctx.rebuild_thread_context()
                     kv_chunk.trace_ctx.trace_slice_start(
@@ -2186,11 +2291,15 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                 )
 
     def start_prefill_thread(self):
+        recv = self._make_worker_recv(self.server_socket)
+
         def bootstrap_thread():
             """This thread recvs pre-alloc notification from the decode engine"""
             # KVPoll.Bootstrapping -> KVPoll.WaitingForInput
-            while True:
-                waiting_req_bytes = self.server_socket.recv_multipart()
+            while not self._stopped:
+                waiting_req_bytes = recv()
+                if waiting_req_bytes is None:
+                    continue
                 room = waiting_req_bytes[0].decode("ascii")
                 # Staging: decode reports consumption watermark back to prefill
                 if room == "WATERMARK":
@@ -2321,12 +2430,18 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                         )
                         self.update_status(room, KVPoll.WaitingForInput)
 
-        threading.Thread(target=bootstrap_thread).start()
+        t = threading.Thread(target=bootstrap_thread, daemon=True)
+        t.start()
+        self._worker_threads.append(t)
 
     def start_decode_thread(self):
+        recv = self._make_worker_recv(self.server_socket)
+
         def decode_thread():
-            while True:
-                msg = self.server_socket.recv_multipart()
+            while not self._stopped:
+                msg = recv()
+                if msg is None:
+                    continue
                 if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
                     self._handle_aux_data(msg)
                     continue
@@ -2379,8 +2494,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                     failure_reason=reason,
                 )
 
-        threading.Thread(target=decode_thread).start()
-        self._start_heartbeat_checker_thread()
+        t = threading.Thread(target=decode_thread, daemon=True)
+        t.start()
+        self._worker_threads.append(t)
+        t = self._start_heartbeat_checker_thread()
+        self._worker_threads.append(t)
 
     def add_transfer_request(
         self,

python/sglang/srt/disaggregation/mori/conn.py

@@ -293,6 +293,8 @@ class TransferTarget:
 
 class MoriKVManager(CommonKVManager):
     AUX_DATA_HEADER = b"AUX_DATA"
+    # Implements teardown() below, so runtime PD role switching is supported.
+    supports_role_switch = True
 
     # The bootstrap socket carries several message kinds, so the status message
     # is tagged. Mori has always shipped the failure reason with it.
@@ -315,6 +317,9 @@ class MoriKVManager(CommonKVManager):
         self.transfer_lock = threading.Lock()
         self._zmq_ctx = zmq.Context()
         self._socket_local = threading.local()
+        # Set by teardown() to make worker threads exit (PoC: P<->D role switch).
+        self._stopped = False
+        self._worker_threads: List[threading.Thread] = []
         self._send_aux_rdma = envs.SGLANG_MORI_SEND_AUX_RDMA.get()
         self._register_local_buffers()
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
@@ -325,7 +330,10 @@ class MoriKVManager(CommonKVManager):
             self._wait_poll_ms = envs.SGLANG_MORI_WAIT_POLL_MS.get()
             self._transfer_timeout_ms = envs.SGLANG_MORI_TRANSFER_TIMEOUT_MS.get()
             for shard, queue in enumerate(self._transfer_queues):
-                threading.Thread(
+                # Track the thread so teardown() can join it: otherwise every
+                # P->D->P flip that re-enters PREFILL leaks _num_shards threads
+                # (each parked forever in FastQueue.get()).
+                t = threading.Thread(
                     target=self._transfer_worker,
                     args=(queue,),
                     daemon=True,
@@ -333,7 +341,9 @@ class MoriKVManager(CommonKVManager):
                         f"mori-xfer-dp{self.system_dp_rank}-"
                         f"tp{self.attn_tp_rank}-s{shard}"
                     ),
-                ).start()
+                )
+                t.start()
+                self._worker_threads.append(t)
             self._start_bootstrap_thread()
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self._start_decode_thread()
@@ -420,6 +430,11 @@ class MoriKVManager(CommonKVManager):
     def _transfer_worker(self, queue: FastQueue) -> None:
         while True:
             kv_chunk = queue.get()
+            # teardown() pushes a None sentinel to unblock get() and stop the
+            # worker: FastQueue.get() blocks indefinitely, so checking _stopped
+            # alone can never wake a parked worker during a role switch.
+            if kv_chunk is None:
+                break
             try:
                 self._process_transfer_chunk(kv_chunk)
             except Exception as exc:
@@ -683,11 +698,13 @@ class MoriKVManager(CommonKVManager):
         logger.debug("Room %s marked Failed via ABORT from decode", bootstrap_room)
 
     def _start_bootstrap_thread(self) -> None:
+        recv = self._make_worker_recv(self.server_socket)
+
         def bootstrap_worker():
-            while True:
+            while not self._stopped:
                 try:
-                    msg = self.server_socket.recv_multipart()
-                    if not msg:
+                    msg = recv()
+                    if msg is None:
                         continue
 
                     tag = msg[0]
@@ -705,15 +722,23 @@ class MoriKVManager(CommonKVManager):
                     else:
                         self._handle_transfer_message(payload)
                 except Exception:
+                    if self._stopped:
+                        break
                     logger.exception("Bootstrap worker failed")
 
-        threading.Thread(target=bootstrap_worker, daemon=True).start()
+        t = threading.Thread(target=bootstrap_worker, daemon=True)
+        t.start()
+        self._worker_threads.append(t)
 
     def _start_decode_thread(self) -> None:
+        recv = self._make_worker_recv(self.server_socket)
+
         def decode_worker():
-            while True:
+            while not self._stopped:
                 try:
-                    msg = self.server_socket.recv_multipart()
+                    msg = recv()
+                    if msg is None:
+                        continue
                     if msg and msg[0] == MoriKVManager.AUX_DATA_HEADER:
                         self._handle_aux_data(msg)
                         continue
@@ -732,9 +757,69 @@ class MoriKVManager(CommonKVManager):
                         failure_reason=reason,
                     )
                 except Exception:
+                    if self._stopped:
+                        break
                     logger.exception("Decode status worker failed")
 
-        threading.Thread(target=decode_worker, daemon=True).start()
+        t = threading.Thread(target=decode_worker, daemon=True)
+        t.start()
+        self._worker_threads.append(t)
+
+    def teardown(self) -> None:
+        """Stop worker threads and release transport resources so this
+        KVManager can be discarded during a P<->D role switch.
+
+        The KV cache pool memory is owned by the scheduler and is NOT freed
+        here; only mori-side registrations / sockets / engine are released.
+        """
+        self._stopped = True
+        # Transfer workers (PREFILL role) park in FastQueue.get(), which has no
+        # timeout; push a None sentinel per shard to wake and stop them so the
+        # join below returns instead of leaking the thread.
+        for queue in getattr(self, "_transfer_queues", []):
+            try:
+                queue.put(None)
+            except Exception:
+                logger.exception("Failed to signal mori transfer worker on teardown")
+        # Join workers before touching their sockets: ZMQ sockets aren't
+        # thread-safe, so don't close server_socket while a worker may poll it.
+        for t in self._worker_threads:
+            t.join(timeout=3.0)
+        self._worker_threads = []
+        # Drop the queues so their buffered tasks/senders are released too.
+        self._transfer_queues = []
+        try:
+            self.server_socket.close(linger=0)
+        except Exception:
+            logger.exception("Failed to close mori server_socket during teardown")
+        # destroy() force-closes every socket in the context (incl. per-thread
+        # cached PUSH sockets); plain term() would block waiting on them.
+        try:
+            self._zmq_ctx.destroy(linger=0)
+        except Exception:
+            logger.exception("Failed to destroy mori zmq context during teardown")
+        # Deregister RDMA memory and drop the IOEngine reference.
+        try:
+            for descs in (self.kv_mem_descs, self.aux_mem_descs):
+                for desc in descs:
+                    try:
+                        self.engine.deregister_memory(desc)
+                    except Exception:
+                        pass
+            for component_descs in self.state_mem_descs:
+                for desc in component_descs:
+                    try:
+                        self.engine.deregister_memory(desc)
+                    except Exception:
+                        pass
+        finally:
+            self.kv_mem_descs = []
+            self.aux_mem_descs = []
+            self.state_mem_descs = []
+            self.engine = None
+        logger.info(
+            "MoriKVManager torn down (was role=%s)", self.disaggregation_mode.value
+        )
 
     def _add_remote_peer(self, register_info: KVArgsRegisterInfo) -> None:
         engine_key = register_info.engine_key

python/sglang/srt/disaggregation/role_switch.py

@@ -0,0 +1,284 @@
+"""Runtime prefill<->decode role switching for PD disaggregation.
+
+The token KV pool is role-independent and never reallocated; only the
+role-specific disaggregation structures are torn down and rebuilt on a flip.
+Kept out of scheduler.py to avoid growing it further.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Callable, Optional, Tuple
+
+from sglang.srt.disaggregation.common.conn import CommonKVReceiver
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import PdRoleSwitchReqInput, PdRoleSwitchReqOutput
+from sglang.srt.runtime_context import get_context, get_disagg
+from sglang.srt.utils import get_available_gpu_memory
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
+logger = logging.getLogger(__name__)
+
+
+class PdRoleSwitchRestart(Exception):
+    """Break out of the current role's event loop after a successful switch."""
+
+
+def run_event_loop_supervisor(
+    scheduler: Scheduler, dispatch_once: Callable[[Scheduler], None]
+) -> None:
+    """Re-dispatch the scheduler event loop after each runtime role switch."""
+    while True:
+        try:
+            return dispatch_once(scheduler)
+        except PdRoleSwitchRestart:
+            logger.info(
+                "Re-dispatching event loop after PD role switch -> %s",
+                scheduler.disaggregation_mode.value,
+            )
+
+
+def handle_pd_role_switch(
+    scheduler: Scheduler, recv_req: PdRoleSwitchReqInput
+) -> PdRoleSwitchReqOutput:
+    """Flip the scheduler's disaggregation role at runtime. The instance must be
+    idle; rebuild failure is fatal to the instance (no in-place rollback)."""
+    old_role = scheduler.disaggregation_mode.value
+    new_role = (recv_req.new_role or "").lower()
+
+    def _fail(msg: str, safe_to_restore: bool = False) -> PdRoleSwitchReqOutput:
+        logger.warning(
+            "PD role switch rejected (%s -> %s): %s", old_role, new_role, msg
+        )
+        return PdRoleSwitchReqOutput(
+            success=False,
+            message=msg,
+            old_role=old_role,
+            new_role=new_role,
+            safe_to_restore=safe_to_restore,
+        )
+
+    rejection = _reject_reason(scheduler, new_role)
+    if rejection is not None:
+        return _fail(*rejection)
+    if new_role == old_role:
+        return PdRoleSwitchReqOutput(
+            success=True,
+            message="already in target role",
+            old_role=old_role,
+            new_role=new_role,
+        )
+    if not scheduler.is_fully_idle():
+        return _fail(
+            "instance is not idle; drain all requests before switching",
+            safe_to_restore=True,
+        )
+
+    required_graph_gb = recv_req.decode_cuda_graph_memory_gb
+    # Same condition ensure_decode_cuda_graphs skips on, so the check cannot be
+    # bypassed while the capture still runs.
+    will_capture_graphs = (
+        new_role == "decode" and not scheduler.tp_worker.get_decode_cuda_graph_bs()
+    )
+    if will_capture_graphs and required_graph_gb is None:
+        return _fail(
+            "decode_cuda_graph_memory_gb is required before capturing decode graphs",
+            safe_to_restore=True,
+        )
+    if will_capture_graphs and required_graph_gb is not None:
+        if required_graph_gb < 0:
+            return _fail(
+                "decode_cuda_graph_memory_gb must be non-negative",
+                safe_to_restore=True,
+            )
+        try:
+            available_graph_gb = get_available_gpu_memory(
+                scheduler.device, scheduler.ps.gpu_id
+            )
+        except Exception as e:
+            return _fail(
+                f"failed to check decode CUDA graph headroom: {e}",
+                safe_to_restore=True,
+            )
+        if available_graph_gb < required_graph_gb:
+            return _fail(
+                "insufficient decode CUDA graph headroom: "
+                f"required={required_graph_gb:.3f} GB, "
+                f"available={available_graph_gb:.3f} GB",
+                safe_to_restore=True,
+            )
+
+    scheduler._pd_role_switch_in_progress = True
+    try:
+        # Teardown + role flip + rebuild are one logical atomic step. If any of
+        # them raises, the instance is left half-torn-down (old role released,
+        # new role not up) and isn't safe to serve, so mark it unhealthy. There
+        # is no in-place rollback.
+        try:
+            teardown_disaggregation(scheduler)
+            get_context().override("role_switch.flip", disaggregation_mode=new_role)
+            scheduler.init_disaggregation()
+            scheduler._sync_disaggregation_mode_to_subcomponents()
+        except Exception as e:
+            scheduler._pd_role_switch_unhealthy = True
+            logger.critical(
+                "PD role switch (%s -> %s) failed during teardown/rebuild; "
+                "instance unhealthy: %s",
+                old_role,
+                new_role,
+                e,
+            )
+            return _fail(
+                f"role switch failed; instance unhealthy, restart required: {e}"
+            )
+
+        if new_role == "decode":
+            # Best-effort deferred capture; a failure only degrades to eager.
+            try:
+                scheduler.tp_worker.ensure_decode_cuda_graphs(
+                    recv_req.decode_cuda_graph_bs
+                )
+            except Exception:
+                logger.exception("Decode CUDA graph capture on role switch failed")
+
+        # Break out of the old-role event loop so the supervisor re-dispatches.
+        scheduler._event_loop_should_restart = True
+        logger.info("PD role switch succeeded: %s -> %s", old_role, new_role)
+        return PdRoleSwitchReqOutput(
+            success=True, message="ok", old_role=old_role, new_role=new_role
+        )
+    except Exception as e:
+        logger.exception("PD role switch failed")
+        return _fail(f"role switch raised: {e}")
+    finally:
+        scheduler._pd_role_switch_in_progress = False
+
+
+def _reject_reason(scheduler: Scheduler, new_role: str) -> Optional[Tuple[str, bool]]:
+    """Why the switch must be rejected before draining, or None to proceed.
+
+    Table-driven: the first failing precondition's message is returned.
+    """
+    sa = scheduler.server_args
+    km = _current_kv_manager(scheduler)
+    # (failed?, safe to restore routing?, lazy message)
+    checks = (
+        (
+            not sa.enable_pd_role_switch,
+            True,
+            lambda: "--enable-pd-role-switch is not set on this instance",
+        ),
+        (
+            scheduler._pd_role_switch_unhealthy,
+            False,
+            lambda: (
+                "instance is unhealthy after a failed role switch; restart required"
+            ),
+        ),
+        (
+            scheduler._pd_role_switch_in_progress,
+            False,
+            lambda: "another role switch is already in progress",
+        ),
+        (
+            new_role not in ("prefill", "decode"),
+            True,
+            lambda: f"invalid new_role={new_role!r}",
+        ),
+        (
+            scheduler.disaggregation_mode == DisaggregationMode.NULL,
+            True,
+            lambda: "instance is not running in PD disaggregation mode",
+        ),
+        (
+            km is not None and not km.supports_role_switch,
+            True,
+            lambda: (
+                f"transfer backend {get_disagg().disaggregation_transfer_backend!r} "
+                "does not support runtime role switch"
+            ),
+        ),
+        (
+            getattr(km, "enable_staging", False),
+            True,
+            lambda: (
+                "staging buffer (SGLANG_DISAGG_STAGING_BUFFER) is not "
+                "supported with runtime role switch"
+            ),
+        ),
+    )
+    return next(
+        ((msg(), safe_to_restore) for failed, safe_to_restore, msg in checks if failed),
+        None,
+    )
+
+
+def _current_kv_manager(scheduler: Scheduler):
+    """The KV manager of the current role's disaggregation queue, or None."""
+    if scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
+        q = getattr(scheduler, "disagg_prefill_bootstrap_queue", None)
+    elif scheduler.disaggregation_mode == DisaggregationMode.DECODE:
+        q = getattr(scheduler, "disagg_decode_prealloc_queue", None)
+    else:
+        q = None
+    return getattr(q, "kv_manager", None) if q is not None else None
+
+
+def teardown_disaggregation(scheduler: Scheduler) -> None:
+    """Release the current role's disaggregation structures (queues, metadata
+    buffers, KV transfer manager) so the other role can be rebuilt."""
+    mode = scheduler.disaggregation_mode
+    if mode == DisaggregationMode.PREFILL:
+        q = getattr(scheduler, "disagg_prefill_bootstrap_queue", None)
+        if q is not None:
+            km = getattr(q, "kv_manager", None)
+            if km is not None:
+                km.teardown()
+            scheduler.disagg_prefill_bootstrap_queue = None
+        scheduler.disagg_prefill_inflight_queue = []
+    elif mode == DisaggregationMode.DECODE:
+        q = getattr(scheduler, "disagg_decode_prealloc_queue", None)
+        if q is not None:
+            km = getattr(q, "kv_manager", None)
+            if km is not None:
+                km.teardown()
+            scheduler.disagg_decode_prealloc_queue = None
+        scheduler.disagg_decode_transfer_queue = None
+        # clear socket ctx in CommonKVReceiver
+        CommonKVReceiver.close_all_sockets()
+    scheduler.disagg_metadata_buffers = None
+    scheduler.req_to_metadata_buffer_idx_allocator = None
+    _release_prefix_cache_for_role_switch(scheduler)
+
+
+def _release_prefix_cache_for_role_switch(scheduler: Scheduler) -> None:
+    """Release the prefix (radix/hicache) cache so a flip works with radix ON.
+
+    With radix disabled (ChunkCache) the flip needs nothing here: ChunkCache
+    keeps no persistent prefixes and, since the instance is idle before the
+    switch, the allocator is already empty. This is the historical
+    ``--disable-radix-cache`` path, left untouched by the guard below.
+
+    With radix (or hicache) enabled, finished prefixes stay in the tree and keep
+    their KV-pool slots *locked* even while idle. Carried across a role switch
+    that means (a) the new role would match against stale prefixes whose KV no
+    longer means what it did (corruption) and (b) those locked slots would leak
+    on every flip. Reset mirrors ``Scheduler.flush_cache``'s cache-release block
+    (the instance is already fully idle, checked before teardown) and, for
+    hicache, best-effort clears the storage backend so it is released completely.
+    """
+    if scheduler.disable_radix_cache:
+        return
+    tree_cache = scheduler.tree_cache
+    if tree_cache is not None:
+        clear_storage = getattr(tree_cache, "clear_storage_backend", None)
+        if callable(clear_storage):
+            try:
+                clear_storage()
+            except Exception:
+                logger.exception("hicache storage release on role switch failed")
+        tree_cache.reset()
+    scheduler.req_to_token_pool.clear()
+    scheduler.token_to_kv_pool_allocator.clear()

python/sglang/srt/entrypoints/engine.py

@@ -1407,6 +1407,7 @@ class Engine(EngineScoreMixin, EngineBase):
             "load_format": tm.config_value("load_format"),
             "reasoning_parser": tm.config_value("reasoning_parser"),
             "tool_call_parser": tm.config_value("tool_call_parser"),
+            "disaggregation_mode": tm.config_value("disaggregation_mode"),
         }
 
     def init_weights_update_group(

python/sglang/srt/entrypoints/grpc_bridge.py

@@ -411,6 +411,9 @@ class RuntimeHandle:
             "load_format": self.tokenizer_manager.config_value("load_format"),
             "reasoning_parser": self.tokenizer_manager.config_value("reasoning_parser"),
             "tool_call_parser": self.tokenizer_manager.config_value("tool_call_parser"),
+            "disaggregation_mode": self.tokenizer_manager.config_value(
+                "disaggregation_mode"
+            ),
             "model_type": getattr(model_config.hf_config, "model_type", None),
             "architectures": getattr(model_config.hf_config, "architectures", None),
         }

python/sglang/srt/entrypoints/http_server.py

@@ -133,6 +133,7 @@ from sglang.srt.managers.io_struct import (
     OpenSessionReqInput,
     ParseFunctionCallReq,
     PauseGenerationReqInput,
+    PdRoleSwitchReqInput,
     ProfileReq,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
@@ -779,6 +780,9 @@ async def model_info():
         "tool_call_parser": _global_state.tokenizer_manager.config_value(
             "tool_call_parser"
         ),
+        "disaggregation_mode": _global_state.tokenizer_manager.config_value(
+            "disaggregation_mode"
+        ),
         "has_image_understanding": model_config.is_image_understandable_model,
         "has_audio_understanding": model_config.is_audio_understandable_model,
         "model_type": getattr(model_config.hf_config, "model_type", None),
@@ -1561,6 +1565,23 @@ async def slow_down(obj: Annotated[SlowDownReqInput, Body()], request: Request):
         return _create_error_response(e)
 
 
+@app.api_route("/pd_role_switch", methods=["POST"])
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def pd_role_switch(
+    obj: Annotated[PdRoleSwitchReqInput, Body()], request: Request
+):
+    """Switch this instance's PD disaggregation role (prefill<->decode) at runtime.
+    Requires --enable-pd-role-switch; the instance must be idle."""
+    try:
+        result = await _global_state.tokenizer_manager.pd_role_switch(obj, request)
+    except Exception as e:
+        return _create_error_response(e)
+    return ORJSONResponse(
+        msgspec_to_builtins(result),
+        status_code=HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST,
+    )
+
+
 @app.api_route("/load_lora_adapter", methods=["POST"])
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def load_lora_adapter(

python/sglang/srt/managers/disagg_service.py

@@ -1,5 +1,6 @@
 """Start bootstrap/kv-store-related server"""
 
+import logging
 import os
 
 from sglang.srt.disaggregation.utils import (
@@ -14,14 +15,30 @@ from sglang.srt.runtime_context import (
     get_serving,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def start_disagg_service():
     # Start kv bootstrap server on prefill
     disagg_mode = DisaggregationMode(get_disagg().disaggregation_mode)
     transfer_backend = TransferBackend(get_disagg().disaggregation_transfer_backend)
 
-    if disagg_mode == DisaggregationMode.PREFILL:
-        # only start bootstrap server on prefill tm
+    # With role switching, run bootstrap on every instance (not just prefill) so
+    # one flipped to prefill already has it; it isn't rebuilt on flip.
+    start_bootstrap = disagg_mode == DisaggregationMode.PREFILL or (
+        get_disagg().enable_pd_role_switch and disagg_mode != DisaggregationMode.NULL
+    )
+
+    if start_bootstrap and get_disagg().enable_pd_role_switch:
+        logger.warning(
+            "Role switch starts a bootstrap server on this instance at %s:%d. "
+            "If another PD instance runs on the same host, give each one a "
+            "distinct --disaggregation-bootstrap-port or the bind will conflict.",
+            get_serving().host,
+            get_disagg().disaggregation_bootstrap_port,
+        )
+
+    if start_bootstrap:
         kv_bootstrap_server_class = get_kv_class(
             transfer_backend, KVClassType.BOOTSTRAP_SERVER
         )

python/sglang/srt/managers/io_struct.py

@@ -2049,6 +2049,24 @@ class SlowDownReqOutput(BaseReq, kw_only=True):
     pass
 
 
+class PdRoleSwitchReqInput(BaseReq, kw_only=True):
+    # Target role; "" is an invalid sentinel rejected by the handler.
+    new_role: Literal["prefill", "decode", ""] = ""
+    # Optional decode bs to capture on a flip to decode (capture-to-fit);
+    # None uses the server's configured decode bs list.
+    decode_cuda_graph_bs: Optional[List[int]] = None
+    # Measured graph footprint from a matching decode peer.
+    decode_cuda_graph_memory_gb: Optional[float] = None
+
+
+class PdRoleSwitchReqOutput(BaseReq, kw_only=True):
+    success: bool = False
+    message: str = ""
+    old_role: str = ""
+    new_role: str = ""
+    safe_to_restore: bool = False
+
+
 class AbortReq(BaseReq, kw_only=True):
     # Whether to abort all requests
     abort_all: bool = False

python/sglang/srt/managers/schedule_batch.py

@@ -2225,6 +2225,7 @@ def release_req(
     # Callers that will recompute the KV instead (PD true-retraction rebootstrap)
     # pass offload_kv=False to skip the wasteful device->host copy.
     backup_saved = True
+    # The config bag reflects role flips; server_args keeps the launch role.
     if get_disagg().disaggregation_mode == "decode" and offload_kv:
         backup_saved = retraction_backup(
             req,

python/sglang/srt/managers/scheduler.py

@@ -74,6 +74,7 @@ from sglang.srt.configs.model_config import (
 )
 from sglang.srt.constrained.grammar_manager import GrammarManager
 from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix
+from sglang.srt.disaggregation import role_switch
 from sglang.srt.disaggregation.checksum import KvChecksumComputer
 from sglang.srt.disaggregation.decode import (
     DecodePreallocQueue,
@@ -159,6 +160,7 @@ from sglang.srt.managers.io_struct import (
     MMInputsProcessError,
     OpenSessionReqInput,
     PauseGenerationReqInput,
+    PdRoleSwitchReqInput,
     ProfileReq,
     ReleaseMemoryOccupationReqInput,
     RemoveExternalCorpusReqInput,
@@ -1284,6 +1286,13 @@ class Scheduler(
         self.hisparse_coordinator.set_decode_producer_stream(self.forward_stream)
 
     def init_running_status(self):
+        # Set by a runtime PD role switch to break out of the current event loop.
+        self._event_loop_should_restart = False
+        # Guards against concurrent/re-entrant PD role switches.
+        self._pd_role_switch_in_progress = False
+        # Set if a role switch tore down the old role but failed to rebuild
+        # either the new or the old role; the instance can no longer serve.
+        self._pd_role_switch_unhealthy = False
         # Set by the ShutdownReq handler to break the event loop for graceful shutdown.
         self.gracefully_exit = False
         self.waiting_queue: List[Req] = []
@@ -1804,6 +1813,7 @@ class Scheduler(
                     self.weight_updater.check_weights,
                 ),
                 (SlowDownReqInput, self.slow_down),
+                (PdRoleSwitchReqInput, self.handle_pd_role_switch),
                 (
                     ProfileReq,
                     lambda req: self.profiler_manager._profile(req),
@@ -2133,6 +2143,13 @@ class Scheduler(
         if self.external_corpus_manager is not None:
             self.external_corpus_manager.check_pending_load()
 
+        # A runtime PD role switch rebuilt the disaggregation structures for a new
+        # role. The response has already been sent above; now break out of the
+        # current (old-role) event loop so the supervisor can re-dispatch.
+        if get_disagg().enable_pd_role_switch and self._event_loop_should_restart:
+            self._event_loop_should_restart = False
+            raise role_switch.PdRoleSwitchRestart()
+
     @staticmethod
     def _tokenized_requests(recv_req):
         if isinstance(
@@ -5151,7 +5168,7 @@ class Scheduler(
         draft_graph_memory_usage = (
             None if self.draft_worker is None else self.draft_worker.graph_memory_usage
         )
-        ret["memory_usage"] = build_memory_usage(
+        memory_usage = build_memory_usage(
             weight_gb=self.tp_worker.model_runner.weight_load_mem_usage,
             kv_cache_gb=self.token_to_kv_pool_allocator.get_kvcache().mem_usage,
             startup_available_gb=self.startup_available_gpu_memory_gb,
@@ -5160,8 +5177,29 @@ class Scheduler(
             target_graph_memory_usage=self.tp_worker.graph_memory_usage,
             draft_graph_memory_usage=draft_graph_memory_usage,
         )
+        ret["memory_usage"] = memory_usage
         ret["startup_time"] = self.startup_time
         ret["effective_max_running_requests_per_dp"] = self.max_running_requests
+        # PD role switch: report this instance's role and the decode CUDA graph
+        # batch sizes it captured, which a router feeds back as
+        # PdRoleSwitchReqInput.decode_cuda_graph_bs. Unset until
+        # init_disaggregation runs, which also re-derives it on every flip.
+        disaggregation_mode = getattr(self, "disaggregation_mode", None)
+        if disaggregation_mode is not None:
+            ret["disaggregation_mode"] = disaggregation_mode.value
+            ret["decode_cuda_graph_bs"] = self.tp_worker.get_decode_cuda_graph_bs()
+            ret["decode_cuda_graph_memory_gb"] = round(
+                sum(
+                    memory_usage["graph"][phase]
+                    for phase in (
+                        "decode",
+                        "target_verify",
+                        "draft_decode",
+                        "draft_extend",
+                    )
+                ),
+                3,
+            )
 
         if get_exec().moe.elastic_ep_backend is not None:
             from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
@@ -5739,6 +5777,24 @@ class Scheduler(
         self.forward_sleep_time = t
         return SlowDownReqOutput()
 
+    def handle_pd_role_switch(self, recv_req: PdRoleSwitchReqInput):
+        return role_switch.handle_pd_role_switch(self, recv_req)
+
+    def _sync_disaggregation_mode_to_subcomponents(self):
+        # Push the (possibly flipped) mode into sub-components that cache it.
+        # object.__setattr__ because some are frozen dataclasses.
+        for name in (
+            "invariant_checker",
+            "load_inquirer",
+            "output_streamer",
+            "batch_result_processor",
+        ):
+            comp = getattr(self, name, None)
+            if comp is not None and hasattr(comp, "disaggregation_mode"):
+                object.__setattr__(
+                    comp, "disaggregation_mode", self.disaggregation_mode
+                )
+
     def expert_distribution_handle(self, recv_req: ExpertDistributionReq):
         action = recv_req.action
         if action == ExpertDistributionReqType.START_RECORD:
@@ -5818,6 +5874,15 @@ class Scheduler(
 
 
 def dispatch_event_loop(scheduler: Scheduler):
+    if scheduler.server_args.enable_pd_role_switch:
+        return role_switch.run_event_loop_supervisor(
+            scheduler,
+            _dispatch_event_loop_once,
+        )
+    return _dispatch_event_loop_once(scheduler)
+
+
+def _dispatch_event_loop_once(scheduler: Scheduler):
     # The live PP property asserts before torch.distributed init (MLX stub).
     disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
     if disaggregation_mode == DisaggregationMode.NULL:

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -424,12 +424,13 @@ class SchedulerInvariantChecker:
             )
 
     def _check_req_pool(self):
-        if self.disaggregation_mode == DisaggregationMode.DECODE:
-            req_total_size = (
-                self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
-            )
-        else:
-            req_total_size = self.req_to_token_pool.size
+        # Account for pre-alloc headroom whenever the pool has it. The decode
+        # pool always does; with runtime P<->D role switching a prefill instance
+        # may also hold a decode-flavored (pre-alloc) pool, so key off the pool
+        # itself rather than the current role.
+        req_total_size = (
+            self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
+        )
 
         session_req_count = self.pool_stats_observer.session_held_req_count()
         if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:

python/sglang/srt/managers/tokenizer_control_mixin.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import fastapi
 
+from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.managers.communicator import FanOutCommunicator
 from sglang.srt.managers.io_struct import (
     AddExternalCorpusReqInput,
@@ -48,6 +49,8 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqOutput,
     LoRAUpdateOutput,
     OpenSessionReqInput,
+    PdRoleSwitchReqInput,
+    PdRoleSwitchReqOutput,
     ProfileReq,
     ProfileReqOutput,
     ProfileReqType,
@@ -77,6 +80,7 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.load_snapshot import LoadSnapshot
 from sglang.srt.runtime_context import (
+    get_disagg,
     get_lora,
     get_parallel,
     get_serving,
@@ -115,6 +119,7 @@ _COMMUNICATOR_SPECS = [
     ("resume_memory_occupation", ResumeMemoryOccupationReqOutput),
     ("check_weights", CheckWeightsReqOutput),
     ("slow_down", SlowDownReqOutput),
+    ("pd_role_switch", PdRoleSwitchReqOutput),
     ("flush_cache", FlushCacheReqOutput),
     ("add_external_corpus", AddExternalCorpusReqOutput),
     ("remove_external_corpus", RemoveExternalCorpusReqOutput),
@@ -840,6 +845,42 @@ class TokenizerControlMixin:
         self.auto_create_handle_loop()
         await self.slow_down_communicator(obj)
 
+    async def pd_role_switch(
+        self: TokenizerManager,
+        obj: PdRoleSwitchReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> PdRoleSwitchReqOutput:
+        self.auto_create_handle_loop()
+        if not self.server_args.enable_pd_role_switch:
+            return PdRoleSwitchReqOutput(
+                success=False,
+                message="--enable-pd-role-switch is not set on this server",
+                old_role=get_disagg().disaggregation_mode,
+                new_role=obj.new_role,
+                safe_to_restore=True,
+            )
+        results = await self.pd_role_switch_communicator(obj)
+        all_success = all(r.success for r in results)
+        safe_to_restore = bool(results) and all(r.safe_to_restore for r in results)
+        if all_success:
+            # Keep the tokenizer-manager's view of the role in sync so future
+            # control ops and bootstrap routing behave consistently.
+            self.record_config_updates(
+                "tokenizer.pd_role_switch", disaggregation_mode=obj.new_role
+            )
+            self.disaggregation_mode = DisaggregationMode(obj.new_role)
+            msg = "ok"
+        else:
+            # Surface only the failing workers' messages.
+            msg = "; ".join(r.message for r in results if not r.success)
+        return PdRoleSwitchReqOutput(
+            success=all_success,
+            message=msg,
+            old_role=results[0].old_role if results else "",
+            new_role=obj.new_role,
+            safe_to_restore=safe_to_restore,
+        )
+
     async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]:
         self.auto_create_handle_loop()
         req = GetInternalStateReq()

python/sglang/srt/managers/tp_worker.py

@@ -471,6 +471,18 @@ class TpModelWorker(BaseTpWorker):
         for mr in self.model_runner_list[1:]:
             mr.init_cuda_graphs(capture_decode_cuda_graph=capture_decode_cuda_graph)
 
+    def ensure_decode_cuda_graphs(self, capture_bs: Optional[List[int]] = None):
+        """Idempotently capture decode cuda graphs for all model runners (used
+        for the on-flip capture during a runtime PD role switch)."""
+        self.model_runner.ensure_decode_cuda_graphs(capture_bs)
+        for mr in self.model_runner_list[1:]:
+            mr.ensure_decode_cuda_graphs(capture_bs)
+
+    def get_decode_cuda_graph_bs(self) -> List[int]:
+        """Decode bs captured as CUDA graphs (empty on a not-yet-flipped prefill,
+        or on a runner that never allocates a KV pool, e.g. the MLX stub)."""
+        return list(getattr(self.model_runner, "decode_cuda_graph_capture_bs", []))
+
     def start_startup_weight_load(self) -> None:
         """Start deferred checkpoint prefetching for all model runners."""
         self.model_runner.start_startup_weight_load()

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -1009,32 +1009,34 @@ class KVCacheConfigurator:
     def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool:
         extra_max_context_len = get_req_to_token_extra_context_len()
 
-        if get_disagg().disaggregation_mode == "decode":
-            # Extra slots for pre-allocated requests
-            pre_alloc_size = get_disagg().disaggregation_decode_extra_slots
+        disagg = get_disagg()
+        if disagg.disaggregation_mode == "decode" or disagg.enable_pd_role_switch:
+            # A flip-capable prefill needs the decode pool shape, and the extra-slot
+            # default is only computed for a decode launch.
+            pre_alloc_size = disagg.disaggregation_decode_extra_slots
+            if disagg.enable_pd_role_switch:
+                pre_alloc_size = pre_alloc_size or 0
             if self.mambaish_config:
-                req_to_token_pool = self._build_hybrid_mamba_decode_req_pool(
+                return self._build_hybrid_mamba_decode_req_pool(
                     max_num_reqs=max_num_reqs,
                     extra_max_context_len=extra_max_context_len,
                     pre_alloc_size=pre_alloc_size,
                 )
-            else:
-                req_to_token_pool = self._build_decode_req_pool(
-                    max_num_reqs=max_num_reqs,
-                    extra_max_context_len=extra_max_context_len,
-                    pre_alloc_size=pre_alloc_size,
-                )
-        elif self.mambaish_config:
-            req_to_token_pool = self._build_hybrid_req_pool(
+            return self._build_decode_req_pool(
+                max_num_reqs=max_num_reqs,
+                extra_max_context_len=extra_max_context_len,
+                pre_alloc_size=pre_alloc_size,
+            )
+
+        if self.mambaish_config:
+            return self._build_hybrid_req_pool(
                 max_num_reqs=max_num_reqs,
                 extra_max_context_len=extra_max_context_len,
             )
-        else:
-            req_to_token_pool = self._build_default_req_pool(
-                max_num_reqs=max_num_reqs,
-                extra_max_context_len=extra_max_context_len,
-            )
-        return req_to_token_pool
+        return self._build_default_req_pool(
+            max_num_reqs=max_num_reqs,
+            extra_max_context_len=extra_max_context_len,
+        )
 
     def _get_mamba_layer_ids_for_req_pool(self) -> list:
         mamba_layer_ids = [

python/sglang/srt/mem_cache/memory_pool.py

@@ -273,6 +273,10 @@ class ReqToTokenPool:
     """A memory pool that maps a request to its token locations."""
 
     enable_mamba_extra_buffer_lazy: bool = False
+    # Extra pre-allocation headroom (reserved for in-transfer decode requests).
+    # 0 for a plain pool; the decode-flavored pool (DecodeReqToTokenPool) sets a
+    # positive value. Declared here so callers can read it without getattr.
+    pre_alloc_size: int = 0
     # Class default: some decode pools borrow another __init__ (see
     # DecodeReqToTokenPool) but inherit alloc_rows.
     _on_alloc_rows: Optional[Callable[[List[int]], None]] = None

python/sglang/srt/model_executor/model_runner.py

@@ -92,6 +92,7 @@ from sglang.srt.mem_cache.kv_cache_configurator import (
 from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
     cuda_graph_fully_disabled,
 )
 from sglang.srt.model_executor.forward_batch_info import (
@@ -107,6 +108,7 @@ from sglang.srt.model_executor.graph_memory_usage import (
     replace_graph_memory_usage,
     replace_graph_time_usage,
 )
+from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput
 from sglang.srt.model_executor.model_runner_components import misc_utils
 from sglang.srt.model_executor.model_runner_components.attention_backend_setup import (
     build_attention_backends,
@@ -936,6 +938,11 @@ class ModelRunner:
         self.init_indexer_capturer()
 
         self.graph_shared_output = None
+        # Set once real decode CUDA graphs are captured (makes on-flip role-switch
+        # capture idempotent).
+        self.decode_cuda_graph_captured = False
+        # Captured decode bs; exposed via /get_server_info for role-switch queries.
+        self.decode_cuda_graph_capture_bs: list[int] = []
 
     def maybe_init_hisparse_coordinator(self):
         if not self.enable_hisparse:
@@ -1492,6 +1499,51 @@ class ModelRunner:
             capture.time_usage,
             phases=("decode", "target_verify", "draft_decode"),
         )
+        # Bookkeeping for the PD role switch: mark the graphs as captured (makes
+        # the on-flip capture idempotent) and record the captured bs so it can be
+        # queried via /get_server_info.
+        self.decode_cuda_graph_captured = self.decode_cuda_graph_runner is not None
+        self.decode_cuda_graph_capture_bs = list(
+            getattr(self.decode_cuda_graph_runner, "capture_bs", []) or []
+        )
+
+    def ensure_decode_cuda_graphs(self, capture_bs: Optional[list[int]] = None):
+        """Idempotently capture decode CUDA graphs after startup.
+
+        Used by the PD role switch: an instance launched as prefill runs fully
+        eager (decode CUDA graph disabled). On the first flip to decode we
+        enable the decode CUDA graph and capture it here, so the flipped
+        instance replays decode graphs instead of running eager.
+        """
+        if self.decode_cuda_graph_captured:
+            logger.info("Decode CUDA graphs already captured; skipping re-capture.")
+            return
+
+        cfg = get_exec().graph.cuda_graph_config
+        was_disabled = cfg is not None and cfg.decode.backend == Backend.DISABLED
+        if was_disabled:
+            # Prefill was launched with the decode CUDA graph disabled; enable it
+            # for the decode role.
+            logger.info(
+                "Enabling decode CUDA graph on role switch (was disabled at startup)."
+            )
+            cfg.decode.backend = Backend.FULL
+            get_context().override(
+                "model_runner.ensure_decode_cuda_graphs", disable_cuda_graph=False
+            )
+
+        if capture_bs:
+            # Capture-to-fit: only the requested (router-sized) batch sizes.
+            filtered_bs = sorted({int(b) for b in capture_bs if int(b) > 0})
+            if filtered_bs:
+                cfg.decode.bs = filtered_bs
+
+        if was_disabled:
+            # graph_shared_output is skipped at startup when decode is disabled,
+            # so build it now (before the decode runner reads its logits buffer).
+            self.graph_shared_output = GraphSharedOutput.create_for_model_runner(self)
+
+        self.init_decode_cuda_graph()
 
     def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
         self.prefill_cuda_graph_runner = None

python/sglang/srt/observability/metrics_collector.py

@@ -1124,8 +1124,11 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
         )
         collector: Optional[SchedulerMetricsCollector] = None
         if enable_metrics:
-            engine_type = DisaggregationMode.to_engine_type(
-                get_disagg().disaggregation_mode
+            # Keep one metric series across role flips.
+            engine_type = (
+                "dynamic"
+                if get_disagg().enable_pd_role_switch
+                else DisaggregationMode.to_engine_type(get_disagg().disaggregation_mode)
             )
             labels = {
                 "model_name": get_serving().served_model_name,