chore: import upstream snapshot with attribution

python/sglang/srt/managers/async_dynamic_batch_tokenizer.py

@@ -0,0 +1,171 @@
+"""
+Asynchronous dynamic batch tokenizer for SGLang.
+
+This module provides an async tokenizer with dynamic batching capabilities
+to reduce tokenization overhead when multiple requests arrive concurrently.
+"""
+
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncDynamicbatchTokenizer:
+    """Asynchronous tokenizer with dynamic batching for single string prompts.
+
+    Dynamically batches pending encode requests from a queue to reduce overhead.
+    Only handles single string prompts - regular batch processing of multiple
+    strings per request should be handled at a higher level.
+    A single-thread ThreadPoolExecutor is used so the event loop stays responsive.
+
+    Note: Uses lazy initialization for asyncio components because this class
+    is instantiated in TokenizerManager.__init__() before the event loop starts.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        max_batch_size: int = 32,
+        batch_wait_timeout_s: float = 0.002,
+    ) -> None:
+        self.tokenizer = tokenizer
+        self.max_batch_size = max_batch_size
+        self.batch_wait_timeout_s = batch_wait_timeout_s
+
+        # Single queue for all encode requests - initialized lazily
+        self._queue: Optional[asyncio.Queue] = None
+        self._batcher_task: Optional[asyncio.Task] = None
+
+        # Single-thread executor for blocking tokenizer calls
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._initialized = False
+
+    def _ensure_initialized(self):
+        """Lazy initialization of event loop dependent components."""
+        if not self._initialized:
+            self._queue = asyncio.Queue()
+            self._batcher_task = asyncio.create_task(self._dynamic_batch_loop())
+            self._initialized = True
+
+    async def __call__(self, prompt: str, **kwargs) -> Any:
+        """Encode a single prompt."""
+        return await self.encode(prompt, **kwargs)
+
+    async def encode(self, prompt: str, **kwargs) -> Any:
+        """Encode a single prompt."""
+        self._ensure_initialized()
+        result_future: asyncio.Future = asyncio.get_running_loop().create_future()
+        await self._queue.put((prompt, kwargs, result_future))
+        return await result_future
+
+    async def _dynamic_batch_loop(self):
+        """Dynamically batch incoming encode requests for efficiency."""
+        while True:
+            try:
+                # Get the first request
+                prompt, kwargs, result_future = await self._queue.get()
+
+                # Collect requests into dynamic batch
+                prompts = [prompt]
+                kwargs_list = [kwargs]
+                result_futures = [result_future]
+
+                # Check if there are more items immediately available in the queue
+                # If queue is empty, process single item immediately without timeout
+                if self._queue.empty():
+                    # No other requests waiting, process immediately
+                    pass
+                else:
+                    # There might be more requests, wait for dynamic batching opportunity
+                    start_time = asyncio.get_running_loop().time()
+
+                    # Collect more requests up to max_batch_size or batch_wait_timeout_s
+                    while len(prompts) < self.max_batch_size:
+                        elapsed = asyncio.get_running_loop().time() - start_time
+                        if elapsed >= self.batch_wait_timeout_s:
+                            break
+
+                        remaining_time = self.batch_wait_timeout_s - elapsed
+                        try:
+                            prompt, kwargs, result_future = await asyncio.wait_for(
+                                self._queue.get(), remaining_time
+                            )
+                            prompts.append(prompt)
+                            kwargs_list.append(kwargs)
+                            result_futures.append(result_future)
+                        except asyncio.TimeoutError:
+                            break
+
+                # Log dynamic batch information
+                logger.debug(
+                    f"AsyncDynamicbatchTokenizer: Processing dynamic batch of size {len(prompts)}"
+                )
+
+                # Process the dynamic batch
+                await self._process_dynamic_batch(prompts, kwargs_list, result_futures)
+
+            except Exception as e:
+                logger.error(f"Error in dynamic batch loop: {e}")
+                # Continue the loop to handle other requests
+
+    async def _process_dynamic_batch(
+        self,
+        prompts: List[str],
+        kwargs_list: List[Dict],
+        result_futures: List[asyncio.Future],
+    ) -> None:
+        """Process a dynamic batch of encode requests for single string prompts."""
+        # Check if all kwargs are identical for efficient batch processing
+        first_kw = kwargs_list[0]
+        can_batch = all(kw == first_kw for kw in kwargs_list[1:])
+        kwargs = first_kw if can_batch else None
+
+        try:
+            # If every request uses identical kwargs we can run a single
+            # batch tokenizer call for a big speed-up.
+            if can_batch and len(prompts) > 1:
+                encode_fn = partial(self.tokenizer, prompts, **kwargs)
+                results = await asyncio.get_running_loop().run_in_executor(
+                    self._executor, encode_fn
+                )
+
+                for i, fut in enumerate(result_futures):
+                    if not fut.done():
+                        data = {k: v[i] for k, v in results.items()}
+                        fut.set_result(data)
+            else:
+                # Process each request individually due to different kwargs
+                if len(prompts) > 1 and not can_batch:
+                    logger.warning(
+                        f"AsyncDynamicbatchTokenizer: Dynamic batching disabled for batch of {len(prompts)} "
+                        f"requests due to differing kwargs. This reduces performance benefits. "
+                        f"Consider using consistent tokenization parameters across requests."
+                    )
+
+                encode_fn = lambda prompts=prompts, kwargs=kwargs_list: [
+                    self.tokenizer(p, **kw) for p, kw in zip(prompts, kwargs_list)
+                ]
+                results = await asyncio.get_running_loop().run_in_executor(
+                    self._executor, encode_fn
+                )
+
+                for fut, res in zip(result_futures, results):
+                    if not fut.done():
+                        fut.set_result(res)
+        except Exception as e:
+            logger.error(f"Error in dynamic batch processing: {e}")
+            for fut in result_futures:
+                if not fut.done():
+                    fut.set_exception(e)
+
+    def __del__(self):
+        """Clean up background tasks."""
+        if hasattr(self, "_batcher_task") and self._batcher_task:
+            if not self._batcher_task.done():
+                self._batcher_task.cancel()
+        if hasattr(self, "_executor"):
+            self._executor.shutdown(wait=False)

python/sglang/srt/managers/communicator.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import asyncio
+import copy
+from typing import Callable, Generic, List, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class FanOutCommunicator(Generic[T]):
+    """Fan-out request + collect response primitive over zmq.
+
+    One send is fanned out to `fan_out` recipients; the caller awaits until
+    all `fan_out` responses are collected. Supports two modes:
+    - "queueing": requests are serialized; concurrent callers wait in a FIFO queue.
+    - "watching": concurrent callers share a single in-flight request and all
+      receive the same result when it completes.
+
+    Only one request is in-flight at any time in either mode.
+    """
+
+    def __init__(
+        self,
+        send: Callable[[T], None],
+        fan_out: int,
+        mode: str = "queueing",
+    ):
+        self._send = send
+        self._fan_out = fan_out
+        self._mode = mode
+        self._result_event: Optional[asyncio.Event] = None
+        self._result_values: Optional[List[T]] = None
+        self._queueing_lock = asyncio.Lock()
+
+        assert mode in ["queueing", "watching"]
+
+    async def queueing_call(self, obj: T):
+        # asyncio.Lock is FIFO-fair: a new caller cannot acquire while earlier
+        # callers are still waiting, so requests are strictly serialized in
+        # arrival order. It also releases on exception/cancellation, so a
+        # failed caller never blocks the callers queued behind it.
+        async with self._queueing_lock:
+            if obj is not None:
+                self._send(obj)
+
+            self._result_event = asyncio.Event()
+            self._result_values = []
+            await self._result_event.wait()
+            result_values = self._result_values
+            self._result_event = self._result_values = None
+
+            return result_values
+
+    async def watching_call(self, obj):
+        if self._result_event is None:
+            assert self._result_values is None
+            self._result_values = []
+            self._result_event = asyncio.Event()
+
+            if obj is not None:
+                self._send(obj)
+
+        # Capture local refs before await -- after event fires, the first
+        # awakened coroutine clears shared state; later awaiters use local refs.
+        values = self._result_values
+        event = self._result_event
+        await event.wait()
+
+        result_values = copy.deepcopy(values)
+        if self._result_event is event:
+            self._result_event = self._result_values = None
+        return result_values
+
+    async def __call__(self, obj):
+        if self._mode == "queueing":
+            return await self.queueing_call(obj)
+        else:
+            return await self.watching_call(obj)
+
+    def handle_recv(self, recv_obj: T):
+        self._result_values.append(recv_obj)
+        if len(self._result_values) == self._fan_out:
+            self._result_event.set()
+
+    @staticmethod
+    def merge_results(results):
+        all_success = all([r.success for r in results])
+        all_message = [r.message for r in results]
+        all_message = " | ".join(all_message)
+        return all_success, all_message

python/sglang/srt/managers/configure_logging.py

@@ -0,0 +1,70 @@
+"""
+Copyright 2023-2025 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""
+Configure the logging settings of a server.
+
+Usage:
+python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000
+"""
+
+import argparse
+
+import requests
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="http://localhost:30000")
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default=None,
+        choices=["debug", "info", "warning", "error", "critical"],
+        help="Set runtime log level",
+    )
+    parser.add_argument("--log-requests", action="store_true")
+    parser.add_argument("--log-requests-level", type=int, default=3)
+    parser.add_argument(
+        "--dump-requests-folder", type=str, default="/tmp/sglang_request_dump"
+    )
+    parser.add_argument("--dump-requests-threshold", type=int, default=1000)
+    parser.add_argument(
+        "--dump-requests-exclude-meta-keys",
+        type=str,
+        default=None,
+        help=(
+            "Comma-separated meta_info keys to strip from each dumped request "
+            "(e.g. 'routed_experts,hidden_states'). Pass an empty string to "
+            "keep all keys. If not set, the server default is used."
+        ),
+    )
+    args = parser.parse_args()
+
+    payload = {
+        "log_requests": args.log_requests,
+        "log_requests_level": args.log_requests_level,  # Log full requests
+        "dump_requests_folder": args.dump_requests_folder,
+        "dump_requests_threshold": args.dump_requests_threshold,
+        "log_level": args.log_level,
+    }
+    if args.dump_requests_exclude_meta_keys is not None:
+        payload["dump_requests_exclude_meta_keys"] = [
+            k.strip()
+            for k in args.dump_requests_exclude_meta_keys.split(",")
+            if k.strip()
+        ]
+
+    response = requests.post(args.url + "/configure_logging", json=payload)
+    assert response.status_code == 200

python/sglang/srt/managers/data_parallel_controller.py

@@ -0,0 +1,715 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A controller that dispatches requests to multiple data parallel workers."""
+
+import faulthandler
+import logging
+import multiprocessing as mp
+import signal
+import threading
+import time
+from enum import Enum, auto
+from typing import Callable, List, Optional
+
+import psutil
+import setproctitle
+import zmq
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
+from sglang.srt.managers.io_struct import (
+    ActiveRanksOutput,
+    BatchTokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    BlockReqInput,
+    ProfileReq,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    sock_recv,
+    sock_send,
+    unwrap_from_pickle,
+    wrap_as_pickle,
+)
+from sglang.srt.managers.load_snapshot import create_load_snapshot_reader
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.scheduler import run_scheduler_process
+from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
+from sglang.srt.observability.req_time_stats import DPControllerReqTimeStats
+from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
+from sglang.srt.server_args import (
+    DP_ATTENTION_HANDSHAKE_PORT_DELTA,
+    PortArgs,
+    ServerArgs,
+)
+from sglang.srt.utils import numa_utils
+from sglang.srt.utils.common import (
+    configure_logger,
+    kill_itself_when_parent_died,
+    maybe_reindex_device_id,
+)
+from sglang.srt.utils.network import (
+    NetworkAddress,
+    bind_port,
+    get_zmq_socket,
+    get_zmq_socket_on_host,
+)
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+from sglang.srt.utils.watchdog import Watchdog
+from sglang.utils import TypeBasedDispatcher, get_exception_traceback
+
+logger = logging.getLogger(__name__)
+
+SCHEDULER_PIDS_ARG = "scheduler_pids"
+
+
+class LoadBalanceMethod(Enum):
+    """Load balance method."""
+
+    ROUND_ROBIN = auto()
+    FOLLOW_BOOTSTRAP_ROOM = auto()
+    TOTAL_REQUESTS = auto()
+    TOTAL_TOKENS = auto()
+
+    @classmethod
+    def from_str(cls, method: str):
+        method = method.upper()
+        try:
+            return cls[method]
+        except KeyError as exc:
+            raise ValueError(f"Invalid load balance method: {method}") from exc
+
+
+class DPBudget:
+    def __init__(self, dp_size: int):
+        self.dp_size = dp_size
+        self.total_requests = [0] * dp_size
+        self.total_tokens = [0] * dp_size
+        self.last_timestamp = [0.0] * dp_size
+
+    def update_budget(self, loads):
+        """Update budget from shm snapshots, skipping stale reads."""
+        for load in loads:
+            if load.timestamp == self.last_timestamp[load.dp_rank]:
+                continue
+            self.last_timestamp[load.dp_rank] = load.timestamp
+            self.total_requests[load.dp_rank] = (
+                load.num_running_reqs + load.num_waiting_reqs
+            )
+            self.total_tokens[load.dp_rank] = load.num_total_tokens
+
+    def dispatch(self, method: LoadBalanceMethod, estimated_tokens: int = 0):
+        if method == LoadBalanceMethod.TOTAL_REQUESTS:
+            target_rank = self.total_requests.index(min(self.total_requests))
+        elif method == LoadBalanceMethod.TOTAL_TOKENS:
+            # Use total_requests as a tie-breaker when total_tokens are equal
+            target_rank = min(
+                range(self.dp_size),
+                key=lambda i: (self.total_tokens[i], self.total_requests[i]),
+            )
+        else:
+            return None
+
+        # Increment the load of that worker by one as a heuristic
+        self.total_requests[target_rank] += 1
+        self.total_tokens[target_rank] += estimated_tokens
+        return target_rank
+
+
+class DataParallelController:
+    """A controller that dispatches requests to multiple data parallel workers."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        run_scheduler_process_func: Callable,
+    ) -> None:
+        # Parse args
+        self.server_args = server_args
+        self.port_args = port_args
+        self.load_balance_method = LoadBalanceMethod.from_str(
+            server_args.load_balance_method
+        )
+        self.run_scheduler_process_func = run_scheduler_process_func
+
+        # Init inter-process communication
+        self.context = zmq.Context(1 + server_args.dp_size)
+        if server_args.node_rank == 0:
+            self.recv_from_tokenizer = get_zmq_socket(
+                self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
+            )
+
+        # Dispatch method
+        self.round_robin_counter = 0
+        dispatch_lookup = {
+            LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
+            LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM: self.follow_bootstrap_room_scheduler,
+            LoadBalanceMethod.TOTAL_REQUESTS: self.total_requests_scheduler,
+            LoadBalanceMethod.TOTAL_TOKENS: self.total_tokens_scheduler,
+        }
+        self.dispatching = dispatch_lookup[self.load_balance_method]
+        self.refresh_load_budget_on_dispatch = self.load_balance_method in (
+            LoadBalanceMethod.TOTAL_REQUESTS,
+            LoadBalanceMethod.TOTAL_TOKENS,
+        )
+
+        # Load balance budget
+        self.dp_budget = DPBudget(server_args.dp_size)
+        self.load_snapshot_reader = create_load_snapshot_reader(
+            server_args,
+            port_args,
+            caller="DataParallelController",
+        )
+        self._last_refresh_time = 0.0
+
+        # To protect changing env vars to set CUDA_VISIBLE_DEVICES.
+        self.env_lock = threading.Lock()
+
+        # Launch data parallel workers
+        self.scheduler_procs = []
+        self.workers: List[zmq.Socket] = [None] * server_args.dp_size
+        self.status: List[bool] = [True] * server_args.dp_size
+
+        if server_args.enable_dp_attention:
+            self.launch_dp_attention_schedulers(server_args, port_args)
+            # When local control broadcast is enabled, send control messages to
+            # every DP group leader (attn_tp_rank=0) so each leader broadcasts
+            # within its own attn_tp_group instead of the full tp_group.
+            # Otherwise fall back to the original behaviour: send to only the
+            # first leader, which then broadcasts over the full tp_group.
+            local_ctrl = server_args.enable_dp_attention_local_control_broadcast
+            self.control_message_step = 1 if local_ctrl else server_args.tp_size
+        else:
+            self.launch_dp_schedulers(server_args, port_args)
+            self.control_message_step = 1
+
+        self.init_dispatcher()
+
+        self.soft_watchdog = Watchdog.create(
+            debug_name="DataParallelController",
+            watchdog_timeout=server_args.soft_watchdog_timeout,
+            soft=True,
+            test_stuck_time=envs.SGLANG_TEST_STUCK_DP_CONTROLLER.get(),
+        )
+
+        if server_args.enable_metrics:
+            start_cpu_monitor_thread("data_parallel_controller")
+
+    def send_to_all_workers(self, obj):
+        for i, worker in enumerate(self.workers):
+            if self.status[i]:
+                sock_send(worker, obj)
+
+    def send_control_message(self, obj):
+        # Send control messages to first worker of tp group
+        for worker in self.workers[:: self.control_message_step]:
+            sock_send(worker, obj)
+
+    def update_active_ranks(self, ranks: ActiveRanksOutput):
+        self.status = ranks.status
+
+    def refresh_load_budget(self):
+        # Throttle to at most once per 20ms.  When a burst of requests
+        # arrives, dispatching_with_trace() calls this before every
+        # dispatch.  Each call reads the latest scheduler snapshot and
+        # overwrites the speculative +1 increments that DPBudget.dispatch()
+        # added for previously dispatched requests in this burst.  Without
+        # throttling, the budget resets to the (stale) scheduler-reported
+        # value on every request, causing the entire burst to land on a
+        # single DP rank.  The 20ms interval lets the burst complete
+        # using speculative counters, then refreshes from the real
+        # scheduler load for the next batch.
+        now = time.perf_counter()
+        if now - self._last_refresh_time < 0.02:
+            return
+        self._last_refresh_time = now
+        self.dp_budget.update_budget(self.load_snapshot_reader.read_all())
+
+    def dispatching_with_trace(self, req: Req, refresh_load_budget: bool = True):
+        if refresh_load_budget and self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+
+        time_stats = DPControllerReqTimeStats.new_from_obj(
+            unwrap_from_pickle(req.time_stats)
+        )
+
+        time_stats.set_dp_dispatch_time()
+        req.time_stats = wrap_as_pickle(time_stats)
+        self.dispatching(req)
+        req.time_stats = time_stats
+        req.time_stats.set_dp_dispatch_finish_time()
+
+    def dispatch_batch_generate(self, batch_req: BatchTokenizedGenerateReqInput):
+        if self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+        for req in batch_req:
+            self.dispatching_with_trace(req, refresh_load_budget=False)
+
+    def dispatch_batch_embedding(self, batch_req: BatchTokenizedEmbeddingReqInput):
+        if self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+        for req in batch_req:
+            self.dispatching_with_trace(req, refresh_load_budget=False)
+
+    def init_dispatcher(self):
+        self._request_dispatcher = TypeBasedDispatcher(
+            [
+                (TokenizedGenerateReqInput, self.dispatching_with_trace),
+                (TokenizedEmbeddingReqInput, self.dispatching_with_trace),
+                (BatchTokenizedGenerateReqInput, self.dispatch_batch_generate),
+                (BatchTokenizedEmbeddingReqInput, self.dispatch_batch_embedding),
+                (BlockReqInput, self.send_to_all_workers),
+                (ProfileReq, self.send_to_all_workers),
+                (ActiveRanksOutput, self.update_active_ranks),
+            ]
+        )
+        self._request_dispatcher.add_fallback_fn(self.send_control_message)
+
+    def launch_dp_schedulers(self, server_args, port_args):
+        base_gpu_id = 0
+
+        threads = []
+        sockets = []
+        ready_events = []
+        for dp_rank in range(server_args.dp_size):
+            tmp_port_args = PortArgs.init_new(server_args)
+            tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
+            tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
+            tmp_port_args.instance_id = port_args.instance_id
+
+            # This port is checked free in PortArgs.init_new.
+            # We hold it first so that the next dp worker gets a different port
+            sockets.append(bind_port(tmp_port_args.nccl_port))
+
+            ready_event = threading.Event()
+            ready_events.append(ready_event)
+
+            # Create a thread for each worker
+            thread = threading.Thread(
+                target=self.launch_tensor_parallel_group_thread,
+                args=(server_args, tmp_port_args, base_gpu_id, dp_rank, ready_event),
+            )
+            threads.append(thread)
+            base_gpu_id += (
+                server_args.tp_size * server_args.pp_size * server_args.gpu_id_step
+            )
+
+            if server_args.node_rank == 0:
+                self.workers[dp_rank] = get_zmq_socket(
+                    self.context,
+                    zmq.PUSH,
+                    tmp_port_args.scheduler_input_ipc_name,
+                    True,
+                )
+
+        # Free all sockets before starting the threads to launch TP workers
+        for sock in sockets:
+            sock.close()
+
+        # Start all threads
+        for thread in threads:
+            thread.start()
+        for event in ready_events:
+            event.wait()
+
+    def launch_tensor_parallel_group_thread(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        base_gpu_id: int,
+        dp_rank: int,
+        ready_event: threading.Event,
+    ):
+        self.launch_tensor_parallel_group(server_args, port_args, base_gpu_id, dp_rank)
+        ready_event.set()
+
+        # This thread cannot be closed because otherwise the `kill_itself_when_parent_died`
+        # function in scheduler.py will kill the scheduler.
+        while True:
+            time.sleep(30 * 24 * 3600)
+
+    def _broadcast_worker_ports(
+        self, server_args: ServerArgs, worker_ports: Optional[List[int]] = None
+    ) -> List[int]:
+        """Broadcast worker ports from node 0 to all other nodes.
+
+        Node 0 acts as the server, waiting for all other nodes to connect and
+        sending them the pre-allocated worker ports. Other nodes act as clients,
+        connecting to node 0 to receive their copy of the worker ports.
+
+        Args:
+            server_args: Server arguments containing node configuration.
+            worker_ports: Pre-allocated worker ports to broadcast.
+
+        Returns:
+            List of worker ports (same on all nodes after broadcast).
+        """
+        # Determine the endpoint for inter-node communication
+        if server_args.dist_init_addr is None:
+            na = NetworkAddress(
+                server_args.host or "127.0.0.1",
+                server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA,
+            )
+        else:
+            na = NetworkAddress.parse(server_args.dist_init_addr)
+            na = NetworkAddress(na.host, na.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA)
+        endpoint = na.to_tcp()
+
+        if server_args.node_rank == 0:
+            # Node 0: Broadcast worker ports to all other nodes
+            return self._broadcast_ports_as_server(
+                endpoint, server_args.nnodes - 1, worker_ports
+            )
+        else:
+            # Other nodes: Receive worker ports from node 0
+            return self._receive_ports_as_client(endpoint, server_args.node_rank)
+
+    def _broadcast_ports_as_server(
+        self, endpoint: str, expected_clients: int, worker_ports: List[int]
+    ) -> List[int]:
+        """Broadcast worker ports to all client nodes."""
+        logger.debug(f"Broadcasting worker ports to {expected_clients} client nodes")
+        logger.debug(f"Worker ports: {worker_ports}")
+
+        rep_socket = get_zmq_socket(self.context, zmq.REP, endpoint, True)
+
+        try:
+            connected_clients = 0
+            while connected_clients < expected_clients:
+                # Wait for client handshake
+                client_rank = sock_recv(rep_socket)
+                logger.debug(f"Received handshake from node {client_rank}")
+
+                # Send worker ports to client
+                sock_send(rep_socket, wrap_as_pickle(worker_ports))
+                connected_clients += 1
+                logger.debug(
+                    f"Sent worker ports to {connected_clients}/{expected_clients} nodes"
+                )
+
+            logger.debug("Worker port broadcast completed")
+            return worker_ports
+        finally:
+            if self.server_args.elastic_ep_backend is None:
+                rep_socket.close()
+            else:
+                threading.Thread(
+                    target=self._reply_ports_as_server,
+                    args=(rep_socket, worker_ports),
+                    daemon=True,
+                ).start()
+
+    def _reply_ports_as_server(self, rep_socket: zmq.Socket, worker_ports: List[int]):
+        """
+        Runs as a background thread to broadcast worker ports for recovered EP ranks
+        """
+        while True:
+            # Wait for client handshake
+            try:
+                client_rank = sock_recv(rep_socket)
+            except Exception:
+                logger.exception(
+                    "Failed to recv/decode handshake in reply thread; continue"
+                )
+                continue
+            logger.debug(f"Received handshake from node {client_rank}")
+
+            # Send worker ports to client
+            sock_send(rep_socket, wrap_as_pickle(worker_ports))
+            logger.debug(f"Sent worker ports to node {client_rank}")
+
+    def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]:
+        """Receive worker ports from the server node."""
+        logger.debug(f"Connecting to node 0 to receive worker ports")
+
+        req_socket = get_zmq_socket(self.context, zmq.REQ, endpoint, False)
+        req_socket.setsockopt(zmq.RCVTIMEO, 600 * 1000)  # 10 minute timeout
+        req_socket.setsockopt(zmq.SNDTIMEO, 600 * 1000)
+
+        try:
+            # Send handshake with our node rank
+            sock_send(req_socket, wrap_as_pickle(str(node_rank)))
+
+            # Receive worker ports
+            worker_ports = sock_recv(req_socket)
+            logger.debug(f"Received {len(worker_ports)} worker ports from node 0")
+            return worker_ports
+        except zmq.Again:
+            logger.error("Timeout waiting for worker ports from node 0")
+            raise RuntimeError(
+                "Failed to receive worker ports from node 0 within timeout"
+            )
+        finally:
+            req_socket.close()
+
+    def launch_dp_attention_schedulers(
+        self, server_args: ServerArgs, port_args: PortArgs
+    ):
+        if server_args.dist_init_addr is None:
+            bind_host = "127.0.0.1"
+        else:
+            bind_host = NetworkAddress.parse(server_args.dist_init_addr).host
+
+        # Pre-allocate worker ports on node 0 to avoid conflicts
+        worker_ports = []
+        if server_args.node_rank == 0:
+            for dp_rank in range(server_args.dp_size):
+                worker_port, worker_socket = get_zmq_socket_on_host(
+                    self.context, zmq.PUSH, host=bind_host
+                )
+                worker_ports.append(worker_port)
+                self.workers[dp_rank] = worker_socket
+                logger.debug(
+                    "Assigned port %s to worker %s on host %s",
+                    worker_port,
+                    dp_rank,
+                    bind_host,
+                )
+
+        broadcasted_ports = self._broadcast_worker_ports(
+            server_args, worker_ports if worker_ports else None
+        )
+        self.launch_tensor_parallel_group(
+            server_args, port_args, 0, None, broadcasted_ports
+        )
+
+    def launch_tensor_parallel_group(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        base_gpu_id: int,
+        dp_rank: Optional[int],
+        worker_ports: Optional[List[int]] = None,
+    ):
+        if not server_args.enable_dp_attention:
+            logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
+
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=server_args.enable_memory_saver
+        )
+
+        scheduler_pipe_readers = []
+
+        pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1)
+        nnodes_per_pp_rank = max(server_args.nnodes // server_args.pp_size, 1)
+        pp_rank_range = range(
+            pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
+            pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank + 1),
+        )
+
+        nnodes_per_tp_group = nnodes_per_pp_rank
+        tp_size_per_node = server_args.tp_size // nnodes_per_tp_group
+        tp_rank_range = range(
+            tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group),
+            tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
+        )
+
+        attn_cp_rank = 0
+        moe_dp_rank = 0
+        for pp_rank in pp_rank_range:
+            for tp_rank in tp_rank_range:
+                rank_port_args = port_args
+
+                if server_args.enable_dp_attention:
+                    # dp attention has different sharding logic
+                    _, _, dp_rank, _ = compute_dp_attention_world_info(
+                        server_args.enable_dp_attention,
+                        tp_rank,
+                        server_args.tp_size,
+                        server_args.dp_size,
+                        server_args.attn_cp_size,
+                    )
+                    # compute zmq ports for this dp rank
+                    rank_port_args = PortArgs.init_new(
+                        server_args, dp_rank, worker_ports
+                    )
+                    # Data parallelism reuses the tensor parallelism group,
+                    # so all dp ranks should use the same nccl port.
+                    rank_port_args.nccl_port = port_args.nccl_port
+                    rank_port_args.instance_id = port_args.instance_id
+
+                reader, writer = mp.Pipe(duplex=False)
+                gpu_id = (
+                    server_args.base_gpu_id
+                    + base_gpu_id
+                    + ((pp_rank % pp_size_per_node) * tp_size_per_node)
+                    + (tp_rank % tp_size_per_node) * server_args.gpu_id_step
+                )
+                attn_dp_size = (
+                    server_args.dp_size if server_args.enable_dp_attention else 1
+                )
+
+                # Parallelism hierarchy (outermost to innermost):
+                # - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
+                # - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
+                attn_tp_size = (
+                    server_args.tp_size // attn_dp_size // server_args.attn_cp_size
+                )
+                attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
+                moe_dp_rank = tp_rank // (
+                    server_args.tp_size // server_args.moe_dp_size
+                )
+                moe_ep_rank = (
+                    tp_rank
+                    % (server_args.tp_size // server_args.moe_dp_size)
+                    // (
+                        server_args.tp_size
+                        // server_args.moe_dp_size
+                        // server_args.ep_size
+                    )
+                )
+
+                with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
+                    proc = mp.Process(
+                        target=self.run_scheduler_process_func,
+                        args=(
+                            server_args,
+                            rank_port_args,
+                            gpu_id,
+                            tp_rank,
+                            attn_cp_rank,
+                            moe_dp_rank,
+                            moe_ep_rank,
+                            pp_rank,
+                            dp_rank,
+                            writer,
+                        ),
+                    )
+                    with (
+                        memory_saver_adapter.configure_subprocess(),
+                        numa_utils.configure_subprocess(server_args, gpu_id),
+                    ):
+                        proc.start()
+                self.scheduler_procs.append(proc)
+                scheduler_pipe_readers.append(reader)
+
+        # Wait for model to finish loading
+        scheduler_info = []
+        for i in range(len(scheduler_pipe_readers)):
+            scheduler_info.append(scheduler_pipe_readers[i].recv())
+
+        self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
+        self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
+
+    def maybe_external_dp_rank_routing(self, req: Req):
+        if req.routed_dp_rank is not None:
+            logger.debug(f"Direct routing to DP rank {req.routed_dp_rank}")
+            sock_send(self.workers[req.routed_dp_rank], req)
+            return True
+        return False
+
+    def round_robin_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
+        while True:
+            if self.status[self.round_robin_counter]:
+                logger.debug(f"Choose worker {self.round_robin_counter}")
+                sock_send(self.workers[self.round_robin_counter], req)
+                self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                    self.workers
+                )
+                break
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                self.workers
+            )
+
+    def follow_bootstrap_room_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
+        assert req.bootstrap_room is not None, (
+            "req.bootstrap_room should not be None. Do not send requests directly to "
+            "prefill or decode instances; send to the router instead."
+        )
+        target_rank = req.bootstrap_room % len(self.workers)
+        sock_send(self.workers[target_rank], req)
+
+    def total_requests_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+        target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_REQUESTS)
+        sock_send(self.workers[target_worker], req)
+
+    def total_tokens_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+        estimated_tokens = len(req.input_ids)
+        target_worker = self.dp_budget.dispatch(
+            LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=estimated_tokens
+        )
+        sock_send(self.workers[target_worker], req)
+
+    def event_loop(self):
+        while True:
+            while True:
+                self.soft_watchdog.feed()
+                try:
+                    recv_req = sock_recv(self.recv_from_tokenizer, flags=zmq.NOBLOCK)
+                except zmq.ZMQError:
+                    break
+                self._request_dispatcher(recv_req)
+
+
+def run_data_parallel_controller_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+    run_scheduler_process_func: Callable = run_scheduler_process,
+):
+    setproctitle.setproctitle("sglang::data_parallel_controller")
+    faulthandler.enable()
+    kill_itself_when_parent_died()
+    parent_process = psutil.Process().parent()
+
+    configure_logger(server_args)
+    if server_args.enable_trace:
+        process_tracing_init(
+            server_args.otlp_traces_endpoint,
+            "sglang",
+            trace_modules=server_args.trace_modules,
+        )
+        thread_label = "DP Controller"
+        if server_args.disaggregation_mode == "prefill":
+            thread_label = "Prefill DP Controller"
+        elif server_args.disaggregation_mode == "decode":
+            thread_label = "Decode DP Controller"
+        trace_set_thread_info(thread_label)
+
+    try:
+        controller = DataParallelController(
+            server_args, port_args, run_scheduler_process_func
+        )
+        scheduler_pids = [
+            proc.pid for proc in controller.scheduler_procs if proc is not None
+        ]
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "max_total_num_tokens": controller.max_total_num_tokens,
+                "max_req_input_len": controller.max_req_input_len,
+                SCHEDULER_PIDS_ARG: scheduler_pids,
+            }
+        )
+        if server_args.node_rank == 0:
+            controller.event_loop()
+        for proc in controller.scheduler_procs:
+            proc.join()
+            logger.error(
+                f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
+            )
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"DataParallelController hit an exception: {traceback}")
+        parent_process.send_signal(signal.SIGQUIT)

python/sglang/srt/managers/detokenizer_manager.py

@@ -0,0 +1,508 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""DetokenizerManager is a process that detokenizes the token ids."""
+
+import dataclasses
+import logging
+import os
+import signal
+from collections import OrderedDict, defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import psutil
+import pybase64
+import setproctitle
+import torch
+import zmq
+
+from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import (
+    BatchEmbeddingOutput,
+    BatchStrOutput,
+    BatchTokenIDOutput,
+    ConfigureLoggingReq,
+    FreezeGCReq,
+    sock_recv,
+    sock_send,
+)
+from sglang.srt.managers.multi_tokenizer_mixin import MultiHttpWorkerDetokenizerMixin
+from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import configure_logger, freeze_gc, kill_itself_when_parent_died
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.srt.utils.network import get_zmq_socket
+from sglang.srt.utils.patch_tokenizer import decode_without_hf_kwargs
+from sglang.srt.utils.watchdog import Watchdog
+from sglang.utils import (
+    TypeBasedDispatcher,
+    find_printable_text,
+    get_exception_traceback,
+)
+
+logger = logging.getLogger(__name__)
+
+# Maximum number of request states that detokenizer can hold. When exceeded,
+# oldest request states will be evicted. Default: 65536 (1<<16).
+# For more details, see: https://github.com/sgl-project/sglang/issues/2812
+# Use power of 2 values for better memory allocation.
+DETOKENIZER_MAX_STATES = int(os.environ.get("SGLANG_DETOKENIZER_MAX_STATES", 1 << 16))
+
+
+@dataclasses.dataclass
+class DecodeStatus:
+    """Store the status of incremental decoding."""
+
+    decoded_text: str
+    decode_ids: List[int]
+    surr_offset: int
+    read_offset: int
+    # Offset that's sent to tokenizer for incremental update.
+    sent_offset: int = 0
+    decoded_text_len: int = dataclasses.field(init=False)
+    decoded_text_chunks: List[str] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        self.decoded_text_len = len(self.decoded_text)
+
+    def append_decoded_text(self, text: str):
+        if text:
+            self.decoded_text_chunks.append(text)
+            self.decoded_text_len += len(text)
+
+    def get_decoded_text(self) -> str:
+        if self.decoded_text_chunks:
+            self.decoded_text += "".join(self.decoded_text_chunks)
+            self.decoded_text_chunks.clear()
+        return self.decoded_text
+
+
+class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
+    """DetokenizerManager is a process that detokenizes the token ids."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        # Init inter-process communication
+        self.init_ipc_channels(port_args, server_args)
+
+        # Init tokenizer
+        self.init_tokenizer(server_args)
+
+        # Init running status
+        self.init_running_status(server_args)
+
+        # Init dispatcher
+        self.init_request_dispatcher()
+
+    def init_ipc_channels(self, port_args: PortArgs, server_args: ServerArgs):
+        context = zmq.Context(2)
+        self.recv_from_scheduler = get_zmq_socket(
+            context, zmq.PULL, port_args.detokenizer_ipc_name, True
+        )
+        # In multi-tokenizer mode, results are pushed back to each TokenizerWorker
+        # directly via SocketMapping inside multi_http_worker_event_loop, so the
+        # single send_to_tokenizer socket is unused.
+        if server_args.tokenizer_worker_num == 1:
+            self.send_to_tokenizer = get_zmq_socket(
+                context, zmq.PUSH, port_args.tokenizer_ipc_name, False
+            )
+
+    def init_tokenizer(self, server_args: ServerArgs):
+        if server_args.skip_tokenizer_init:
+            self.tokenizer = None
+        else:
+            self.tokenizer = get_tokenizer(
+                server_args.tokenizer_path,
+                tokenizer_mode=server_args.tokenizer_mode,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+                tokenizer_backend=server_args.tokenizer_backend,
+            )
+
+    def init_running_status(self, server_args: ServerArgs):
+        self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES)
+        self.disable_tokenizer_batch_decode = server_args.disable_tokenizer_batch_decode
+        self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss"
+
+        self.soft_watchdog = Watchdog.create(
+            debug_name="DetokenizerManager",
+            watchdog_timeout=server_args.soft_watchdog_timeout,
+            soft=True,
+            test_stuck_time=envs.SGLANG_TEST_STUCK_DETOKENIZER.get(),
+        )
+
+        if server_args.enable_metrics:
+            start_cpu_monitor_thread("detokenizer")
+
+    def init_request_dispatcher(self):
+        self._request_dispatcher = TypeBasedDispatcher(
+            [
+                (BatchEmbeddingOutput, self.handle_batch_embedding_out),
+                (BatchTokenIDOutput, self.handle_batch_token_id_out),
+                (FreezeGCReq, self.handle_freeze_gc_req),
+                (ConfigureLoggingReq, self.handle_configure_logging_req),
+            ]
+        )
+
+    def event_loop(self):
+        """The event loop that handles requests"""
+        while True:
+            with self.soft_watchdog.disable():
+                recv_obj = sock_recv(self.recv_from_scheduler)
+            output = self._request_dispatcher(recv_obj)
+            if output is not None:
+                sock_send(self.send_to_tokenizer, output)
+            self.soft_watchdog.feed()
+
+    def trim_matched_stop(
+        self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
+    ):
+        if not finished_reason:
+            return output
+
+        matched = finished_reason.get("matched", None)
+        if not matched:
+            return output
+
+        # TODO(lmzheng): handle the case where multiple stop strs are hit
+
+        # Trim stop str.
+        if isinstance(matched, str) and isinstance(output, str):
+            pos = output.find(matched)
+            if pos == -1:
+                return output
+            end = pos + len(matched)
+            return output[:end] if no_stop_trim else output[:pos]
+
+        # Trim stop token.
+        if isinstance(matched, int) and isinstance(output, list):
+            if no_stop_trim:
+                return output
+            # 200012 <|call|> is the tool call token and one of eos tokens for gpt-oss model
+            if output[-1] == 200012 and self.is_tool_call_parser_gpt_oss:
+                return output
+            assert len(output) > 0
+            # NOTE: We can always assume the last token is the matched stop token
+            return output[:-1]
+        return output
+
+    def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOutput):
+        # If it is embedding model, no detokenization is needed.
+        return recv_obj
+
+    def _grouped_batch_decode(
+        self,
+        ids_list: List[List[int]],
+        skip_list: List[bool],
+        space_list: List[bool],
+    ) -> List[str]:
+        """Batch decode with grouping by (skip_special_tokens, spaces_between_special_tokens)."""
+        n = len(ids_list)
+        if n == 0:
+            return []
+
+        # Empty token spans decode to "" but tokenizer.batch_decode (and the
+        # slow per-row decode_without_hf_kwargs path) still pays per-row
+        # overhead; under high-concurrency streaming this adds up. Filter
+        # empties out, decode the rest, then scatter back.
+        keep_idx: Optional[List[int]] = None
+        if not all(ids_list):
+            keep_idx = [i for i, ids in enumerate(ids_list) if ids]
+            if not keep_idx:
+                return [""] * n
+            ids_list = [ids_list[i] for i in keep_idx]
+            skip_list = [skip_list[i] for i in keep_idx]
+            space_list = [space_list[i] for i in keep_idx]
+
+        if not getattr(self.tokenizer, "is_fast", False):
+            decoded = [
+                decode_without_hf_kwargs(self.tokenizer, ids, skip)
+                for ids, skip in zip(ids_list, skip_list)
+            ]
+        else:
+            # fast path: all rows share the same (skip, space) flags.
+            first_skip, first_space = skip_list[0], space_list[0]
+            if all(
+                s == first_skip and sp == first_space
+                for s, sp in zip(skip_list, space_list)
+            ):
+                decoded = self.tokenizer.batch_decode(
+                    ids_list,
+                    skip_special_tokens=first_skip,
+                    spaces_between_special_tokens=first_space,
+                )
+            else:
+                # Group indices by (skip, space) tuple and decode each group.
+                groups: Dict[Tuple[bool, bool], List[int]] = defaultdict(list)
+                for idx, (skip, space) in enumerate(zip(skip_list, space_list)):
+                    groups[(skip, space)].append(idx)
+
+                decoded = [""] * len(ids_list)
+                for (skip, space), indices in groups.items():
+                    group_decoded = self.tokenizer.batch_decode(
+                        [ids_list[idx] for idx in indices],
+                        skip_special_tokens=skip,
+                        spaces_between_special_tokens=space,
+                    )
+                    for idx, text in zip(indices, group_decoded):
+                        decoded[idx] = text
+
+        if keep_idx is None:
+            return decoded
+        results = [""] * n
+        for i, text in zip(keep_idx, decoded):
+            results[i] = text
+        return results
+
+    def _decode_batch_token_id_output(self, recv_obj: BatchTokenIDOutput):
+        bs = len(recv_obj.rids)
+
+        # Initialize decode status
+        read_ids, surr_ids = [], []
+        for i in range(bs):
+            rid = recv_obj.rids[i]
+            if rid not in self.decode_status:
+                s = DecodeStatus(
+                    decoded_text=recv_obj.decoded_texts[i],
+                    decode_ids=list(recv_obj.decode_ids[i]),
+                    surr_offset=0,
+                    read_offset=recv_obj.read_offsets[i],
+                )
+                self.decode_status[rid] = s
+            else:
+                s = self.decode_status[rid]
+                s.decode_ids.extend(recv_obj.decode_ids[i])
+
+            read_ids.append(
+                self.trim_matched_stop(
+                    s.decode_ids[s.surr_offset :],
+                    recv_obj.finished_reasons[i],
+                    recv_obj.no_stop_trim[i],
+                )
+            )
+            surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
+
+        # Decode token ids to strings
+        if not self.disable_tokenizer_batch_decode:
+            surr_texts = self._grouped_batch_decode(
+                surr_ids,
+                recv_obj.skip_special_tokens,
+                recv_obj.spaces_between_special_tokens,
+            )
+            read_texts = self._grouped_batch_decode(
+                read_ids,
+                recv_obj.skip_special_tokens,
+                recv_obj.spaces_between_special_tokens,
+            )
+        else:
+            # Do not use batch decode to prevent some detokenization edge cases (e.g., gpt-oss).
+            surr_texts = [
+                self.tokenizer.decode(
+                    surr, skip_special_tokens=skip, spaces_between_special_tokens=space
+                )
+                for surr, skip, space in zip(
+                    surr_ids,
+                    recv_obj.skip_special_tokens,
+                    recv_obj.spaces_between_special_tokens,
+                )
+            ]
+            read_texts = [
+                self.tokenizer.decode(
+                    read, skip_special_tokens=skip, spaces_between_special_tokens=space
+                )
+                for read, skip, space in zip(
+                    read_ids,
+                    recv_obj.skip_special_tokens,
+                    recv_obj.spaces_between_special_tokens,
+                )
+            ]
+
+        # Incremental decoding
+        output_strs = []
+        for i in range(bs):
+            rid = recv_obj.rids[i]
+            try:
+                s = self.decode_status[rid]
+            except KeyError:
+                raise RuntimeError(
+                    f"Decode status not found for request {rid}. "
+                    "It may be due to the request being evicted from the decode status due to memory pressure. "
+                    "Please increase the maximum number of requests by setting "
+                    "the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. "
+                    f"The current value is {DETOKENIZER_MAX_STATES}. "
+                    "For more details, see: https://github.com/sgl-project/sglang/issues/2812"
+                )
+            new_text = read_texts[i][len(surr_texts[i]) :]
+            if recv_obj.finished_reasons[i] is None:
+                # Streaming. Invariant: sent_offset >= decoded_text_len. The
+                # gap (`pending`) is "printable but uncommitted" text emitted
+                # in a prior "�" recovery step; we skip it from this step's
+                # emission so we don't double-send.
+                pending = s.sent_offset - s.decoded_text_len
+                if new_text and not new_text.endswith("�"):
+                    # Clean text: commit to decoded_text and advance offsets.
+                    s.append_decoded_text(new_text)
+                    s.surr_offset = s.read_offset
+                    s.read_offset = len(s.decode_ids)
+                    s.sent_offset = s.decoded_text_len
+                    output_strs.append(new_text[pending:] if pending else new_text)
+                else:
+                    # Incomplete UTF-8: emit the printable prefix only; do not
+                    # commit (token offsets stay so the next iteration retries
+                    # with more tokens).
+                    printable = find_printable_text(new_text)
+                    s.sent_offset = s.decoded_text_len + len(printable)
+                    output_strs.append(printable[pending:] if pending else printable)
+                continue
+
+            if rid in self.decode_status:
+                del self.decode_status[rid]
+
+            # Finished: materialize once, trim the matched stop, emit the tail.
+            output_str = self.trim_matched_stop(
+                s.get_decoded_text() + new_text,
+                recv_obj.finished_reasons[i],
+                recv_obj.no_stop_trim[i],
+            )
+            incremental_output = output_str[s.sent_offset :]
+            s.sent_offset = len(output_str)
+            output_strs.append(incremental_output)
+
+        return output_strs
+
+    @staticmethod
+    def _b64_encode_per_request(
+        data_list: Optional[List[Optional[torch.Tensor]]],
+    ) -> Optional[List[Optional[str]]]:
+        """Encode a per-request list of tensors as base64 strings, off the
+        tokenizer hot path. Returns None when the input is None; per-item None
+        stays None.
+        """
+        if data_list is None:
+            return None
+        return [
+            (
+                pybase64.b64encode(item.numpy().tobytes()).decode("utf-8")
+                if item is not None
+                else None
+            )
+            for item in data_list
+        ]
+
+    def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
+        # If handling idle batch, set output_strs to [].
+        output_strs = (
+            self._decode_batch_token_id_output(recv_obj)
+            if len(recv_obj.rids) > 0
+            else []
+        )
+        routed_experts = self._b64_encode_per_request(recv_obj.routed_experts)
+        indexer_topk = self._b64_encode_per_request(recv_obj.indexer_topk)
+        return BatchStrOutput(
+            rids=recv_obj.rids,
+            http_worker_ipcs=recv_obj.http_worker_ipcs,
+            finished_reasons=recv_obj.finished_reasons,
+            output_strs=output_strs,
+            output_ids=recv_obj.output_ids,
+            prompt_tokens=recv_obj.prompt_tokens,
+            reasoning_tokens=recv_obj.reasoning_tokens,
+            completion_tokens=recv_obj.completion_tokens,
+            cached_tokens=recv_obj.cached_tokens,
+            cached_tokens_details=recv_obj.cached_tokens_details,
+            image_tokens=recv_obj.image_tokens,
+            audio_tokens=recv_obj.audio_tokens,
+            video_tokens=recv_obj.video_tokens,
+            spec_verify_ct=recv_obj.spec_verify_ct,
+            spec_num_correct_drafts=recv_obj.spec_num_correct_drafts,
+            spec_num_block_accept_tokens=recv_obj.spec_num_block_accept_tokens,
+            spec_num_cap_tokens=recv_obj.spec_num_cap_tokens,
+            spec_correct_drafts_histogram=recv_obj.spec_correct_drafts_histogram,
+            spec_cap_lens_histogram=recv_obj.spec_cap_lens_histogram,
+            input_token_logprobs_val=recv_obj.input_token_logprobs_val,
+            input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
+            output_token_logprobs_val=recv_obj.output_token_logprobs_val,
+            output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
+            input_top_logprobs_val=recv_obj.input_top_logprobs_val,
+            input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
+            output_top_logprobs_val=recv_obj.output_top_logprobs_val,
+            output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
+            input_token_ids_logprobs_val=recv_obj.input_token_ids_logprobs_val,
+            input_token_ids_logprobs_idx=recv_obj.input_token_ids_logprobs_idx,
+            output_token_ids_logprobs_val=recv_obj.output_token_ids_logprobs_val,
+            output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx,
+            output_token_entropy_val=recv_obj.output_token_entropy_val,
+            output_hidden_states=recv_obj.output_hidden_states,
+            routed_experts=routed_experts,
+            indexer_topk=indexer_topk,
+            customized_info=recv_obj.customized_info,
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
+            retraction_counts=recv_obj.retraction_counts,
+            token_steps=recv_obj.token_steps,
+            dp_ranks=recv_obj.dp_ranks,
+            time_stats=recv_obj.time_stats,
+        )
+
+    def handle_freeze_gc_req(self, recv_req: FreezeGCReq):
+        freeze_gc("Detokenizer Manager")
+        return None
+
+    def handle_configure_logging_req(self, recv_req: ConfigureLoggingReq):
+        if recv_req.log_level is not None:
+            logging.getLogger().setLevel(recv_req.log_level.upper())
+
+
+def is_health_check_request(rid: Optional[str]) -> bool:
+    return isinstance(rid, str) and rid.startswith(HEALTH_CHECK_RID_PREFIX)
+
+
+class LimitedCapacityDict(OrderedDict):
+    def __init__(self, capacity: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.capacity = capacity
+
+    def __setitem__(self, key, value):
+        if len(self) >= self.capacity:
+            # Remove the oldest element (first item in the dict)
+            self.popitem(last=False)
+        # Set the new item
+        super().__setitem__(key, value)
+
+
+def run_detokenizer_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    detokenizer_manager_class=DetokenizerManager,
+):
+    kill_itself_when_parent_died()
+    setproctitle.setproctitle("sglang::detokenizer")
+    configure_logger(server_args)
+    parent_process = psutil.Process().parent()
+
+    manager = None
+    try:
+        manager = detokenizer_manager_class(server_args, port_args)
+        if server_args.tokenizer_worker_num == 1:
+            manager.event_loop()
+        else:
+            manager.multi_http_worker_event_loop()
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"DetokenizerManager hit an exception: {traceback}")
+        if manager is not None:
+            manager.maybe_clear_socket_mapping()
+        parent_process.send_signal(signal.SIGQUIT)

python/sglang/srt/managers/disagg_service.py

@@ -0,0 +1,44 @@
+"""Start bootstrap/kv-store-related server"""
+
+import os
+
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    KVClassType,
+    TransferBackend,
+    get_kv_class,
+)
+from sglang.srt.server_args import ServerArgs
+
+
+def start_disagg_service(
+    server_args: ServerArgs,
+):
+    # Start kv bootstrap server on prefill
+    disagg_mode = DisaggregationMode(server_args.disaggregation_mode)
+    transfer_backend = TransferBackend(server_args.disaggregation_transfer_backend)
+
+    if disagg_mode == DisaggregationMode.PREFILL:
+        # only start bootstrap server on prefill tm
+        kv_bootstrap_server_class = get_kv_class(
+            transfer_backend, KVClassType.BOOTSTRAP_SERVER
+        )
+        bootstrap_server = kv_bootstrap_server_class(
+            host=server_args.host,
+            port=server_args.disaggregation_bootstrap_port,
+        )
+        is_create_store = (
+            server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND
+        )
+        if is_create_store:
+            try:
+                from memfabric_hybrid import create_config_store
+
+                ascend_url = os.getenv("ASCEND_MF_STORE_URL")
+                create_config_store(ascend_url)
+            except Exception as e:
+                error_message = f"Failed create mf store, invalid ascend_url."
+                error_message += f" With exception {e}"
+                raise error_message
+
+        return bootstrap_server

python/sglang/srt/managers/embed_types.py

@@ -0,0 +1,58 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Structs for embedding injection.
+
+These are placed in a separate module to avoid circular imports between
+io_struct.py and schedule_batch.py.
+"""
+
+from typing import List
+
+import msgspec
+import torch
+
+
+class PositionalEmbeds(msgspec.Struct, array_like=True):
+    """Embeddings to place at specific token positions.
+
+    Accepts either a list of [1, hidden_dim] tensors or a pre-stacked [N, hidden_dim] tensor.
+    In both cases, __post_init__ stacks into a single [N, hidden_dim] tensor to reduce
+    ZMQ serialization overhead.
+
+    Attributes:
+        embeds: Stacked tensor of shape [N, hidden_dim] after __post_init__.
+        positions: List of positions where embeddings should be injected.
+    """
+
+    embeds: torch.Tensor
+    positions: List[int]
+
+    def __post_init__(self):
+        # Normalize list of tensors into a single [N, hidden_dim] tensor.
+        # Dispatch by element rank to avoid a per-element unsqueeze.
+        if isinstance(self.embeds, list):
+            if not self.embeds:
+                self.embeds = torch.cat(self.embeds, dim=0)  # raises — empty is invalid
+            elif self.embeds[0].dim() == 1:
+                # [hidden_dim] elements → stack adds the leading dim.
+                self.embeds = torch.stack(self.embeds, dim=0)
+            else:
+                # [1, hidden_dim] (already has the leading dim) → plain concat.
+                self.embeds = torch.cat(self.embeds, dim=0)
+        if self.embeds.shape[0] != len(self.positions):
+            raise ValueError(
+                f"embeds length ({self.embeds.shape[0]}) != "
+                f"positions length ({len(self.positions)})"
+            )

python/sglang/srt/managers/hisparse_coordinator.py

@@ -0,0 +1,843 @@
+# to be combined with the sparse coordinator class and sparse algorithm family
+
+import logging
+from typing import List, NamedTuple, Union
+
+import torch
+
+from sglang.jit_kernel.hisparse import (
+    load_cache_to_device_buffer_dsv4_mla,
+    load_cache_to_device_buffer_mla,
+)
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.mem_cache.allocator.hisparse import (
+    DeepSeekV4HiSparseTokenToKVPoolAllocator,
+    HiSparseTokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.hisparse_memory_pool import (
+    HiSparseDSATokenToKVPool,
+)
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.memory_pool_host import (
+    DeepSeekV4PagedHostPool,
+    MLATokenToKVPoolHost,
+)
+from sglang.srt.utils import get_device_module, is_hip
+
+device_module = get_device_module()
+
+_is_hip = is_hip()
+
+logger = logging.getLogger(__name__)
+
+
+class HiSparseAct(NamedTuple):
+    start_event: device_module.Event
+    finish_event: device_module.Event
+    req: Req
+
+
+class HiSparseTokenStats(NamedTuple):
+    device_tokens: int
+    device_token_usage: float
+    host_tokens: int
+    host_token_usage: float
+
+
+class HiSparseCoordinator:
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: Union[
+            HiSparseTokenToKVPoolAllocator,
+            DeepSeekV4HiSparseTokenToKVPoolAllocator,
+        ],
+        top_k: int,
+        device_buffer_size: int,
+        device: str,
+        tp_group,
+        host_to_device_ratio: int = 2,
+        swap_in_block_size: int = 960,
+    ):
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.top_k = top_k
+        self.device_buffer_size = device_buffer_size
+        self.device = device
+        self.swap_in_block_size = swap_in_block_size
+        self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio
+
+        self.is_dsv4_hisparse = isinstance(
+            self.token_to_kv_pool_allocator, DeepSeekV4HiSparseTokenToKVPoolAllocator
+        )
+        if self.is_dsv4_hisparse:
+            self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache
+            page_size = self.mem_pool_device.page_size
+            num_host_pages = (
+                self.token_to_kv_pool_allocator.size_full // self.compress_ratio
+                + page_size
+                - 1
+            ) // page_size
+            self.mem_pool_host = DeepSeekV4PagedHostPool(
+                pool_name="dsv4_hisparse_c4",
+                device_buffers=self.mem_pool_device.kv_buffer,
+                item_bytes=self.mem_pool_device.bytes_per_page_padded,
+                num_host_pages=num_host_pages,
+                slot_page_size=page_size,
+                layout="layer_first",
+            )
+            self.item_size_bytes = (
+                self.mem_pool_device.kv_cache_total_dim
+                * self.mem_pool_device.store_dtype.itemsize
+            )
+        else:
+            assert isinstance(
+                self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator
+            )
+            self.mem_pool_device: HiSparseDSATokenToKVPool = (
+                self.token_to_kv_pool_allocator.get_kvcache()
+            )
+            self.mem_pool_host = MLATokenToKVPoolHost(
+                device_pool=self.mem_pool_device,
+                host_to_device_ratio=host_to_device_ratio,
+                host_size=0,
+                page_size=self.mem_pool_device.page_size,
+                layout="layer_first",
+                override_kv_cache_dim=self.mem_pool_device.kv_cache_dim,
+            )
+            self.item_size_bytes = self.mem_pool_host.token_stride_size
+        self.page_size = self.mem_pool_device.page_size
+
+        max_num_req_slots = req_to_token_pool.req_to_token.shape[0]
+        max_context_len = req_to_token_pool.max_context_len
+        max_compressed_context_len = (
+            max_context_len + self.compress_ratio - 1
+        ) // self.compress_ratio
+
+        # to have an extra page for new tokens
+        self.padded_buffer_size = (
+            self.device_buffer_size + self.mem_pool_device.page_size
+        )
+
+        self.req_to_device_buffer = torch.zeros(
+            (max_num_req_slots, self.padded_buffer_size),
+            dtype=torch.int64,
+            device=device,
+        )
+        self.req_device_buffer_size = torch.zeros(
+            max_num_req_slots, dtype=torch.int64, device="cpu"
+        )
+        self.req_to_host_pool = torch.full(
+            (max_num_req_slots, max_compressed_context_len + self.page_size),
+            -1,
+            dtype=torch.int64,
+            device=device,
+        )
+        self.req_to_host_pool_allocated_len = torch.zeros(
+            max_num_req_slots, dtype=torch.int64, device="cpu"
+        )
+
+        self.write_staging_stream = device_module.Stream()
+        self.decode_backup_stream = device_module.Stream()
+        self.ack_staging_queue: List[HiSparseAct] = []
+        self.decode_producer_stream = None
+        self._backup_done_event = device_module.Event()
+        self._has_pending_backup = False
+
+        self.tp_group = tp_group
+        self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
+
+        # initialize data structures for swap-in kernel
+        layer_num = self.mem_pool_device.layer_num
+        self.req_device_buffer_tokens = torch.full(
+            (layer_num, max_num_req_slots, self.padded_buffer_size),
+            -1,
+            dtype=torch.int32,
+            device=device,
+        )
+        self.req_device_buffer_token_locs = torch.full(
+            (layer_num, max_num_req_slots, self.padded_buffer_size),
+            -1,
+            dtype=torch.int32,
+            device=device,
+        )
+        self._lru_init = torch.arange(
+            self.device_buffer_size, dtype=torch.int16, device=device
+        )
+        self.lru_slots = (
+            self._lru_init.view(1, 1, -1)
+            .repeat(layer_num, max_num_req_slots, 1)
+            .contiguous()
+        )
+        self._device_buffer_arange_i32 = torch.arange(
+            self.device_buffer_size, dtype=torch.int32, device=device
+        )
+
+        # Pre-allocated output buffer for swap_in_selected_pages (CUDA-graph safe)
+        self.top_k_device_locs_buffer = torch.full(
+            (max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
+        )
+        self.raw_indices_buffer = torch.full(
+            (max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
+        )
+        # Scalar tensor: number of real (non-padded) requests in the batch.
+        # Updated before each graph replay so padded blocks early-return.
+        self.num_real_reqs = torch.zeros(1, dtype=torch.int32, device=device)
+
+        # CPU flag: True means "skip backup on the next decode step" because
+        # staging already backed up all prefill tokens.  Cleared after one step.
+        self._skip_first_backup = [False] * max_num_req_slots
+
+    def set_decode_producer_stream(self, stream) -> None:
+        self.decode_producer_stream = stream
+
+    def destroy(self) -> None:
+        # Drain in-flight transfers so the buffer is idle, then unregister it.
+        # See HostKVCache.destroy for why the explicit unregister matters.
+        self.write_staging_stream.synchronize()
+        self.decode_backup_stream.synchronize()
+        self.mem_pool_host.destroy()
+
+    def get_token_stats(self) -> HiSparseTokenStats:
+        device_allocator = self.token_to_kv_pool_allocator.hisparse_attn_allocator
+        device_capacity = device_allocator.size
+        device_tokens = device_capacity - device_allocator.available_size()
+        host_capacity = self.mem_pool_host.size
+        host_tokens = host_capacity - self.mem_pool_host.available_size()
+        return HiSparseTokenStats(
+            device_tokens=device_tokens,
+            device_token_usage=(
+                device_tokens / device_capacity if device_capacity > 0 else 0.0
+            ),
+            host_tokens=host_tokens,
+            host_token_usage=(
+                host_tokens / host_capacity if host_capacity > 0 else 0.0
+            ),
+        )
+
+    def admit_request_into_staging(self, req: Req) -> None:
+        req.hisparse_staging = True
+
+        full_kv_indices = self.req_to_token_pool.req_to_token[
+            req.req_pool_idx, : req.extend_range.end
+        ].to(dtype=torch.int64, copy=True)
+        device_indices = (
+            self.mem_pool_device.translate_loc_from_full_to_hisparse_device(
+                full_kv_indices
+            )
+        )
+
+        prefill_len = len(device_indices)
+        host_indices = self.mem_pool_host.alloc_paged_token_slots(
+            self.req_to_host_pool,
+            self.req_to_host_pool_allocated_len,
+            req.req_pool_idx,
+            0,
+            prefill_len,
+        )
+
+        start_event = device_module.Event()
+        finish_event = device_module.Event()
+        start_event.record()
+        with device_module.stream(self.write_staging_stream):
+            start_event.wait(self.write_staging_stream)
+            self.mem_pool_host.backup_from_device_all_layer(
+                self.mem_pool_device,
+                host_indices,
+                device_indices,
+                io_backend="kernel",
+            )
+            finish_event.record()
+            if host_indices.is_cuda:
+                host_indices.record_stream(self.write_staging_stream)
+            if device_indices.is_cuda:
+                device_indices.record_stream(self.write_staging_stream)
+
+        self.ack_staging_queue.append(HiSparseAct(start_event, finish_event, req))
+
+    def admit_request_direct(self, req: Req) -> None:
+        """Direct-to-host path: KV data already resides in host pool via RDMA.
+
+        Skips staging DMA entirely. Only allocates a small device buffer
+        (4KB) for decode-time swap-in, then marks the request as ready.
+        Host indices were already written to req_to_host_pool.
+
+        Metadata fixups after alloc_device_buffer():
+        - alloc_device_buffer() sets device_buffer_tokens = [0, 1, ..., buf_size-1],
+          which tells the swap-in kernel that those tokens are cached in the device
+          buffer.  In the staging path this is correct (prefill filled the buffer),
+          but here the buffer is empty.
+        """
+        self.alloc_device_buffer(req)
+
+        host_len = self.host_token_len(req.kv_allocated_len)
+        if host_len <= self.device_buffer_size:
+            # Short sequences (seq_len <= device_buffer_size): the kernel fast path
+            # returns device_buffer_locs directly without any host loading, so we
+            # must preload all tokens from host pool into the device buffer
+            # TODO(hzh0425): Optimize this.
+            self._preload_to_device_buffer(req)
+        else:
+            # Long sequence: reset device_buffer_tokens to -1 so the kernel
+            # sees all slots as empty -> every top-k lookup is a miss -> host load.
+            self.req_device_buffer_tokens[
+                :, req.req_pool_idx, : self.device_buffer_size
+            ] = -1
+
+        req.hisparse_staging = False
+        self._skip_first_backup[req.req_pool_idx] = True
+        logger.debug("HiSparse: admitting request %s directly", req.rid)
+
+    def host_token_len(self, kv_allocated_len: int) -> int:
+        if self.is_dsv4_hisparse:
+            return kv_allocated_len // self.compress_ratio
+        return kv_allocated_len
+
+    def _preload_to_device_buffer(self, req: Req) -> None:
+        """Preload all tokens from host pool into the device buffer."""
+        n = self.host_token_len(req.kv_allocated_len)
+        host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
+        device_locs = self.req_to_device_buffer[req.req_pool_idx, :n]
+
+        for layer_id in range(self.mem_pool_device.layer_num):
+            self.mem_pool_host.load_to_device_per_layer(
+                self.mem_pool_device,
+                host_indices,
+                device_locs,
+                layer_id,
+                io_backend="kernel",
+            )
+
+    def alloc_device_buffer(self, req: Req) -> None:
+        if self.is_dsv4_hisparse:
+            allocated_len = req.extend_range.end
+            alloc_size = self.padded_buffer_size
+        else:
+            allocated_len = req.kv_allocated_len
+            page_size = self.mem_pool_device.page_size
+            # Allocate only enough for current tokens (page-aligned).
+            # When prefill already fills device_buffer_size, include the reserved page.
+            alloc_size = min(
+                ((allocated_len + page_size - 1) // page_size) * page_size,
+                self.device_buffer_size,
+            )
+            if alloc_size == self.device_buffer_size:
+                alloc_size = self.padded_buffer_size
+
+        compressed_logical_indices = (
+            self.mem_pool_device.translate_loc_from_full_to_compressed(
+                self.req_to_token_pool.req_to_token[req.req_pool_idx, :allocated_len]
+            )
+        )
+        compressed_len = len(compressed_logical_indices)
+
+        buffer_indices = self.token_to_kv_pool_allocator.alloc_device_buffer(
+            compressed_logical_indices, alloc_size
+        )
+        if buffer_indices is None:
+            logger.error(
+                "HiSparse: alloc_device_buffer failed for req %s "
+                "(compressed_len=%d, alloc_size=%d)",
+                req.rid,
+                compressed_len,
+                alloc_size,
+            )
+            raise RuntimeError("HiSparse alloc_device_buffer returned None")
+
+        buffer_indices = buffer_indices.to(torch.int32)
+        self.req_to_device_buffer[req.req_pool_idx, :alloc_size] = buffer_indices
+        self.req_device_buffer_size[req.req_pool_idx] = alloc_size
+
+        self.req_device_buffer_tokens[
+            :, req.req_pool_idx, : self.device_buffer_size
+        ] = self._device_buffer_arange_i32
+        self.req_device_buffer_token_locs[:, req.req_pool_idx, :alloc_size] = (
+            buffer_indices[:alloc_size]
+        )
+
+    def _grow_device_buffers(
+        self,
+        seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        req_pool_indices_cpu: torch.Tensor,
+    ) -> torch.Tensor:
+        """Grow device buffers for requests whose sequence length exceeds current capacity."""
+        current_caps = self.req_device_buffer_size[req_pool_indices_cpu]
+        short_reqs_cpu = seq_lens_cpu <= self.device_buffer_size
+        needs_grow_cpu = short_reqs_cpu & (seq_lens_cpu > current_caps)
+
+        if torch.any(needs_grow_cpu):
+            page_size = self.mem_pool_device.page_size
+            grow_indices = torch.where(needs_grow_cpu)[0]
+
+            # Compute all grow sizes on CPU, then do a single bulk allocation
+            req_idxs = []
+            old_caps = []
+            new_caps = []
+            grow_sizes = []
+            total_grow = 0
+            for i in grow_indices.tolist():
+                req_idx = int(req_pool_indices_cpu[i])
+                current_cap = int(current_caps[i])
+                seq_len = int(seq_lens_cpu[i])
+
+                new_cap = min(
+                    ((seq_len + page_size - 1) // page_size) * page_size,
+                    self.device_buffer_size,
+                )
+                if new_cap == self.device_buffer_size:
+                    new_cap = self.padded_buffer_size
+                grow_size = new_cap - current_cap
+                if grow_size <= 0:
+                    continue
+                req_idxs.append(req_idx)
+                old_caps.append(current_cap)
+                new_caps.append(new_cap)
+                grow_sizes.append(grow_size)
+                total_grow += grow_size
+
+            if total_grow > 0:
+                all_new_indices = (
+                    self.token_to_kv_pool_allocator.hisparse_attn_allocator.alloc(
+                        total_grow
+                    )
+                )
+                if all_new_indices is None:
+                    logger.error(
+                        "HiSparse: _grow_device_buffers bulk alloc failed "
+                        "(total_grow=%d)",
+                        total_grow,
+                    )
+                    raise RuntimeError(
+                        f"HiSparse _grow_device_buffers failed (total_grow={total_grow})"
+                    )
+
+                offset = 0
+                for req_idx, current_cap, new_cap, grow_size in zip(
+                    req_idxs, old_caps, new_caps, grow_sizes
+                ):
+                    chunk = all_new_indices[offset : offset + grow_size]
+                    offset += grow_size
+                    self.req_to_device_buffer[req_idx, current_cap:new_cap] = chunk
+                    self.req_device_buffer_token_locs[
+                        :, req_idx, current_cap:new_cap
+                    ] = chunk
+                    self.req_device_buffer_size[req_idx] = new_cap
+
+        reserved_positions = (seq_lens - 1).clamp(max=self.device_buffer_size)
+        return self.req_to_device_buffer[req_pool_indices, reserved_positions]
+
+    def has_ongoing_staging(self) -> bool:
+        return len(self.ack_staging_queue) > 0
+
+    def collect_ready_reqs(self) -> List[Req]:
+        ready_reqs: List[Req] = []
+        if len(self.ack_staging_queue) == 0:
+            return ready_reqs
+
+        finish_count = 0
+        for _, finish_event, _ in self.ack_staging_queue:
+            if not finish_event.query():
+                break
+            finish_count += 1
+        queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
+        if self.tp_world_size > 1:
+            # synchronize TP workers to make sure the same update to scheduler
+            torch.distributed.all_reduce(
+                queue_size,
+                op=torch.distributed.ReduceOp.MIN,
+                group=self.tp_group,
+            )
+        finish_count = int(queue_size.item())
+        while finish_count > 0:
+            _, _, req = self.ack_staging_queue.pop(0)
+            # prepare device buffer and update req
+            self.alloc_device_buffer(req)
+            self._skip_first_backup[req.req_pool_idx] = True
+            req.hisparse_staging = False
+            finish_count -= 1
+            ready_reqs.append(req)
+        return ready_reqs
+
+    def map_last_loc_to_buffer(
+        self,
+        seq_lens: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        req_pool_indices_cpu: torch.Tensor,
+    ) -> None:
+        self._eager_backup_previous_token(
+            seq_lens, req_pool_indices, seq_lens_cpu, req_pool_indices_cpu
+        )
+
+        if not self.is_dsv4_hisparse:
+            # Grow device buffers if needed and resolve the latest-token slot.
+            reserved_buffer_loc = self._grow_device_buffers(
+                seq_lens, req_pool_indices, seq_lens_cpu, req_pool_indices_cpu
+            )
+            self.req_device_buffer_token_locs[
+                :, req_pool_indices, self.device_buffer_size
+            ] = reserved_buffer_loc.to(torch.int32)
+
+            compressed_locs = self.token_to_kv_pool_allocator.get_last_loc_compressed(
+                out_cache_loc
+            )
+            # ROCm: the decode remap creates a temporary hisparse device slot per
+            # new token (via the page_size==1 allocator path). Free the stale
+            # slot before pointing the mapping at the reserved device-buffer slot,
+            # otherwise the temporary slots leak and corrupt later swap-in lookups.
+            # CUDA keeps the original behavior: the swap-in kernel consumes only
+            # top_k_device_locs, so stale mapping entries are harmless there.
+            if _is_hip:
+                previous_locs = self.mem_pool_device._translate_loc_to_hisparse_device(
+                    compressed_locs
+                )
+                stale_locs = previous_locs[
+                    (previous_locs > 0) & (previous_locs != reserved_buffer_loc)
+                ]
+                if stale_locs.numel() > 0:
+                    self.token_to_kv_pool_allocator.free_hisparse_indices(stale_locs)
+
+            self.mem_pool_device.full_to_hisparse_device_index_mapping[
+                compressed_locs
+            ] = reserved_buffer_loc
+            return
+
+        active_reqs = seq_lens % self.compress_ratio == 0
+        if not torch.any(active_reqs):
+            return
+
+        active_seq_lens = seq_lens[active_reqs]
+        active_out_cache_loc = out_cache_loc[active_reqs]
+        active_req_pool_indices = req_pool_indices[active_reqs]
+
+        compressed_seq_lens = active_seq_lens // self.compress_ratio
+        reserved_positions = (compressed_seq_lens - 1).clamp(
+            max=self.device_buffer_size
+        )
+        reserved_buffer_loc = self.req_to_device_buffer[
+            active_req_pool_indices, reserved_positions
+        ]
+
+        self.req_device_buffer_token_locs[
+            :, active_req_pool_indices, self.device_buffer_size
+        ] = reserved_buffer_loc.to(torch.int32)
+
+        compressed_locs = self.token_to_kv_pool_allocator.get_last_loc_compressed(
+            active_out_cache_loc
+        )
+        self.mem_pool_device.full_to_hisparse_device_index_mapping[compressed_locs] = (
+            reserved_buffer_loc
+        )
+
+    def _eager_backup_previous_token(
+        self,
+        seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        req_pool_indices_cpu: torch.Tensor,
+    ) -> None:
+        """Back up the previous compressed token to host memory.
+
+        Each newly produced compressed token (one per `compress_ratio` decode
+        steps) must be backed up to host so the swap-in kernel can later
+        recover it.
+
+        Two cases are skipped:
+        - The first decode step right after staging: all prefill tokens were
+          already backed up during staging, so there is nothing new to save.
+        - Steps where `(seq_len - 1) % compress_ratio != 0`: no new compressed
+          token was produced this step.
+        """
+        # Build the list of batch positions that need a host backup.
+        # Skip the first decode step after staging (prefill already backed up),
+        # and skip non-aligned steps that did not produce a new compressed token.
+        backup_indices = []
+        for i in range(len(seq_lens_cpu)):
+            req_idx = int(req_pool_indices_cpu[i])
+            if self._skip_first_backup[req_idx]:
+                self._skip_first_backup[req_idx] = False
+                continue
+            if (int(seq_lens_cpu[i]) - 1) % self.compress_ratio == 0:
+                backup_indices.append(i)
+
+        if not backup_indices:
+            return
+
+        backup_indices_gpu = torch.tensor(
+            backup_indices, dtype=torch.int64, device=self.device
+        )
+        backup_req_indices = req_pool_indices[backup_indices_gpu]
+
+        # The previous compressed token's position and its device buffer slot:
+        #  compressed_pos = (seq_len - 1) // compress_ratio - 1
+        #  - short: slot = compressed_pos          (within the regular buffer)
+        #  - long:  slot = device_buffer_size      (the reserved slot)
+        prev_seq_lens = seq_lens[backup_indices_gpu] - 1
+        compressed_prev_seq_lens = prev_seq_lens // self.compress_ratio
+        actual_compressed_pos = compressed_prev_seq_lens - 1
+
+        buffer_slot = actual_compressed_pos.clamp(max=self.device_buffer_size)
+
+        device_locs = self.req_to_device_buffer[backup_req_indices, buffer_slot]
+
+        host_locs_list = []
+        for i in backup_indices:
+            req_idx = int(req_pool_indices_cpu[i])
+            start_pos = (int(seq_lens_cpu[i]) - 1) // self.compress_ratio - 1
+            host_locs = self.mem_pool_host.alloc_paged_token_slots(
+                self.req_to_host_pool,
+                self.req_to_host_pool_allocated_len,
+                req_idx,
+                start_pos,
+                1,
+            )
+            host_locs_list.append(host_locs)
+        host_locs = torch.cat(host_locs_list)
+
+        self.wait_for_pending_backup()
+        schedule_stream = device_module.current_stream()
+        with device_module.stream(self.decode_backup_stream):
+            self.decode_backup_stream.wait_stream(schedule_stream)
+            if self.decode_producer_stream is not None:
+                self.decode_backup_stream.wait_stream(self.decode_producer_stream)
+            self.mem_pool_host.backup_from_device_all_layer(
+                self.mem_pool_device,
+                host_locs,
+                device_locs,
+                io_backend="kernel",
+            )
+            self._backup_done_event.record()
+            if host_locs.is_cuda:
+                host_locs.record_stream(self.decode_backup_stream)
+            if backup_req_indices.is_cuda:
+                backup_req_indices.record_stream(self.decode_backup_stream)
+            if actual_compressed_pos.is_cuda:
+                actual_compressed_pos.record_stream(self.decode_backup_stream)
+            if device_locs.is_cuda:
+                device_locs.record_stream(self.decode_backup_stream)
+        self._has_pending_backup = True
+
+    def wait_for_pending_backup(self) -> None:
+        if not self._has_pending_backup:
+            return
+        self._backup_done_event.wait(device_module.current_stream())
+        self._has_pending_backup = False
+
+    def naive_load_topk(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        top_k_tokens: torch.Tensor,
+        layer_id: int,
+    ) -> torch.Tensor:
+        """Load top-k selected tokens into device memory and return their device indices.
+
+        This is a naive per-request loop implementation for debugging/validation.
+        Production code uses swap_in_selected_pages (JIT CUDA kernel) instead.
+
+        Note: dsv4 hisparse is not supported — DeepSeekV4SingleKVPoolHost has no
+        load_to_device_per_layer and indices live in compressed space. Currently
+        only used as a kernel oracle in test_hisparse_unit.py (non-dsv4 path).
+
+        Args:
+            req_pool_indices: Pool indices for each request.  Shape: (num_reqs,)
+            seq_lens: Sequence lengths for each request.  Shape: (num_reqs,)
+            top_k_tokens: Selected token positions per request.  Shape: (num_reqs, top_k)
+            layer_id: The layer to load KV cache for.
+
+        Returns:
+            Device KV cache indices for the selected tokens.  Shape: (num_reqs, top_k)
+        """
+        assert (
+            not self.is_dsv4_hisparse
+        ), "naive_load_topk is not implemented for dsv4 hisparse"
+        num_reqs = req_pool_indices.size(0)
+        top_k_indices = torch.full(
+            (num_reqs, self.top_k), -1, dtype=torch.int32, device=self.device
+        )
+
+        for i in range(num_reqs):
+            seq_len = int(seq_lens[i].item())
+            top_n = min(seq_len, self.top_k)
+            if top_n == 0:
+                continue
+
+            req_idx = int(req_pool_indices[i].item())
+            selected_tokens = top_k_tokens[i, :top_n].to(dtype=torch.int64)
+
+            assert torch.all(
+                selected_tokens >= 0
+            ), f"Req {req_idx}: selected tokens contain negative positions"
+            assert torch.all(selected_tokens < seq_len), (
+                f"Req {req_idx}: selected tokens {selected_tokens.tolist()} "
+                f"out of range for seq_len={seq_len}"
+            )
+
+            if seq_len <= self.device_buffer_size:
+                device_indices = self.req_to_device_buffer[req_idx, selected_tokens]
+            else:
+                device_indices = torch.empty(
+                    top_n, dtype=torch.int64, device=self.device
+                )
+
+                is_latest_token = selected_tokens == (seq_len - 1)
+                needs_host_load = ~is_latest_token
+
+                device_indices[is_latest_token] = self.req_to_device_buffer[
+                    req_idx, self.device_buffer_size
+                ]
+
+                num_to_load = int(needs_host_load.sum().item())
+                if num_to_load > 0:
+                    tokens_to_load = selected_tokens[needs_host_load]
+                    host_locs = self.req_to_host_pool[req_idx, tokens_to_load]
+
+                    invalid_mask = host_locs < 0
+                    if torch.any(invalid_mask):
+                        bad_positions = tokens_to_load[invalid_mask].tolist()
+                        raise AssertionError(
+                            f"Req {req_idx} (seq_len={seq_len}, layer={layer_id}): "
+                            f"missing host backup at token positions {bad_positions}"
+                        )
+
+                    buffer_locs = self.req_to_device_buffer[req_idx, :num_to_load]
+                    device_indices[needs_host_load] = buffer_locs
+
+                    self.mem_pool_host.load_to_device_per_layer(
+                        self.mem_pool_device,
+                        host_locs,
+                        buffer_locs,
+                        layer_id,
+                        io_backend="kernel",
+                    )
+
+            top_k_indices[i, :top_n] = device_indices.to(torch.int32)
+
+        return top_k_indices
+
+    def abort_staging_request(self, req: Req) -> None:
+        """Remove a request from the staging queue and free its host + device resources.
+
+        Must be called when aborting a request that has been admitted into staging
+        but has not yet completed (i.e. req.hisparse_staging is True).
+        """
+        # Remove from staging queue
+        self.ack_staging_queue = [
+            act for act in self.ack_staging_queue if act.req is not req
+        ]
+        # Wait for any in-flight staging DMA to complete before freeing
+        self.write_staging_stream.synchronize()
+
+        prefill_len = req.extend_range.end
+        allocated_locs = self.req_to_token_pool.req_to_token[
+            req.req_pool_idx, :prefill_len
+        ]
+        self.token_to_kv_pool_allocator.free_hisparse(allocated_locs)
+
+        # Free host memory that was allocated during admit_request_into_staging
+        host_indices = self.mem_pool_host.allocated_host_indices(
+            self.req_to_host_pool,
+            req.req_pool_idx,
+            self.req_to_host_pool_allocated_len[req.req_pool_idx],
+        )
+        if host_indices.numel() > 0:
+            self.mem_pool_host.free(host_indices)
+        self.req_to_host_pool[req.req_pool_idx, :] = -1
+        self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
+        self._skip_first_backup[req.req_pool_idx] = False
+        req.hisparse_staging = False
+
+    def retract_req(self, req: Req) -> None:
+        if req.hisparse_staging:
+            self.abort_staging_request(req)
+        else:
+            self.request_finished(req)
+
+    def request_finished(self, req: Req):
+        # release resources only after the execution of a potential overlapped batch
+        if self.decode_producer_stream is not None:
+            device_module.current_stream().wait_stream(self.decode_producer_stream)
+        self.wait_for_pending_backup()
+
+        # Use kv_allocated_len (not seqlen): under speculative decoding the
+        # allocator can over-allocate beyond the committed seqlen, and those
+        # extra slots may carry stale mapping entries pointing at buffer slots
+        # we just freed via free_hisparse_indices(all_hi). If left set, the
+        # subsequent release_kv_cache -> allocator.free -> free_hisparse path
+        # re-frees them (double-free into the page allocator's free list).
+        allocated_len = req.kv_allocated_len
+
+        # release memory -- only free actually-allocated buffer indices
+        current_cap = int(self.req_device_buffer_size[req.req_pool_idx])
+        if current_cap > 0:
+            side_buf_hi = self.req_to_device_buffer[req.req_pool_idx, :current_cap]
+            all_hi = torch.unique(side_buf_hi[side_buf_hi > 0])
+            if all_hi.numel() > 0:
+                self.token_to_kv_pool_allocator.free_hisparse_indices(all_hi)
+
+        allocated_locs = self.req_to_token_pool.req_to_token[
+            req.req_pool_idx, :allocated_len
+        ]
+        compressed_locs = self.mem_pool_device.translate_loc_from_full_to_compressed(
+            allocated_locs
+        )
+        self.mem_pool_device.full_to_hisparse_device_index_mapping[compressed_locs] = 0
+
+        host_indices = self.mem_pool_host.allocated_host_indices(
+            self.req_to_host_pool,
+            req.req_pool_idx,
+            self.req_to_host_pool_allocated_len[req.req_pool_idx],
+        )
+        if host_indices.numel() > 0:
+            self.mem_pool_host.free(host_indices)
+
+        # clear req info
+        self.req_device_buffer_tokens[:, req.req_pool_idx, :] = -1
+        self.req_device_buffer_token_locs[:, req.req_pool_idx, :] = -1
+        self.req_to_device_buffer[req.req_pool_idx, :] = 0
+        self.req_device_buffer_size[req.req_pool_idx] = 0
+        self.req_to_host_pool[req.req_pool_idx, :] = -1
+        self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
+        self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init)
+        self._skip_first_backup[req.req_pool_idx] = False
+
+    def swap_in_selected_pages(
+        self,
+        req_pool_indices: torch.Tensor,
+        compressed_seq_lens: torch.Tensor,
+        top_k_result: torch.Tensor,
+        layer_id: int,
+    ) -> torch.Tensor:
+        """Swap selected top-k tokens into device memory and return their indices."""
+        num_reqs = req_pool_indices.size(0)
+
+        top_k_indices = self.top_k_device_locs_buffer[:num_reqs]
+        top_k_indices.fill_(-1)
+
+        swap_in_fn = (
+            load_cache_to_device_buffer_dsv4_mla
+            if self.is_dsv4_hisparse
+            else load_cache_to_device_buffer_mla
+        )
+        swap_in_fn(
+            top_k_tokens=top_k_result,
+            device_buffer_tokens=self.req_device_buffer_tokens[layer_id],
+            host_cache_locs=self.req_to_host_pool,
+            device_buffer_locs=self.req_device_buffer_token_locs[layer_id],
+            host_cache=self.mem_pool_host.kv_buffer[layer_id],
+            device_buffer=self.mem_pool_device.kv_buffer[layer_id],
+            top_k_device_locs=top_k_indices,
+            req_pool_indices=req_pool_indices,
+            seq_lens=compressed_seq_lens,
+            lru_slots=self.lru_slots[layer_id],
+            item_size_bytes=self.item_size_bytes,
+            num_top_k=self.top_k,
+            hot_buffer_size=self.device_buffer_size,
+            page_size=1,
+            block_size=self.swap_in_block_size,
+            num_real_reqs=self.num_real_reqs,
+        )
+        return top_k_indices

python/sglang/srt/managers/load_snapshot.py

@@ -0,0 +1,654 @@
+"""Load snapshot: publish scheduler load metrics for DP balancing and /v1/loads.
+
+Architecture
+------------
+
+Each scheduler periodically publishes a ``LoadSnapshot`` containing its
+current load metrics (running reqs, tokens, throughput, ...).  Two
+transport backends are supported:
+
+**SHM mode** (single-node, default)::
+
+    Scheduler  ──ShmLoadSnapshotWriter──▶  /dev/shm mmap file
+                                               ▲
+    TokenizerManager  ──ShmLoadSnapshotReader───┘  (for /v1/loads)
+    DataParallelController  ──ShmLoadSnapshotReader─┘  (for dispatch)
+
+**ZMQ mode** (multi-node DP attention, or ``SGLANG_LOAD_SNAPSHOT_USE_ZMQ=1``)::
+
+    Scheduler (any node)  ──ZmqLoadSnapshotWriter (PUSH)──▶  network
+                                                               │
+    ZmqShmLoadSnapshotReader (PULL, node 0)  ◀─────────────────┘
+        │  drains zmq, writes to SHM
+        ▼
+    /dev/shm mmap file (node 0)
+        ▲
+    TokenizerManager / DataParallelController  ──ShmLoadSnapshotReader──┘
+
+Shared memory does not work across nodes, so multi-node DP attention
+requires the ZMQ transport.  The ``ZmqShmLoadSnapshotReader`` on node 0
+receives snapshots from all schedulers via zmq PUSH/PULL and writes them
+into the local SHM file.  All readers (TokenizerManager,
+DataParallelController) on
+node 0 then read from SHM.
+
+``zmq_reader_owner()`` decides which process on node 0 binds the zmq
+PULL socket (only one can bind); the other reads plain SHM.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import logging
+import mmap
+import os
+import struct
+from contextlib import contextmanager
+from typing import Optional
+
+import msgspec
+import msgspec.msgpack
+import msgspec.structs
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.network import is_zmq_endpoint_ipv6
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def should_use_zmq(server_args) -> bool:
+    """Whether to use zmq PUSH/PULL instead of shared memory for load snapshots.
+
+    Shared memory (mmap) only works within a single node.  When schedulers
+    run on multiple nodes (multi-node DP attention), they cannot write to
+    the SHM file on node 0, so we fall back to zmq transport.  The env var
+    ``SGLANG_LOAD_SNAPSHOT_USE_ZMQ`` forces zmq mode for testing.
+    """
+    return (
+        server_args.enable_dp_attention and server_args.nnodes > 1
+    ) or envs.SGLANG_LOAD_SNAPSHOT_USE_ZMQ.get()
+
+
+_LOAD_AWARE_METHODS = frozenset({"total_requests", "total_tokens"})
+
+
+def _tokenizer_load_snapshot_owner_caller(server_args) -> str:
+    """The caller that plays the tokenizer-side zmq owner role.
+
+    In multi-tokenizer mode (``tokenizer_worker_num > 1``) there are N
+    independent ``TokenizerWorker`` processes that would all try to bind the
+    same zmq PULL endpoint.  Instead, the single ``MultiTokenizerRouter``
+    process owns the socket (polls zmq -> SHM) and every worker reads SHM.
+    """
+    if server_args.tokenizer_worker_num > 1:
+        return "MultiTokenizerRouter"
+    return "TokenizerManager"
+
+
+def zmq_reader_owner(server_args, caller: str) -> bool:
+    """Decide which process owns the zmq PULL socket.
+
+    Exactly one of ``"DataParallelController"``, ``"TokenizerManager"``, or
+    ``"MultiTokenizerRouter"`` must return True when zmq mode is active.  The
+    owner polls zmq -> SHM; the others read SHM.
+
+    Rules:
+      - Non-zero node_rank: no TokenizerManager, DataParallelController only
+        launches schedulers and waits -> nobody owns it.
+      - dp_size == 1: no DataParallelController exists -> tokenizer-side owner
+        owns it.
+      - dp_size > 1, load-aware method: DataParallelController polls on every
+        dispatch via refresh_load_budget() -> DataParallelController owns it.
+      - dp_size > 1, round-robin / other: DataParallelController never reads
+        load data -> tokenizer-side owner owns it (polls on /v1/loads calls).
+
+    The tokenizer-side owner is the ``"MultiTokenizerRouter"`` caller in
+    multi-tokenizer mode, otherwise the ``"TokenizerManager"`` caller.
+    """
+    if not should_use_zmq(server_args):
+        return False
+    if server_args.node_rank != 0:
+        return False
+    tokenizer_owner = _tokenizer_load_snapshot_owner_caller(server_args)
+    if server_args.dp_size == 1:
+        return caller == tokenizer_owner
+    if server_args.load_balance_method.lower() in _LOAD_AWARE_METHODS:
+        return caller == "DataParallelController"
+    return caller == tokenizer_owner
+
+
+# ---------------------------------------------------------------------------
+# LoadSnapshot data class
+# ---------------------------------------------------------------------------
+
+
+class MemoryMetrics(msgspec.Struct, array_like=True):
+    """Memory breakdown metrics."""
+
+    weight_gb: float
+    kv_cache_gb: float
+    graph_gb: float
+    token_capacity: int
+
+
+class SpeculativeMetrics(msgspec.Struct, array_like=True):
+    """Speculative decoding metrics."""
+
+    accept_length: float
+    accept_rate: float
+
+
+class LoRAMetrics(msgspec.Struct, array_like=True):
+    """LoRA adapter pool metrics."""
+
+    slots_used: int
+    slots_total: int
+    utilization: float
+
+
+class DisaggregationMetrics(msgspec.Struct, array_like=True):
+    """PD disaggregation metrics."""
+
+    mode: str  # "prefill", "decode", or "null"
+    prefill_bootstrap_queue_reqs: int = 0
+    prefill_inflight_queue_reqs: int = 0
+    decode_prealloc_queue_reqs: int = 0
+    decode_transfer_queue_reqs: int = 0
+    decode_retracted_queue_reqs: int = 0
+    kv_transfer_speed_gb_s: float = 0.0
+    kv_transfer_latency_ms: float = 0.0
+
+
+class QueueMetrics(msgspec.Struct, array_like=True):
+    """Detailed queue info breakdown."""
+
+    waiting: int
+    grammar: int
+    paused: int
+    retracted: int
+
+
+_CORE_KEYS = (
+    "timestamp",
+    "dp_rank",
+    "num_running_reqs",
+    "num_waiting_reqs",
+    "num_waiting_uncached_tokens",
+    "num_used_tokens",
+    "num_total_tokens",
+    "max_total_num_tokens",
+    "max_running_requests",
+    "token_usage",
+    "gen_throughput",
+    "cache_hit_rate",
+    "utilization",
+)
+
+
+class LoadSnapshot(msgspec.Struct, omit_defaults=True):
+    """Per-DP-rank load metrics: the SHM/zmq wire format and the /v1/loads source."""
+
+    timestamp: float = 0.0
+    dp_rank: int = 0
+    num_running_reqs: int = 0
+    num_waiting_reqs: int = 0
+    num_waiting_uncached_tokens: int = 0
+    num_used_tokens: int = 0
+    num_total_tokens: int = 0
+    max_total_num_tokens: int = 0
+    max_running_requests: int = 0
+    token_usage: float = 0.0
+    gen_throughput: float = 0.0
+    cache_hit_rate: float = 0.0
+    utilization: float = 0.0
+
+    memory: Optional[MemoryMetrics] = None
+    speculative: Optional[SpeculativeMetrics] = None
+    lora: Optional[LoRAMetrics] = None
+    disaggregation: Optional[DisaggregationMetrics] = None
+    queues: Optional[QueueMetrics] = None
+
+    VALID_SECTIONS = frozenset(
+        {"core", "memory", "spec", "lora", "disagg", "queues", "all"}
+    )
+
+    def to_dict(self, include: Optional[set[str]] = None) -> dict:
+        load = {key: getattr(self, key) for key in _CORE_KEYS}
+
+        if include is None or "all" in include:
+            include_all = True
+        else:
+            if not (include <= self.VALID_SECTIONS):
+                raise ValueError(
+                    f"Invalid include sections: {include - self.VALID_SECTIONS}. "
+                    f"Valid options: {sorted(self.VALID_SECTIONS)}"
+                )
+            if include == {"core"}:
+                return load
+            include_all = False
+
+        for field, include_name, section in (
+            ("memory", "memory", self.memory),
+            ("speculative", "spec", self.speculative),
+            ("lora", "lora", self.lora),
+            ("disaggregation", "disagg", self.disaggregation),
+            ("queues", "queues", self.queues),
+        ):
+            if section is None or (not include_all and include_name not in include):
+                continue
+            load[field] = msgspec.structs.asdict(section)
+
+        return load
+
+
+def _enc_hook(obj):
+    """Coerce numpy scalars to native Python; msgpack has no numpy types."""
+    to_item = getattr(obj, "item", None)
+    if to_item is not None:
+        return to_item()
+    raise NotImplementedError(f"cannot encode {type(obj).__name__} in load snapshot")
+
+
+snapshot_encoder = msgspec.msgpack.Encoder(enc_hook=_enc_hook)
+snapshot_decoder = msgspec.msgpack.Decoder(LoadSnapshot)
+
+
+# ---------------------------------------------------------------------------
+# SHM file layout utilities
+# ---------------------------------------------------------------------------
+
+MAGIC = b"SLNS"
+VERSION = 2
+HEADER_STRUCT = struct.Struct("<4sHHI")
+SLOT_LEN_STRUCT = struct.Struct("<I")
+SLOT_SIZE = 16 * 1024
+
+
+@contextmanager
+def file_lock(fd: int, lock_type: int):
+    fcntl.flock(fd, lock_type)
+    try:
+        yield
+    finally:
+        fcntl.flock(fd, fcntl.LOCK_UN)
+
+
+def shm_path_for(ipc_name: str) -> str:
+    name = os.path.basename(ipc_name.rstrip("/")) or "default"
+    safe_name = "".join(c if c.isalnum() or c in "._-" else "_" for c in name)
+    digest = hashlib.blake2s(ipc_name.encode(), digest_size=4).hexdigest()
+    return f"/dev/shm/sglang_loads_{safe_name}_{digest}.shm"
+
+
+def file_size(dp_size: int, slot_size: int = SLOT_SIZE) -> int:
+    return HEADER_STRUCT.size + dp_size * slot_size
+
+
+def slot_offset(dp_rank: int, slot_size: int = SLOT_SIZE) -> int:
+    return HEADER_STRUCT.size + dp_rank * slot_size
+
+
+# ---------------------------------------------------------------------------
+# Writers
+# ---------------------------------------------------------------------------
+
+
+class ShmLoadSnapshotWriter:
+    def __init__(
+        self, path: str, dp_size: int, dp_rank: int, publish_interval: int = 1
+    ):
+        if dp_rank < 0 or dp_rank >= dp_size:
+            raise ValueError(f"invalid dp_rank={dp_rank} for dp_size={dp_size}")
+        self.publish_interval = max(1, publish_interval)
+        self.publish_counter = 0
+
+        self.path = path
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+        self.slot_size = SLOT_SIZE
+        self.fd = -1
+        size = file_size(dp_size, self.slot_size)
+
+        self.fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
+        try:
+            with file_lock(self.fd, fcntl.LOCK_EX):
+                os.ftruncate(self.fd, size)
+                self.mmap = mmap.mmap(self.fd, size, access=mmap.ACCESS_WRITE)
+                HEADER_STRUCT.pack_into(
+                    self.mmap, 0, MAGIC, VERSION, dp_size, self.slot_size
+                )
+                self._write_payload(LoadSnapshot(dp_rank=dp_rank))
+        except Exception:
+            if self.fd >= 0:
+                os.close(self.fd)
+            raise
+
+    def write(self, snapshot: LoadSnapshot) -> None:
+        if snapshot.dp_rank != self.dp_rank:
+            raise ValueError(
+                f"snapshot dp_rank={snapshot.dp_rank} does not match writer dp_rank={self.dp_rank}"
+            )
+
+        with file_lock(self.fd, fcntl.LOCK_EX):
+            self._write_payload(snapshot)
+
+    def _write_payload(self, snapshot: LoadSnapshot) -> None:
+        payload = snapshot_encoder.encode(snapshot)
+        max_payload_size = self.slot_size - SLOT_LEN_STRUCT.size
+        if len(payload) > max_payload_size:
+            raise ValueError(
+                f"load snapshot payload size {len(payload)} exceeds slot payload "
+                f"capacity {max_payload_size}"
+            )
+
+        offset = slot_offset(self.dp_rank, self.slot_size)
+        payload_start = offset + SLOT_LEN_STRUCT.size
+        payload_end = payload_start + len(payload)
+        slot_end = offset + self.slot_size
+
+        SLOT_LEN_STRUCT.pack_into(self.mmap, offset, 0)
+        self.mmap[payload_start:payload_end] = payload
+        self.mmap[payload_end:slot_end] = b"\0" * (slot_end - payload_end)
+        SLOT_LEN_STRUCT.pack_into(self.mmap, offset, len(payload))
+
+    def close(self) -> None:
+        self.mmap.close()
+        os.close(self.fd)
+
+
+class ZmqLoadSnapshotWriter:
+    """Sends load snapshots via zmq PUSH to a ZmqShmLoadSnapshotReader.
+
+    CONFLATE is set so only the latest message is kept in the send
+    buffer when the reader is slower than the writer.
+    """
+
+    def __init__(
+        self, endpoint: str, dp_size: int, dp_rank: int, publish_interval: int = 1
+    ):
+        import zmq as _zmq
+
+        if dp_rank < 0 or dp_rank >= dp_size:
+            raise ValueError(f"invalid dp_rank={dp_rank} for dp_size={dp_size}")
+        self.publish_interval = max(1, publish_interval)
+        self.publish_counter = 0
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+
+        self._zmq = _zmq
+        self._ctx = _zmq.Context.instance()
+        self._socket = self._ctx.socket(_zmq.PUSH)
+        if is_zmq_endpoint_ipv6(endpoint):
+            self._socket.setsockopt(_zmq.IPV6, 1)
+        self._socket.setsockopt(_zmq.LINGER, 0)
+        self._socket.setsockopt(_zmq.CONFLATE, 1)
+        self._socket.connect(endpoint)
+
+    def write(self, snapshot: LoadSnapshot) -> None:
+        if snapshot.dp_rank != self.dp_rank:
+            raise ValueError(
+                f"snapshot dp_rank={snapshot.dp_rank} does not match "
+                f"writer dp_rank={self.dp_rank}"
+            )
+        try:
+            self._socket.send(snapshot_encoder.encode(snapshot), self._zmq.NOBLOCK)
+        except self._zmq.Again:
+            pass
+
+    def close(self) -> None:
+        self._socket.close()
+
+
+# ---------------------------------------------------------------------------
+# Readers
+# ---------------------------------------------------------------------------
+
+
+class ShmLoadSnapshotReader:
+    def __init__(self, path: str, dp_size: int):
+        self.path = path
+        self.dp_size = dp_size
+        self.mmap: Optional[mmap.mmap] = None
+        self.fd: Optional[int] = None
+        self.slot_size = SLOT_SIZE
+        self._header_warning_logged = False
+        self._attach()
+
+    def _attach(self) -> bool:
+        if self.mmap is not None:
+            return True
+
+        try:
+            fd = os.open(self.path, os.O_RDONLY)
+        except FileNotFoundError:
+            return False
+
+        size = os.fstat(fd).st_size
+        if size < HEADER_STRUCT.size:
+            os.close(fd)
+            return False
+
+        try:
+            with file_lock(fd, fcntl.LOCK_SH):
+                mapped = mmap.mmap(fd, size, access=mmap.ACCESS_READ)
+                magic, version, dp_size, slot_size = HEADER_STRUCT.unpack_from(
+                    mapped, 0
+                )
+        except (OSError, ValueError):
+            os.close(fd)
+            return False
+
+        if (
+            magic != MAGIC
+            or version != VERSION
+            or dp_size != self.dp_size
+            or slot_size < SLOT_LEN_STRUCT.size
+            or size < file_size(self.dp_size, slot_size)
+        ):
+            mapped.close()
+            os.close(fd)
+            if not self._header_warning_logged:
+                logger.warning("load shm header mismatch at %s", self.path)
+                self._header_warning_logged = True
+            return False
+
+        self.mmap = mapped
+        self.fd = fd
+        self.slot_size = slot_size
+        return True
+
+    def read(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        if dp_rank < 0 or dp_rank >= self.dp_size:
+            return None
+        if not self._attach():
+            return None
+
+        assert self.fd is not None
+        with file_lock(self.fd, fcntl.LOCK_SH):
+            return self._read_slot(dp_rank)
+
+    def _read_slot(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        assert self.mmap is not None
+        offset = slot_offset(dp_rank, self.slot_size)
+        (payload_len,) = SLOT_LEN_STRUCT.unpack_from(self.mmap, offset)
+        max_payload_size = self.slot_size - SLOT_LEN_STRUCT.size
+        if payload_len == 0 or payload_len > max_payload_size:
+            return None
+
+        payload_start = offset + SLOT_LEN_STRUCT.size
+        payload_end = payload_start + payload_len
+        try:
+            return snapshot_decoder.decode(self.mmap[payload_start:payload_end])
+        except Exception as e:
+            logger.debug("load snapshot decode failed for rank %s: %s", dp_rank, e)
+            return None
+
+    def read_all(self) -> list[LoadSnapshot]:
+        if not self._attach():
+            return []
+
+        assert self.fd is not None
+        with file_lock(self.fd, fcntl.LOCK_SH):
+            loads = []
+            for r in range(self.dp_size):
+                load = self._read_slot(r)
+                if load is not None:
+                    loads.append(load)
+            return loads
+
+    def close(self) -> None:
+        if self.mmap is not None:
+            self.mmap.close()
+            self.mmap = None
+        if self.fd is not None:
+            os.close(self.fd)
+            self.fd = None
+
+
+class ZmqShmLoadSnapshotReader:
+    """Receives snapshots via zmq PULL from writers, writes to SHM, reads from SHM.
+
+    Transparently wraps a ShmLoadSnapshotReader.  Every read() / read_all()
+    first drains the PULL socket into SHM so callers always see fresh data.
+    """
+
+    def __init__(self, endpoint: str, shm_path: str, dp_size: int):
+        import zmq as _zmq
+
+        self._zmq = _zmq
+        self._ctx = _zmq.Context.instance()
+        self._socket = self._ctx.socket(_zmq.PULL)
+        if is_zmq_endpoint_ipv6(endpoint):
+            self._socket.setsockopt(_zmq.IPV6, 1)
+        self._socket.setsockopt(_zmq.LINGER, 0)
+        self._socket.setsockopt(_zmq.CONFLATE, 1)
+        self._socket.bind(endpoint)
+
+        self._endpoint = endpoint
+        self._shm_path = shm_path
+        self.dp_size = dp_size
+        self._shm_reader = ShmLoadSnapshotReader(shm_path, dp_size)
+        self._shm_writers: dict[int, ShmLoadSnapshotWriter] = {}
+
+    def _poll(self) -> None:
+        """Drain zmq messages and write latest per dp_rank to SHM."""
+        latest: dict[int, LoadSnapshot] = {}
+        while True:
+            try:
+                data = self._socket.recv(self._zmq.NOBLOCK)
+            except self._zmq.Again:
+                break
+            try:
+                snapshot = snapshot_decoder.decode(data)
+                if 0 <= snapshot.dp_rank < self.dp_size:
+                    latest[snapshot.dp_rank] = snapshot
+            except Exception as e:
+                logger.warning("load snapshot zmq decode failed: %s", e)
+
+        for dp_rank, snapshot in latest.items():
+            if dp_rank not in self._shm_writers:
+                self._shm_writers[dp_rank] = ShmLoadSnapshotWriter(
+                    self._shm_path, self.dp_size, dp_rank
+                )
+            try:
+                self._shm_writers[dp_rank].write(snapshot)
+            except Exception as e:
+                logger.warning(
+                    "load snapshot shm write failed for rank %d: %s", dp_rank, e
+                )
+
+    def fileno(self) -> int:
+        """Edge-triggered fd that becomes readable when zmq messages arrive.
+
+        Lets an owner process register the reader with an event loop and drain
+        it via ``poll()`` instead of polling on a timer.
+        """
+        return self._socket.getsockopt(self._zmq.FD)
+
+    def poll(self) -> None:
+        """Drain the zmq PULL socket into SHM.
+
+        Public entry point so an owner process (e.g. MultiTokenizerRouter) can
+        keep SHM fresh without touching internals.
+        """
+        self._poll()
+
+    def read(self, dp_rank: int) -> Optional[LoadSnapshot]:
+        self._poll()
+        return self._shm_reader.read(dp_rank)
+
+    def read_all(self) -> list[LoadSnapshot]:
+        self._poll()
+        return self._shm_reader.read_all()
+
+    def close(self) -> None:
+        for w in self._shm_writers.values():
+            w.close()
+        self._shm_writers.clear()
+        self._shm_reader.close()
+        self._socket.close()
+        if self._endpoint.startswith("ipc://"):
+            try:
+                os.unlink(self._endpoint[len("ipc://") :])
+            except OSError:
+                pass
+
+
+# ---------------------------------------------------------------------------
+# Factory functions
+# ---------------------------------------------------------------------------
+
+
+def _zmq_addr_for(port_args) -> str:
+    """Return the zmq PUSH/PULL address from PortArgs.
+
+    For dp_attention (TCP mode), uses the ``load_collector_ipc_name`` field
+    stored in PortArgs.  For single-node IPC (env-var override), derives
+    a deterministic IPC path from ``instance_id``.
+    """
+    ipc_name = getattr(port_args, "load_collector_ipc_name", "")
+    if ipc_name:
+        return ipc_name
+    safe = "".join(
+        c if c.isalnum() or c in "._-" else "_" for c in port_args.instance_id
+    )
+    digest = hashlib.blake2s(port_args.instance_id.encode(), digest_size=4).hexdigest()
+    return f"ipc:///tmp/sglang_load_collector_{safe}_{digest}.sock"
+
+
+def create_load_snapshot_writer(
+    server_args,
+    port_args,
+    dp_size: int,
+    dp_rank: int,
+    publish_interval: int = 1,
+):
+    """Return a SHM or ZMQ writer based on server configuration."""
+    if should_use_zmq(server_args):
+        return ZmqLoadSnapshotWriter(
+            _zmq_addr_for(port_args), dp_size, dp_rank, publish_interval
+        )
+    return ShmLoadSnapshotWriter(
+        shm_path_for(port_args.instance_id), dp_size, dp_rank, publish_interval
+    )
+
+
+def create_load_snapshot_reader(server_args, port_args, caller: str):
+    """Create a load snapshot reader.
+
+    Args:
+        caller: ``"DataParallelController"``, ``"TokenizerManager"``, or
+            ``"MultiTokenizerRouter"`` -- determines who binds the zmq PULL
+            socket when zmq mode is active.
+    """
+    dp_size = server_args.dp_size
+    if zmq_reader_owner(server_args, caller):
+        return ZmqShmLoadSnapshotReader(
+            _zmq_addr_for(port_args), shm_path_for(port_args.instance_id), dp_size
+        )
+    return ShmLoadSnapshotReader(shm_path_for(port_args.instance_id), dp_size)

python/sglang/srt/managers/min_free_slots_delayer.py

@@ -0,0 +1,41 @@
+from typing import Optional
+
+
+def resolve_min_free_slots(
+    user_value: Optional[int],
+    max_running_requests: int,
+    is_dflash_family: bool = False,
+) -> Optional[int]:
+    """Resolve the min-free-slots threshold (None = disabled).
+
+    A user value (>1) is capped to the DFlash formula so the trigger never
+    delays more aggressively than the legacy heuristic. When unset, DFlash
+    workloads fall back to the formula (preserving the always-on behavior);
+    other workloads stay disabled. Also disabled when max_running_requests < 8.
+    """
+    max_running_requests = max(0, int(max_running_requests))
+    formula = min(4, max(2, (max_running_requests + 5) // 6))
+    if user_value is None:
+        user_value = formula if is_dflash_family else None
+
+    if user_value is None or user_value <= 1:
+        return None
+    if max_running_requests < 8:
+        return None
+    return min(user_value, formula)
+
+
+class MinFreeSlotsDelayer:
+    """Delay fresh prefill admissions until at least ``min_free_slots`` running-
+    request slots free up, batching them into one admission instead of one at a
+    time. Useful when each admission is expensive (e.g. DFlash's draft prefill).
+
+    Per-rank local: running-batch slots are private to each DP rank, so a rank
+    with free slots does not wait for a congested peer.
+    """
+
+    def __init__(self, min_free_slots: int):
+        self._min_free_slots = min_free_slots
+
+    def should_delay(self, *, running_bs: int, num_allocatable_reqs: int) -> bool:
+        return running_bs > 0 and num_allocatable_reqs < self._min_free_slots

python/sglang/srt/managers/multi_tokenizer_mixin.py

@@ -0,0 +1,782 @@
+from __future__ import annotations
+
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""
+Mixin classes and utils for multi-http-worker mode
+This file uses multiple processes to handle requests and tokenization, reducing the overhead of python and http server.
+"""
+
+import asyncio
+import logging
+import multiprocessing as multiprocessing
+import os
+import pickle
+import signal
+import sys
+import threading
+import zlib
+from multiprocessing import shared_memory
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
+
+import psutil
+import setproctitle
+import zmq
+import zmq.asyncio
+
+from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend
+from sglang.srt.managers.disagg_service import start_disagg_service
+from sglang.srt.managers.io_struct import (
+    BaseBatchReq,
+    BaseReq,
+    BatchEmbeddingOutput,
+    BatchStrOutput,
+    BatchTokenIDOutput,
+    ContinueGenerationReqInput,
+    FreezeGCReq,
+    PauseContinueBroadcastReq,
+    PauseGenerationReqInput,
+    TokenizerWorkerRegistrationReq,
+    async_sock_recv,
+    async_sock_send,
+    sock_recv,
+    sock_send,
+    unwrap_from_pickle,
+    wrap_as_pickle,
+)
+from sglang.srt.managers.load_snapshot import (
+    create_load_snapshot_reader,
+    zmq_reader_owner,
+)
+from sglang.srt.managers.tokenizer_manager import TokenizerManager
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import (
+    configure_logger,
+    kill_itself_when_parent_died,
+    kill_process_tree,
+)
+from sglang.srt.utils.network import get_zmq_socket
+from sglang.utils import get_exception_traceback
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.detokenizer_manager import DetokenizerManager
+
+logger = logging.getLogger(__name__)
+
+
+class SocketMapping:
+    def __init__(self):
+        self._zmq_context = zmq.Context()
+        self._mapping: Dict[str, zmq.Socket] = {}
+
+    def clear_all_sockets(self):
+        for socket in self._mapping.values():
+            socket.close()
+        self._mapping.clear()
+
+    def _register_ipc_mapping(self, ipc_name: str, is_tokenizer: bool):
+        type_str = "tokenizer" if is_tokenizer else "detokenizer"
+        if ipc_name in self._mapping:
+            logger.warning(f"{type_str} already registered {ipc_name=}, skipping...")
+            return
+        logger.info(f"Registering {type_str} {ipc_name=} in SocketMapping...")
+        socket = get_zmq_socket(self._zmq_context, zmq.PUSH, ipc_name, False)
+        self._mapping[ipc_name] = socket
+
+    def send_output(self, ipc_name: str, output: Any, is_tokenizer: bool = False):
+        if ipc_name is None:
+            # Some unhandled cases
+            logger.warning(f"IPC name is None, output type={type(output)}, skipping...")
+            return
+
+        if ipc_name not in self._mapping:
+            self._register_ipc_mapping(ipc_name, is_tokenizer=is_tokenizer)
+        sock_send(self._mapping[ipc_name], output)
+
+
+def _extract_field_by_index(
+    output: Any, field_name: str, index: int, check_length: bool = True
+) -> Any:
+    """Extract a field value from output by index, handling None and length checks.
+
+    Args:
+        output: The output object containing the field
+        field_name: The name of the field to extract
+        index: The index to access in the field list
+        check_length: If True, check both field existence and length. If False, only check field existence.
+
+    Returns:
+        A list containing the field value at index, or None if not available.
+    """
+    field = getattr(output, field_name, None)
+    if field is None:
+        return None
+
+    should_wrap_result = field_name in ("customized_info", "time_stats")
+    if should_wrap_result:
+        field = unwrap_from_pickle(field)
+        if field is None:
+            return None
+
+    if isinstance(field, dict):
+        new_field = {}
+        for k, v in field.items():
+            if len(v) > index:
+                new_field[k] = [v[index]] if should_wrap_result else v[index]
+            else:
+                new_field[k] = [None] if should_wrap_result else None
+        if should_wrap_result:
+            return wrap_as_pickle(new_field) if new_field else None
+        return new_field
+
+    if check_length:
+        if len(field) <= index:
+            return None
+
+    new_field = [field[index]]
+    return wrap_as_pickle(new_field) if should_wrap_result else new_field
+
+
+def _handle_output_by_index(output, i):
+    """NOTE: A maintainable method is better here."""
+    if isinstance(output, BatchTokenIDOutput):
+        new_output = BatchTokenIDOutput(
+            rids=[output.rids[i]],
+            spec_verify_ct=_extract_field_by_index(output, "spec_verify_ct", i),
+            spec_num_correct_drafts=_extract_field_by_index(
+                output, "spec_num_correct_drafts", i
+            ),
+            spec_correct_drafts_histogram=_extract_field_by_index(
+                output, "spec_correct_drafts_histogram", i
+            ),
+            spec_num_block_accept_tokens=_extract_field_by_index(
+                output, "spec_num_block_accept_tokens", i
+            ),
+            spec_num_cap_tokens=_extract_field_by_index(
+                output, "spec_num_cap_tokens", i
+            ),
+            spec_cap_lens_histogram=_extract_field_by_index(
+                output, "spec_cap_lens_histogram", i
+            ),
+            time_stats=_extract_field_by_index(output, "time_stats", i),
+            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
+            decoded_texts=_extract_field_by_index(output, "decoded_texts", i),
+            decode_ids=_extract_field_by_index(output, "decode_ids", i),
+            read_offsets=_extract_field_by_index(output, "read_offsets", i),
+            output_ids=_extract_field_by_index(output, "output_ids", i),
+            skip_special_tokens=_extract_field_by_index(
+                output, "skip_special_tokens", i
+            ),
+            spaces_between_special_tokens=_extract_field_by_index(
+                output, "spaces_between_special_tokens", i
+            ),
+            no_stop_trim=_extract_field_by_index(output, "no_stop_trim", i),
+            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
+            completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
+            reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
+            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
+            cached_tokens_details=_extract_field_by_index(
+                output, "cached_tokens_details", i
+            ),
+            image_tokens=_extract_field_by_index(output, "image_tokens", i),
+            audio_tokens=_extract_field_by_index(output, "audio_tokens", i),
+            video_tokens=_extract_field_by_index(output, "video_tokens", i),
+            input_token_logprobs_val=_extract_field_by_index(
+                output, "input_token_logprobs_val", i, check_length=False
+            ),
+            input_token_logprobs_idx=_extract_field_by_index(
+                output, "input_token_logprobs_idx", i, check_length=False
+            ),
+            output_token_logprobs_val=_extract_field_by_index(
+                output, "output_token_logprobs_val", i, check_length=False
+            ),
+            output_token_logprobs_idx=_extract_field_by_index(
+                output, "output_token_logprobs_idx", i, check_length=False
+            ),
+            input_top_logprobs_val=_extract_field_by_index(
+                output, "input_top_logprobs_val", i, check_length=False
+            ),
+            input_top_logprobs_idx=_extract_field_by_index(
+                output, "input_top_logprobs_idx", i, check_length=False
+            ),
+            output_top_logprobs_val=_extract_field_by_index(
+                output, "output_top_logprobs_val", i, check_length=False
+            ),
+            output_top_logprobs_idx=_extract_field_by_index(
+                output, "output_top_logprobs_idx", i, check_length=False
+            ),
+            input_token_ids_logprobs_val=_extract_field_by_index(
+                output, "input_token_ids_logprobs_val", i, check_length=False
+            ),
+            input_token_ids_logprobs_idx=_extract_field_by_index(
+                output, "input_token_ids_logprobs_idx", i, check_length=False
+            ),
+            output_token_ids_logprobs_val=_extract_field_by_index(
+                output, "output_token_ids_logprobs_val", i, check_length=False
+            ),
+            output_token_ids_logprobs_idx=_extract_field_by_index(
+                output, "output_token_ids_logprobs_idx", i, check_length=False
+            ),
+            output_token_entropy_val=_extract_field_by_index(
+                output, "output_token_entropy_val", i, check_length=False
+            ),
+            output_hidden_states=_extract_field_by_index(
+                output, "output_hidden_states", i, check_length=False
+            ),
+            routed_experts=_extract_field_by_index(
+                output, "routed_experts", i, check_length=False
+            ),
+            indexer_topk=_extract_field_by_index(
+                output, "indexer_topk", i, check_length=False
+            ),
+            retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
+            token_steps=_extract_field_by_index(
+                output, "token_steps", i, check_length=False
+            ),
+            customized_info=_extract_field_by_index(
+                output, "customized_info", i, check_length=False
+            ),
+            dp_ranks=_extract_field_by_index(output, "dp_ranks", i, check_length=False),
+        )
+    elif isinstance(output, BatchEmbeddingOutput):
+        new_output = BatchEmbeddingOutput(
+            rids=[output.rids[i]],
+            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
+            embeddings=_extract_field_by_index(output, "embeddings", i),
+            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
+            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
+        )
+    elif isinstance(output, BatchStrOutput):
+        new_output = BatchStrOutput(
+            rids=[output.rids[i]],
+            spec_verify_ct=_extract_field_by_index(output, "spec_verify_ct", i),
+            spec_num_correct_drafts=_extract_field_by_index(
+                output, "spec_num_correct_drafts", i
+            ),
+            spec_correct_drafts_histogram=_extract_field_by_index(
+                output, "spec_correct_drafts_histogram", i
+            ),
+            spec_num_block_accept_tokens=_extract_field_by_index(
+                output, "spec_num_block_accept_tokens", i
+            ),
+            spec_num_cap_tokens=_extract_field_by_index(
+                output, "spec_num_cap_tokens", i
+            ),
+            spec_cap_lens_histogram=_extract_field_by_index(
+                output, "spec_cap_lens_histogram", i
+            ),
+            time_stats=_extract_field_by_index(output, "time_stats", i),
+            finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
+            output_strs=_extract_field_by_index(output, "output_strs", i),
+            output_ids=_extract_field_by_index(output, "output_ids", i),
+            prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i),
+            completion_tokens=_extract_field_by_index(output, "completion_tokens", i),
+            reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i),
+            cached_tokens=_extract_field_by_index(output, "cached_tokens", i),
+            cached_tokens_details=_extract_field_by_index(
+                output, "cached_tokens_details", i
+            ),
+            image_tokens=_extract_field_by_index(output, "image_tokens", i),
+            audio_tokens=_extract_field_by_index(output, "audio_tokens", i),
+            video_tokens=_extract_field_by_index(output, "video_tokens", i),
+            input_token_logprobs_val=_extract_field_by_index(
+                output, "input_token_logprobs_val", i, check_length=False
+            ),
+            input_token_logprobs_idx=_extract_field_by_index(
+                output, "input_token_logprobs_idx", i, check_length=False
+            ),
+            output_token_logprobs_val=_extract_field_by_index(
+                output, "output_token_logprobs_val", i, check_length=False
+            ),
+            output_token_logprobs_idx=_extract_field_by_index(
+                output, "output_token_logprobs_idx", i, check_length=False
+            ),
+            input_top_logprobs_val=_extract_field_by_index(
+                output, "input_top_logprobs_val", i, check_length=False
+            ),
+            input_top_logprobs_idx=_extract_field_by_index(
+                output, "input_top_logprobs_idx", i, check_length=False
+            ),
+            output_top_logprobs_val=_extract_field_by_index(
+                output, "output_top_logprobs_val", i, check_length=False
+            ),
+            output_top_logprobs_idx=_extract_field_by_index(
+                output, "output_top_logprobs_idx", i, check_length=False
+            ),
+            input_token_ids_logprobs_val=_extract_field_by_index(
+                output, "input_token_ids_logprobs_val", i, check_length=False
+            ),
+            input_token_ids_logprobs_idx=_extract_field_by_index(
+                output, "input_token_ids_logprobs_idx", i, check_length=False
+            ),
+            output_token_ids_logprobs_val=_extract_field_by_index(
+                output, "output_token_ids_logprobs_val", i, check_length=False
+            ),
+            output_token_ids_logprobs_idx=_extract_field_by_index(
+                output, "output_token_ids_logprobs_idx", i, check_length=False
+            ),
+            output_token_entropy_val=_extract_field_by_index(
+                output, "output_token_entropy_val", i, check_length=False
+            ),
+            output_hidden_states=_extract_field_by_index(
+                output, "output_hidden_states", i, check_length=False
+            ),
+            routed_experts=_extract_field_by_index(
+                output, "routed_experts", i, check_length=False
+            ),
+            indexer_topk=_extract_field_by_index(
+                output, "indexer_topk", i, check_length=False
+            ),
+            customized_info=_extract_field_by_index(
+                output, "customized_info", i, check_length=False
+            ),
+            dp_ranks=_extract_field_by_index(output, "dp_ranks", i, check_length=False),
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
+            retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
+            token_steps=_extract_field_by_index(
+                output, "token_steps", i, check_length=False
+            ),
+        )
+    else:
+        new_output = output
+    return new_output
+
+
+class MultiHttpWorkerDetokenizerMixin:
+    """Mixin class for DetokenizerManager"""
+
+    def maybe_clear_socket_mapping(self: DetokenizerManager):
+        if hasattr(self, "socket_mapping"):
+            self.socket_mapping.clear_all_sockets()
+
+    def multi_http_worker_event_loop(self: DetokenizerManager):
+        """The event loop that handles requests, for multi multi-http-worker mode"""
+        self.socket_mapping = SocketMapping()
+        while True:
+            recv_obj = sock_recv(self.recv_from_scheduler)
+            output = self._request_dispatcher(recv_obj)
+            if output is None:
+                continue
+
+            # Fan out the output back to the originating tokenizer worker(s).
+            # In multi-detokenizer mode the upstream MultiDetokenizerRouter may
+            # forward either batched or single requests, so handle both shapes.
+            if isinstance(recv_obj, BaseBatchReq):
+                for i, ipc_name in enumerate(recv_obj.http_worker_ipcs):
+                    new_output = _handle_output_by_index(output, i)
+                    self.socket_mapping.send_output(
+                        ipc_name, new_output, is_tokenizer=True
+                    )
+            elif isinstance(recv_obj, BaseReq):
+                self.socket_mapping.send_output(
+                    recv_obj.http_worker_ipc, output, is_tokenizer=True
+                )
+            else:
+                raise ValueError(
+                    f"multi_http_worker_event_loop got unexpected req type {type(recv_obj)}"
+                )
+
+
+class MultiTokenizerRouter:
+    """A router between tokenizer managers and the scheduler/detokenizer manager.
+
+    Forward: tokenizer managers → router → scheduler.
+    Backward: detokenizer manager → router → tokenizer managers.
+    Also broadcasts pause/continue to all tokenizer managers for consistent is_pause state.
+    """
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        self.server_args = server_args
+        context = zmq.asyncio.Context(3)
+        self.recv_from_detokenizer = get_zmq_socket(
+            context, zmq.PULL, port_args.tokenizer_ipc_name, True
+        )
+        self.send_to_scheduler = get_zmq_socket(
+            context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
+        )
+        self.receive_from_worker = get_zmq_socket(
+            context, zmq.PULL, port_args.tokenizer_worker_ipc_name, True
+        )
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(target=self._run_loop, daemon=True)
+        self._thread.start()
+        self._task = asyncio.run_coroutine_threadsafe(
+            self.router_worker_obj(), self._loop
+        )
+        self._handle_task = asyncio.run_coroutine_threadsafe(
+            print_exception_wrapper(self.handle_loop), self._loop
+        )
+
+        # In multi-tokenizer mode the N TokenizerWorker processes cannot each
+        # bind the zmq PULL socket used for load snapshots, so the single
+        # MultiTokenizerRouter process owns it (zmq -> SHM) and the workers
+        # read SHM only. Drain it event-driven via the socket's fd instead of
+        # polling on a timer.
+        self.load_snapshot_reader = None
+        if zmq_reader_owner(server_args, "MultiTokenizerRouter"):
+            self.load_snapshot_reader = create_load_snapshot_reader(
+                server_args, port_args, caller="MultiTokenizerRouter"
+            )
+            self._loop.call_soon_threadsafe(self._register_load_snapshot_reader)
+
+        self.disaggregation_bootstrap_server = start_disagg_service(self.server_args)
+
+        # Worker IPC names for pause/continue broadcasting
+        self.all_worker_ipcs: set[str] = set()
+        # Shared socket mapping (both coroutines run on self._loop, so safe)
+        self.socket_mapping = SocketMapping()
+
+    def _run_loop(self):
+        self._loop.run_forever()
+
+    def _register_load_snapshot_reader(self):
+        """Drain zmq load snapshots into SHM whenever the PULL socket is readable.
+
+        zmq exposes an edge-triggered fd; ``poll()`` drains it until empty, which
+        also re-arms the fd, so TokenizerWorkers reading SHM stay up to date
+        without any timer.
+        """
+        assert self.load_snapshot_reader is not None
+        self._loop.add_reader(
+            self.load_snapshot_reader.fileno(), self.load_snapshot_reader.poll
+        )
+        # Drain anything already queued before the fd was registered.
+        self.load_snapshot_reader.poll()
+
+    async def router_worker_obj(self):
+        """Forward path: workers → scheduler, with pause/continue broadcast."""
+        while True:
+            recv_obj = await async_sock_recv(self.receive_from_worker)
+
+            if isinstance(recv_obj, TokenizerWorkerRegistrationReq):
+                if recv_obj.worker_ipc_name not in self.all_worker_ipcs:
+                    self.all_worker_ipcs.add(recv_obj.worker_ipc_name)
+                    logger.info(
+                        f"Router registered worker IPC: {recv_obj.worker_ipc_name} "
+                        f"(total: {len(self.all_worker_ipcs)})"
+                    )
+                continue
+
+            if isinstance(
+                recv_obj, (PauseGenerationReqInput, ContinueGenerationReqInput)
+            ):
+                # Broadcast to ALL workers so every worker's is_pause is set
+                is_pause = isinstance(recv_obj, PauseGenerationReqInput)
+                broadcast = PauseContinueBroadcastReq(is_pause=is_pause)
+                for ipc_name in self.all_worker_ipcs:
+                    self.socket_mapping.send_output(ipc_name, broadcast)
+                # Forward to scheduler rank 0 (it broadcasts to all TP/PP/DP
+                # ranks internally). Skip for abort mode which drains via polling.
+                if not (
+                    isinstance(recv_obj, PauseGenerationReqInput)
+                    and recv_obj.mode == "abort"
+                ):
+                    await async_sock_send(self.send_to_scheduler, recv_obj)
+                continue
+
+            await async_sock_send(self.send_to_scheduler, recv_obj)
+
+    async def handle_loop(self):
+        """Backward path: detokenizer → route results to correct worker."""
+        while True:
+            recv_obj = await async_sock_recv(self.recv_from_detokenizer)
+            await self._distribute_result_to_workers(recv_obj)
+
+    async def _distribute_result_to_workers(self, recv_obj):
+        if isinstance(recv_obj, BaseReq):
+            ipc_names = [recv_obj.http_worker_ipc]
+        elif isinstance(recv_obj, BaseBatchReq):
+            ipc_names = recv_obj.http_worker_ipcs
+        else:
+            raise ValueError(f"Unknown recv_obj type: {type(recv_obj)}")
+
+        for i, ipc_name in enumerate(ipc_names):
+            new_recv_obj = _handle_output_by_index(recv_obj, i)
+            self.socket_mapping.send_output(ipc_name, new_recv_obj)
+
+
+class MultiDetokenizerRouter:
+    """Route scheduler outputs to one of N DetokenizerManager workers.
+
+    Each request is pinned to a worker by hashing its ``http_worker_ipc`` with
+    ``zlib.crc32`` (deterministic across runs), so all outputs of the same rid
+    always land on the same detokenizer and ``decode_status`` stays consistent.
+    """
+
+    def __init__(self, ipc_name_list: List[str], port_args: PortArgs):
+        self.ipc_name_list = ipc_name_list
+        self.num_workers = len(ipc_name_list)
+        self.socket_mapping = SocketMapping()
+        context = zmq.Context(2)
+        self.recv_from_scheduler = get_zmq_socket(
+            context, zmq.PULL, port_args.detokenizer_ipc_name, True
+        )
+
+    def _pick(self, key: str) -> str:
+        return self.ipc_name_list[zlib.crc32(key.encode()) % self.num_workers]
+
+    def _send(self, ipc_name: str, obj: Any) -> None:
+        self.socket_mapping.send_output(ipc_name, obj, is_tokenizer=False)
+
+    def event_loop(self):
+        while True:
+            recv_obj = sock_recv(self.recv_from_scheduler)
+
+            # FreezeGCReq must freeze every detokenizer process.
+            if isinstance(recv_obj, FreezeGCReq):
+                for ipc in self.ipc_name_list:
+                    self._send(ipc, recv_obj)
+                continue
+
+            # Single request: route by its own http_worker_ipc.
+            if isinstance(recv_obj, BaseReq):
+                assert (
+                    recv_obj.http_worker_ipc is not None
+                ), f"Single req {recv_obj.rid=} missing http_worker_ipc"
+                self._send(self._pick(recv_obj.http_worker_ipc), recv_obj)
+                continue
+
+            # Batch request.
+            if isinstance(recv_obj, BaseBatchReq):
+                # Idle/no-op batch (rids=[]): broadcast to all detokenizers
+                if not recv_obj.rids:
+                    for ipc in self.ipc_name_list:
+                        self._send(ipc, recv_obj)
+                    continue
+
+                ipcs = recv_obj.http_worker_ipcs
+                assert (
+                    ipcs is not None
+                    and len(ipcs) == len(recv_obj.rids)
+                    and all(x is not None for x in ipcs)
+                ), f"Batch req {recv_obj.rids=} has invalid http_worker_ipcs"
+
+                # Split per-item and route each by its own ipc.
+                for i, ipc_key in enumerate(ipcs):
+                    one = _handle_output_by_index(recv_obj, i)
+                    if one is recv_obj:
+                        raise TypeError(f"Cannot split {type(recv_obj)}")
+                    one.http_worker_ipcs = [ipc_key]
+                    self._send(self._pick(ipc_key), one)
+                continue
+
+            raise ValueError(
+                f"MultiDetokenizerRouter got unsupported type {type(recv_obj)}"
+            )
+
+
+def run_multi_detokenizer_router_process(
+    ipc_name_list: List[str],
+    server_args: ServerArgs,
+    port_args: PortArgs,
+):
+    kill_itself_when_parent_died()
+    setproctitle.setproctitle("sglang::detokenizer_router")
+    configure_logger(server_args)
+    parent_process = psutil.Process().parent()
+
+    router = None
+    try:
+        router = MultiDetokenizerRouter(ipc_name_list, port_args)
+        router.event_loop()
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"MultiDetokenizerRouter hit an exception: {traceback}")
+        if router is not None:
+            router.socket_mapping.clear_all_sockets()
+        parent_process.send_signal(signal.SIGQUIT)
+
+
+class TokenizerWorker(TokenizerManager):
+    """Tokenizer Worker in multi-http-worker mode"""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        setproctitle.setproctitle(f"sglang::tokenizer_worker:{os.getpid()}")
+        import torch
+
+        torch.set_num_threads(1)
+        # prevent init prefill bootstrapserver again
+        disaggregation_mode = server_args.disaggregation_mode
+        server_args.override(
+            "tokenizer_worker.suppress_bootstrap", disaggregation_mode="null"
+        )
+        super().__init__(server_args, port_args)
+
+        self.worker_id = os.getpid()
+        self.tokenizer_ipc_name = port_args.tokenizer_ipc_name
+
+        # For PD disaggregtion
+        self.server_args.override(
+            "tokenizer_worker.restore_disaggregation_mode",
+            disaggregation_mode=disaggregation_mode,
+        )
+        self.disaggregation_mode = DisaggregationMode(
+            self.server_args.disaggregation_mode
+        )
+        self.disaggregation_transfer_backend = TransferBackend(
+            self.server_args.disaggregation_transfer_backend
+        )
+
+        # Register this worker with the router for pause/continue broadcasting
+        reg = TokenizerWorkerRegistrationReq(worker_ipc_name=self.tokenizer_ipc_name)
+        self._dispatch_to_scheduler(reg)
+
+        # Future for awaiting pause/continue broadcast confirmation
+        self._pause_continue_future: Optional[asyncio.Future] = None
+
+        # Register PauseContinueBroadcastReq in the result dispatcher so
+        # handle_loop routes it to _handle_pause_continue_broadcast
+        from sglang.utils import TypeBasedDispatcher
+
+        self._result_dispatcher += TypeBasedDispatcher(
+            [(PauseContinueBroadcastReq, self._handle_pause_continue_broadcast)]
+        )
+
+    async def pause_generation(self, obj: PauseGenerationReqInput):
+        loop = asyncio.get_event_loop()
+        self._pause_continue_future = loop.create_future()
+        # Send to router which will broadcast to all workers
+        # (router also handles forwarding to scheduler for non-abort modes)
+        self._dispatch_to_scheduler(obj)
+        await self._pause_continue_future
+
+        if obj.mode == "abort":
+            # Abort polling: only the originator checks its own lock state
+            while True:
+                self.abort_request(abort_all=True)
+                is_locked = await self.model_update_lock.is_locked()
+                if not is_locked:
+                    break
+                await asyncio.sleep(1.0)
+
+    async def continue_generation(self, obj: ContinueGenerationReqInput):
+        loop = asyncio.get_event_loop()
+        self._pause_continue_future = loop.create_future()
+        self._dispatch_to_scheduler(obj)
+        await self._pause_continue_future
+
+    def _handle_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq):
+        """Called from handle_loop when a broadcast arrives from the router."""
+        loop = asyncio.get_event_loop()
+        loop.create_task(self._apply_pause_continue_broadcast(obj))
+
+    async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq):
+        """Apply pause/continue state under the condition lock."""
+        async with self.is_pause_cond:
+            if obj.is_pause:
+                self.is_pause = True
+            else:
+                self.is_pause = False
+                self.is_pause_cond.notify_all()
+
+        # Resolve the pending future if this worker initiated the pause/continue
+        if self._pause_continue_future and not self._pause_continue_future.done():
+            self._pause_continue_future.set_result(True)
+            self._pause_continue_future = None
+
+
+def get_tokenizer_worker_class(server_args: ServerArgs) -> Type[TokenizerWorker]:
+    worker_class = server_args.get_tokenizer_worker_class()
+    if not isinstance(worker_class, type) or not issubclass(
+        worker_class, TokenizerWorker
+    ):
+        raise TypeError(
+            "ServerArgs.get_tokenizer_worker_class() must return a TokenizerWorker "
+            f"subclass, got {worker_class!r}"
+        )
+
+    return worker_class
+
+
+async def print_exception_wrapper(func):
+    """
+    Sometimes an asyncio function does not print exception.
+    We do another wrapper to handle the exception.
+    """
+    try:
+        await func()
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"MultiTokenizerRouter hit an exception: {traceback}")
+        if hasattr(func, "__self__") and isinstance(
+            func.__self__, MultiTokenizerRouter
+        ):
+            func.__self__.dump_requests_before_crash()
+        kill_process_tree(os.getpid(), include_parent=True)
+        sys.exit(1)
+
+
+def get_main_process_id() -> int:
+    """Get the main process ID."""
+    return multiprocessing.current_process()._parent_pid
+
+
+def write_to_shared_memory(obj, name: str) -> shared_memory.SharedMemory:
+    """Write data to shared memory"""
+    serialized = pickle.dumps(obj)
+    size = len(serialized)
+    try:
+        # Try to open existing shared memory
+        shm = shared_memory.SharedMemory(name=name)
+        # If size is insufficient, close and recreate
+        if shm.size < size:
+            shm.close()
+            shm.unlink()
+            shm = shared_memory.SharedMemory(create=True, size=size, name=name)
+    except FileNotFoundError:
+        # If not present, create new shared memory
+        shm = shared_memory.SharedMemory(create=True, size=size, name=name)
+
+    shm.buf[:size] = serialized
+    return shm
+
+
+def read_from_shared_memory(name: str) -> Any:
+    """Read data from shared memory"""
+    try:
+        shm = shared_memory.SharedMemory(name=name)
+        data = pickle.loads(bytes(shm.buf))
+        shm.close()
+        return data
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Shared memory {name} not found")
+
+
+def write_data_for_multi_tokenizer(
+    port_args: PortArgs, server_args: ServerArgs, scheduler_info: Dict
+):
+    """Write args information to share memory for multi-tokenizer"""
+    # get main process ID
+    main_pid = get_main_process_id()
+    current_pid = os.getpid()
+    logger.info(f"main process ID: {main_pid}, current process ID: {current_pid}")
+    args = (port_args, server_args, scheduler_info)
+    args_shm = write_to_shared_memory(args, f"multi_tokenizer_args_{current_pid}")
+    args_shm.close()
+
+    return args_shm

python/sglang/srt/managers/multimodal_processor.py

@@ -0,0 +1,83 @@
+# TODO: also move pad_input_ids into this module
+import importlib
+import inspect
+import logging
+import pkgutil
+
+from sglang.srt.configs.model_config import ModelImpl
+from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
+from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+PROCESSOR_MAPPING = {}
+
+
+def import_processors(package_name: str, overwrite: bool = False):
+    package = importlib.import_module(package_name)
+    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
+        if not ispkg:
+            try:
+                module = importlib.import_module(name)
+            except Exception as e:
+                logger.warning(f"Ignore import error when loading {name}: {e}")
+                continue
+            all_members = inspect.getmembers(module, inspect.isclass)
+            classes = [
+                member
+                for name, member in all_members
+                if member.__module__ == module.__name__
+            ]
+            for cls in (
+                cls for cls in classes if issubclass(cls, BaseMultimodalProcessor)
+            ):
+                assert hasattr(cls, "models")
+                for arch in getattr(cls, "models"):
+                    if overwrite:
+                        for model_cls, processor_cls in PROCESSOR_MAPPING.items():
+                            if model_cls.__name__ == arch.__name__:
+                                del PROCESSOR_MAPPING[model_cls]
+                                break
+                    PROCESSOR_MAPPING[arch] = cls
+
+
+def get_mm_processor(
+    hf_config,
+    server_args: ServerArgs,
+    processor,
+    transport_mode,
+    model_config=None,
+    **kwargs,
+) -> BaseMultimodalProcessor:
+    model_impl = str(getattr(server_args, "model_impl", "auto")).lower()
+    uses_transformers_backend = model_impl == "transformers"
+    if model_impl == "auto" and model_config is not None:
+        from sglang.srt.model_loader.utils import get_resolved_model_impl
+
+        uses_transformers_backend = (
+            get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
+        )
+
+    for model_cls, processor_cls in PROCESSOR_MAPPING.items():
+        if model_cls.__name__ not in hf_config.architectures:
+            continue
+        if not uses_transformers_backend or getattr(
+            processor_cls, "supports_transformers_backend", False
+        ):
+            return processor_cls(
+                hf_config, server_args, processor, transport_mode, **kwargs
+            )
+
+    if uses_transformers_backend:
+        from sglang.srt.multimodal.processors.transformers_auto import (
+            TransformersAutoMultimodalProcessor,
+        )
+
+        return TransformersAutoMultimodalProcessor(
+            hf_config, server_args, processor, transport_mode, **kwargs
+        )
+
+    raise ValueError(
+        f"No processor registered for architecture: {hf_config.architectures}.\n"
+        f"Registered architectures: {[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}"
+    )

python/sglang/srt/managers/overlap_utils.py

@@ -0,0 +1,527 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional, Sequence
+
+import msgspec
+import torch
+
+from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras
+from sglang.srt.environ import envs
+from sglang.srt.utils import is_cuda, is_hip, is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.eagle_info import EagleDraftInput
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+
+def decide_needs_cpu_seq_lens(
+    server_args: ServerArgs,
+    attn_backends: Sequence[AttentionBackend],
+) -> bool:
+    """Whether FutureMap must publish seq_lens_cpu / sum.
+
+    OR over per-backend needs_cpu_seq_lens; force True under TBO (it reads the
+    CPU mirror outside the backend layer to split the batch) or ngram (its
+    USE_FULL_MASK verify path reads the host mirror regardless of backend).
+    """
+    # Local import: keep overlap_utils' module-level deps leaf-only so it stays
+    # importable everywhere; spec_info pulls in the spec/schedule_batch graph.
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+    if server_args.enable_two_batch_overlap:
+        # FIXME: support TBO without seq lens cpu value
+        return True
+    algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+    if algo.is_ngram():
+        # ngram's USE_FULL_MASK verify path reads seq_lens_cpu per req to size
+        # the tree mask, regardless of the attn backend (e.g. Triton opts out).
+        return True
+    # Skip unset slots (e.g. draft_extend_attn_backend on some spec configs);
+    # missing flag -> True so undeclared backends stay on the legacy path.
+    return any(
+        getattr(b, "needs_cpu_seq_lens", True) for b in attn_backends if b is not None
+    )
+
+
+def decide_needs_confidence_relay(server_args: ServerArgs) -> bool:
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+    algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
+    if not algo.is_dspark():
+        return False
+    return read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
+
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
+
+# Token-buf consume tracking: init to -1, assert non-negative on gather,
+# write -1 back. Catches "gather without intermediate stash" bugs. CI enables
+# via the existing SGLANG_IS_IN_CI; off in production.
+_DEBUG_ASSERT = envs.SGLANG_IS_IN_CI.get()
+
+
+@torch.compile(dynamic=True, disable=_is_npu)
+def _assert_nonneg_and_invalidate(
+    values: torch.Tensor, buf: torch.Tensor, indices: torch.Tensor
+) -> None:
+    """Fused: assert all `values >= 0` and scatter -1 into `buf[indices]`.
+    Compiled so the reduction + assert + scatter run as one kernel launch."""
+    torch._assert_async((values >= 0).all())
+    buf[indices] = -1
+
+
+def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None:
+    """Materialize input_ids at forward entry. Two sources:
+
+    - Prefill: H2D copy from pinned CPU staging (prefill_input_ids_cpu).
+    - Decode/spec_v2: gather from FutureMap (last iter's sampled token).
+    """
+    if batch.prefill_input_ids_cpu is not None:
+        prefill_gpu = batch.prefill_input_ids_cpu.to(batch.device, non_blocking=True)
+        if batch.mix_running_indices is not None:
+            decode_gpu = future_map.output_tokens_buf[batch.mix_running_indices]
+            if _DEBUG_ASSERT:
+                _assert_nonneg_and_invalidate(
+                    decode_gpu,
+                    future_map.output_tokens_buf,
+                    batch.mix_running_indices,
+                )
+            batch.input_ids = torch.cat([prefill_gpu, decode_gpu])
+        else:
+            batch.input_ids = prefill_gpu
+        batch.prefill_input_ids_cpu = None
+        batch.mix_running_indices = None
+    elif batch.input_ids is None and future_map.spec_algo.is_none():
+        batch.input_ids = future_map.output_tokens_buf[batch.req_pool_indices]
+        if _DEBUG_ASSERT:
+            _assert_nonneg_and_invalidate(
+                batch.input_ids, future_map.output_tokens_buf, batch.req_pool_indices
+            )
+
+    # Only the overlap path relays spec extras through the future_map; the
+    # synchronous (non-overlap) V2 path installs next_draft_input directly.
+    if batch.enable_overlap and not batch.spec_algorithm.is_none():
+        future_map._resolve_spec_extras(batch)
+
+
+CONFIDENCE_RELAY_RING_LAG: int = 2
+CONFIDENCE_RELAY_RING_DEPTH: int = CONFIDENCE_RELAY_RING_LAG + 1
+
+
+class ResolvedConfidence(msgspec.Struct):
+
+    confidence: torch.Tensor
+    generation: torch.Tensor
+
+
+@dataclass
+class RelayPayload:
+    """Per-iteration stash payload for the FutureMap bufs. Non-spec fills only
+    `bonus_tokens`; which spec extras get relayed is decided by
+    `FutureMap.spec_algo`, not by this payload's shape."""
+
+    bonus_tokens: torch.Tensor
+    topk_p: Optional[torch.Tensor] = None
+    topk_index: Optional[torch.Tensor] = None
+    hidden_states: Optional[torch.Tensor] = None
+    draft_probs: Optional[torch.Tensor] = None
+    dsa_topk_indices: Optional[torch.Tensor] = None
+
+    @classmethod
+    def from_draft_input(cls, draft_input: EagleDraftInput) -> RelayPayload:
+        return cls(
+            bonus_tokens=draft_input.bonus_tokens,
+            topk_p=draft_input.topk_p,
+            topk_index=draft_input.topk_index,
+            hidden_states=draft_input.hidden_states,
+            draft_probs=getattr(draft_input, "draft_probs", None),
+            dsa_topk_indices=getattr(draft_input, "dsa_topk_indices", None),
+        )
+
+
+class ConfidenceRelay(msgspec.Struct):
+
+    device: torch.device
+    req_pool_size: int
+    pool: Any
+    confidence_buf: Optional[torch.Tensor] = None
+    conf_ring: Optional[torch.Tensor] = None
+    gen_ring: Optional[torch.Tensor] = None
+    copy_done: Optional[list] = None
+    ring_pos: int = 0
+    initialized: bool = False
+
+    def _lazy_init(self, confidence: torch.Tensor) -> None:
+        self.initialized = True
+        gamma = confidence.shape[-1]
+        self.confidence_buf = torch.empty(
+            (self.req_pool_size, gamma), dtype=torch.float32, device=self.device
+        )
+        if _is_cuda:
+            depth = CONFIDENCE_RELAY_RING_DEPTH
+            self.conf_ring = torch.empty(
+                (depth, self.req_pool_size, gamma),
+                dtype=torch.float32,
+                pin_memory=True,
+            )
+            self.gen_ring = torch.zeros((depth, self.req_pool_size), dtype=torch.int64)
+            self.copy_done = [
+                torch.get_device_module(self.device).Event() for _ in range(depth)
+            ]
+
+    def scatter(self, indices: torch.Tensor, confidence: torch.Tensor) -> None:
+        if not self.initialized:
+            self._lazy_init(confidence)
+        self.confidence_buf[indices] = confidence.to(self.confidence_buf.dtype)
+
+    def issue_ring_copy(self, *, stream, publish_ready) -> None:
+        if not self.initialized or stream is None or publish_ready is None:
+            return
+        slot = self.ring_pos % CONFIDENCE_RELAY_RING_DEPTH
+        stream.wait_event(publish_ready)
+        with torch.get_device_module(self.device).stream(stream):
+            self.conf_ring[slot].copy_(self.confidence_buf, non_blocking=True)
+            self.copy_done[slot].record()
+        self.gen_ring[slot].copy_(self.pool.req_generation)
+        self.ring_pos += 1
+
+    def resolve(
+        self, batch: ScheduleBatch, *, stream, publish_ready
+    ) -> Optional[ResolvedConfidence]:
+        if not self.initialized:
+            return None
+        draft_input = batch.spec_info
+        if draft_input is None:
+            return None
+        fi = draft_input.future_indices
+        if fi is None or fi.shape[0] == 0:
+            return None
+
+        if stream is None or publish_ready is None:
+            idx = batch.req_pool_indices
+            idx_cpu = batch.req_pool_indices_cpu
+            return ResolvedConfidence(
+                confidence=self.confidence_buf[idx].cpu(),
+                generation=self.pool.req_generation[idx_cpu].clone(),
+            )
+
+        if self.ring_pos < CONFIDENCE_RELAY_RING_LAG:
+            return None
+        slot = (self.ring_pos - CONFIDENCE_RELAY_RING_LAG) % CONFIDENCE_RELAY_RING_DEPTH
+        if not self.copy_done[slot].query():
+            return None
+
+        idx_cpu = batch.req_pool_indices_cpu
+        return ResolvedConfidence(
+            confidence=self.conf_ring[slot][idx_cpu],
+            generation=self.gen_ring[slot][idx_cpu],
+        )
+
+
+class FutureMap:
+    """Always-on pool-indexed relay for cross-iter values. Forward writes via
+    publish/stash; next iter reads via resolve_forward_inputs / resolve_seq_lens_cpu.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        spec_algo: SpeculativeAlgorithm,
+        req_to_token_pool: ReqToTokenPool,
+        needs_cpu_seq_lens: bool = True,
+        needs_confidence_relay: bool = False,
+    ):
+        # Bufs indexed by req_pool_idx; slot 0 mirrors KV padding row so
+        # CUDA-graph padded batches (req_pool_idx == 0) are harmless.
+        self.device = device
+        self.spec_algo = spec_algo
+        # Computed by decide_needs_cpu_seq_lens(); see that helper for the
+        # full decision (per-backend flag + TBO / piecewise CG overrides).
+        self.needs_cpu_seq_lens = needs_cpu_seq_lens
+        self.needs_confidence_relay = needs_confidence_relay
+        self.req_pool_size = req_to_token_pool.req_to_token.shape[0]
+
+        if _DEBUG_ASSERT:
+            # Poisoned init: every row must be written before its first gather.
+            self.output_tokens_buf = torch.full(
+                (self.req_pool_size,), -1, dtype=torch.int64, device=self.device
+            )
+            self.new_seq_lens_buf = torch.full(
+                (self.req_pool_size,), -1, dtype=torch.int64, device=self.device
+            )
+        else:
+            self.output_tokens_buf = torch.empty(
+                (self.req_pool_size,), dtype=torch.int64, device=self.device
+            )
+            self.new_seq_lens_buf = torch.empty(
+                (self.req_pool_size,), dtype=torch.int64, device=self.device
+            )
+        # Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare
+        # D2H pulls (gated only on publish, off the schedule stream). CUDA-only:
+        # recovers occupancy lost to the WAR barrier (also CUDA-only); other
+        # platforms have no barrier and use the plain .cpu() bootstrap path.
+        if _is_cuda:
+            self.new_seq_lens_cpu_pinned = torch.empty(
+                (self.req_pool_size,), dtype=torch.int64, pin_memory=True
+            )
+            self.fwd_prepare_d2h_stream = torch.get_device_module(self.device).Stream()
+        else:
+            self.new_seq_lens_cpu_pinned = None
+            self.fwd_prepare_d2h_stream = None
+        # Lazy-inited on the first non-empty stash (peeks tensor shapes); non-spec's is a no-op.
+        self._forward_buf_initialized = False
+
+        self.publish_ready = None  # lazy device.Event(); only spec_v2 needs it
+        # Debug consume-once state: armed by a recording publish, consumed by
+        # resolve; arm/consume strictly alternate across all batch interleavings.
+        self._publish_fresh = False
+
+        self.confidence_relay = ConfidenceRelay(
+            device=self.device,
+            req_pool_size=self.req_pool_size,
+            pool=req_to_token_pool,
+        )
+
+    def _lazy_init_forward_buf(self, payload: RelayPayload):
+        # Local import (see decide_needs_cpu_seq_lens): keep module-level deps leaf.
+        from sglang.srt.speculative.spec_utils import spec_need_hidden_states
+
+        self._forward_buf_initialized = True
+
+        # Spec extras are gated by spec_algo, not by the payload's shape, so a
+        # non-spec stash allocates no extra bufs (only output_tokens_buf).
+        self.need_topk = self.spec_algo.is_some() and self.spec_algo.need_topk()
+        self.need_hidden_states = (
+            self.spec_algo.is_some()
+            and spec_need_hidden_states()
+            and payload.hidden_states is not None
+        )
+
+        if self.need_topk:
+            topk_p0 = payload.topk_p[0]
+            topk_index0 = payload.topk_index[0]
+            self.topk_p_buf = torch.empty(
+                (self.req_pool_size, *topk_p0.shape),
+                dtype=topk_p0.dtype,
+                device=self.device,
+            )
+            self.topk_index_buf = torch.empty(
+                (self.req_pool_size, *topk_index0.shape),
+                dtype=topk_index0.dtype,
+                device=self.device,
+            )
+        if self.need_hidden_states:
+            hidden_states0 = payload.hidden_states[0]
+            self.hidden_states_buf = torch.empty(
+                (self.req_pool_size, *hidden_states0.shape),
+                dtype=hidden_states0.dtype,
+                device=self.device,
+            )
+
+        self.draft_probs_buf = None
+        if payload.draft_probs is not None:
+            draft_probs0 = payload.draft_probs[0]
+            self.draft_probs_buf = torch.empty(
+                (self.req_pool_size, *draft_probs0.shape),
+                dtype=draft_probs0.dtype,
+                device=self.device,
+            )
+
+        self.dsa_topk_indices_buf = None
+        if payload.dsa_topk_indices is not None:
+            seed0 = payload.dsa_topk_indices[0]
+            self.dsa_topk_indices_buf = torch.empty(
+                (self.req_pool_size, *seed0.shape),
+                dtype=payload.dsa_topk_indices.dtype,
+                device=self.device,
+            )
+
+    def resolve_confidence_cpu(
+        self, batch: ScheduleBatch
+    ) -> Optional[ResolvedConfidence]:
+        if not self.needs_confidence_relay:
+            return None
+        return self.confidence_relay.resolve(
+            batch,
+            stream=self.fwd_prepare_d2h_stream,
+            publish_ready=self.publish_ready,
+        )
+
+    def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
+        if self.spec_algo.is_ngram():
+            # FIXME: remove once precomputed draft is supported.
+            return
+        draft_input: EagleDraftInput = batch.spec_info
+        if draft_input is None:
+            # FIXME(lsyin): only prefill; not compatible with mixed mode
+            return
+        indices = draft_input.future_indices
+        if indices.shape[0] == 0:
+            return
+        # FIXME: indices = batch.req_pool_indices, pinned 2 iters via
+        # record_batch_in_overlap; record_stream here is redundant.
+        indices.record_stream(torch.get_device_module(self.device).current_stream())
+        if self.need_topk:
+            hidden_states_buf = (
+                self.hidden_states_buf if self.need_hidden_states else None
+            )
+            (
+                draft_input.topk_p,
+                draft_input.topk_index,
+                bonus_tokens,
+                hidden_states,
+            ) = gather_spec_extras(
+                indices,
+                self.topk_p_buf,
+                self.topk_index_buf,
+                self.output_tokens_buf,
+                hidden_states_buf,
+            )
+            draft_input.bonus_tokens = bonus_tokens
+            if hidden_states is not None:
+                draft_input.hidden_states = hidden_states
+            if self.draft_probs_buf is not None and draft_input.draft_probs is not None:
+                draft_input.draft_probs = self.draft_probs_buf[indices]
+        else:
+            draft_input.bonus_tokens = self.output_tokens_buf[indices]
+        if self.need_hidden_states and not self.need_topk:
+            draft_input.hidden_states = self.hidden_states_buf[indices]
+        if self.dsa_topk_indices_buf is not None:
+            draft_input.dsa_topk_indices = self.dsa_topk_indices_buf[indices]
+        if _DEBUG_ASSERT:
+            _assert_nonneg_and_invalidate(
+                draft_input.bonus_tokens, self.output_tokens_buf, indices
+            )
+
+    def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None:
+        # Lazy pull from new_seq_lens_buf for spec_v2 (accept_lens not known to
+        # schedule). The CPU mirror is gated by needs_cpu_seq_lens; backends that
+        # opt out take the GPU-only path below. A private D2H stream overlaps the copy.
+        draft_input = batch.spec_info
+        if draft_input is None:
+            return
+
+        fi = draft_input.future_indices
+        if fi is None:
+            return
+        if self.publish_ready is not None:
+            if _DEBUG_ASSERT:
+                # Consume-once: every event wait must be re-armed by a fresh
+                # forward publish; a stale consume means a publish went missing.
+                assert self._publish_fresh, "resolve without a fresh forward publish"
+                self._publish_fresh = False
+            if _is_hip:
+                # Temporary workaround: Event.wait() regresses TPOT on AMD MI355.
+                self.publish_ready.synchronize()
+            else:
+                self.publish_ready.wait()
+        batch.seq_lens = self.new_seq_lens_buf[fi]
+
+        if not self.needs_cpu_seq_lens:
+            # GPU gather above is kept (SB.seq_lens must advance each verify);
+            # skip the .cpu() D2H. Downstream takes the GPU-only path.
+            batch.seq_lens_cpu = None
+            batch.seq_lens_sum = None
+            if _DEBUG_ASSERT:
+                # Poison consumed rows: each row must be re-published/seeded
+                # before the next resolve gathers it (safe here: the forward's
+                # re-publish is fenced behind this stream via wait_stream).
+                _assert_nonneg_and_invalidate(batch.seq_lens, self.new_seq_lens_buf, fi)
+            return
+
+        if self.fwd_prepare_d2h_stream is None or self.publish_ready is None:
+            batch.seq_lens_cpu = batch.seq_lens.cpu()  # bootstrap / non-CUDA
+            batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
+            if _DEBUG_ASSERT:
+                _assert_nonneg_and_invalidate(batch.seq_lens, self.new_seq_lens_buf, fi)
+            return
+
+        # Mechanism: don't sync the schedule stream; gate a private stream on the
+        # publish event and copy into the static pinned buffer.
+        self.fwd_prepare_d2h_stream.wait_event(self.publish_ready)
+        with torch.get_device_module(self.device).stream(self.fwd_prepare_d2h_stream):
+            self.new_seq_lens_cpu_pinned.copy_(self.new_seq_lens_buf, non_blocking=True)
+        self.fwd_prepare_d2h_stream.synchronize()
+
+        # FIXME: fi == batch.req_pool_indices; unify future_indices and req_pool_indices.
+        batch.seq_lens_cpu = self.new_seq_lens_cpu_pinned[batch.req_pool_indices_cpu]
+        batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
+        if _DEBUG_ASSERT:
+            # After the D2H copy completed (synchronize above), so the pinned
+            # mirror is not poisoned.
+            _assert_nonneg_and_invalidate(batch.seq_lens, self.new_seq_lens_buf, fi)
+
+    def publish(
+        self,
+        future_indices: torch.Tensor,
+        new_seq_lens: torch.Tensor,
+        confidence: Optional[torch.Tensor] = None,
+    ) -> None:
+        indices = future_indices
+        if indices.shape[0] == 0:
+            return  # DP idle
+        self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype)
+        publish_confidence = self.needs_confidence_relay and confidence is not None
+        if publish_confidence:
+            self.confidence_relay.scatter(indices, confidence)
+        # Only spec_v2 needs the event; it gates the seq_lens D2H on the private stream.
+        if self.spec_algo.is_some():
+            device_module = torch.get_device_module(self.device)
+            if self.publish_ready is None:
+                self.publish_ready = device_module.Event()
+            else:
+                # Chain the records: event fire implies every prior publish is
+                # visible, so an off-forward-stream publish (PD-decode prebuilt
+                # seeding) cannot drop the in-flight forward's fence.
+                device_module.current_stream().wait_event(self.publish_ready)
+            self.publish_ready.record()
+            self._publish_fresh = True
+        if publish_confidence:
+            self.confidence_relay.issue_ring_copy(
+                stream=self.fwd_prepare_d2h_stream,
+                publish_ready=self.publish_ready,
+            )
+
+    def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None:
+        if self.spec_algo.is_ngram():
+            # FIXME: remove once precomputed draft is supported.
+            return
+        indices = future_indices
+        if indices.shape[0] == 0:
+            # DP idle: payload is empty stub; lazy-init shape peek would IndexError.
+            return
+        if not self._forward_buf_initialized:
+            self._lazy_init_forward_buf(payload)
+        self.output_tokens_buf[indices] = payload.bonus_tokens.to(
+            self.output_tokens_buf.dtype
+        )
+
+        if self.need_topk:
+            self.topk_p_buf[indices] = payload.topk_p.to(self.topk_p_buf.dtype)
+            self.topk_index_buf[indices] = payload.topk_index.to(
+                self.topk_index_buf.dtype
+            )
+        if self.need_hidden_states:
+            self.hidden_states_buf[indices] = payload.hidden_states.to(
+                self.hidden_states_buf.dtype
+            )
+        if self.draft_probs_buf is not None and payload.draft_probs is not None:
+            self.draft_probs_buf[indices] = payload.draft_probs
+        if (
+            self.dsa_topk_indices_buf is not None
+            and payload.dsa_topk_indices is not None
+        ):
+            self.dsa_topk_indices_buf[indices] = payload.dsa_topk_indices.to(
+                self.dsa_topk_indices_buf.dtype
+            )

python/sglang/srt/managers/prefill_delayer.py

@@ -0,0 +1,406 @@
+import dataclasses
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, NamedTuple, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_bool_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector
+
+_DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _State:
+    delayed_count: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+
+    def bump_delayed_count(self) -> "_State":
+        return dataclasses.replace(self, delayed_count=self.delayed_count + 1)
+
+
+class _NegotiateOutput(NamedTuple):
+    next_state: Optional[_State]
+    input_estimation: str
+    output_allow: bool
+    output_reason: str
+    num_prefillable: int
+    num_token_watermark_force_allow: int
+    # Accumulated wait of the prefill being released on this pass. Carried
+    # explicitly because `next_state` is None on every release path and thus
+    # cannot convey it to the metrics observation.
+    wait_forward_passes: int = 0
+    wait_seconds: float = 0.0
+
+
+class PrefillDelayer:
+    def __init__(
+        self,
+        dp_size: int,
+        attn_tp_size: int,
+        cpu_group,
+        server_args,
+        max_delay_passes: int,
+        token_usage_low_watermark: Optional[float],
+        metrics_collector: Optional["SchedulerMetricsCollector"] = None,
+        device: Optional["torch.device"] = "cpu",
+        device_group=None,
+    ):
+        self._max_delay_passes = max_delay_passes
+        self._token_usage_low_watermark = token_usage_low_watermark
+        # Queue-based trigger is opt-in: activates only when queue_min_ratio
+        # is explicitly set. Additive with the slot-based trigger.
+        self._queue_min_ratio = server_args.prefill_delayer_queue_min_ratio
+        # Fall back to 5000ms if unset; this is a local safety cap, not a
+        # semantic default, so we don't surface it via ServerArgs.
+        self._max_delay_ms = server_args.prefill_delayer_max_delay_ms
+        if self._max_delay_ms is None:
+            self._max_delay_ms = 5000.0
+        self._queue_trigger_enabled = self._queue_min_ratio is not None
+        logger.info(
+            f"PrefillDelayer initialized with "
+            f"max_delay_passes={self._max_delay_passes} "
+            f"token_usage_low_watermark={self._token_usage_low_watermark} "
+            f"queue_min_ratio={self._queue_min_ratio} "
+            f"max_delay_ms={self._max_delay_ms} "
+            f"queue_trigger_enabled={self._queue_trigger_enabled}"
+        )
+        self.dp_size = dp_size
+        self.enable_dp_attention = server_args.enable_dp_attention
+        dp_size_dim = dp_size if self.enable_dp_attention else 1
+
+        # Mirror scheduler_dp_attn_mixin's NCCL all-gather path: when the
+        # env flag is on (or overlap scheduling is disabled), ride the NCCL
+        # device group on `device` instead of gloo on CPU.
+        use_nccl = (
+            server_args.disable_overlap_schedule
+            or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get()
+        )
+        if use_nccl:
+            assert (
+                device_group is not None
+            ), "device_group is required when using NCCL for PrefillDelayer all-gather"
+            self._gather_group = device_group
+            self._gather_device = device
+        else:
+            self._gather_group = cpu_group
+            self._gather_device = "cpu"
+
+        # Fields packed per rank into the all-gather tensor: prefillable,
+        # token_watermark_force_allow, running_batch, max_prefill_bs,
+        # waiting_queue_len.
+        self._global_info_buffer = torch.empty(
+            (dp_size_dim, attn_tp_size, 5),
+            dtype=torch.int64,
+            device=self._gather_device,
+        )
+
+        self._metrics_collector = metrics_collector
+
+        self._curr_state: Optional[_State] = None
+        self.skip_first_delayer = True
+
+        assert (
+            not server_args.disable_overlap_schedule
+        ), "To use PrefillDelayer, disable_overlap_schedule must be False."
+
+    def _negotiate_should_allow_prefill(
+        self,
+        local_prefillable: bool,
+        token_usage: float,
+        running_batch: int = 0,
+        max_prefill_bs: int = 0,
+        max_running_requests: int = 0,
+        waiting_queue_len: int = 0,
+    ) -> _NegotiateOutput:
+        out = self._negotiate_should_allow_prefill_pure(
+            prev_state=self._curr_state,
+            local_prefillable=local_prefillable,
+            token_usage=token_usage,
+            running_batch=running_batch,
+            max_prefill_bs=max_prefill_bs,
+            max_running_requests=max_running_requests,
+            waiting_queue_len=waiting_queue_len,
+        )
+        self._curr_state = out.next_state
+        return out
+
+    # (Almost) pure function, do not modify self state
+    def _negotiate_should_allow_prefill_pure(
+        self,
+        prev_state: Optional[_State],
+        local_prefillable: bool,
+        token_usage: float,
+        running_batch: int = 0,
+        max_prefill_bs: int = 0,
+        max_running_requests: int = 0,
+        waiting_queue_len: int = 0,
+    ) -> _NegotiateOutput:
+        # Compute local states
+        local_token_watermark_force_allow = (
+            local_prefillable
+            and ((x := self._token_usage_low_watermark) is not None)
+            and (token_usage < x)
+        )
+
+        # Gather global states
+        tp0_info = self._gather_info(
+            local_prefillable=local_prefillable,
+            local_token_watermark_force_allow=local_token_watermark_force_allow,
+            running_batch=running_batch,
+            max_prefill_bs=max_prefill_bs,
+            waiting_queue_len=waiting_queue_len,
+        )
+        global_prefillable = tp0_info[:, 0]
+        global_token_watermark_force_allow = tp0_info[:, 1]
+        global_running_batch = tp0_info[:, 2]
+        global_max_prefill_bs = tp0_info[:, 3]
+        global_waiting_queue_len = tp0_info[:, 4]
+
+        # Compute derived global states
+        if global_prefillable.min().item() > 0:
+            prefillable_status = "all"
+        elif global_prefillable.max().item() == 0:
+            prefillable_status = "none"
+        else:
+            prefillable_status = "mixed"
+        global_exists_token_watermark_force_allow = (
+            global_token_watermark_force_allow.max().item() > 0
+        )
+        debug_info = dict(
+            input_estimation=prefillable_status,
+            num_prefillable=global_prefillable.sum().item(),
+            num_token_watermark_force_allow=global_token_watermark_force_allow.sum().item(),
+        )
+
+        # Wait accumulated so far, taken from prev_state. Release paths attach
+        # this so the wait histograms observe the real value; delay paths leave
+        # the defaults (0) since the wait isn't finished and isn't observed.
+        wait_info = dict(
+            wait_forward_passes=prev_state.delayed_count if prev_state else 0,
+            wait_seconds=(
+                (time.perf_counter() - prev_state.start_time) if prev_state else 0.0
+            ),
+        )
+
+        # Compute outputs
+        if prefillable_status == "all":
+            # Safety valve: low KV usage means GPU is underutilized, skip
+            # delay. Mirrors the check in the "mixed" branch.
+            if global_exists_token_watermark_force_allow:
+                return _NegotiateOutput(
+                    next_state=None,
+                    output_allow=True,
+                    output_reason="token_watermark",
+                    **debug_info,
+                    **wait_info,
+                )
+
+            if not self.enable_dp_attention:
+                max_running_requests = (
+                    max_running_requests + self.dp_size - 1
+                ) // self.dp_size
+
+            global_running_batch_max = int(global_running_batch.max().item())
+            global_max_prefill_bs_max = int(global_max_prefill_bs.max().item())
+            global_waiting_queue_max = int(global_waiting_queue_len.max().item())
+
+            # Queue-based trigger: delay prefill until the waiting queue
+            # reaches queue_min = min(running_req * ratio, max_prefill_bs),
+            # capped by a wall-clock timeout to bound worst-case TTFT.
+            # Targets workloads where decode requests finish one-at-a-time
+            # and fragment prefill into many tiny batches.
+            queue_condition = False
+            if self._queue_trigger_enabled and global_running_batch_max > 0:
+                queue_min_effective = min(
+                    int(global_running_batch_max * self._queue_min_ratio),
+                    global_max_prefill_bs_max,
+                )
+                queue_condition = (
+                    queue_min_effective > 0
+                    and global_waiting_queue_max < queue_min_effective
+                )
+                if queue_condition and prev_state is not None:
+                    elapsed_ms = (time.perf_counter() - prev_state.start_time) * 1000.0
+                    if elapsed_ms >= self._max_delay_ms:
+                        queue_condition = False
+
+            slot_condition = (
+                max_running_requests - global_running_batch_max
+                < global_max_prefill_bs_max
+            )
+
+            if slot_condition or queue_condition:
+                # When the "max_decode_bs - running_bs < max_prefill_bs" condition is met,
+                # the first merge_batch causes the decoding to fail to reach the maximum batch size.
+                if self.skip_first_delayer:
+                    self.skip_first_delayer = False
+                    pass
+                else:
+                    next_state = prev_state or _State()
+                    next_state = next_state.bump_delayed_count()
+                    return _NegotiateOutput(
+                        next_state=next_state,
+                        output_allow=False,
+                        output_reason="delay",
+                        **debug_info,
+                    )
+            exist_previous_wait = prev_state is not None
+            return _NegotiateOutput(
+                next_state=None,
+                output_allow=True,
+                output_reason="wait_success" if exist_previous_wait else "no_wait",
+                **debug_info,
+                **wait_info,
+            )
+        elif prefillable_status == "none":
+            return _NegotiateOutput(
+                next_state=None,
+                # It does not matter whether we allow or not, thus we allow for simplicity
+                output_allow=True,
+                output_reason="",
+                **debug_info,
+                **wait_info,
+            )
+        elif prefillable_status == "mixed":
+            if global_exists_token_watermark_force_allow:
+                return _NegotiateOutput(
+                    next_state=None,
+                    output_allow=True,
+                    output_reason="token_watermark",
+                    **debug_info,
+                    **wait_info,
+                )
+
+            prev_delayed_count = prev_state.delayed_count if prev_state else 0
+            if prev_delayed_count < self._max_delay_passes - 1:
+                next_state = prev_state or _State()
+                next_state = next_state.bump_delayed_count()
+                return _NegotiateOutput(
+                    next_state=next_state,
+                    output_allow=False,
+                    output_reason="delay",
+                    **debug_info,
+                )
+            else:
+                return _NegotiateOutput(
+                    next_state=None,
+                    output_allow=True,
+                    output_reason="wait_timeout",
+                    **debug_info,
+                    **wait_info,
+                )
+        else:
+            raise NotImplementedError
+
+    def _gather_info(
+        self,
+        local_prefillable: bool,
+        local_token_watermark_force_allow: bool,
+        running_batch: int = 0,
+        max_prefill_bs: int = 0,
+        waiting_queue_len: int = 0,
+    ):
+        local_info = torch.tensor(
+            [
+                int(local_prefillable),
+                int(local_token_watermark_force_allow),
+                running_batch,
+                max_prefill_bs,
+                waiting_queue_len,
+            ],
+            device=self._gather_device,
+            dtype=torch.int64,
+        )
+        torch.distributed.all_gather_into_tensor(
+            self._global_info_buffer.flatten(),
+            local_info,
+            group=self._gather_group,
+        )
+        tp0_info = self._global_info_buffer[:, 0, :]
+        return tp0_info
+
+
+class PrefillDelayerSinglePassExecutor:
+    def __init__(self, prefill_delayer: PrefillDelayer, token_usage: float):
+        self._prefill_delayer = prefill_delayer
+        self._token_usage = token_usage
+        self._result: Optional[_NegotiateOutput] = None
+
+    @property
+    def _called(self) -> bool:
+        return self._result is not None
+
+    def finalize(self, *, actual_prefill: bool):
+        if not self._called:
+            self.negotiate_should_allow_prefill(local_prefillable=False)
+
+        _record_single_pass_result(
+            actual_execution=actual_prefill,
+            output=self._result,
+            metrics_collector=self._prefill_delayer._metrics_collector,
+        )
+
+    def negotiate_should_allow_prefill(
+        self,
+        local_prefillable: bool,
+        running_batch: int = 0,
+        max_prefill_bs: int = 0,
+        max_running_requests: int = 0,
+        waiting_queue_len: int = 0,
+    ) -> bool:
+        if not self._called:
+            self._result = self._prefill_delayer._negotiate_should_allow_prefill(
+                local_prefillable=local_prefillable,
+                token_usage=self._token_usage,
+                running_batch=running_batch,
+                max_prefill_bs=max_prefill_bs,
+                max_running_requests=max_running_requests,
+                waiting_queue_len=waiting_queue_len,
+            )
+        return self._result.output_allow
+
+
+def _record_single_pass_result(
+    actual_execution: bool,
+    output: _NegotiateOutput,
+    metrics_collector: Optional["SchedulerMetricsCollector"],
+) -> None:
+    if _DEBUG_LOG:
+        if output.output_allow and (output.output_reason == "wait_timeout"):
+            logger.info(
+                f"PrefillDelayer timeout thus not forbid prefill "
+                f"(num_prefillable={output.num_prefillable}, "
+                f"actual_execution={actual_execution})"
+            )
+        elif output.output_allow and (output.output_reason == "token_watermark"):
+            logger.info(
+                f"PrefillDelayer force allow prefill due to low watermark. "
+                f"(num_prefillable={output.num_prefillable}, "
+                f"num_token_watermark_force_allow={output.num_token_watermark_force_allow}, "
+                f"actual_execution={actual_execution})"
+            )
+        else:
+            assert output.output_reason in {
+                "",
+                "wait_success",
+                "no_wait",
+                "delay",
+            }
+
+    if metrics_collector is not None:
+        metrics_collector.observe_prefill_delayer_outcome(
+            forward_passes=output.wait_forward_passes,
+            wait_seconds=output.wait_seconds,
+            input_estimation=output.input_estimation,
+            output_allow=output.output_allow,
+            output_reason=output.output_reason,
+            actual_execution=actual_execution,
+        )

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

@@ -0,0 +1,982 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.environ import envs
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.schedule_batch import (
+    FINISH_ABORT,
+    FINISH_MATCHED_TOKEN,
+    Req,
+    ScheduleBatch,
+)
+from sglang.srt.mem_cache.common import (
+    maybe_cache_unfinished_req,
+    release_kv_cache,
+)
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
+from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
+        DecodeKVCacheOffloadManager,
+    )
+    from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+    from sglang.srt.managers.scheduler_components.logprob_result_processor import (
+        SchedulerLogprobResultProcessor,
+    )
+    from sglang.srt.managers.scheduler_components.metrics_reporter import (
+        SchedulerMetricsReporter,
+    )
+    from sglang.srt.managers.scheduler_components.output_streamer import (
+        SchedulerOutputStreamer,
+    )
+    from sglang.srt.managers.tp_worker import BaseTpWorker
+    from sglang.srt.managers.utils import (
+        EmbeddingBatchResult,
+        GenerationBatchResult,
+    )
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerBatchResultProcessor:
+    is_generation: bool
+    disaggregation_mode: DisaggregationMode
+    enable_overlap: bool
+    enable_overlap_mlx: bool
+    server_args: ServerArgs
+    model_config: ModelConfig
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    tree_cache: BasePrefixCache
+    hisparse_coordinator: Optional[HiSparseCoordinator]
+    req_to_token_pool: ReqToTokenPool
+    decode_offload_manager: Optional[DecodeKVCacheOffloadManager]
+    metrics_collector: SchedulerMetricsCollector
+    metrics_reporter: SchedulerMetricsReporter
+    draft_worker: BaseTpWorker
+    model_worker: BaseTpWorker
+    logprob_result_processor: SchedulerLogprobResultProcessor
+    output_streamer: SchedulerOutputStreamer
+    abort_request: Callable
+
+    def process_batch_result_prebuilt(self, batch: ScheduleBatch):
+        assert self.disaggregation_mode == DisaggregationMode.DECODE
+        use_free_group = self.server_args.disaggregation_decode_enable_radix_cache
+        if use_free_group:
+            self.token_to_kv_pool_allocator.free_group_begin()
+        for req in batch.reqs:
+            req.time_stats.set_decode_prebuilt_finish_time()
+            req.update_finish_state()
+            if req.finished():
+                req.time_stats.set_quick_finish_time()
+                if self.server_args.enable_hisparse:
+                    self.hisparse_coordinator.request_finished(req)
+                release_kv_cache(req, self.tree_cache)
+
+        # Note: Logprobs should be handled on the prefill engine.
+        self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
+        if use_free_group:
+            self.token_to_kv_pool_allocator.free_group_end()
+
+    def _maybe_collect_routed_experts(self, req: Req):
+        """Collect routed experts for a finished request.
+
+        Returns immediately if `return_routed_experts` was not set on the
+        request, so non-opted-in reqs don't pay the host-gather cost.
+
+        Honors the caller's absolute start so the response covers
+        `[start_len, seqlen - 1)`. The default start_len is 0, which returns
+        the full sequence.
+
+        Logs a soft warning if the resulting tensor's row count differs from
+        the expected `seqlen - 1 - start_len`, to catch silent regressions.
+        """
+        if not req.return_routed_experts:
+            return
+        capturer = get_global_experts_capturer()
+        if capturer is None:
+            return
+        start_len = req.routed_experts_start_len
+        seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
+        req.routed_experts = capturer.get_topk(
+            req_pool_idx=req.req_pool_idx,
+            seqlen=seqlen,
+            req_to_token_pool=self.req_to_token_pool,
+            start_len=start_len,
+        )
+
+        expected_rows = max(0, seqlen - 1 - start_len)
+        if (
+            req.routed_experts is not None
+            and req.routed_experts.shape[0] != expected_rows
+        ):
+            logger.warning(
+                "routed_experts row-count mismatch for req %s: got %d, expected %d "
+                "(seqlen=%d, raw_seqlen=%d, cached_tokens=%d, start_len=%s). "
+                "This indicates a silent bug.",
+                req.rid,
+                req.routed_experts.shape[0],
+                expected_rows,
+                seqlen,
+                req.seqlen,
+                req.cached_tokens,
+                req.routed_experts_start_len,
+            )
+
+    def _maybe_collect_indexer_topk(self, req: Req):
+        capturer = get_global_indexer_capturer()
+        if capturer is None:
+            return
+        seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
+        req.indexer_topk = capturer.get_topk(
+            req_pool_idx=req.req_pool_idx,
+            seqlen=seqlen,
+            req_to_token_pool=self.req_to_token_pool,
+        )
+
+    def _maybe_collect_customized_info(
+        self,
+        i: int,
+        req: Req,
+        logits_output: LogitsProcessorOutput,
+    ):
+        if logits_output is not None and logits_output.customized_info is not None:
+            if req.customized_info is None:
+                req.customized_info = {}
+            for k, v in logits_output.customized_info.items():
+                if k not in req.customized_info:
+                    req.customized_info[k] = []
+                # Copy the element so it doesn't retain the entire batch
+                # tensor/array via a view reference.
+                elem = v[i]
+                if isinstance(elem, torch.Tensor):
+                    elem = elem.clone()
+                elif hasattr(elem, "copy") and callable(elem.copy):
+                    elem = elem.copy()
+                req.customized_info[k].append(elem)
+
+    def process_batch_result_prefill(
+        self,
+        batch: ScheduleBatch,
+        result: Union[GenerationBatchResult, EmbeddingBatchResult],
+    ):
+        skip_stream_req = None
+
+        if self.is_generation:
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+            if result.routed_experts_output is not None:
+                result.routed_experts_output.finalize()
+                result.routed_experts_output = None
+            if result.indexer_topk_output is not None:
+                result.indexer_topk_output.finalize()
+                result.indexer_topk_output = None
+
+            (
+                logits_output,
+                next_token_ids,
+                extend_input_len_per_req,
+                extend_logprob_start_len_per_req,
+            ) = (
+                result.logits_output,
+                result.next_token_ids,
+                result.extend_input_len_per_req,
+                result.extend_logprob_start_len_per_req,
+            )
+
+            # Move next_token_ids and logprobs to cpu
+            next_token_ids = next_token_ids.tolist()
+            self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
+
+            self._validate_pp_skip_output_comm(batch, result)
+
+            hidden_state_offset = 0
+
+            # Check finish conditions
+            logprob_pt = 0
+
+            for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
+                if (
+                    req.finished() and req.inflight_middle_chunks <= 0
+                ) or req.is_retracted:
+                    # Decode req in a mixed batch, or a retracted req. Keep an
+                    # aborted middle chunk in the chunked branch long enough to
+                    # drain its accounting without streaming it.
+                    continue
+
+                if req.inflight_middle_chunks <= 0:
+                    req.time_stats.set_prefill_finished_time()
+
+                    # req output_ids are set here
+                    req.output_ids.append(next_token_id)
+
+                    self._maybe_update_reasoning_tokens(req, next_token_id)
+
+                    req.update_finish_state()
+                    if req.finished():
+                        self._maybe_collect_routed_experts(req)
+                        self._maybe_collect_indexer_topk(req)
+                        release_kv_cache(req, self.tree_cache)
+                        req.time_stats.set_completion_time()
+                    elif not batch.decoding_reqs or req not in batch.decoding_reqs:
+                        maybe_cache_unfinished_req(req, self.tree_cache)
+                        if self.server_args.enable_hisparse:
+                            self.hisparse_coordinator.admit_request_into_staging(req)
+
+                    self._maybe_collect_customized_info(i, req, logits_output)
+
+                    if batch.return_logprob:
+                        logprob_pt = self._apply_prefill_logprobs(
+                            req=req,
+                            i=i,
+                            logits_output=logits_output,
+                            extend_input_len_per_req=extend_input_len_per_req,
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
+                            next_token_ids=next_token_ids,
+                            logprob_pt=logprob_pt,
+                        )
+
+                    if (
+                        req.return_hidden_states
+                        and logits_output.hidden_states is not None
+                    ):
+                        hidden_state_offset = self._append_prefill_hidden_states(
+                            req=req,
+                            logits_output=logits_output,
+                            hidden_state_offset=hidden_state_offset,
+                        )
+
+                    if req.grammar is not None:
+                        self._apply_prefill_grammar(
+                            req=req, next_token_id=next_token_id
+                        )
+
+                else:
+                    # being chunked reqs' prefill is not finished
+                    req.inflight_middle_chunks -= 1
+                    # There is only at most one request being currently chunked.
+                    # Because this request does not finish prefill,
+                    # we don't want to stream the request currently being chunked.
+                    skip_stream_req = req
+
+                    # Incrementally update input logprobs.
+                    if batch.return_logprob:
+                        logprob_pt = self._apply_chunked_prefill_logprobs(
+                            req=req,
+                            i=i,
+                            logits_output=logits_output,
+                            extend_input_len_per_req=extend_input_len_per_req,
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
+                            logprob_pt=logprob_pt,
+                        )
+
+                    req.time_stats.set_last_chunked_prefill_finish_time()
+
+        else:  # embedding or reward model
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+
+            embeddings = self._convert_embeddings(result=result)
+            phs = result.pooled_hidden_states
+
+            if phs is not None:
+                if isinstance(phs, list):
+                    phs = [t.cpu().detach() for t in phs]
+                else:
+                    phs = phs.cpu().detach()
+
+            # Check finish conditions
+            for i, req in enumerate(batch.reqs):
+                if req.is_retracted:
+                    continue
+
+                req.embedding = embeddings[i]
+                if req.return_pooled_hidden_states and phs is not None:
+                    req.pooled_hidden_state = phs[i]
+                if req.inflight_middle_chunks <= 0:
+                    req.time_stats.set_prefill_finished_time()
+                    # Dummy output token for embedding models
+                    req.output_ids.append(0)
+                    req.update_finish_state()
+
+                    if req.finished():
+                        release_kv_cache(req, self.tree_cache)
+                        req.time_stats.set_completion_time()
+                    else:
+                        maybe_cache_unfinished_req(req, self.tree_cache)
+                else:
+                    # being chunked reqs' prefill is not finished
+                    req.inflight_middle_chunks -= 1
+                    req.time_stats.set_last_chunked_prefill_finish_time()
+
+        self.output_streamer.stream_output(
+            batch.reqs, batch.return_logprob, skip_stream_req
+        )
+
+        can_run_cuda_graph = result.can_run_cuda_graph
+        self.metrics_reporter.report_prefill_stats(
+            batch=batch,
+            prefill_stats=batch.prefill_stats,
+            can_run_cuda_graph=can_run_cuda_graph,
+            dp_cooperation_info=batch.dp_cooperation_info,
+        )
+
+    def _convert_embeddings(self, *, result: EmbeddingBatchResult) -> list:
+        is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
+
+        embeddings = result.embeddings
+
+        if is_sparse:
+            batch_ids, token_ids = embeddings.indices()
+            values = embeddings.values()
+
+            embeddings = [{} for _ in range(embeddings.size(0))]
+            for i in range(batch_ids.shape[0]):
+                embeddings[batch_ids[i].item()][token_ids[i].item()] = values[i].item()
+        else:
+            if isinstance(embeddings, torch.Tensor):
+                embeddings = embeddings.tolist()
+            else:
+                embeddings = [tensor.tolist() for tensor in embeddings]
+        return embeddings
+
+    def move_logprobs_to_cpu(
+        self,
+        *,
+        batch: ScheduleBatch,
+        logits_output: LogitsProcessorOutput,
+    ) -> None:
+        if batch.return_logprob:
+            if logits_output.next_token_logprobs is not None:
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.tolist()
+                )
+            if logits_output.input_token_logprobs is not None:
+                logits_output.input_token_logprobs = tuple(
+                    logits_output.input_token_logprobs.tolist()
+                )
+            if logits_output.next_token_top_logprobs_val:
+                logits_output.next_token_top_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_top_logprobs_val
+                ]
+                logits_output.next_token_top_logprobs_idx = [
+                    x.tolist() for x in logits_output.next_token_top_logprobs_idx
+                ]
+            if logits_output.next_token_token_ids_logprobs_val:
+                logits_output.next_token_token_ids_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
+                ]
+
+    def _apply_prefill_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+        extend_input_len_per_req: Optional[List[int]],
+        extend_logprob_start_len_per_req: Optional[List[int]],
+        next_token_ids: List[int],
+        logprob_pt: int,
+    ) -> int:
+        assert extend_logprob_start_len_per_req is not None
+        assert extend_input_len_per_req is not None
+        extend_logprob_start_len = extend_logprob_start_len_per_req[i]
+        extend_input_len = extend_input_len_per_req[i]
+
+        num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs(
+            req,
+            extend_input_len,
+            extend_logprob_start_len,
+        )
+
+        if req.return_logprob:
+            self.logprob_result_processor.add_logprob_return_values(
+                i,
+                req,
+                logprob_pt,
+                next_token_ids,
+                num_input_logprobs,
+                logits_output,
+            )
+        logprob_pt += num_input_logprobs
+        return logprob_pt
+
+    @staticmethod
+    def _validate_pp_skip_output_comm(
+        batch: ScheduleBatch,
+        result: Union[GenerationBatchResult, EmbeddingBatchResult],
+    ):
+        """Validate PP skip output comm correctness.
+
+        - When skip=True: all reqs must be middle chunks (inflight_middle_chunks > 0)
+          so placeholder zeros are never consumed via req.output_ids.append().
+        - When skip=False: at least one req should consume next_token_ids
+          (inflight_middle_chunks <= 0), otherwise warn.
+        """
+        if not envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get():
+            return
+
+        if not getattr(result, "skipped_output_comm", False):
+            if batch.forward_mode.is_extend() and not batch.forward_mode.is_prebuilt():
+                has_consumed_output = any(
+                    req.inflight_middle_chunks <= 0
+                    for req in batch.reqs
+                    if not req.finished() and not req.is_retracted
+                )
+                if not has_consumed_output and len(batch.reqs) > 0:
+                    chunks = list([r.inflight_middle_chunks for r in batch.reqs])
+                    logger.warning(
+                        f"PP non-skip output comm: no req consumed next_token_ids. "
+                        f"contains_last_prefill_chunk={batch.contains_last_prefill_chunk}, "
+                        f"num_reqs={len(batch.reqs)}, all inflight_middle_chunks={chunks}"
+                    )
+            return
+
+        for req in batch.reqs:
+            if not req.finished() and not req.is_retracted:
+                assert req.inflight_middle_chunks > 0, (
+                    f"PP skip output comm invariant violated: req {req.rid} "
+                    f"has inflight_middle_chunks={req.inflight_middle_chunks} "
+                    f"but output was skipped (contains_last_prefill_chunk="
+                    f"{batch.contains_last_prefill_chunk}). "
+                    f"Placeholder zeros would be appended to output_ids."
+                )
+
+    def _append_prefill_hidden_states(
+        self,
+        *,
+        req: Req,
+        logits_output: LogitsProcessorOutput,
+        hidden_state_offset: int,
+    ) -> int:
+        req.hidden_states.append(
+            logits_output.hidden_states[
+                hidden_state_offset : (
+                    hidden_state_offset := hidden_state_offset
+                    + len(req.origin_input_ids)
+                )
+            ]
+            .cpu()
+            .clone()
+            .tolist()
+        )
+        return hidden_state_offset
+
+    def _apply_prefill_grammar(self, *, req: Req, next_token_id: int) -> None:
+        # FIXME: this try-except block is for handling unexpected xgrammar issue.
+        try:
+            req.grammar.accept_token(next_token_id)
+        except ValueError as e:
+            # Grammar accept_token can raise ValueError if the token is not in the grammar.
+            # This can happen if the grammar is not set correctly or the token is invalid.
+            logger.error(
+                f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
+            )
+            req.to_finish = FINISH_ABORT()
+        req.grammar.finished = req.finished()
+
+    def _apply_chunked_prefill_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+        extend_input_len_per_req: Optional[List[int]],
+        extend_logprob_start_len_per_req: Optional[List[int]],
+        logprob_pt: int,
+    ) -> int:
+        extend_logprob_start_len = extend_logprob_start_len_per_req[i]
+        extend_input_len = extend_input_len_per_req[i]
+        if extend_logprob_start_len < extend_input_len:
+            # Update input logprobs.
+            num_input_logprobs = (
+                self.logprob_result_processor.calculate_num_input_logprobs(
+                    req,
+                    extend_input_len,
+                    extend_logprob_start_len,
+                )
+            )
+            if req.return_logprob:
+                self.logprob_result_processor.add_input_logprob_return_values(
+                    i,
+                    req,
+                    logits_output,
+                    logprob_pt,
+                    num_input_logprobs,
+                    last_prefill_chunk=False,
+                )
+            logprob_pt += num_input_logprobs
+        return logprob_pt
+
+    def _resolve_spec_v2_tokens(
+        self,
+        result: GenerationBatchResult,
+        batch: ScheduleBatch,
+    ) -> List[List[int]]:
+        """Resolve the padded next token ids for spec-v2 (overlap and non-overlap)."""
+        assert result.next_token_ids.is_cpu
+        assert result.accept_lens.is_cpu
+
+        next_token_ids = result.next_token_ids.tolist()
+        accept_lens = result.accept_lens.tolist()
+        result.num_correct_drafts = sum(accept_lens) - len(batch.reqs)
+        result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens]
+
+        block_accept_lens = (
+            result.block_accept_lens.tolist()
+            if result.block_accept_lens is not None
+            else None
+        )
+        result.num_block_accept_tokens = (
+            sum(block_accept_lens) if block_accept_lens else 0
+        )
+        cap_lens = result.cap_lens.tolist() if result.cap_lens is not None else None
+        result.num_cap_tokens = sum(cap_lens) if cap_lens else 0
+
+        # Feed the adaptive controller now that accept_lens is on CPU,
+        # instead of doing a synchronous GPU→CPU copy in the worker hot path.
+        # BaseSpecWorker provides a no-op default for non-adaptive workers.
+        self.model_worker.on_verify_complete_cpu(
+            result.num_correct_drafts_per_req_cpu, batch_size=len(batch.reqs)
+        )
+
+        predict_tokens = []
+        # In adaptive spec-v2, the worker state may already have switched when this
+        # delayed result is processed. Use the draft token count recorded on result.
+        stride = result.speculative_num_draft_tokens
+        assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
+
+        for i, req in enumerate(batch.reqs):
+            accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
+
+            if req.is_retracted or req.finished():
+                # Nothing to settle: no worker pre-claims the bonus, so
+                # kv_committed_len already holds the committed prefix.
+                pass
+            else:
+                if req.grammar is not None:
+                    # Stop accepting once the grammar terminates, so the
+                    # over-drafted suffix is never committed to KV nor emitted.
+                    # This advances the grammar FSM; the result loop only syncs
+                    # grammar.finished.
+                    accept_tokens = self._accept_grammar_tokens(req, accept_tokens)
+
+                # Commit the full accepted run (drafts + bonus).
+                num_accept_tokens = len(accept_tokens)
+                req.kv_committed_len += num_accept_tokens
+                req.spec_verify_ct += 1
+
+                num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
+                req.spec_num_correct_drafts += num_correct_drafts
+                req.update_spec_correct_drafts_histogram(num_correct_drafts)
+
+                if block_accept_lens is not None:
+                    req.spec_num_block_accept_tokens += block_accept_lens[i]
+                if cap_lens is not None:
+                    req.spec_num_cap_tokens += cap_lens[i]
+                    req.update_spec_cap_lens_histogram(cap_lens[i])
+
+            predict_tokens.append(accept_tokens)
+
+        return predict_tokens
+
+    def _accept_grammar_tokens(
+        self, req: Req, tokens: Union[int, List[int]]
+    ) -> List[int]:
+        """Advance the grammar over the accepted token(s), stopping at the token
+        that terminates it.
+
+        ``tokens`` is a single sampled token (normal decode) or the whole
+        verified run (spec decode). Returns the retained prefix; for spec the
+        suffix past grammar completion is dropped so it is never committed to KV
+        nor emitted. Advances the grammar FSM only -- ``grammar.finished`` is
+        synced by the caller once the finish state is updated.
+        """
+        if isinstance(tokens, int):
+            tokens = [tokens]
+        retained = []
+        try:
+            for token_id in tokens:
+                req.grammar.accept_token(token_id)
+                retained.append(token_id)
+                if req.grammar.is_terminated():
+                    break
+        except ValueError as e:
+            # accept_token raises ValueError if the token is not in the grammar
+            # (misconfigured grammar or invalid token); abort the request.
+            logger.error(
+                f"Grammar accept_token failed for req {req.rid} with token "
+                f"{tokens}: {e}"
+            )
+            req.to_finish = FINISH_ABORT()
+        return retained
+
+    def process_batch_result_idle(
+        self,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
+        self.output_streamer._stream_output_generation(
+            batch.reqs, batch.return_logprob, is_idle_batch=True
+        )
+
+    def process_batch_result_decode(
+        self,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+        if result.routed_experts_output is not None:
+            result.routed_experts_output.finalize()
+            result.routed_experts_output = None
+        if result.indexer_topk_output is not None:
+            result.indexer_topk_output.finalize()
+            result.indexer_topk_output = None
+
+        logits_output, next_token_ids, can_run_cuda_graph = (
+            result.logits_output,
+            result.next_token_ids,
+            result.can_run_cuda_graph,
+        )
+
+        next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
+            batch=batch,
+            result=result,
+            logits_output=logits_output,
+            next_token_ids=next_token_ids,
+        )
+
+        self.metrics_reporter.num_generated_tokens += len(batch.reqs)
+        if not batch.spec_algorithm.is_none():
+            self.metrics_reporter.update_spec_metrics(
+                batch.batch_size(),
+                result.num_correct_drafts,
+                num_block_accept_tokens=result.num_block_accept_tokens,
+                num_cap_tokens=result.num_cap_tokens,
+            )
+        if self.server_args.enable_metrics:
+            self.metrics_collector.increment_decode_cuda_graph_pass(
+                value=can_run_cuda_graph
+            )
+
+        self.token_to_kv_pool_allocator.free_group_begin()
+
+        for i, req in enumerate(batch.reqs):
+            req: Req
+
+            if (self.enable_overlap or self.enable_overlap_mlx) and (
+                req.finished() or req.is_retracted
+            ):
+                # NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled.
+                # And all the over-allocated tokens will be freed in `release_kv_cache`.
+                continue
+
+            # next_token_id is a per-req list: 1 token for non-spec, the verified
+            # run for spec (already grammar-truncated in _resolve_spec_v2_tokens).
+            next_token_id = next_token_ids[i]
+            is_spec = not batch.spec_algorithm.is_none()
+
+            req.output_ids.extend(next_token_id)
+            new_accept_len = len(next_token_id)
+
+            self._maybe_update_reasoning_tokens(req, next_token_id)
+            req.time_stats.set_last_decode_finish_time()
+            req.update_finish_state(new_accept_len)
+
+            self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
+
+            if req.return_logprob:
+                self._apply_decode_logprobs(
+                    req=req,
+                    i=i,
+                    batch=batch,
+                    next_token_id=next_token_id,
+                    next_token_logprobs=next_token_logprobs,
+                    logits_output=logits_output,
+                )
+
+            if req.return_hidden_states and logits_output.hidden_states is not None:
+                # hidden_states is [bs * stride, hidden_dim], one row per emitted
+                # token; stride = speculative_num_draft_tokens for spec, 1 for non-spec.
+                stride = result.speculative_num_draft_tokens or 1
+                accept_len = len(next_token_id)
+                start = i * stride
+                req.hidden_states.extend(
+                    logits_output.hidden_states[start : start + accept_len]
+                    .cpu()
+                    .tolist()
+                )
+
+            if req.grammar is not None:
+                if not is_spec:
+                    # Normal decode advances the grammar for its single token
+                    # here; spec already advanced it in _resolve_spec_v2_tokens.
+                    self._accept_grammar_tokens(req, next_token_id)
+                req.grammar.finished = req.finished()
+
+        self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
+        self.token_to_kv_pool_allocator.free_group_end()
+
+        self.metrics_reporter.forward_ct_decode = (
+            self.metrics_reporter.forward_ct_decode + 1
+        ) % (1 << 30)
+        self.metrics_reporter.report_decode_stats(
+            can_run_cuda_graph,
+            running_batch=batch,
+            num_correct_drafts=result.num_correct_drafts,
+        )
+
+    def _normalize_decode_outputs(
+        self,
+        *,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        logits_output: LogitsProcessorOutput,
+        next_token_ids: Union[torch.Tensor, List[int]],
+    ) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]:
+        next_token_logprobs = None
+        # Normalize to a uniform per-req list of accepted tokens (List[List[int]]):
+        # spec unpacks the padded verify output; non-spec wraps its single token.
+        if not batch.spec_algorithm.is_none():
+            next_token_ids = self._resolve_spec_v2_tokens(result, batch)
+        else:
+            # CUDA workers return a device tensor, MLX a host list[int]; both -> list.
+            ids = (
+                next_token_ids.tolist()
+                if torch.is_tensor(next_token_ids)
+                else next_token_ids
+            )
+            next_token_ids = [[t] for t in ids]
+
+        if batch.return_logprob:
+            next_token_logprobs = logits_output.next_token_logprobs.tolist()
+            if logits_output.next_token_top_logprobs_val:
+                logits_output.next_token_top_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_top_logprobs_val
+                ]
+                logits_output.next_token_top_logprobs_idx = [
+                    x.tolist() for x in logits_output.next_token_top_logprobs_idx
+                ]
+
+            if logits_output.next_token_token_ids_logprobs_val:
+                logits_output.next_token_token_ids_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
+                ]
+        return next_token_ids, next_token_logprobs
+
+    def _apply_decode_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        batch: ScheduleBatch,
+        next_token_id: Union[int, List[int]],
+        next_token_logprobs: list,
+        logits_output: LogitsProcessorOutput,
+    ) -> None:
+        # accepted_ids is already a per-req list; non-spec logprobs are flat, so
+        # the scalar logprob still needs wrapping.
+        if not batch.spec_algorithm.is_none():
+            accepted_logprobs = next_token_logprobs[i]
+            accepted_ids = next_token_id
+            max_accept = len(accepted_logprobs)
+        else:
+            accepted_logprobs = [next_token_logprobs[i]]
+            accepted_ids = next_token_id
+            max_accept = 1
+
+        for j, tok_id in enumerate(accepted_ids):
+            req.logprob.output_token_logprobs_val.append(accepted_logprobs[j])
+            req.logprob.output_token_logprobs_idx.append(tok_id)
+            if req.logprob.top_logprobs_num > 0:
+                flat_idx = i * max_accept + j
+                req.logprob.output_top_logprobs_val.append(
+                    logits_output.next_token_top_logprobs_val[flat_idx]
+                )
+                req.logprob.output_top_logprobs_idx.append(
+                    logits_output.next_token_top_logprobs_idx[flat_idx]
+                )
+            if req.logprob.token_ids_logprob is not None:
+                flat_idx = i * max_accept + j
+                req.logprob.output_token_ids_logprobs_val.append(
+                    logits_output.next_token_token_ids_logprobs_val[flat_idx]
+                )
+                req.logprob.output_token_ids_logprobs_idx.append(
+                    logits_output.next_token_token_ids_logprobs_idx[flat_idx]
+                )
+
+    def _handle_finish_state_updated_req(
+        self,
+        req: Req,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+    ):
+        # Called here (after update_finish_state) so req.finished() is valid
+        # for mamba_lazy_post_decode_at_boundary inside.
+        self._mamba_prefix_cache_update(req, batch, result, i)
+
+        if (
+            self.server_args.disaggregation_decode_enable_offload_kvcache
+            and not req.finished()
+        ):
+            self.decode_offload_manager.offload_kv_cache(req)
+
+        if req.finished():
+            # isinstance narrowing: create_worker may also return plain
+            # TpModelWorker-based drafts, which carry no spec-worker hooks.
+            if isinstance(self.draft_worker, BaseSpecWorker):
+                self.draft_worker.note_request_finished(
+                    rid=req.rid,
+                    natural_stop=isinstance(req.finished_reason, FINISH_MATCHED_TOKEN),
+                )
+
+            # delete feature to save memory
+            if req.multimodal_inputs is not None and req.session is None:
+                req.multimodal_inputs.release_features()
+            self._maybe_collect_routed_experts(req)
+            self._maybe_collect_indexer_topk(req)
+
+            if self.server_args.disaggregation_decode_enable_offload_kvcache:
+                # Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
+                if not self.decode_offload_manager.offload_kv_cache(req):
+                    self.decode_offload_manager.finalize_release_on_finish(req)
+            else:
+                if self.server_args.enable_hisparse:
+                    self.hisparse_coordinator.request_finished(req)
+                prepare_release = getattr(
+                    self.model_worker, "prepare_for_kv_cache_release", None
+                )
+                if callable(prepare_release):
+                    prepare_release(req)
+                is_insert = (
+                    req.mamba_lazy_is_insert
+                    if get_server_args().enable_mamba_extra_buffer_lazy()
+                    else True
+                )
+                release_kv_cache(req, self.tree_cache, is_insert=is_insert)
+
+            req.time_stats.set_completion_time()
+
+        self._maybe_collect_customized_info(i, req, logits_output)
+
+    def _maybe_update_reasoning_tokens(
+        self,
+        req: Req,
+        next_token_id: Union[int, List[int]],
+    ):
+        think_end_id = self.model_config.think_end_id
+        if req.require_reasoning and think_end_id is not None:
+            req.update_reasoning_tokens(next_token_id, think_end_id)
+
+    def _mamba_prefix_cache_update(
+        self,
+        req: Req,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        i: int,
+    ) -> None:
+        """Update mamba track state at ping-pong boundaries.
+
+        Non-lazy: swap the ping-pong index so the next forward writes to
+        the alternate slot.
+        Lazy: keep the same index (prealloc handles the swap) and run
+        post-decode cleanup to free the temporary second slot.
+        """
+        if req.mamba_ping_pong_track_buffer is None:
+            return
+
+        lazy = get_server_args().enable_mamba_extra_buffer_lazy()
+        at_boundary, track_seqlen = self._mamba_check_track_boundary(
+            req, batch, result, i
+        )
+
+        if not at_boundary:
+            return
+
+        req.mamba_last_track_seqlen = track_seqlen
+        if lazy:
+            self.mamba_lazy_post_decode_at_boundary(req, batch)
+        else:
+            req.mamba_next_track_idx = (
+                batch.req_to_token_pool.get_mamba_ping_pong_other_idx(
+                    req.mamba_next_track_idx
+                )
+            )
+
+    def _mamba_check_track_boundary(self, req, batch, result, i):
+        """Check if this decode step crosses a mamba track interval boundary.
+
+        Returns (at_boundary, track_seqlen).  The boundary condition
+        matches what the forward's tracking mask used:
+        ``prepare_for_decode`` increments both ``seq_lens_cpu`` and
+        ``kv_committed_len`` by 1, then checks
+        ``seq_lens_cpu % interval == 0``.  Using ``kv_committed_len``
+        here reproduces that check exactly, and the value is always a
+        multiple of ``interval`` (hence page-aligned).
+
+        For spec decode, the boundary is detected by comparing the
+        accepted seq_len range against interval boundaries.
+        """
+        interval = get_server_args().mamba_track_interval
+
+        if batch.spec_algorithm.is_none():
+            if req.kv_committed_len % interval == 0:
+                return True, req.kv_committed_len
+        elif result.num_correct_drafts_per_req_cpu is not None:
+            cur = req.seqlen - 1
+            prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1
+            if cur // interval != prev // interval:
+                return True, cur // interval * interval
+
+        return False, 0
+
+    def mamba_lazy_post_decode_at_boundary(self, req: Req, batch: ScheduleBatch):
+        """Post-decode cleanup at a lazy-mode track boundary.
+
+        Finished reqs: if prealloc failed (other slot is -1), the forward
+        overwrote the only slot with corrupted state, so mark
+        is_insert=False to skip the cache insert.  If the other slot is
+        occupied (stale prealloc from an overlap extra forward), free it
+        so the prealloc assert in the next prepare_for_decode holds.
+
+        Running reqs: free the old ping-pong slot so we go back to
+        holding only 1 slot until the next boundary.
+        """
+        other_idx = 1 - req.mamba_next_track_idx
+        other_val = req.mamba_ping_pong_track_buffer[other_idx].item()
+        if other_val != -1:
+            pool = batch.req_to_token_pool
+            pool.mamba_allocator.free(
+                req.mamba_ping_pong_track_buffer[other_idx].unsqueeze(0)
+            )
+            pool.set_mamba_ping_pong_slot(req, other_idx, -1)
+        elif req.finished():
+            req.mamba_lazy_is_insert = False

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

@@ -0,0 +1,322 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed.parallel_state import get_tp_group
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.environ import envs
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+    cuda_graph_fully_disabled,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.observability.metrics_collector import DPCooperationInfo
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.utils.common import require_mlp_tp_gather
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+
+_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
+
+
+@dataclass
+class MLPSyncBatchInfo:
+    dp_size: int
+    tp_size: int
+    cp_size: int
+
+    num_tokens: int
+    num_tokens_for_logprob: int
+    can_cuda_graph: bool
+    is_extend_in_batch: bool
+    local_can_run_tbo: bool
+    local_forward_mode: int
+    can_run_breakable_cuda_graph: bool
+
+    # some gathered elements
+    tp0_info: torch.Tensor = None
+    global_num_tokens: list[int] = None
+    global_num_tokens_for_logprob: list[int] = None
+    tbo_split_seq_index: torch.Tensor = None
+    global_forward_mode: int = None
+    dp_cooperation_info: Optional[DPCooperationInfo] = None
+
+    def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
+        return torch.tensor(
+            [
+                self.num_tokens,
+                self.num_tokens_for_logprob,
+                int(self.can_cuda_graph),
+                int(self.is_extend_in_batch),
+                int(self.local_can_run_tbo),
+                self.local_forward_mode,
+                int(self.can_run_breakable_cuda_graph),
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
+        return torch.tensor(
+            [
+                0,  # num_tokens
+                0,  # num_tokens_for_logprob
+                1,  # can_cuda_graph
+                0,  # is_extend_in_batch
+                1,  # local_can_run_tbo
+                ForwardMode.IDLE.value,  # local_forward_mode
+                0,  # can_run_breakable_cuda_graph
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    def all_gather(self, device, group: torch.distributed.ProcessGroup):
+        local_info_tensor = self._get_local_tensor(device=device)
+        global_info_tensor = torch.empty(
+            (self.dp_size, self.tp_size * self.cp_size, 7),
+            dtype=torch.int64,
+            device=device,
+        )
+
+        torch.distributed.all_gather_into_tensor(
+            global_info_tensor.flatten(),
+            local_info_tensor,
+            group=group,
+        )
+        if device == "cpu":
+            tp_active_ranks = get_tp_group().active_ranks_cpu
+        else:
+            tp_active_ranks = get_tp_group().active_ranks
+
+        # Set fallback values for inactive ranks
+        tp_info = global_info_tensor.view(self.dp_size * self.tp_size * self.cp_size, 7)
+        tp_info[tp_active_ranks == 0] = self._get_fallback_tensor(device=device)
+
+        tp0_info = global_info_tensor[:, 0, :]
+        self.tp0_info = tp0_info
+        # Perform only one Device-to-Host (D2H) memory copy
+        cpu_data = tp0_info[:, :2].cpu()
+        self.global_num_tokens = cpu_data[:, 0].tolist()
+        self.global_num_tokens_for_logprob = cpu_data[:, 1].tolist()
+        self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
+        self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
+        self.can_run_breakable_cuda_graph = bool(tp0_info[:, 6].min().item())
+        if _ENABLE_METRICS_DP_ATTENTION:
+            self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
+
+
+def _update_gather_batch(
+    batch: ScheduleBatch,
+    mlp_sync_info: MLPSyncBatchInfo,
+    require_mlp_tp_gather: bool,
+    skip_all_gather=False,
+):
+    # TODO: handle the case when moe_dense_tp_size != 1
+    if not require_mlp_tp_gather:
+        batch.global_num_tokens = [mlp_sync_info.num_tokens]
+        batch.global_num_tokens_for_logprob = [mlp_sync_info.num_tokens_for_logprob]
+    else:
+        batch.global_num_tokens = mlp_sync_info.global_num_tokens
+        batch.global_num_tokens_for_logprob = (
+            mlp_sync_info.global_num_tokens_for_logprob
+        )
+    if not skip_all_gather:
+        batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
+        batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
+        batch.global_forward_mode = mlp_sync_info.global_forward_mode
+
+    # Check forward mode for cuda graph
+    batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
+    batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_breakable_cuda_graph
+
+
+def prepare_mlp_sync_batch_raw(
+    local_batch: ScheduleBatch,
+    dp_size: int,
+    attn_tp_size: int,
+    attn_cp_size: int,
+    tp_group: GroupCoordinator,
+    get_idle_batch: Callable[[], ScheduleBatch],
+    disable_cuda_graph: bool,
+    require_mlp_tp_gather: bool,
+    disable_overlap_schedule: bool,
+    offload_tags: set[str],
+):
+    # Check if other DP workers have running batches
+    if (
+        local_batch is None
+        or local_batch.forward_mode.is_prebuilt()
+        or local_batch.forward_mode.is_idle()
+    ):
+        num_tokens = 0
+        num_tokens_for_logprob = 0
+    elif local_batch.forward_mode.is_decode():
+        num_tokens = local_batch.batch_size()
+        num_tokens_for_logprob = num_tokens
+    else:
+        num_tokens = local_batch.extend_num_tokens
+        num_tokens_for_logprob = sum(
+            # We should have at least 1 token for sample in every case.
+            max(extend_len - logprob_start_len, 1)
+            for logprob_start_len, extend_len in zip(
+                local_batch.extend_logprob_start_lens,
+                local_batch.extend_lens,
+            )
+        )
+        assert (
+            local_batch.return_logprob
+            or num_tokens_for_logprob == local_batch.batch_size()
+        )
+
+    skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()
+    can_cuda_graph = (
+        local_batch is None
+        or local_batch.forward_mode.is_decode_or_idle()
+        or local_batch.forward_mode.is_prebuilt()
+    ) and not disable_cuda_graph
+    # Idle/None ranks are permissive (like can_cuda_graph): the all-gather
+    # min()-reduces this across DP ranks, so a prefill batch with idle ranks
+    # still resolves to True (idle ranks become a padded dummy extend).
+    can_run_breakable_cuda_graph = (
+        local_batch is None
+        or local_batch.forward_mode.is_idle()
+        or local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
+    ) and check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
+
+    is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
+    if local_batch is not None:
+        local_batch.is_extend_in_batch = is_extend_in_batch
+
+    tbo_preparer = TboDPAttentionPreparer()
+    if len(offload_tags) == 0 and (
+        disable_overlap_schedule
+        or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get()
+    ):
+        group = tp_group.device_group
+        device = tp_group.device
+    else:
+        group = tp_group.cpu_group
+        device = "cpu"
+
+    local_can_run_tbo, local_forward_mode = tbo_preparer.prepare_all_gather(local_batch)
+
+    mlp_sync_info = MLPSyncBatchInfo(
+        dp_size=dp_size,
+        tp_size=attn_tp_size,
+        cp_size=attn_cp_size,
+        num_tokens=num_tokens,
+        num_tokens_for_logprob=num_tokens_for_logprob,
+        can_cuda_graph=can_cuda_graph,
+        is_extend_in_batch=is_extend_in_batch,
+        local_can_run_tbo=local_can_run_tbo,
+        local_forward_mode=local_forward_mode,
+        can_run_breakable_cuda_graph=can_run_breakable_cuda_graph,
+    )
+
+    if not skip_all_gather:
+        mlp_sync_info.all_gather(device=device, group=group)
+
+        mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = (
+            tbo_preparer.compute_output(
+                mlp_sync_info.tp0_info[:, 4:6],
+            )
+        )
+
+    # Decide whether to emit idle batch
+    if skip_all_gather:
+        # Skip idle batch when attn-dp=1
+        need_idle_batch = dp_size > 1
+    else:
+        need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0
+
+    batch_to_gather = local_batch
+    if need_idle_batch:
+        if local_batch is None:
+            batch_to_gather = local_batch = get_idle_batch()
+        elif local_batch.forward_mode.is_prebuilt():
+            # NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
+            batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
+
+    if batch_to_gather is not None:
+        _update_gather_batch(
+            batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
+        )
+
+    if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
+        local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
+
+    return local_batch
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerDPAttnAdapter:
+    tp_group: GroupCoordinator
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    tree_cache: BasePrefixCache
+    offload_tags: set[str]
+    ps: ParallelState
+    server_args: ServerArgs
+    model_config: ModelConfig
+    enable_overlap: bool
+    spec_algorithm: SpeculativeAlgorithm
+    get_require_mlp_sync: Callable[[], bool]
+
+    def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch):
+        return prepare_mlp_sync_batch_raw(
+            local_batch,
+            dp_size=self.server_args.dp_size,
+            attn_tp_size=self.ps.attn_tp_size,
+            attn_cp_size=self.ps.attn_cp_size,
+            tp_group=self.tp_group,
+            get_idle_batch=self.get_idle_batch,
+            disable_cuda_graph=cuda_graph_fully_disabled(),
+            require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
+            disable_overlap_schedule=self.server_args.disable_overlap_schedule,
+            offload_tags=self.offload_tags,
+        )
+
+    def maybe_prepare_mlp_sync_batch(
+        self,
+        batch: Optional[ScheduleBatch],
+        need_sync: Optional[bool] = None,
+    ) -> Optional[ScheduleBatch]:
+        """
+        Helper to prepare MLP sync batch for DP attention.
+        Should be called after get_new_batch_prefill().
+
+        Args:
+            batch: The batch to process
+            need_sync: If specified, overrides self.get_require_mlp_sync() for prepare_mlp_sync_batch decision
+        """
+        if need_sync if need_sync is not None else self.get_require_mlp_sync():
+            batch = self.prepare_mlp_sync_batch(batch)
+        return batch
+
+    def get_idle_batch(self) -> ScheduleBatch:
+        idle_batch = ScheduleBatch.init_new(
+            [],
+            self.req_to_token_pool,
+            self.token_to_kv_pool_allocator,
+            self.tree_cache,
+            self.model_config,
+            self.enable_overlap,
+            self.spec_algorithm,
+        )
+        idle_batch.prepare_for_idle()
+        return idle_batch

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

@@ -0,0 +1,65 @@
+import logging
+import time
+from typing import Callable, Optional, Tuple
+
+from sglang.srt.managers.io_struct import FlushCacheReqInput, FlushCacheReqOutput
+from sglang.srt.managers.scheduler_components.ipc_channels import (
+    SchedulerIpcChannels,
+)
+
+
+class SchedulerFlushWrapper:
+    def __init__(
+        self,
+        *,
+        flush_cache: Callable[[], bool],
+        is_fully_idle: Callable[[], bool],
+        ipc_channels: SchedulerIpcChannels,
+    ) -> None:
+        self._flush_cache = flush_cache
+        self._is_fully_idle = is_fully_idle
+        self._ipc_channels = ipc_channels
+        self._pending: Optional[Tuple[FlushCacheReqInput, float]] = None
+
+    def handle(self, recv_req: FlushCacheReqInput) -> Optional[FlushCacheReqOutput]:
+        if self._pending is not None:
+            return FlushCacheReqOutput(
+                success=False,
+                message="Another flush_cache is already in progress.",
+            )
+
+        timeout_s = float(recv_req.timeout_s or 0.0)
+        if timeout_s <= 0.0:
+            return FlushCacheReqOutput(success=self._flush_cache())
+
+        if self._is_fully_idle():
+            return FlushCacheReqOutput(success=self._flush_cache())
+
+        self._pending = (recv_req, time.monotonic() + timeout_s)
+        return None
+
+    def check_pending(self) -> None:
+        if self._pending is None:
+            return
+
+        pending_req, deadline = self._pending
+
+        if self._is_fully_idle():
+            success = self._flush_cache()
+            self._pending = None
+            self._ipc_channels.send_to_tokenizer.send_output(
+                FlushCacheReqOutput(success=success), pending_req
+            )
+            return
+
+        if time.monotonic() >= deadline:
+            logging.warning(
+                "Deferred flush_cache timed out while waiting for idle state."
+            )
+            self._pending = None
+            self._ipc_channels.send_to_tokenizer.send_output(
+                FlushCacheReqOutput(
+                    success=False, message="Timed out waiting for idle state."
+                ),
+                pending_req,
+            )

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

@@ -0,0 +1,35 @@
+import zmq
+
+from sglang.srt.environ import envs
+from sglang.srt.observability.req_time_stats import real_time
+from sglang.srt.platforms import current_platform
+
+
+class IdleSleeper:
+    """
+    In setups which have long inactivity periods it is desirable to reduce
+    system power consumption when sglang does nothing. This would lead not only
+    to power savings, but also to more CPU thermal headroom when a request
+    eventually comes. This is important in cases when multiple GPUs are connected
+    as each GPU would otherwise pin one thread at 100% CPU usage.
+
+    The simplest solution is to use zmq.Poller on all sockets that may receive
+    data that needs handling immediately.
+    """
+
+    def __init__(self, sockets):
+        self.poller = zmq.Poller()
+        self.last_empty_time = real_time()
+        for s in sockets:
+            self.poller.register(s, zmq.POLLIN)
+
+        self.empty_cache_interval = envs.SGLANG_EMPTY_CACHE_INTERVAL.get()
+
+    def maybe_sleep(self):
+        self.poller.poll(1000)
+        if (
+            self.empty_cache_interval > 0
+            and real_time() - self.last_empty_time > self.empty_cache_interval
+        ):
+            self.last_empty_time = real_time()
+            current_platform.empty_cache()

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

@@ -0,0 +1,486 @@
+from __future__ import annotations
+
+import logging
+from collections import deque
+from dataclasses import dataclass, field
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Deque,
+    List,
+    Optional,
+    Tuple,
+)
+
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.environ import envs
+from sglang.srt.managers.scheduler_components.pool_stats_observer import (
+    PoolStats,
+    SchedulerPoolStatsObserver,
+)
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.common import (
+    ceil_align,
+    raise_error_or_warn,
+)
+from sglang.srt.utils.watchdog import WatchdogRaw
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+logger = logging.getLogger(__name__)
+
+# Number of recent busy-check messages buffered for the level-1 dump-on-leak path.
+BUSY_MEM_CHECK_LOG_RING_SIZE = 1000
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerInvariantChecker:
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    disaggregation_mode: DisaggregationMode
+    page_size: int
+    full_tokens_per_layer: Optional[int]
+    swa_tokens_per_layer: Optional[int]
+    max_total_num_tokens: int
+    server_args: ServerArgs
+    tree_cache: BasePrefixCache
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    req_to_token_pool: ReqToTokenPool
+    pool_stats_observer: SchedulerPoolStatsObserver
+    get_last_batch: Callable
+    get_running_batch: Callable
+    count_req_pool_leak_warnings: int = 0
+    count_memory_leak_warnings: int = 0
+    recent_busy_msgs: Deque[str] = field(
+        default_factory=lambda: deque(maxlen=BUSY_MEM_CHECK_LOG_RING_SIZE)
+    )
+
+    @staticmethod
+    def _check_pool_invariant(
+        pool_name: str,
+        available: int,
+        evictable: int,
+        protected: int,
+        session_held: int,
+        total: int,
+        uncached: int = 0,
+    ) -> Tuple[bool, str]:
+        """Check: available + evictable + protected + session_held + uncached == total."""
+        total_accounted = available + evictable + protected + session_held + uncached
+        leak = total_accounted != total
+        msg = (
+            f"[{pool_name}] {total=}, {available=}, {evictable=}, "
+            f"{protected=}, {session_held=}, {uncached=}"
+        )
+        return leak, msg
+
+    def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
+        if self.is_hybrid_swa and not self.full_tokens_per_layer:
+            return False, ""
+        if self.is_hybrid_swa:
+            protected = self.tree_cache.full_protected_size()
+            session_held = self.pool_stats_observer.session_held_full_tokens()
+            total = self.full_tokens_per_layer
+        elif self.is_hybrid_ssm:
+            # Branch on cache type for the protected accessor (MambaRadixCache
+            # splits full/mamba; ChunkCache only has the single protected_size).
+            # Use the allocator's `.size` for `total`: static max_total_num_tokens for
+            # non-unified pools, the dynamic byte-coordinated cap (matching
+            # `available_size`) for the unified pool.
+            if self.tree_cache.supports_mamba():
+                protected = self.tree_cache.full_protected_size()
+            else:
+                protected = self.tree_cache.protected_size()
+            session_held = self.pool_stats_observer.session_held_tokens()
+            total = self.token_to_kv_pool_allocator.size
+        else:
+            protected = self.tree_cache.protected_size()
+            session_held = self.pool_stats_observer.session_held_tokens()
+            total = self.max_total_num_tokens
+        full_evictable_size = ps.full_evictable_size
+        allocator = self.token_to_kv_pool_allocator
+        if getattr(self.server_args, "dcp_size", 1) > 1 and allocator.page_size > 1:
+            # DCP stores logical tokens in widened physical pages.  Prefix cache
+            # counters are logical-token based, while the allocator frees whole
+            # physical pages, so round cached tokens up to physical page units.
+            full_evictable_size = (
+                (full_evictable_size + allocator.page_size - 1)
+                // allocator.page_size
+                * allocator.page_size
+            )
+        leak, msg = self._check_pool_invariant(
+            "full",
+            ps.full_available_size,
+            full_evictable_size,
+            protected,
+            session_held,
+            total,
+            uncached,
+        )
+        if (
+            leak
+            and getattr(self.server_args, "dcp_size", 1) > 1
+            and allocator.page_size > 1
+        ):
+            # Radix/Mamba cache accounting is logical-token based while DCP full
+            # KV allocation is physical-page based. Partial physical pages can
+            # leave a small page-level slack even when all pages are owned by
+            # either the allocator or the prefix cache.
+            return False, f"{msg}, dcp_physical_page_slack_allowed=True"
+        return leak, msg
+
+    def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
+        return self._check_pool_invariant(
+            "swa",
+            ps.swa_available_size,
+            ps.swa_evictable_size,
+            self.tree_cache.swa_protected_size(),
+            self.pool_stats_observer.session_held_swa_tokens(),
+            self.swa_tokens_per_layer,
+            uncached,
+        )
+
+    def _check_mamba_pool(self, ps: PoolStats) -> Tuple[bool, str]:
+        ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
+        if ckpt_pool is not None:
+            return self._check_mamba_pool_with_int8(ps, ckpt_pool)
+        leak, msg = self._check_pool_invariant(
+            "mamba",
+            ps.mamba_available_size,
+            ps.mamba_evictable_size,
+            self.tree_cache.mamba_protected_size(),
+            self.pool_stats_observer.session_held_mamba_slots(),
+            self.req_to_token_pool.mamba_pool.size,
+        )
+        if leak:
+            # Page-level leak diagnosis for mamba
+            free_full_pages = set(
+                self.token_to_kv_pool_allocator.free_pages.tolist()
+                + self.token_to_kv_pool_allocator.release_pages.tolist()
+            )
+            cached_full_pages = set(self.tree_cache.all_values_flatten().tolist())
+            expected_full_pages = set(
+                range(1, self.token_to_kv_pool_allocator.size + 1)
+            )
+            leaked_full_pages = (
+                expected_full_pages - free_full_pages - cached_full_pages
+            )
+            mamba_allocator = self.req_to_token_pool.mamba_allocator
+            free_mamba_pages = set(mamba_allocator.free_slots.tolist())
+            cached_mamba_pages = set(
+                self.tree_cache.all_mamba_values_flatten().tolist()
+            )
+            expected_mamba_pages = set(range(1, mamba_allocator.size + 1))
+            leaked_mamba_pages = (
+                expected_mamba_pages - free_mamba_pages - cached_mamba_pages
+            )
+            msg += (
+                f", leaked_full_pages={leaked_full_pages or None}"
+                f", leaked_mamba_pages={leaked_mamba_pages or None}"
+            )
+        return leak, msg
+
+    def _check_mamba_pool_with_int8(self, ps: PoolStats, ckpt_pool) -> Tuple[bool, str]:
+        """Two-pool invariant for int8 mamba checkpoints.
+
+        The radix-cached states live in the int8 checkpoint pool, NOT the active
+        bf16 pool. So the single-pool equation (active.available + radix_cached ==
+        active.size) is wrong -- it double-counts the radix states against a pool
+        that does not hold them. Instead check the two pools independently:
+
+          * active bf16 pool: backs running requests only; the radix owns ZERO
+            active slots. Checked at idle (in-flight == 0) -> available == total.
+          * int8 checkpoint pool: backs the radix-cached states; its occupancy is
+            exactly the radix evictable + protected counts.
+        """
+        active_leak, active_msg = self._check_pool_invariant(
+            "mamba-active",
+            ps.mamba_available_size,
+            ps.mamba_evictable_size,  # 0 in int8 mode (radix owns no active slots)
+            0,
+            self.pool_stats_observer.session_held_mamba_slots(),
+            self.req_to_token_pool.mamba_pool.size,
+        )
+        int8_leak, int8_msg = self._check_pool_invariant(
+            "mamba-int8",
+            ckpt_pool.available_size(),
+            self.tree_cache.mamba_evictable_size(),
+            self.tree_cache.mamba_protected_size(),
+            0,
+            ckpt_pool.num_slots,
+        )
+        return active_leak or int8_leak, active_msg + "\n" + int8_msg
+
+    def _get_total_uncached_sizes(
+        self,
+    ) -> Tuple[int, int]:
+        """Sum uncached tokens for full and SWA pools across all active batches.
+
+        Returns (full_uncached, swa_uncached). For non-SWA models, swa_uncached is 0.
+
+        For full pool: uncached = allocated - cache_protected_len
+        For SWA pool:  uncached = allocated - max(cache_protected_len, swa_evicted_seqlen)
+        """
+        # After decode: running_batch IS last_batch (same object), count once.
+        # After prefill: they differ, both hold uncached tokens.
+        # Use identity (is / is not), not membership or ==: ScheduleBatch's
+        # dataclass __eq__ compares tensor fields and raises on ambiguous bools.
+        last_batch = self.get_last_batch()
+        running_batch = self.get_running_batch()
+        batches = [last_batch]
+        if (
+            running_batch is not None
+            and running_batch is not last_batch
+            and not running_batch.is_empty()
+        ):
+            batches.append(running_batch)
+
+        full_uncached = 0
+        swa_uncached = 0
+        for batch in batches:
+            for req in batch.reqs:
+                assert req.kv_committed_freed == req.kv_overallocated_freed
+                if req.kv_committed_freed or req.req_pool_idx is None:
+                    continue
+
+                allocated_len = req.kv_allocated_len
+                if self.page_size > 1:
+                    allocated_len = ceil_align(allocated_len, self.page_size)
+                    assert req.cache_protected_len % self.page_size == 0
+
+                full_uncached += allocated_len - req.cache_protected_len
+                if self.is_hybrid_swa:
+                    swa_uncached += allocated_len - max(
+                        req.cache_protected_len, req.swa_evicted_seqlen
+                    )
+
+        return full_uncached, swa_uncached
+
+    def self_check_during_busy(self):
+        if self.get_last_batch() is None:
+            return
+
+        ps = self.pool_stats_observer.get_pool_stats()
+        full_uncached, swa_uncached = self._get_total_uncached_sizes()
+
+        full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached)
+
+        swa_leak, swa_msg = False, ""
+        if self.is_hybrid_swa:
+            swa_leak, swa_msg = self._check_swa_pool(ps, uncached=swa_uncached)
+
+        level = envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get()
+        full_line = f"[Mem Check (BUSY)] {full_msg}"
+        swa_line = f"[Mem Check (BUSY)] {swa_msg}" if swa_msg else None
+
+        if level > 1:
+            # Verbose: log every iteration.
+            logger.info(full_line)
+            if swa_line:
+                logger.info(swa_line)
+        elif level == 1:
+            # Quiet: buffer and stay silent; flush the recent ones only on a leak.
+            self.recent_busy_msgs.append(full_line)
+            if swa_line:
+                self.recent_busy_msgs.append(swa_line)
+            if full_leak or swa_leak:
+                for msg in self.recent_busy_msgs:
+                    logger.info(msg)
+
+        assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}"
+        assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}"
+
+        if envs.SGLANG_CHECK_KV_PAGE_INVARIANTS.get():
+            self._check_kv_page_invariants()
+
+    def _check_kv_page_invariants(self):
+        """committed<=allocated for every req/slot, and no double free:
+          A. no owner references a page that is in the free pool (use-after-free).
+          B. the free pool has no duplicate pages (two owners freed the same page).
+        All heavy work runs on GPU to avoid per-token device->host sync."""
+        rtt = self.req_to_token_pool.req_to_token
+        row_width = rtt.shape[1]
+
+        def _add_owner(req_or_slot, label, rpi, committed, allocated):
+            assert 0 <= committed <= allocated <= row_width
+            owners.append((label, rpi, allocated))
+
+        owners: list[tuple[str, Optional[int], int]] = []
+        batch = self.get_last_batch()
+        if batch is not None:
+            for req in batch.reqs:
+                _add_owner(
+                    req,
+                    f"req {req.rid}",
+                    req.req_pool_idx,
+                    req.kv_committed_len,
+                    req.kv_allocated_len,
+                )
+        sess = getattr(self.tree_cache, "slots", None)
+        if sess:
+            for sid, slot in sess.items():
+                if getattr(slot, "is_holding_kv", False):
+                    _add_owner(
+                        slot,
+                        f"slot {sid[:8]}",
+                        slot.req_pool_idx,
+                        slot.kv_committed_len,
+                        slot.kv_allocated_len,
+                    )
+
+        active = [
+            (label, rpi, al) for label, rpi, al in owners if rpi is not None and al > 0
+        ]
+        if not active:
+            return
+
+        idx = torch.as_tensor([rpi for _, rpi, _ in active], device=rtt.device)
+        allocs = torch.as_tensor([al for _, _, al in active], device=rtt.device)
+        mask = torch.arange(row_width, device=rtt.device)[None, :] < allocs[:, None]
+        owner_pages = rtt[idx][mask] // self.page_size
+
+        # Sub-allocators to check: a flat allocator is its own single sub; a
+        # hybrid-SWA wrapper exposes full_attn_allocator + swa_attn_allocator.
+        alloc = self.token_to_kv_pool_allocator
+        sub_allocs = (
+            [alloc]
+            if getattr(alloc, "free_pages", None) is not None
+            else [
+                sub
+                for n in ("full_attn_allocator", "swa_attn_allocator")
+                if (sub := getattr(alloc, n, None)) is not None
+                and getattr(sub, "free_pages", None) is not None
+            ]
+        )
+        if not sub_allocs:
+            return
+
+        def _free_pages(a):
+            free = a.free_pages
+            release = getattr(a, "release_pages", None)
+            return (
+                torch.cat((free, release))
+                if release is not None and len(release) > 0
+                else free
+            )
+
+        # Check B: every sub-pool's free set has no duplicate pages.
+        for i, sub in enumerate(sub_allocs):
+            free = _free_pages(sub)
+            uniq = torch.unique(free)
+            if uniq.numel() != free.numel():
+                raise_error_or_warn(
+                    self,
+                    envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+                    "count_memory_leak_warnings",
+                    f"KV double free: sub-pool {i} has {free.numel() - uniq.numel()} duplicate pages.",
+                )
+
+        # Check A: owner pages (full-pool indices) must not be in the full free
+        # set (sub_allocs[0] is the full pool, even on hybrid-SWA).
+        full_unique = torch.unique(_free_pages(sub_allocs[0]))
+        stale = owner_pages[torch.isin(owner_pages, full_unique)]
+        if stale.numel() > 0:
+            raise_error_or_warn(
+                self,
+                envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+                "count_memory_leak_warnings",
+                f"KV page use-after-free: {stale.numel()} owner page refs are in "
+                f"the free pool, sample pages={torch.unique(stale)[:8].tolist()}.",
+            )
+
+    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
+
+        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:
+            msg = (
+                "req_to_token_pool memory leak detected!"
+                f"available_size={len(self.req_to_token_pool.free_slots)}, "
+                f"session_held={session_req_count}, "
+                f"total_size={self.req_to_token_pool.size}\n"
+            )
+            raise_error_or_warn(
+                self,
+                envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+                "count_req_pool_leak_warnings",
+                msg,
+            )
+
+    def _report_leak(self, pool_name: str, token_msg: str):
+        msg = f"{pool_name} memory leak detected! {token_msg}"
+        raise_error_or_warn(
+            self,
+            envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
+            "count_memory_leak_warnings",
+            msg,
+        )
+
+    def _check_all_pools(
+        self, ps: PoolStats, uncached: int = 0
+    ) -> Tuple[bool, List[str]]:
+        """Check memory invariant across all pools. Returns (has_leak, messages)."""
+        has_leak = False
+        messages = []
+
+        full_leak, full_msg = self._check_full_pool(ps, uncached=uncached)
+        has_leak |= full_leak
+        messages.append(full_msg)
+
+        if self.is_hybrid_swa:
+            swa_leak, swa_msg = self._check_swa_pool(ps)
+            has_leak |= swa_leak
+            messages.append(swa_msg)
+
+        if self.is_hybrid_ssm and self.tree_cache.supports_mamba():
+            mamba_leak, mamba_msg = self._check_mamba_pool(ps)
+            has_leak |= mamba_leak
+            messages.append(mamba_msg)
+
+        return has_leak, messages
+
+    def _check_tree_cache(self):
+        if (
+            self.tree_cache.is_tree_cache()
+            and (self.is_hybrid_swa and self.tree_cache.supports_swa())
+            or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
+        ):
+            self.tree_cache.sanity_check()
+
+
+def create_scheduler_watchdog(
+    scheduler: Scheduler, watchdog_timeout: float, soft: bool = False
+) -> WatchdogRaw:
+    def dump_info() -> str:
+        if scheduler.is_initializing:
+            return ""
+        _, messages = scheduler.invariant_checker._check_all_pools(
+            scheduler.pool_stats_observer.get_pool_stats(),
+        )
+        return (
+            f"{scheduler.cur_batch_for_debug.batch_size()=}\n"
+            f"{scheduler.cur_batch_for_debug.reqs=}\n" + "\n".join(messages)
+        )
+
+    return WatchdogRaw(
+        debug_name="Scheduler",
+        get_counter=lambda: scheduler.forward_ct,
+        is_active=lambda: (
+            scheduler.is_initializing or scheduler.cur_batch_for_debug is not None
+        ),
+        watchdog_timeout=watchdog_timeout,
+        soft=soft,
+        dump_info=dump_info,
+    )

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

@@ -0,0 +1,87 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import zmq
+
+from sglang.srt.managers.scheduler_components.output_sender import SenderWrapper
+from sglang.srt.server_args import PortArgs
+from sglang.srt.utils.network import get_zmq_socket
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SchedulerIpcChannels:
+    recv_from_tokenizer: Union[zmq.Socket, "ScriptedTokenizerRecvProxy"]
+    recv_from_rpc: Optional[zmq.Socket]
+    send_to_tokenizer: SenderWrapper
+    send_to_detokenizer: SenderWrapper
+    send_metrics_from_scheduler: Optional[zmq.Socket]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        port_args: PortArgs,
+        is_rank_zero: bool,
+        skip_tokenizer_init: bool,
+        metrics_enabled: bool,
+        enable_scripted_runtime: bool,
+    ) -> "SchedulerIpcChannels":
+        context = zmq.Context(2)
+
+        if is_rank_zero:
+            recv_from_tokenizer = get_zmq_socket(
+                context, zmq.PULL, port_args.scheduler_input_ipc_name, False
+            )
+            if enable_scripted_runtime:
+                from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+                    ScriptedTokenizerRecvProxy,
+                )
+
+                recv_from_tokenizer = ScriptedTokenizerRecvProxy(
+                    underlying=recv_from_tokenizer
+                )
+            recv_from_rpc = get_zmq_socket(
+                context, zmq.DEALER, port_args.rpc_ipc_name, False
+            )
+
+            send_to_tokenizer_raw = get_zmq_socket(
+                context, zmq.PUSH, port_args.tokenizer_ipc_name, False
+            )
+            if skip_tokenizer_init:
+                # Directly send to the TokenizerManager
+                send_to_detokenizer_raw = get_zmq_socket(
+                    context, zmq.PUSH, port_args.tokenizer_ipc_name, False
+                )
+            else:
+                # Send to the DetokenizerManager
+                send_to_detokenizer_raw = get_zmq_socket(
+                    context, zmq.PUSH, port_args.detokenizer_ipc_name, False
+                )
+
+            send_to_tokenizer = SenderWrapper(send_to_tokenizer_raw)
+            send_to_detokenizer = SenderWrapper(send_to_detokenizer_raw)
+        else:
+            recv_from_tokenizer = None
+            recv_from_rpc = None
+            send_to_tokenizer = SenderWrapper(None)
+            send_to_detokenizer = SenderWrapper(None)
+
+        if metrics_enabled:
+            send_metrics_from_scheduler = get_zmq_socket(
+                context, zmq.PUSH, port_args.metrics_ipc_name, False
+            )
+        else:
+            send_metrics_from_scheduler = None
+
+        return cls(
+            recv_from_tokenizer=recv_from_tokenizer,
+            recv_from_rpc=recv_from_rpc,
+            send_to_tokenizer=send_to_tokenizer,
+            send_to_detokenizer=send_to_detokenizer,
+            send_metrics_from_scheduler=send_metrics_from_scheduler,
+        )

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

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Optional,
+)
+
+import msgspec
+import zmq
+
+from sglang.srt.disaggregation.kv_events import (
+    EventPublisherFactory,
+    KVEventBatch,
+    select_kv_publisher_dp_rank,
+)
+from sglang.srt.managers.io_struct import hook_custom_types, sock_send
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+
+
+class SchedulerStats: ...  # type: ignore[no-redef]
+
+
+class KvMetrics(msgspec.Struct, tag=True, kw_only=True, array_like=True):
+    request_active_slots: int = 0
+    request_total_slots: int = 0
+    kv_active_blocks: int = 0
+    kv_total_blocks: int = 0
+    num_requests_waiting: int = 0
+    gpu_cache_usage_perc: float = 0.0
+    gpu_prefix_cache_hit_rate: float = 0.0
+    data_parallel_rank: int = 0
+
+
+hook_custom_types(KvMetrics)
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerKvEventsPublisher:
+    kv_events_config: Optional[str]
+    ps: ParallelState
+    attn_tp_rank: int
+    attn_cp_rank: int
+    attn_dp_rank: int
+    dp_rank: Optional[int]
+    tree_cache: BasePrefixCache
+    send_metrics_from_scheduler: Optional[zmq.Socket]
+    max_running_requests: int
+    max_total_num_tokens: int
+    get_stats: Callable
+    enable_kv_cache_events: bool = False
+    kv_event_publisher: Any = None
+
+    def __post_init__(self) -> None:
+        self.init_kv_events(self.kv_events_config)
+
+    def init_kv_events(self, kv_events_config: Optional[str]):
+        self.enable_kv_cache_events = bool(
+            kv_events_config
+            and self.ps.pp_rank == 0
+            and self.ps.attn_tp_rank == 0
+            and self.ps.attn_cp_rank == 0
+        )
+
+        if self.enable_kv_cache_events:
+            self.kv_event_publisher = EventPublisherFactory.create(
+                kv_events_config,
+                select_kv_publisher_dp_rank(
+                    self.ps.attn_dp_size, self.ps.attn_dp_rank, self.ps.dp_rank
+                ),
+            )
+
+    def emit_kv_metrics(self):
+        if not self.enable_kv_cache_events:
+            return
+
+        kv_metrics = KvMetrics()
+        kv_metrics.request_active_slots = self.get_stats().num_running_reqs.total
+        kv_metrics.request_total_slots = self.max_running_requests
+        kv_metrics.kv_active_blocks = int(
+            self.get_stats().token_usage * self.max_total_num_tokens
+        )
+        kv_metrics.kv_total_blocks = self.max_total_num_tokens
+        kv_metrics.num_requests_waiting = self.get_stats().num_queue_reqs.total
+        kv_metrics.gpu_cache_usage_perc = self.get_stats().token_usage
+        kv_metrics.gpu_prefix_cache_hit_rate = self.get_stats().cache_hit_rate
+        kv_metrics.data_parallel_rank = (
+            self.ps.dp_rank if self.ps.dp_rank is not None else 0
+        )
+
+        if not self.send_metrics_from_scheduler.closed:
+            sock_send(self.send_metrics_from_scheduler, kv_metrics)
+
+    def publish_kv_events(self):
+        if not self.enable_kv_cache_events:
+            return
+
+        events = self.tree_cache.take_events()
+        if events:
+            batch = KVEventBatch(ts=time.time(), events=events)
+            self.kv_event_publisher.publish(batch)

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

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.load_snapshot import (
+    DisaggregationMetrics,
+    LoadSnapshot,
+    LoRAMetrics,
+    MemoryMetrics,
+    QueueMetrics,
+    SpeculativeMetrics,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.managers.scheduler_components.pool_stats_observer import (
+        SchedulerPoolStatsObserver,
+    )
+    from sglang.srt.managers.tp_worker import BaseTpWorker
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerLoadInquirer:
+    disaggregation_mode: DisaggregationMode
+    ps: ParallelState
+    server_args: ServerArgs
+    max_total_num_tokens: int
+    max_running_requests: int
+    pool_stats_observer: SchedulerPoolStatsObserver
+    tp_worker: BaseTpWorker
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    spec_algorithm: SpeculativeAlgorithm
+    get_running_batch: Callable
+    get_waiting_queue: Callable
+    get_stats: Callable
+    get_chunked_req: Callable
+    get_disagg_prefill_bootstrap_queue: Callable
+    get_disagg_prefill_inflight_queue: Callable
+    get_disagg_decode_prealloc_queue: Callable
+    get_disagg_decode_transfer_queue: Callable
+    get_spec_total_num_accept_tokens: Callable
+    get_spec_total_num_forward_ct: Callable
+
+    def _get_num_pending_tokens(self, chunk_deduct: int = 0) -> int:
+        """Get the total number of tokens pending prefill.
+
+        This includes tokens from waiting queue requests plus remaining tokens
+        from the currently chunked request.
+
+        Args:
+            chunk_deduct: extra tokens to subtract from the chunked request's
+                remaining count. At batch-scheduling time the current chunk
+                has been planned but ``prefix_indices`` does not yet include it,
+                so callers pass ``extend_input_len`` here. At load-reporting
+                time ``prefix_indices`` is already up-to-date, so the default
+                0 is correct.
+        """
+        num_pending_tokens = sum(req.seqlen for req in self.get_waiting_queue())
+        if self.get_chunked_req() is not None:
+            req = self.get_chunked_req()
+            num_pending_tokens += req.seqlen - len(req.prefix_indices) - chunk_deduct
+        return num_pending_tokens
+
+    def get_num_waiting_uncached_tokens(self) -> int:
+        """Get uncached input tokens waiting for prefill compute."""
+        if self.disaggregation_mode == DisaggregationMode.DECODE:
+            return 0
+        num_tokens = 0
+        for req in self.get_waiting_queue():
+            # if match-in-waiting-queue disabled, this metric returns seq_lens
+            num_tokens += max(0, req.seqlen - req.num_matched_prefix_tokens)
+        cr = self.get_chunked_req()
+        if cr is not None:
+            num_tokens += max(0, cr.seqlen - len(cr.prefix_indices))
+        return num_tokens
+
+    def get_loads(self) -> LoadSnapshot:
+        """Build the per-DP-rank load snapshot for DP balancing and /v1/loads."""
+        stats = self.get_stats()
+        num_running_reqs = len(self.get_running_batch().reqs)
+
+        waiting_queues = [self.get_waiting_queue()]
+        pending_token_queues = [self.get_waiting_queue()]
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            prefill_bootstrap_queue = self.get_disagg_prefill_bootstrap_queue().queue
+            waiting_queues.append(prefill_bootstrap_queue)
+            pending_token_queues.append(prefill_bootstrap_queue)
+        elif self.disaggregation_mode == DisaggregationMode.DECODE:
+            decode_prealloc_queue = self.get_disagg_decode_prealloc_queue().queue
+            decode_transfer_queue = self.get_disagg_decode_transfer_queue().queue
+            decode_retracted_queue = (
+                self.get_disagg_decode_prealloc_queue().retracted_queue
+            )
+            waiting_queues.append(decode_prealloc_queue)
+            waiting_queues.append(decode_transfer_queue)
+            waiting_queues.append(decode_retracted_queue)
+            # In disaggregated decode, transfer-queue requests and transferred
+            # waiting-queue requests have already pre-allocated decode-side KV
+            # slots, so they are already included in num_used_tokens.
+            pending_token_queues = [decode_prealloc_queue, decode_retracted_queue]
+
+        num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
+        num_used_tokens, kv_token_usage = (
+            self.pool_stats_observer.get_pool_stats().get_kv_token_stats()
+        )
+        num_total_tokens = num_used_tokens + sum(
+            req.seqlen for queue in pending_token_queues for req in queue
+        )
+
+        memory = None
+        try:
+            memory = MemoryMetrics(
+                weight_gb=round(self.tp_worker.model_runner.weight_load_mem_usage, 3),
+                kv_cache_gb=round(
+                    self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 3
+                ),
+                graph_gb=round(self.tp_worker.model_runner.graph_mem_usage, 3),
+                token_capacity=int(self.max_total_num_tokens),
+            )
+        except (AttributeError, TypeError) as e:
+            logger.debug(f"Memory metrics not available: {e}")
+
+        speculative = None
+        if (
+            not self.spec_algorithm.is_none()
+            and self.get_spec_total_num_forward_ct() > 0
+        ):
+            speculative = SpeculativeMetrics(
+                accept_length=(
+                    self.get_spec_total_num_accept_tokens()
+                    / self.get_spec_total_num_forward_ct()
+                ),
+                accept_rate=stats.spec_accept_rate,
+            )
+
+        lora = None
+        if self.server_args.enable_lora:
+            lora = LoRAMetrics(
+                slots_used=stats.lora_pool_slots_used,
+                slots_total=stats.lora_pool_slots_total,
+                utilization=stats.lora_pool_utilization,
+            )
+
+        mode_str = "null"
+        prefill_bootstrap = prefill_inflight = 0
+        decode_prealloc = decode_transfer = decode_retracted = 0
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            mode_str = "prefill"
+            prefill_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue)
+            prefill_inflight = len(self.get_disagg_prefill_inflight_queue())
+        elif self.disaggregation_mode == DisaggregationMode.DECODE:
+            mode_str = "decode"
+            decode_prealloc = len(self.get_disagg_decode_prealloc_queue().queue)
+            decode_transfer = len(self.get_disagg_decode_transfer_queue().queue)
+            decode_retracted = len(
+                self.get_disagg_decode_prealloc_queue().retracted_queue
+            )
+        disaggregation = DisaggregationMetrics(
+            mode=mode_str,
+            prefill_bootstrap_queue_reqs=prefill_bootstrap,
+            prefill_inflight_queue_reqs=prefill_inflight,
+            decode_prealloc_queue_reqs=decode_prealloc,
+            decode_transfer_queue_reqs=decode_transfer,
+            decode_retracted_queue_reqs=decode_retracted,
+            kv_transfer_speed_gb_s=stats.kv_transfer_speed_gb_s,
+            kv_transfer_latency_ms=stats.kv_transfer_latency_ms,
+        )
+
+        queues = QueueMetrics(
+            waiting=len(self.get_waiting_queue()),
+            grammar=stats.num_grammar_queue_reqs,
+            paused=stats.num_paused_reqs,
+            retracted=stats.num_retracted_reqs,
+        )
+
+        return LoadSnapshot(
+            dp_rank=int(self.ps.dp_rank) if self.ps.dp_rank is not None else 0,
+            timestamp=time.time(),
+            num_running_reqs=num_running_reqs,
+            num_waiting_reqs=num_waiting_reqs,
+            num_waiting_uncached_tokens=self.get_num_waiting_uncached_tokens(),
+            num_used_tokens=num_used_tokens,
+            num_total_tokens=num_total_tokens,
+            max_total_num_tokens=self.max_total_num_tokens,
+            max_running_requests=self.max_running_requests,
+            token_usage=round(kv_token_usage, 4),
+            gen_throughput=round(stats.gen_throughput, 2),
+            cache_hit_rate=round(stats.cache_hit_rate, 4),
+            utilization=round(stats.utilization, 4),
+            memory=memory,
+            speculative=speculative,
+            lora=lora,
+            disaggregation=disaggregation,
+            queues=queues,
+        )

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

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import (
+    List,
+    Tuple,
+)
+
+import torch
+
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.server_args import (
+    MIS_DELIMITER_TOKEN_ID,
+    ServerArgs,
+)
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerLogprobResultProcessor:
+    server_args: ServerArgs
+    model_config: ModelConfig
+
+    def _process_input_token_logprobs(
+        self, req: Req, input_token_logprobs: List
+    ) -> None:
+        """Process input token logprobs values and indices."""
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Process logprob values - handle multi-item scoring vs regular requests
+        if is_multi_item_scoring:
+            # Multi-item scoring: use all logprobs as-is
+            req.logprob.input_token_logprobs_val = input_token_logprobs
+        else:
+            # Regular request: add None at start, remove last (sampling token)
+            req.logprob.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
+
+        # Process logprob indices based on scoring type
+        if is_multi_item_scoring:
+            # MIS scores come from input_token_ids_logprobs, not input_token_logprobs.
+            # But the shared pipeline requires input_token_logprobs_idx to be the same
+            # length as input_token_logprobs_val (validated at line 816). We fill with
+            # MIS_DELIMITER_TOKEN_ID as a dummy — score_request() ignores this field.
+            delimiter_count = len(req.multi_item_delimiter_indices)
+            input_token_logprobs_idx = [MIS_DELIMITER_TOKEN_ID] * delimiter_count
+        else:
+            # Regular request: include all tokens from logprob_start_len onwards
+            input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
+
+        # Clip padded hash values from image tokens to prevent detokenization errors
+        req.logprob.input_token_logprobs_idx = [
+            x if x < self.model_config.vocab_size - 1 else 0
+            for x in input_token_logprobs_idx
+        ]
+
+    def _process_input_top_logprobs(self, req: Req) -> None:
+        """Process input top logprobs."""
+        if req.logprob.top_logprobs_num <= 0:
+            return
+
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Initialize arrays - multi-item scoring starts empty, others start with None
+        req.logprob.input_top_logprobs_val = [] if is_multi_item_scoring else [None]
+        req.logprob.input_top_logprobs_idx = [] if is_multi_item_scoring else [None]
+
+        # Extend arrays with temp values
+        for val, idx in zip(
+            req.temp_input_top_logprobs_val,
+            req.temp_input_top_logprobs_idx,
+            strict=True,
+        ):
+            req.logprob.input_top_logprobs_val.extend(val)
+            req.logprob.input_top_logprobs_idx.extend(idx)
+
+        # Remove last token (sampling token) for non multi-item scoring requests
+        if not is_multi_item_scoring:
+            req.logprob.input_top_logprobs_val.pop()
+            req.logprob.input_top_logprobs_idx.pop()
+
+        # Clean up temp storage
+        req.temp_input_top_logprobs_idx = None
+        req.temp_input_top_logprobs_val = None
+
+    def _process_input_token_ids_logprobs(self, req: Req) -> None:
+        """Process input token IDs logprobs."""
+        if req.logprob.token_ids_logprob is None:
+            return
+
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        # Initialize arrays - multi-item scoring starts empty, others start with None
+        req.logprob.input_token_ids_logprobs_val = (
+            [] if is_multi_item_scoring else [None]
+        )
+        req.logprob.input_token_ids_logprobs_idx = (
+            [] if is_multi_item_scoring else [None]
+        )
+
+        # Process temp values - convert tensors to lists and extend arrays
+        for val, idx in zip(
+            req.temp_input_token_ids_logprobs_val,
+            req.temp_input_token_ids_logprobs_idx,
+            strict=True,
+        ):
+            val_list = val.tolist() if isinstance(val, torch.Tensor) else val
+            req.logprob.input_token_ids_logprobs_val.extend(
+                val_list if isinstance(val_list, list) else [val_list]
+            )
+            req.logprob.input_token_ids_logprobs_idx.extend(idx)
+
+        # Remove last token (sampling token) for non multi-item scoring requests
+        if not is_multi_item_scoring:
+            req.logprob.input_token_ids_logprobs_val.pop()
+            req.logprob.input_token_ids_logprobs_idx.pop()
+
+        # Clean up temp storage
+        req.temp_input_token_ids_logprobs_idx = None
+        req.temp_input_token_ids_logprobs_val = None
+
+    def _calculate_relevant_tokens_len(self, req: Req) -> int:
+        """Calculate the expected length of logprob arrays based on whether multi-item scoring is enabled.
+
+        For multi-item scoring, only delimiter positions have logprobs.
+        For regular requests, all positions from logprob_start_len onwards have logprobs.
+        """
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        if is_multi_item_scoring:
+            return len(req.multi_item_delimiter_indices)
+        else:
+            return len(req.origin_input_ids[req.logprob_start_len :])
+
+    def calculate_num_input_logprobs(
+        self,
+        req: Req,
+        extend_input_len: int,
+        extend_logprob_start_len: int,
+    ) -> int:
+        """Calculate the number of input logprobs based on whether multi-item scoring is enabled.
+
+        For multi-item scoring, only delimiter positions have logprobs.
+        For regular requests, all positions in the range have logprobs.
+        """
+        is_multi_item_scoring = self._is_multi_item_scoring(req)
+
+        if is_multi_item_scoring:
+            # Count pre-computed delimiter indices within the extend range
+            return sum(
+                1
+                for idx in req.multi_item_delimiter_indices
+                if extend_logprob_start_len <= idx < extend_input_len
+            )
+        else:
+            # Regular request: all tokens in the range
+            return extend_input_len - extend_logprob_start_len
+
+    def _is_multi_item_scoring(self, req: Req) -> bool:
+        """Check if request uses multi-item scoring.
+
+        Multi-item scoring applies to prefill-only requests when a delimiter
+        token is configured. In this mode, only positions containing the
+        delimiter token receive logprobs.
+        """
+        return (
+            self.server_args.enable_mis
+            and req.is_prefill_only
+            and req.multi_item_delimiter_indices is not None
+        )
+
+    def add_input_logprob_return_values(
+        self,
+        i: int,
+        req: Req,
+        output: LogitsProcessorOutput,
+        logprob_pt: int,
+        num_input_logprobs: int,
+        last_prefill_chunk: bool,  # If True, it means prefill is finished.
+    ):
+        """Incrementally add input logprobs to `req`.
+
+        Args:
+            i: The request index in a batch.
+            req: The request. Input logprobs inside req are modified as a
+                consequence of the API
+            logprob_pt: Pointer into the prefill ids processed.
+            output: Logit processor output that's used to compute input logprobs
+            last_prefill_chunk: True if it is the last prefill (when chunked).
+                Some of input logprob operation should only happen at the last
+                prefill (e.g., computing input token logprobs).
+        """
+        assert output.input_token_logprobs is not None
+        if req.input_token_logprobs is None:
+            req.input_token_logprobs = []
+        if req.temp_input_top_logprobs_val is None:
+            req.temp_input_top_logprobs_val = []
+        if req.temp_input_top_logprobs_idx is None:
+            req.temp_input_top_logprobs_idx = []
+        if req.temp_input_token_ids_logprobs_val is None:
+            req.temp_input_token_ids_logprobs_val = []
+        if req.temp_input_token_ids_logprobs_idx is None:
+            req.temp_input_token_ids_logprobs_idx = []
+
+        if req.logprob.input_token_logprobs_val is not None:
+            # The input logprob has been already computed. It only happens
+            # upon retract.
+            if req.logprob.top_logprobs_num > 0:
+                assert req.logprob.input_token_logprobs_val is not None
+            return
+
+        # Important for the performance.
+        assert isinstance(output.input_token_logprobs, tuple)
+        input_token_logprobs: Tuple[int] = output.input_token_logprobs
+        input_token_logprobs = input_token_logprobs[
+            logprob_pt : logprob_pt + num_input_logprobs
+        ]
+        req.input_token_logprobs.extend(input_token_logprobs)
+
+        if req.logprob.top_logprobs_num > 0:
+            req.temp_input_top_logprobs_val.append(output.input_top_logprobs_val[i])
+            req.temp_input_top_logprobs_idx.append(output.input_top_logprobs_idx[i])
+
+        if req.logprob.token_ids_logprob is not None:
+            req.temp_input_token_ids_logprobs_val.append(
+                output.input_token_ids_logprobs_val[i]
+            )
+            req.temp_input_token_ids_logprobs_idx.append(
+                output.input_token_ids_logprobs_idx[i]
+            )
+
+        if last_prefill_chunk:
+            input_token_logprobs = req.input_token_logprobs
+            req.input_token_logprobs = None
+            assert req.logprob.input_token_logprobs_val is None
+            assert req.logprob.input_token_logprobs_idx is None
+            assert req.logprob.input_top_logprobs_val is None
+            assert req.logprob.input_top_logprobs_idx is None
+
+            # Process all input logprob types using helper functions
+            self._process_input_token_logprobs(req, input_token_logprobs)
+            self._process_input_top_logprobs(req)
+
+            self._process_input_token_ids_logprobs(req)
+
+            if req.return_logprob:
+                relevant_tokens_len = self._calculate_relevant_tokens_len(req)
+                assert len(req.logprob.input_token_logprobs_val) == relevant_tokens_len
+                assert len(req.logprob.input_token_logprobs_idx) == relevant_tokens_len
+                if req.logprob.top_logprobs_num > 0:
+                    assert (
+                        len(req.logprob.input_top_logprobs_val) == relevant_tokens_len
+                    )
+                    assert (
+                        len(req.logprob.input_top_logprobs_idx) == relevant_tokens_len
+                    )
+                if req.logprob.token_ids_logprob is not None:
+                    assert (
+                        len(req.logprob.input_token_ids_logprobs_val)
+                        == relevant_tokens_len
+                    )
+                    assert (
+                        len(req.logprob.input_token_ids_logprobs_idx)
+                        == relevant_tokens_len
+                    )
+
+    def add_logprob_return_values(
+        self,
+        i: int,
+        req: Req,
+        pt: int,
+        next_token_ids: List[int],
+        num_input_logprobs: int,
+        output: LogitsProcessorOutput,
+    ):
+        """Attach logprobs to the return values."""
+        if output.next_token_logprobs is not None:
+            req.logprob.output_token_logprobs_val.append(output.next_token_logprobs[i])
+            req.logprob.output_token_logprobs_idx.append(next_token_ids[i])
+
+        # Only add input logprobs if there are input tokens to process
+        # Note: For prefill-only requests with default logprob_start_len, this will be 0,
+        # meaning we only compute output logprobs (which is the intended behavior)
+        if num_input_logprobs > 0:
+            self.add_input_logprob_return_values(
+                i,
+                req,
+                output,
+                pt,
+                num_input_logprobs,
+                last_prefill_chunk=True,
+            )
+        else:
+            self._initialize_empty_logprob_containers(req)
+
+        if req.logprob.top_logprobs_num > 0:
+            req.logprob.output_top_logprobs_val.append(
+                output.next_token_top_logprobs_val[i]
+            )
+            req.logprob.output_top_logprobs_idx.append(
+                output.next_token_top_logprobs_idx[i]
+            )
+
+        if (
+            req.logprob.token_ids_logprob is not None
+            and output.next_token_token_ids_logprobs_val is not None
+        ):
+            # Convert GPU tensor to list if needed
+            logprobs_val = output.next_token_token_ids_logprobs_val[i]
+            if isinstance(logprobs_val, torch.Tensor):
+                logprobs_val = logprobs_val.tolist()
+            req.logprob.output_token_ids_logprobs_val.append(logprobs_val)
+            req.logprob.output_token_ids_logprobs_idx.append(
+                output.next_token_token_ids_logprobs_idx[i]
+            )
+
+        return num_input_logprobs
+
+    def _initialize_empty_logprob_containers(self, req: Req) -> None:
+        """
+        Initialize logprob fields to empty lists if unset.
+
+        This is needed for prefill-only requests where the normal initialization
+        flow might be bypassed, but downstream code expects these fields to be lists.
+        """
+        if req.logprob.input_token_logprobs_val is None:
+            req.logprob.input_token_logprobs_val = []
+        if req.logprob.input_token_logprobs_idx is None:
+            req.logprob.input_token_logprobs_idx = []
+        if req.logprob.input_top_logprobs_val is None:
+            req.logprob.input_top_logprobs_val = []
+        if req.logprob.input_top_logprobs_idx is None:
+            req.logprob.input_top_logprobs_idx = []
+        if req.logprob.input_token_ids_logprobs_val is None:
+            req.logprob.input_token_ids_logprobs_val = []
+        if req.logprob.input_token_ids_logprobs_idx is None:
+            req.logprob.input_token_ids_logprobs_idx = []

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

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Sequence
+
+from sglang.srt.environ import envs
+from sglang.srt.server_args import ServerArgs
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+@dataclass(slots=True, kw_only=True)
+class NewTokenRatioTracker:
+    init: float
+    min: float
+    decay: float
+    current: float
+
+    @classmethod
+    def from_server_args(cls, server_args: ServerArgs) -> NewTokenRatioTracker:
+        init = min(
+            envs.SGLANG_INIT_NEW_TOKEN_RATIO.get()
+            * server_args.schedule_conservativeness,
+            1.0,
+        )
+        min_ratio = min(
+            init * envs.SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR.get(),
+            1.0,
+        )
+        decay = (init - min_ratio) / envs.SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS.get()
+        return cls(init=init, min=min_ratio, decay=decay, current=init)
+
+    def decay_step(self) -> None:
+        self.current = max(self.current - self.decay, self.min)
+
+    def reset(self) -> None:
+        self.current = self.init
+
+    @staticmethod
+    def estimate_new_token_ratio_after_retract(reqs: Sequence[Req]) -> float:
+        total_decoded_tokens = sum(len(r.output_ids) for r in reqs)
+        total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in reqs)
+
+        new_estimate_ratio = (
+            total_decoded_tokens + envs.SGLANG_RETRACT_DECODE_STEPS.get() * len(reqs)
+        ) / (
+            total_max_new_tokens + 1
+        )  # avoid zero division
+        new_estimate_ratio = min(1.0, new_estimate_ratio)
+        return new_estimate_ratio

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

@@ -0,0 +1,28 @@
+from typing import Optional, Union
+
+import zmq
+
+from sglang.srt.managers.io_struct import BaseBatchReq, BaseReq, sock_send
+
+
+class SenderWrapper:
+    def __init__(self, socket: zmq.Socket):
+        self.socket = socket
+
+    def send_output(
+        self,
+        output: Union[BaseReq, BaseBatchReq],
+        recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None,
+    ):
+        if self.socket is None:
+            return
+
+        if (
+            isinstance(recv_obj, BaseReq)
+            and recv_obj.http_worker_ipc is not None
+            and output.http_worker_ipc is None
+        ):
+            # handle communicator reqs for multi-http worker case
+            output.http_worker_ipc = recv_obj.http_worker_ipc
+
+        sock_send(self.socket, output)

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

@@ -0,0 +1,581 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import (
+    Any,
+    Callable,
+    List,
+    Optional,
+)
+
+import torch
+import zmq
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import (
+    BatchEmbeddingOutput,
+    BatchTokenIDOutput,
+    CachedTokensDetails,
+    wrap_as_pickle,
+)
+from sglang.srt.managers.schedule_batch import (
+    BaseFinishReason,
+    Req,
+)
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+logger = logging.getLogger(__name__)
+
+
+DEFAULT_FORCE_STREAM_INTERVAL = envs.SGLANG_FORCE_STREAM_INTERVAL.get()
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerOutputStreamer:
+    send_to_detokenizer: zmq.Socket
+    tree_cache: BasePrefixCache
+    ps: ParallelState
+    server_args: ServerArgs
+    is_generation: bool
+    spec_algorithm: SpeculativeAlgorithm
+    disaggregation_mode: DisaggregationMode
+    enable_hicache_storage: Callable[[], bool]
+    _test_stream_output_count: int = 0
+
+    def _get_storage_backend_type(self) -> str:
+        """Get storage backend type from tree_cache."""
+        storage_backend_type = "none"
+        cache_controller = getattr(self.tree_cache, "cache_controller", None)
+        if cache_controller and hasattr(cache_controller, "storage_backend"):
+            storage_backend = cache_controller.storage_backend
+            if storage_backend is not None:
+                storage_backend_type = type(storage_backend).__name__
+        return storage_backend_type
+
+    def get_cached_tokens_details(self, req: Req) -> Optional[CachedTokensDetails]:
+        """Get detailed cache breakdown for a request, if available.
+
+        Returns:
+            - None if no cached tokens at all
+            - {"device": X, "host": Y} without storage breakdown
+            - {"device": X, "host": Y, "storage": Z} with storage breakdown
+        """
+        if (
+            req.cached_tokens_device > 0
+            or req.cached_tokens_host > 0
+            or req.cached_tokens_storage > 0
+        ):
+            details = {
+                "device": req.cached_tokens_device,
+                "host": req.cached_tokens_host,
+            }
+            # In PD mode the L3 hit is produced on prefill and reported on
+            # decode via metadata, while decode may not have a local storage backend.
+            if req.cached_tokens_storage > 0 or self.enable_hicache_storage():
+                details["storage"] = req.cached_tokens_storage
+            if self.enable_hicache_storage():
+                details["storage_backend"] = self._get_storage_backend_type()
+            return details
+
+        if req.cached_tokens > 0:
+            return {
+                "device": req.cached_tokens,
+                "host": 0,
+            }
+
+        return None
+
+    def stream_output(
+        self,
+        reqs: List[Req],
+        return_logprob: bool,
+        skip_req: Optional[Req] = None,
+    ):
+        """Stream the output to detokenizer."""
+        if self.is_generation:
+            self._stream_output_generation(reqs, return_logprob, skip_req)
+        else:  # embedding or reward model
+            self._stream_output_embedding(reqs)
+
+        if envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get() > 0:
+            self._trigger_crash_for_tests(
+                envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get()
+            )
+
+    def _trigger_crash_for_tests(self, crash_threshold: int):
+        # Crash trigger: crash after stream_output is called N times
+        # This is used for testing purposes.
+        self._test_stream_output_count += 1
+        if self._test_stream_output_count >= crash_threshold:
+            raise RuntimeError(
+                f"Test crash after stream_output called {self._test_stream_output_count} times"
+            )
+
+    def _stream_output_generation(
+        self,
+        reqs: List[Req],
+        return_logprob: bool,
+        skip_req: Optional[Req] = None,
+        is_idle_batch: bool = False,
+    ):
+        return_hidden_states = any(
+            req.return_hidden_states for req in reqs if req is not skip_req
+        )
+        return_routed_experts = any(
+            req.return_routed_experts for req in reqs if req is not skip_req
+        )
+        return_indexer_topk = any(
+            req.return_indexer_topk for req in reqs if req is not skip_req
+        )
+
+        acc = _GenerationStreamAccumulator(
+            return_logprob=return_logprob,
+            return_hidden_states=return_hidden_states,
+            return_routed_experts=return_routed_experts,
+            return_indexer_topk=return_indexer_topk,
+            spec_algorithm=self.spec_algorithm,
+            disaggregation_mode=self.disaggregation_mode,
+            default_stream_interval=self.server_args.stream_interval,
+            default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
+            get_cached_tokens_details=self.get_cached_tokens_details,
+        )
+        for req in reqs:
+            if req is skip_req:
+                continue
+            if req.finished() and req.finished_output:
+                # With the overlap schedule, a request will try to output twice and hit this line twice
+                # because of the one additional delayed token. This "continue" prevented the dummy output.
+                continue
+
+            acc.accept(req=req)
+            self._maybe_log_time_stats(req=req)
+
+        # Send to detokenizer
+        payload = acc.to_payload(
+            dp_rank=self.ps.dp_rank,
+            is_idle_batch=is_idle_batch,
+        )
+        if payload is not None:
+            self.send_to_detokenizer.send_output(payload)
+
+    def _maybe_log_time_stats(self, *, req: Req) -> None:
+        if (
+            req.finished()
+            and self.ps.attn_tp_rank == 0
+            and self.server_args.enable_request_time_stats_logging
+        ):
+            req.log_time_stats()
+
+    def _stream_output_embedding(self, reqs: List[Req]):
+        rids = []
+        http_worker_ipcs = []
+        finished_reasons: List[BaseFinishReason] = []
+
+        embeddings = []
+        prompt_tokens = []
+        cached_tokens = []
+        cached_tokens_details = []  # Detailed breakdown by cache source
+        time_stats = []
+        retraction_counts = []
+        phs_list = []
+        has_phs = False
+        for req in reqs:
+            if req.finished():
+                rids.append(req.rid)
+                http_worker_ipcs.append(req.http_worker_ipc)
+                finished_reasons.append(req.finished_reason.to_json())
+                embeddings.append(req.embedding)
+                prompt_tokens.append(len(req.origin_input_ids))
+                cached_tokens.append(req.cached_tokens)
+
+                # Collect detailed cache breakdown if available
+                cached_tokens_details.append(self.get_cached_tokens_details(req))
+                time_stats.append(req.time_stats)
+                retraction_counts.append(req.retraction_count)
+
+                phs = req.pooled_hidden_state
+                phs_list.append(phs)
+                if phs is not None:
+                    has_phs = True
+
+        # Optimize pooled hidden states (PHS) for IPC serialization.
+        # Two formats, disambiguated on the receiver side by length:
+        #   Stacked:     [stacked_tensor(N, ...)] — len 1, N > 1 requests
+        #   Non-stacked: [tensor_0, tensor_1, ...] — len == N
+        # Stacking reduces N pickle/__reduce_ex__ calls to 1.
+        # Only possible when all entries are non-None and same shape.
+        # See paired receiver logic in tokenizer_manager.py.
+        stacked_phs = None
+        if has_phs:
+            all_have_phs = all(t is not None for t in phs_list)
+            if all_have_phs:
+                if len(phs_list) > 1 and all(
+                    t.shape == phs_list[0].shape for t in phs_list
+                ):
+                    # Stacked: single tensor, wrapped in a list.
+                    stacked_phs = [torch.stack(phs_list)]
+                else:
+                    # Non-stacked: 1 request, mixed shapes, or mixed None.
+                    stacked_phs = phs_list
+            else:
+                # Non-stacked: some requests don't have PHS (None entries).
+                stacked_phs = phs_list
+
+        self.send_to_detokenizer.send_output(
+            BatchEmbeddingOutput(
+                rids=rids,
+                http_worker_ipcs=http_worker_ipcs,
+                time_stats=wrap_as_pickle(time_stats),
+                finished_reasons=finished_reasons,
+                embeddings=embeddings,
+                prompt_tokens=prompt_tokens,
+                cached_tokens=cached_tokens,
+                cached_tokens_details=cached_tokens_details,
+                placeholder_tokens_idx=None,
+                placeholder_tokens_val=None,
+                retraction_counts=retraction_counts,
+                pooled_hidden_states=stacked_phs,
+            )
+        )
+
+
+@dataclass(slots=True, kw_only=True)
+class _GenerationStreamAccumulator:
+    return_logprob: bool
+    return_hidden_states: bool
+    return_routed_experts: bool
+    return_indexer_topk: bool
+    spec_algorithm: Any
+    disaggregation_mode: DisaggregationMode
+    default_stream_interval: int
+    default_force_stream_interval: int
+    get_cached_tokens_details: Callable[[Req], Optional[CachedTokensDetails]]
+
+    rids: list = field(default_factory=list)
+    http_worker_ipcs: list = field(default_factory=list)
+    finished_reasons: list = field(default_factory=list)
+    decoded_texts: list = field(default_factory=list)
+    decode_ids_list: list = field(default_factory=list)
+    read_offsets: list = field(default_factory=list)
+    output_ids: list = field(default_factory=list)
+    skip_special_tokens: list = field(default_factory=list)
+    spaces_between_special_tokens: list = field(default_factory=list)
+    no_stop_trim: list = field(default_factory=list)
+    prompt_tokens: list = field(default_factory=list)
+    reasoning_tokens: list = field(default_factory=list)
+    completion_tokens: list = field(default_factory=list)
+    cached_tokens: list = field(default_factory=list)
+    cached_tokens_details: list = field(
+        default_factory=list
+    )  # Detailed breakdown by cache source
+    image_tokens: list = field(default_factory=list)
+    audio_tokens: list = field(default_factory=list)
+    video_tokens: list = field(default_factory=list)
+    spec_verify_ct: list = field(default_factory=list)
+    spec_num_correct_drafts: list = field(default_factory=list)
+    spec_num_block_accept_tokens: list = field(default_factory=list)
+    spec_num_cap_tokens: list = field(default_factory=list)
+    spec_correct_drafts_histogram: list = field(default_factory=list)
+    spec_cap_lens_histogram: list = field(default_factory=list)
+    retraction_counts: list = field(default_factory=list)
+    output_hidden_states: Optional[list] = None
+    routed_experts: Optional[list] = None
+    indexer_topk: Optional[list] = None
+    customized_info: dict = field(default_factory=dict)
+    time_stats: list = field(default_factory=list)
+    input_token_logprobs_val: Optional[list] = None
+    input_token_logprobs_idx: Optional[list] = None
+    output_token_logprobs_val: Optional[list] = None
+    output_token_logprobs_idx: Optional[list] = None
+    input_top_logprobs_val: Optional[list] = None
+    input_top_logprobs_idx: Optional[list] = None
+    output_top_logprobs_val: Optional[list] = None
+    output_top_logprobs_idx: Optional[list] = None
+    input_token_ids_logprobs_val: Optional[list] = None
+    input_token_ids_logprobs_idx: Optional[list] = None
+    output_token_ids_logprobs_val: Optional[list] = None
+    output_token_ids_logprobs_idx: Optional[list] = None
+
+    def __post_init__(self) -> None:
+        if self.return_hidden_states:
+            self.output_hidden_states = []
+        if self.return_routed_experts:
+            self.routed_experts = []
+        if self.return_indexer_topk:
+            self.indexer_topk = []
+
+        if self.return_logprob:
+            self.input_token_logprobs_val = []
+            self.input_token_logprobs_idx = []
+            self.output_token_logprobs_val = []
+            self.output_token_logprobs_idx = []
+            self.input_top_logprobs_val = []
+            self.input_top_logprobs_idx = []
+            self.output_top_logprobs_val = []
+            self.output_top_logprobs_idx = []
+            self.input_token_ids_logprobs_val = []
+            self.input_token_ids_logprobs_idx = []
+            self.output_token_ids_logprobs_val = []
+            self.output_token_ids_logprobs_idx = []
+
+    def accept(self, *, req: Req) -> None:
+        if req.finished():
+            assert not req.finished_output
+            req.finished_output = True
+            if req.finished_len is None:
+                req.finished_len = len(req.output_ids)
+            should_output = True
+        else:
+            if req.stream:
+                stream_interval = (
+                    req.sampling_params.stream_interval or self.default_stream_interval
+                )
+
+                # origin stream_interval logic
+                should_output = (
+                    len(req.output_ids) % stream_interval == 1
+                    if stream_interval > 1
+                    else len(req.output_ids) % stream_interval == 0
+                )
+
+                if should_output:
+                    # check_match_stop_str_prefix if  tail_str's suffix match stop_str prefix
+                    should_output &= not req.check_match_stop_str_prefix()
+            else:
+                should_output = (
+                    len(req.output_ids) % self.default_force_stream_interval == 0
+                )
+
+        if not should_output:
+            return
+
+        send_token_offset = req.send_token_offset
+        send_output_token_logprobs_offset = req.send_output_token_logprobs_offset
+        self.rids.append(req.rid)
+        self.http_worker_ipcs.append(req.http_worker_ipc)
+        self.finished_reasons.append(
+            req.finished_reason.to_json() if req.finished_reason else None
+        )
+        self.decoded_texts.append(req.decoded_text)
+        decode_ids, read_offset = req.init_incremental_detokenize()
+
+        self.decode_ids_list.append(decode_ids[req.send_decode_id_offset :])
+
+        # Exclude the tokens after stop condition
+        output_ids_ = req.output_ids_through_stop
+
+        req.send_decode_id_offset = len(decode_ids)
+        self.read_offsets.append(read_offset)
+        self.output_ids.append(output_ids_[send_token_offset:])
+        req.send_token_offset = len(output_ids_)
+        self.skip_special_tokens.append(req.sampling_params.skip_special_tokens)
+        self.spaces_between_special_tokens.append(
+            req.sampling_params.spaces_between_special_tokens
+        )
+        self.no_stop_trim.append(req.sampling_params.no_stop_trim)
+        self.prompt_tokens.append(len(req.origin_input_ids))
+        self.reasoning_tokens.append(req.reasoning_tokens)
+        self.completion_tokens.append(len(output_ids_))
+        self.cached_tokens.append(req.cached_tokens)
+
+        # Collect detailed cache breakdown if available
+        self.cached_tokens_details.append(self.get_cached_tokens_details(req))
+
+        # Multimodal prompt token counts. In disagg decode mode the prefill node
+        # already computed these and transferred them via the metadata buffer
+        # (req.mm_*), so prefer the pre-stored values; otherwise compute them
+        # from the request's multimodal items.
+        if req.mm_image_tokens or req.mm_audio_tokens or req.mm_video_tokens:
+            image_t = req.mm_image_tokens
+            audio_t = req.mm_audio_tokens
+            video_t = req.mm_video_tokens
+        elif req.multimodal_inputs:
+            image_t, audio_t, video_t = req.multimodal_inputs.compute_mm_token_counts()
+        else:
+            image_t = audio_t = video_t = 0
+        self.image_tokens.append(image_t)
+        self.audio_tokens.append(audio_t)
+        self.video_tokens.append(video_t)
+
+        self.retraction_counts.append(req.retraction_count)
+
+        self.time_stats.append(req.time_stats)
+
+        if not self.spec_algorithm.is_none():
+            self.spec_verify_ct.append(req.spec_verify_ct)
+            self.spec_num_correct_drafts.append(req.spec_num_correct_drafts)
+            self.spec_num_block_accept_tokens.append(req.spec_num_block_accept_tokens)
+            self.spec_num_cap_tokens.append(req.spec_num_cap_tokens)
+            self.spec_correct_drafts_histogram.append(req.spec_correct_drafts_histogram)
+            self.spec_cap_lens_histogram.append(req.spec_cap_lens_histogram)
+
+        if self.return_logprob:
+            if (
+                req.return_logprob
+                and not req.input_logprob_sent
+                # Decode server does not send input logprobs
+                and self.disaggregation_mode != DisaggregationMode.DECODE
+                # Only send when input logprobs have been computed (after prefill)
+                and req.logprob.input_token_logprobs_val is not None
+            ):
+                self.input_token_logprobs_val.append(
+                    req.logprob.input_token_logprobs_val
+                )
+                self.input_token_logprobs_idx.append(
+                    req.logprob.input_token_logprobs_idx
+                )
+                self.input_top_logprobs_val.append(req.logprob.input_top_logprobs_val)
+                self.input_top_logprobs_idx.append(req.logprob.input_top_logprobs_idx)
+                self.input_token_ids_logprobs_val.append(
+                    req.logprob.input_token_ids_logprobs_val
+                )
+                self.input_token_ids_logprobs_idx.append(
+                    req.logprob.input_token_ids_logprobs_idx
+                )
+                req.input_logprob_sent = True
+            else:
+                self.input_token_logprobs_val.append([])
+                self.input_token_logprobs_idx.append([])
+                self.input_top_logprobs_val.append([])
+                self.input_top_logprobs_idx.append([])
+                self.input_token_ids_logprobs_val.append([])
+                self.input_token_ids_logprobs_idx.append([])
+
+            if req.return_logprob:
+                logprob_end = max(len(output_ids_), 1)
+                self.output_token_logprobs_val.append(
+                    req.logprob.output_token_logprobs_val[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                self.output_token_logprobs_idx.append(
+                    req.logprob.output_token_logprobs_idx[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                self.output_top_logprobs_val.append(
+                    req.logprob.output_top_logprobs_val[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                self.output_top_logprobs_idx.append(
+                    req.logprob.output_top_logprobs_idx[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                self.output_token_ids_logprobs_val.append(
+                    req.logprob.output_token_ids_logprobs_val[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                self.output_token_ids_logprobs_idx.append(
+                    req.logprob.output_token_ids_logprobs_idx[
+                        send_output_token_logprobs_offset:logprob_end
+                    ]
+                )
+                req.send_output_token_logprobs_offset = logprob_end
+            else:
+                self.output_token_logprobs_val.append([])
+                self.output_token_logprobs_idx.append([])
+                self.output_top_logprobs_val.append([])
+                self.output_top_logprobs_idx.append([])
+                self.output_token_ids_logprobs_val.append([])
+                self.output_token_ids_logprobs_idx.append([])
+
+        if self.return_hidden_states:
+            if req.return_hidden_states:
+                # Mirror output_ids_through_stop: spec verify steps can overshoot finished_len.
+                hs = req.hidden_states
+                if req.finished_len is not None:
+                    hs = hs[: req.finished_len]
+                self.output_hidden_states.append(hs)
+            else:
+                self.output_hidden_states.append(None)
+        if self.return_routed_experts:
+            self.routed_experts.append(
+                req.routed_experts if req.return_routed_experts else None
+            )
+        if self.return_indexer_topk:
+            self.indexer_topk.append(
+                req.indexer_topk if req.return_indexer_topk else None
+            )
+
+        current_output_len = len(self.output_ids[-1])
+        if req.customized_info is not None:
+            for key, req_values in req.customized_info.items():
+                if key not in self.customized_info:
+                    self.customized_info[key] = [
+                        [None] * len(prev_output_ids)
+                        for prev_output_ids in self.output_ids[:-1]
+                    ]
+                self.customized_info[key].append(
+                    [None] * current_output_len
+                    if req_values is None
+                    else req_values[send_token_offset : len(output_ids_)]
+                )
+
+        for per_request_values in self.customized_info.values():
+            if len(per_request_values) < len(self.output_ids):
+                per_request_values.append([None] * current_output_len)
+
+    def to_payload(
+        self, *, dp_rank: int, is_idle_batch: bool
+    ) -> Optional[BatchTokenIDOutput]:
+        if not (self.rids or is_idle_batch):
+            return None
+        dp_ranks = [dp_rank] * len(self.rids) if self.rids else None
+        return BatchTokenIDOutput(
+            rids=self.rids,
+            http_worker_ipcs=self.http_worker_ipcs,
+            spec_verify_ct=self.spec_verify_ct,
+            spec_num_correct_drafts=self.spec_num_correct_drafts,
+            spec_num_block_accept_tokens=self.spec_num_block_accept_tokens,
+            spec_num_cap_tokens=self.spec_num_cap_tokens,
+            spec_correct_drafts_histogram=self.spec_correct_drafts_histogram,
+            spec_cap_lens_histogram=self.spec_cap_lens_histogram,
+            time_stats=wrap_as_pickle(self.time_stats),
+            finished_reasons=self.finished_reasons,
+            decoded_texts=self.decoded_texts,
+            decode_ids=self.decode_ids_list,
+            read_offsets=self.read_offsets,
+            output_ids=self.output_ids,
+            skip_special_tokens=self.skip_special_tokens,
+            spaces_between_special_tokens=self.spaces_between_special_tokens,
+            no_stop_trim=self.no_stop_trim,
+            prompt_tokens=self.prompt_tokens,
+            reasoning_tokens=self.reasoning_tokens,
+            completion_tokens=self.completion_tokens,
+            cached_tokens=self.cached_tokens,
+            cached_tokens_details=self.cached_tokens_details,
+            image_tokens=self.image_tokens,
+            audio_tokens=self.audio_tokens,
+            video_tokens=self.video_tokens,
+            input_token_logprobs_val=self.input_token_logprobs_val,
+            input_token_logprobs_idx=self.input_token_logprobs_idx,
+            output_token_logprobs_val=self.output_token_logprobs_val,
+            output_token_logprobs_idx=self.output_token_logprobs_idx,
+            input_top_logprobs_val=self.input_top_logprobs_val,
+            input_top_logprobs_idx=self.input_top_logprobs_idx,
+            output_top_logprobs_val=self.output_top_logprobs_val,
+            output_top_logprobs_idx=self.output_top_logprobs_idx,
+            input_token_ids_logprobs_val=self.input_token_ids_logprobs_val,
+            input_token_ids_logprobs_idx=self.input_token_ids_logprobs_idx,
+            output_token_ids_logprobs_val=self.output_token_ids_logprobs_val,
+            output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx,
+            output_token_entropy_val=None,
+            output_hidden_states=self.output_hidden_states,
+            routed_experts=self.routed_experts,
+            indexer_topk=self.indexer_topk,
+            customized_info=(
+                wrap_as_pickle(self.customized_info) if self.customized_info else None
+            ),
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
+            retraction_counts=self.retraction_counts,
+            dp_ranks=dp_ranks,
+        )

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

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+
+
+class SchedulerStats: ...  # type: ignore[no-redef]
+
+
+@dataclasses.dataclass
+class PoolStats:
+    # For full pools (required)
+    full_num_used: int
+    full_token_usage: float
+    full_available_size: int
+    full_evictable_size: int
+
+    is_hybrid_swa: bool = False
+    is_hybrid_ssm: bool = False
+    is_hisparse: bool = False
+
+    # For hybrid-swa pools
+    swa_num_used: Optional[int] = None
+    swa_token_usage: Optional[float] = None
+    swa_available_size: Optional[int] = None
+    swa_evictable_size: Optional[int] = None
+
+    # For mamba pools
+    mamba_num_used: Optional[int] = None
+    mamba_usage: Optional[float] = None
+    mamba_available_size: Optional[int] = None
+    mamba_evictable_size: Optional[int] = None
+
+    # HiSparse device/host breakdown for decode logs (plain KV pool only)
+    hisparse_device_tokens: Optional[int] = None
+    hisparse_device_token_usage: Optional[float] = None
+    hisparse_host_tokens: Optional[int] = None
+    hisparse_host_token_usage: Optional[float] = None
+
+    def get_kv_token_stats(self) -> Tuple[int, float]:
+        # NOTE: mamba pool is not included in the "token usage" calculation.
+        if self.is_hybrid_swa:
+            num_used = max(self.full_num_used, self.swa_num_used)
+            token_usage = max(self.full_token_usage, self.swa_token_usage)
+        else:
+            num_used = self.full_num_used
+            token_usage = self.full_token_usage
+
+        return num_used, token_usage
+
+    def get_max_pool_usage(self) -> float:
+        usage = self.full_token_usage
+        if self.is_hybrid_swa:
+            usage = max(usage, self.swa_token_usage)
+        if self.is_hybrid_ssm:
+            usage = max(usage, self.mamba_usage)
+        assert usage is not None and usage >= 0, f"{usage=} is not valid"
+        return usage
+
+    def get_prefill_usage_msg_parts(self) -> List[str]:
+        parts = []
+        if self.is_hybrid_swa:
+            parts += [
+                f"full token usage: {self.full_token_usage:.2f}",
+                f"swa token usage: {self.swa_token_usage:.2f}",
+            ]
+        if self.is_hybrid_ssm:
+            if not self.is_hybrid_swa:
+                parts.append(f"full token usage: {self.full_token_usage:.2f}")
+            parts.append(f"mamba usage: {self.mamba_usage:.2f}")
+        if not parts:
+            parts.append(f"token usage: {self.full_token_usage:.2f}")
+        return parts
+
+    def get_decode_usage_msg_parts(self) -> List[str]:
+        parts = []
+        if self.is_hybrid_swa:
+            parts += [
+                f"#full token: {self.full_num_used}",
+                f"full token usage: {self.full_token_usage:.2f}",
+                f"#swa token: {self.swa_num_used}",
+                f"swa token usage: {self.swa_token_usage:.2f}",
+            ]
+        if self.is_hybrid_ssm:
+            if not self.is_hybrid_swa:
+                parts += [
+                    f"#full token: {self.full_num_used}",
+                    f"full token usage: {self.full_token_usage:.2f}",
+                ]
+            parts += [
+                f"mamba num: {self.mamba_num_used}",
+                f"mamba usage: {self.mamba_usage:.2f}",
+            ]
+        if self.is_hisparse:
+            parts += [
+                f"#gpu token: {self.hisparse_device_tokens}",
+                f"gpu token usage: {self.hisparse_device_token_usage:.2f}",
+                f"#cpu token: {self.hisparse_host_tokens}",
+                f"cpu token usage: {self.hisparse_host_token_usage:.2f}",
+            ]
+        if not parts:
+            parts.append(
+                f"#token: {self.full_num_used}, token usage: {self.full_token_usage:.2f}"
+            )
+        return parts
+
+    def update_scheduler_stats(self, stats: SchedulerStats) -> None:
+        """Update pool-related fields on SchedulerStats."""
+        num_used, _ = self.get_kv_token_stats()
+        stats.num_used_tokens = num_used
+        stats.token_usage = round(self.get_max_pool_usage(), 2)
+        stats.full_token_usage = self.full_token_usage
+        if self.is_hybrid_swa:
+            stats.swa_token_usage = self.swa_token_usage
+            stats.swa_available_tokens = self.swa_available_size
+            stats.swa_evictable_tokens = self.swa_evictable_size
+            stats.swa_used_tokens = self.swa_num_used
+        if self.is_hybrid_ssm:
+            stats.mamba_usage = self.mamba_usage
+            stats.mamba_available_tokens = self.mamba_available_size
+            stats.mamba_evictable_tokens = self.mamba_evictable_size
+            stats.mamba_used_tokens = self.mamba_num_used
+        stats.kv_available_tokens = self.full_available_size
+        stats.kv_evictable_tokens = self.full_evictable_size
+        stats.kv_used_tokens = self.full_num_used
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerPoolStatsObserver:
+    tree_cache: BasePrefixCache
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    req_to_token_pool: ReqToTokenPool
+    session_controller: Any
+    hisparse_coordinator: Any
+    is_hybrid_swa: bool
+    is_hybrid_ssm: bool
+    enable_hisparse: bool
+    full_tokens_per_layer: Any
+    swa_tokens_per_layer: Any
+    max_total_num_tokens: int
+    get_last_batch: Callable
+    get_running_batch: Callable
+
+    def streaming_session_count(self) -> int:
+        return sum(
+            1
+            for session in self.session_controller.sessions.values()
+            if session.streaming
+        )
+
+    def active_pool_idxs(self) -> set:
+        """Pool idxs currently owned by reqs in last_batch / running_batch.
+
+        Used to decide which session slots' KV is owned by batch reqs
+        (and thus counted via uncached_size, not session_held).
+        """
+        idxs = set()
+        for batch in [self.get_last_batch(), self.get_running_batch()]:
+            if batch is None or batch.is_empty():
+                continue
+            for req in batch.reqs:
+                if req.req_pool_idx is not None:
+                    idxs.add(req.req_pool_idx)
+        return idxs
+
+    def session_held_tokens(self) -> int:
+        return self.tree_cache.session_held_tokens(self.active_pool_idxs())
+
+    def session_held_full_tokens(self) -> int:
+        return self.tree_cache.session_held_full_tokens(self.active_pool_idxs())
+
+    def session_held_swa_tokens(self) -> int:
+        return self.tree_cache.session_held_swa_tokens(self.active_pool_idxs())
+
+    def session_held_req_count(self) -> int:
+        return self.tree_cache.session_held_req_count()
+
+    def session_held_mamba_slots(self) -> int:
+        return self.tree_cache.session_held_mamba_slots(self.active_pool_idxs())
+
+    def get_pool_stats(self) -> PoolStats:
+        if self.is_hybrid_swa:
+            pool_stats = self._get_swa_token_info()
+        elif self.is_hybrid_ssm:
+            pool_stats = self._get_mamba_token_info()
+        else:
+            pool_stats = self._get_token_info()
+
+        if self.enable_hisparse:
+            pool_stats = self._get_hisparse_token_info(pool_stats)
+
+        # swa + ssm can coexist: overlay mamba fields onto swa stats
+        if self.is_hybrid_ssm:
+            mamba_stats = self._get_mamba_token_info()
+            pool_stats.is_hybrid_ssm = True
+            pool_stats.mamba_num_used = mamba_stats.mamba_num_used
+            pool_stats.mamba_usage = mamba_stats.mamba_usage
+            pool_stats.mamba_available_size = mamba_stats.mamba_available_size
+            pool_stats.mamba_evictable_size = mamba_stats.mamba_evictable_size
+
+        return pool_stats
+
+    def _get_token_info(self) -> PoolStats:
+        available_size = self.token_to_kv_pool_allocator.available_size()
+        evictable_size = self.tree_cache.evictable_size()
+        num_used = self.max_total_num_tokens - (available_size + evictable_size)
+        token_usage = num_used / self.max_total_num_tokens
+        return PoolStats(
+            full_num_used=num_used,
+            full_token_usage=token_usage,
+            full_available_size=available_size,
+            full_evictable_size=evictable_size,
+        )
+
+    def _get_hisparse_token_info(self, pool_stats: PoolStats) -> PoolStats:
+        if self.enable_hisparse and self.hisparse_coordinator is not None:
+            h = self.hisparse_coordinator.get_token_stats()
+            return dataclasses.replace(
+                pool_stats,
+                is_hisparse=True,
+                hisparse_device_tokens=h.device_tokens,
+                hisparse_device_token_usage=h.device_token_usage,
+                hisparse_host_tokens=h.host_tokens,
+                hisparse_host_token_usage=h.host_token_usage,
+            )
+        return pool_stats
+
+    def _get_mamba_token_info(self):
+        is_mamba_radix_cache = (
+            self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache()
+        )
+        full_available_size = self.token_to_kv_pool_allocator.available_size()
+        full_evictable_size = (
+            self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0
+        )
+        mamba_available_size = self.req_to_token_pool.mamba_allocator.available_size()
+        # `mamba_usage`/`mamba_num_used` track the ACTIVE bf16 pool occupancy (running
+        # requests) -- this feeds throttle decisions (get_max_pool_usage) which asserts
+        # usage >= 0. With int8 checkpoints the radix-cached states live in a SEPARATE
+        # int8 pool, so they own ZERO active slots: report evictable=0 against the active
+        # pool (otherwise active.size - (available + radix_cached) goes negative). The
+        # int8 cache pool's own occupancy is validated separately in the invariant check.
+        has_int8_ckpt = (
+            getattr(self.req_to_token_pool, "mamba_ckpt_pool", None) is not None
+        )
+        mamba_evictable_size = (
+            self.tree_cache.mamba_evictable_size()
+            if (is_mamba_radix_cache and not has_int8_ckpt)
+            else 0
+        )
+        full_num_used = self.token_to_kv_pool_allocator.size - (
+            full_available_size + full_evictable_size
+        )
+        mamba_num_used = self.req_to_token_pool.mamba_pool.size - (
+            mamba_available_size + mamba_evictable_size
+        )
+        full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size
+        mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size
+
+        return PoolStats(
+            is_hybrid_ssm=True,
+            full_num_used=full_num_used,
+            full_token_usage=full_token_usage,
+            full_available_size=full_available_size,
+            full_evictable_size=full_evictable_size,
+            mamba_num_used=mamba_num_used,
+            mamba_usage=mamba_usage,
+            mamba_available_size=mamba_available_size,
+            mamba_evictable_size=mamba_evictable_size,
+        )
+
+    def _get_swa_token_info(self) -> PoolStats:
+        full_available_size = self.token_to_kv_pool_allocator.full_available_size()
+        full_evictable_size = self.tree_cache.full_evictable_size()
+        swa_available_size = self.token_to_kv_pool_allocator.swa_available_size()
+        swa_evictable_size = self.tree_cache.swa_evictable_size()
+        full_num_used = self.full_tokens_per_layer - (
+            full_available_size + full_evictable_size
+        )
+        swa_num_used = self.swa_tokens_per_layer - (
+            swa_available_size + swa_evictable_size
+        )
+        # FIXME(hisparse): host-backup transiently over-releases the device pool
+        # counter, producing negative full_num_used / swa_num_used. We clamp to 0
+        # to keep token_usage / leak checks sane, but the underlying accounting
+        # bug should be fixed so the clamp can go away.
+        if self.enable_hisparse:
+            full_num_used = max(0, full_num_used)
+            swa_num_used = max(0, swa_num_used)
+        if not self.full_tokens_per_layer:
+            full_num_used = 0
+            full_available_size = 0
+            full_token_usage = 0.0
+        else:
+            full_token_usage = full_num_used / self.full_tokens_per_layer
+        swa_token_usage = swa_num_used / self.swa_tokens_per_layer
+
+        return PoolStats(
+            is_hybrid_swa=True,
+            full_num_used=full_num_used,
+            full_token_usage=full_token_usage,
+            full_available_size=full_available_size,
+            full_evictable_size=full_evictable_size,
+            swa_num_used=swa_num_used,
+            swa_token_usage=swa_token_usage,
+            swa_available_size=swa_available_size,
+            swa_evictable_size=swa_evictable_size,
+        )

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

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+import logging
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    List,
+    Optional,
+)
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import is_mps, is_npu
+from sglang.srt.utils.profile_merger import ProfileMerger
+from sglang.srt.utils.profile_utils import ProfileManager
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+
+_is_npu = is_npu()
+_is_mps = is_mps()
+if _is_npu:
+    import torch_npu
+
+    patches = [
+        ["profiler.profile", torch_npu.profiler.profile],
+        ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU],
+        ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
+    ]
+    apply_torch_npu_patches(torch_npu, patches)
+elif _is_mps:
+    from sglang.srt.hardware_backend.mlx.profiler import apply_metal_profiler_patches
+
+    apply_metal_profiler_patches()
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True)
+class SchedulerProfilerManager:
+    ps: Any
+    dp_tp_cpu_group: Any
+    get_forward_ct: Callable[[], int]
+
+    def __post_init__(self) -> None:
+        if envs.SGLANG_PROFILE_V2.get():
+            self._profile_manager = ProfileManager(
+                ps=self.ps,
+                cpu_group=self.dp_tp_cpu_group,
+            )
+            return
+
+        self.torch_profiler = None
+        self.torch_profiler_output_dir: Optional[Path] = None
+        self.profiler_activities: Optional[List[str]] = None
+        self.profile_id: Optional[str] = None
+
+        self.profiler_start_forward_ct: Optional[int] = None
+        self.profiler_target_forward_ct: Optional[int] = None
+
+        self.profiler_prefill_ct: Optional[int] = None
+        self.profiler_decode_ct: Optional[int] = None
+        self.profiler_target_prefill_ct: Optional[int] = None
+        self.profiler_target_decode_ct: Optional[int] = None
+
+        self.profile_by_stage: bool = False
+        self.profile_in_progress: bool = False
+        self.merge_profiles = False
+
+        # For ROCM
+        self.rpd_profiler = None
+
+    def _init_profile(
+        self,
+        output_dir: Optional[str],
+        start_step: Optional[int],
+        num_steps: Optional[int],
+        activities: Optional[List[str]],
+        with_stack: Optional[bool],
+        record_shapes: Optional[bool],
+        profile_by_stage: bool,
+        profile_id: str,
+        merge_profiles: bool = False,
+        profile_prefix: str = "",
+        profile_stages: Optional[List[str]] = None,
+    ) -> ProfileReqOutput:
+        if envs.SGLANG_PROFILE_V2.get():
+            return self._profile_manager.configure(
+                output_dir=output_dir,
+                start_step=start_step,
+                num_steps=num_steps,
+                activities=activities,
+                with_stack=with_stack,
+                record_shapes=record_shapes,
+                profile_by_stage=profile_by_stage,
+                profile_id=profile_id,
+                merge_profiles=merge_profiles,
+                profile_prefix=profile_prefix,
+                profile_stages=profile_stages,
+            )
+
+        if self.profile_in_progress:
+            return ProfileReqOutput(
+                success=False,
+                message="Profiling is already in progress. Call /stop_profile first.",
+            )
+
+        self.profile_by_stage = profile_by_stage
+        self.merge_profiles = merge_profiles
+
+        if output_dir is None:
+            output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
+        if activities is None:
+            activities = ["CPU", "GPU"]
+
+        self.torch_profiler_output_dir = Path(output_dir).expanduser()
+        self.torch_profiler_with_stack = with_stack
+        self.torch_profiler_record_shapes = record_shapes
+        self.profiler_activities = activities
+        self.profile_id = profile_id
+        self.profile_prefix = profile_prefix
+
+        if start_step:
+            self.profiler_start_forward_ct = max(start_step, self.get_forward_ct() + 1)
+
+        if num_steps:
+            if self.profile_by_stage:
+                self.profiler_prefill_ct = 0
+                self.profiler_decode_ct = 0
+                self.profiler_target_prefill_ct = num_steps
+                self.profiler_target_decode_ct = num_steps
+            elif start_step:
+                self.profiler_target_forward_ct = (
+                    self.profiler_start_forward_ct + num_steps
+                )
+            else:
+                self.profiler_target_forward_ct = self.get_forward_ct() + num_steps + 1
+            # The caller will be notified when reaching profiler_target_forward_ct
+        else:
+            self.profiler_target_forward_ct = None
+
+        return ProfileReqOutput(success=True, message="Succeeded")
+
+    def _start_profile(
+        self, stage: Optional[ForwardMode] = None
+    ) -> ProfileReqOutput | None:
+        if envs.SGLANG_PROFILE_V2.get():
+            return self._profile_manager.manual_start()
+
+        stage_str = f" for {stage.name}" if stage else ""
+        logger.info(
+            f"Profiling starts{stage_str}. Traces will be saved to: {self.torch_profiler_output_dir} (with profile id: {self.profile_id})",
+        )
+
+        activities = self.profiler_activities
+        with_stack = self.torch_profiler_with_stack
+        record_shapes = self.torch_profiler_record_shapes
+
+        activity_map = {
+            "CPU": torch.profiler.ProfilerActivity.CPU,
+            "GPU": torch.profiler.ProfilerActivity.CUDA,
+        }
+        if hasattr(torch.profiler.ProfilerActivity, "XPU"):
+            activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU
+        torchprof_activities = [
+            activity_map[a] for a in activities if a in activity_map
+        ]
+
+        if "RPD" in activities:  # for ROCM
+            from rpdTracerControl import rpdTracerControl
+
+            rpdTracerControl.skipCreate()
+
+            self.rpd_profile_path = os.path.join(
+                self.torch_profiler_output_dir,
+                "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz",
+            )
+
+            if self.ps.tp_rank == 0:
+                import sqlite3
+
+                from rocpd.schema import RocpdSchema
+
+                if os.path.exists("trace.rpd"):
+                    os.unlink("trace.rpd")
+                schema = RocpdSchema()
+                connection = sqlite3.connect("trace.rpd")
+                schema.writeSchema(connection)
+                connection.commit()
+                del connection
+            torch.distributed.barrier(self.dp_tp_cpu_group)
+
+            self.rpd_profiler = rpdTracerControl()
+            self.rpd_profiler.setPythonTrace(True)
+            self.rpd_profiler.start()
+            self.rpd_profiler.rangePush("", "rpd profile range", "")
+            self.profile_in_progress = True
+        elif torchprof_activities:
+            self.torch_profiler = torch.profiler.profile(
+                activities=torchprof_activities,
+                with_stack=with_stack if with_stack is not None else True,
+                record_shapes=record_shapes if record_shapes is not None else False,
+                on_trace_ready=(
+                    None
+                    if not _is_npu
+                    else torch_npu.profiler.tensorboard_trace_handler(
+                        str(self.torch_profiler_output_dir)
+                    )
+                ),
+                experimental_config=(
+                    None
+                    if not _is_npu
+                    else torch_npu.profiler._ExperimentalConfig(
+                        export_type=torch_npu.profiler.ExportType.Text,
+                        profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+                        msprof_tx=False,
+                        aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
+                        l2_cache=False,
+                        op_attr=False,
+                        data_simplification=False,
+                        record_op_args=False,
+                        gc_detect_threshold=None,
+                    )
+                ),
+            )
+            try:
+                self.torch_profiler.start()
+            except RuntimeError as e:
+                self.torch_profiler = None
+                return ProfileReqOutput(success=False, message=str(e))
+            self.profile_in_progress = True
+
+        if "MEM" in activities:
+            torch.cuda.memory._record_memory_history(max_entries=100000)
+            self.profile_in_progress = True
+
+        if "CUDA_PROFILER" in activities:
+            if self.ps.gpu_id == get_server_args().base_gpu_id:
+                torch.cuda.cudart().cudaProfilerStart()
+            self.profile_in_progress = True
+
+        return ProfileReqOutput(success=True, message="Succeeded")
+
+    def _merge_profile_traces(self) -> str:
+        if not self.merge_profiles:
+            return ""
+
+        if self.ps.tp_rank != 0:
+            return ""
+        if self.ps.dp_size > 1 and self.ps.dp_rank != 0:
+            return ""
+        if self.ps.pp_size > 1 and self.ps.pp_rank != 0:
+            return ""
+        if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0:
+            return ""
+
+        try:
+            logger.info("Starting profile merge...")
+            merger = ProfileMerger(self.torch_profiler_output_dir, self.profile_id)
+            merged_path = merger.merge_chrome_traces()
+
+            summary = merger.get_merge_summary()
+            merge_message = (
+                f" Merged trace: {merged_path} "
+                f"(Events: {summary.get('total_events', '?')}, "
+                f"Files: {summary.get('total_files', '?')})"
+            )
+
+            logger.info(f"Profile merge completed: {merged_path}")
+        except Exception as e:
+            logger.error(f"Failed to merge profiles: {e}", exc_info=True)
+            return f" Merge failed: {e!s}"
+        else:
+            return merge_message
+
+    def _stop_profile(
+        self, stage: Optional[ForwardMode] = None
+    ) -> ProfileReqOutput | None:
+        if envs.SGLANG_PROFILE_V2.get():
+            return self._profile_manager.manual_stop()
+
+        if not self.profile_in_progress:
+            return ProfileReqOutput(
+                success=False,
+                message="Profiling is not in progress. Call /start_profile first.",
+            )
+
+        self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True)
+
+        if self.profile_prefix:
+            stage_prefix = self.profile_prefix + "-"
+        else:
+            stage_prefix = ""
+
+        stage_suffix = f"-{stage.name}" if stage else ""
+        logger.info("Stop profiling" + stage_suffix + "...")
+        if self.torch_profiler is not None:
+            self.torch_profiler.stop()
+            if not _is_npu:
+                # Build filename with only non-zero ranks to maintain backward compatibility
+                filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"]
+
+                # Only add other ranks if parallelism is enabled (size > 1)
+                if self.ps.dp_size > 1:
+                    filename_parts.append(f"DP-{self.ps.dp_rank}")
+                if self.ps.pp_size > 1:
+                    filename_parts.append(f"PP-{self.ps.pp_rank}")
+                if self.ps.moe_ep_size > 1:
+                    filename_parts.append(f"EP-{self.ps.moe_ep_rank}")
+
+                filename = (
+                    stage_prefix
+                    + "-".join(filename_parts)
+                    + stage_suffix
+                    + ".trace.json.gz"
+                )
+
+                self.torch_profiler.export_chrome_trace(
+                    os.path.join(self.torch_profiler_output_dir, filename)
+                )
+            torch.distributed.barrier(self.dp_tp_cpu_group)
+
+        if self.rpd_profiler is not None:
+            self.rpd_profiler.rangePop()
+            self.rpd_profiler.stop()
+            self.rpd_profiler.flush()
+
+            torch.distributed.barrier(self.dp_tp_cpu_group)
+            if self.ps.tp_rank == 0:
+                from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
+
+                rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
+            self.rpd_profiler = None
+            self.rpd_profile_path = None
+
+        if self.profiler_activities is not None and "MEM" in self.profiler_activities:
+            memory_profile_path = os.path.join(
+                self.torch_profiler_output_dir,
+                str(time.time())
+                + f"-TP-{self.ps.tp_rank}-memory"
+                + stage_suffix
+                + ".pickle",
+            )
+            torch.cuda.memory._dump_snapshot(memory_profile_path)
+            torch.cuda.memory._record_memory_history(enabled=None)
+
+        if "CUDA_PROFILER" in self.profiler_activities:
+            if self.ps.gpu_id == get_server_args().base_gpu_id:
+                torch.cuda.cudart().cudaProfilerStop()
+
+        merge_message = self._merge_profile_traces()
+
+        logger.info(
+            "Profiling done. Traces are saved to: %s%s",
+            self.torch_profiler_output_dir,
+            merge_message,
+        )
+        self.torch_profiler = None
+        self.profile_in_progress = False
+        self.profiler_start_forward_ct = None
+
+        return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}")
+
+    def _profile_batch_predicate(self, batch: ScheduleBatch):
+        if envs.SGLANG_PROFILE_V2.get():
+            self._profile_manager.step(forward_mode=batch.forward_mode)
+            return
+
+        if self.profile_by_stage:
+            if batch.forward_mode.is_prefill():
+                if self.profiler_prefill_ct == 0:
+                    self._start_profile(batch.forward_mode)
+                self.profiler_prefill_ct += 1
+                if self.profiler_prefill_ct > self.profiler_target_prefill_ct:
+                    if self.profile_in_progress:
+                        self._stop_profile(stage=ForwardMode.EXTEND)
+            elif batch.forward_mode.is_decode():
+                if self.profiler_decode_ct == 0:
+                    if self.profile_in_progress:
+                        # force trace flush
+                        self._stop_profile(stage=ForwardMode.EXTEND)
+                    self._start_profile(batch.forward_mode)
+                self.profiler_decode_ct += 1
+                if self.profiler_decode_ct > self.profiler_target_decode_ct:
+                    if self.profile_in_progress:
+                        self._stop_profile(stage=ForwardMode.DECODE)
+            elif batch.forward_mode.is_idle():
+                pass
+            else:
+                raise RuntimeError(f"unsupported profile stage: {batch.forward_mode}")
+        else:
+            # Check profiler
+            if (
+                self.profiler_target_forward_ct
+                and self.profiler_target_forward_ct <= self.get_forward_ct()
+            ):
+                self._stop_profile()
+            if (
+                self.profiler_start_forward_ct
+                and self.profiler_start_forward_ct == self.get_forward_ct()
+            ):
+                self._start_profile()
+
+    def _profile(self, recv_req: ProfileReq):
+        if recv_req.req_type == ProfileReqType.START_PROFILE:
+            if recv_req.profile_by_stage or recv_req.start_step:
+                return self._init_profile(
+                    recv_req.output_dir,
+                    recv_req.start_step,
+                    recv_req.num_steps,
+                    recv_req.activities,
+                    recv_req.with_stack,
+                    recv_req.record_shapes,
+                    recv_req.profile_by_stage,
+                    recv_req.profile_id,
+                    recv_req.merge_profiles,
+                    recv_req.profile_prefix,
+                    recv_req.profile_stages,
+                )
+            else:
+                self._init_profile(
+                    recv_req.output_dir,
+                    recv_req.start_step,
+                    recv_req.num_steps,
+                    recv_req.activities,
+                    recv_req.with_stack,
+                    recv_req.record_shapes,
+                    recv_req.profile_by_stage,
+                    recv_req.profile_id,
+                    recv_req.merge_profiles,
+                    recv_req.profile_prefix,
+                )
+                return self._start_profile()
+        else:
+            return self._stop_profile()

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

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from http import HTTPStatus
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    List,
+    Optional,
+    Union,
+)
+
+import zmq
+from torch.distributed import barrier
+
+from sglang.srt.disaggregation.utils import prepare_abort
+from sglang.srt.managers.io_struct import (
+    BatchTokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    sock_recv,
+)
+from sglang.srt.managers.mm_utils import (
+    has_shm_features,
+    unwrap_shm_features,
+)
+from sglang.srt.utils import (
+    broadcast_pyobj,
+    point_to_point_pyobj,
+)
+from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.server_args import ServerArgs
+    from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerRequestReceiver:
+    recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy]
+    recv_from_rpc: Optional[zmq.Socket]
+    recv_skipper: Any
+    input_blocker: Any
+    mm_receiver: Any
+    ps: ParallelState
+    tp_group: Any
+    tp_cpu_group: Any
+    attn_tp_group: Any
+    attn_tp_cpu_group: Any
+    attn_cp_group: Any
+    attn_cp_cpu_group: Any
+    world_group: Any
+    server_args: ServerArgs
+    model_config: ModelConfig
+    max_recv_per_poll: int
+    stream_output: Callable[..., None]
+    get_last_forward_mode: Callable[[], Any]
+    scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None
+
+    def recv_limit_reached(self, num_recv_reqs: int) -> bool:
+        if self.max_recv_per_poll < 0:
+            return False
+        return num_recv_reqs >= self.max_recv_per_poll
+
+    @scheduler_nvtx_method("scheduler.recv_requests")
+    def recv_requests(
+        self,
+    ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]:
+        """Receive results at tp_rank = 0 and broadcast it to all other TP ranks."""
+
+        if self.scripted_scheduler_hook is not None:
+            self.scripted_scheduler_hook.step()
+
+        if self.recv_skipper is not None:
+            if not self.recv_skipper.handle(self.get_last_forward_mode()):
+                return []
+
+        recv_reqs = self._pull_raw_reqs()
+
+        if self.input_blocker is not None:
+            recv_reqs = self.input_blocker.handle(recv_reqs)
+
+        recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs)
+
+        if self.ps.pp_rank == 0:
+            self.unwrap_pickle_wrapper(recv_reqs)
+
+        recv_reqs = self._apply_mm_receiver(recv_reqs)
+
+        self._finalize_shm_features(recv_reqs)
+
+        return recv_reqs
+
+    def _pull_raw_reqs(self) -> Optional[List]:
+        if self.ps.pp_rank == 0:
+            if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
+                recv_reqs = []
+
+                while True:
+                    try:
+                        if self.recv_limit_reached(len(recv_reqs)):
+                            break
+                        recv_req = sock_recv(self.recv_from_tokenizer, zmq.NOBLOCK)
+                    except zmq.ZMQError:
+                        break
+                    recv_reqs.append(recv_req)
+
+                while True:
+                    try:
+                        if self.recv_limit_reached(len(recv_reqs)):
+                            break
+                        recv_rpc = sock_recv(self.recv_from_rpc, zmq.NOBLOCK)
+                    except zmq.ZMQError:
+                        break
+                    recv_reqs.append(recv_rpc)
+            else:
+                recv_reqs = None
+        else:
+            if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
+                dp_offset = (
+                    self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size
+                )
+                recv_reqs = point_to_point_pyobj(
+                    [],
+                    self.ps.pp_rank * self.ps.tp_size + dp_offset,
+                    self.world_group.cpu_group,
+                    (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset,
+                    self.ps.pp_rank * self.ps.tp_size + dp_offset,
+                )
+            else:
+                recv_reqs = None
+        return recv_reqs
+
+    def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List:
+        if self.server_args.enable_dp_attention:
+            if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
+                work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs)
+            else:
+                work_reqs = None
+                control_reqs = None
+
+            if self.ps.attn_tp_size != 1:
+                work_reqs = broadcast_pyobj(
+                    work_reqs,
+                    self.attn_tp_group.rank,
+                    self.attn_tp_cpu_group,
+                    src=self.attn_tp_group.ranks[0],
+                )
+
+            if self.ps.attn_cp_size != 1:
+                work_reqs = broadcast_pyobj(
+                    work_reqs,
+                    self.attn_cp_group.rank,
+                    self.attn_cp_cpu_group,
+                    src=self.attn_cp_group.ranks[0],
+                )
+
+            # When dp_attention_local_control_broadcast is enabled, each DP
+            # group leader already receives control messages from the DP
+            # controller, so we broadcast within attn_tp_group + attn_cp_group
+            # instead of the full tp_group.  This avoids an expensive
+            # all-ranks gloo sync.
+            _local_ctrl = self.server_args.enable_dp_attention_local_control_broadcast
+            if _local_ctrl:
+                if self.ps.attn_tp_size != 1:
+                    control_reqs = broadcast_pyobj(
+                        control_reqs,
+                        self.attn_tp_group.rank,
+                        self.attn_tp_cpu_group,
+                        src=self.attn_tp_group.ranks[0],
+                    )
+                if self.ps.attn_cp_size != 1:
+                    control_reqs = broadcast_pyobj(
+                        control_reqs,
+                        self.attn_cp_group.rank,
+                        self.attn_cp_cpu_group,
+                        src=self.attn_cp_group.ranks[0],
+                    )
+            elif self.ps.tp_size != 1:
+                control_reqs = broadcast_pyobj(
+                    control_reqs,
+                    self.tp_group.rank,
+                    self.tp_cpu_group,
+                    src=self.tp_group.ranks[0],
+                )
+            recv_reqs = work_reqs + control_reqs
+        elif self.ps.tp_size != 1:
+            recv_reqs = broadcast_pyobj(
+                recv_reqs,
+                self.tp_group.rank,
+                self.tp_cpu_group,
+                src=self.tp_group.ranks[0],
+            )
+        return recv_reqs
+
+    def unwrap_pickle_wrapper(self, recv_reqs: Optional[List]) -> None:
+        if not recv_reqs:
+            return
+
+        for req in recv_reqs:
+            if isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)):
+                req.unwrap_pickle_fields()
+            elif isinstance(
+                req, (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput)
+            ):
+                for sub_req in req:
+                    sub_req.unwrap_pickle_fields()
+
+    def _apply_mm_receiver(self, recv_reqs: List) -> List:
+        # Process MM requests under EPD-disaggregation mode
+        if (
+            self.ps.pp_rank == 0
+            and self.server_args.language_only
+            and self.server_args.encoder_transfer_backend
+            in ["zmq_to_scheduler", "mooncake"]
+        ):
+            recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs)
+            for req, error_msg, error_code in abort_reqs:
+                status_code = (
+                    HTTPStatus.BAD_REQUEST
+                    if error_code == 400
+                    else HTTPStatus.INTERNAL_SERVER_ERROR
+                )
+                prepare_abort(req, error_msg, status_code=status_code)
+                self.stream_output([req], req.return_logprob)
+        return recv_reqs
+
+    def _finalize_shm_features(self, recv_reqs: Optional[List]) -> None:
+        # Unwrap shared memory features AFTER all broadcasts complete,
+        # so that ShmPointerMMData metadata (not full tensor data) is what
+        # gets serialized during broadcast_pyobj.
+        if recv_reqs:
+            if self.model_config.is_multimodal and has_shm_features(recv_reqs):
+                # The broadcast source returns with its original objects while
+                # peer ranks may still be unpickling ShmPointerMMData
+                # (-> shm_open).  Synchronize the same CPU groups that carried
+                # SHM-backed work requests before materialize() unlinks them.
+                if self.server_args.enable_dp_attention:
+                    if self.ps.attn_tp_size > 1:
+                        barrier(group=self.attn_tp_cpu_group)
+                    if self.ps.attn_cp_size > 1:
+                        barrier(group=self.attn_cp_cpu_group)
+                elif self.ps.tp_size > 1:
+                    barrier(group=self.tp_cpu_group)
+            for req in recv_reqs:
+                unwrap_shm_features(req)
+
+    def _split_work_and_control_reqs(self, recv_reqs: List):
+        work_reqs = [
+            req
+            for req in recv_reqs
+            if isinstance(
+                req,
+                (
+                    TokenizedGenerateReqInput,
+                    TokenizedEmbeddingReqInput,
+                    BatchTokenizedGenerateReqInput,
+                    BatchTokenizedEmbeddingReqInput,
+                ),
+            )
+        ]
+        control_reqs = [
+            req
+            for req in recv_reqs
+            if not isinstance(
+                req,
+                (
+                    TokenizedGenerateReqInput,
+                    TokenizedEmbeddingReqInput,
+                    BatchTokenizedGenerateReqInput,
+                    BatchTokenizedEmbeddingReqInput,
+                ),
+            )
+        ]
+        return work_reqs, control_reqs

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

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import hashlib
+import logging
+import time
+import traceback
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Iterator, Optional, Tuple
+
+import torch
+
+from sglang.srt.constants import (
+    GPU_MEMORY_ALL_TYPES,
+    GPU_MEMORY_TYPE_CUDA_GRAPH,
+    GPU_MEMORY_TYPE_KV_CACHE,
+    GPU_MEMORY_TYPE_WEIGHTS,
+)
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import (
+    CheckWeightsReqInput,
+    CheckWeightsReqOutput,
+    DestroyWeightsUpdateGroupReqInput,
+    DestroyWeightsUpdateGroupReqOutput,
+    GetWeightsByNameReqInput,
+    GetWeightsByNameReqOutput,
+    InitWeightsUpdateGroupReqInput,
+    InitWeightsUpdateGroupReqOutput,
+    ReleaseMemoryOccupationReqInput,
+    ReleaseMemoryOccupationReqOutput,
+    ResumeMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqOutput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightFromDiskReqOutput,
+    UpdateWeightsFromDistributedReqInput,
+    UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromIPCReqOutput,
+    UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _get_draft_model_runner(draft_worker):
+    # DFlash / FrozenKVMTP workers expose draft_model_runner directly
+    runner = getattr(draft_worker, "draft_model_runner", None)
+    if runner is not None:
+        return runner
+    # EAGLEWorkerV2: _draft_worker.draft_runner
+    inner = getattr(draft_worker, "_draft_worker", None)
+    if inner is not None:
+        runner = getattr(inner, "draft_runner", None)
+        if runner is not None:
+            return runner
+    return None
+
+
+def _merge_checksum_payloads(target: Dict, draft: Dict) -> Dict:
+    merged_checksums = dict(target["checksums"])
+    for name, chk in draft["checksums"].items():
+        merged_checksums[f"draft.{name}"] = chk
+    h = hashlib.sha256()
+    for name in sorted(merged_checksums):
+        h.update(name.encode())
+        h.update(merged_checksums[name].encode())
+    target["checksums"] = merged_checksums
+    target["per_gpu_checksum"] = h.hexdigest()
+    return target
+
+
+@dataclass(kw_only=True, slots=True)
+class SchedulerWeightUpdaterManager:
+    tp_worker: Any
+    draft_worker: Any
+    tp_cpu_group: Any
+    memory_saver_adapter: Any
+    flush_cache: Callable[..., bool]
+    is_fully_idle: Callable[..., bool]
+    scheduler: Optional[Any] = None
+    metrics_collector: Optional[Any] = None
+    offload_tags: set = field(default_factory=set)
+    stashed_model_static_state: Any = None
+
+    @contextmanager
+    def _observe_weight_load(self, source: str) -> Iterator[None]:
+        # Edge-trigger weight_load_duration_seconds at the end of each
+        # update_weights_from_* call. Engine is paused during the update so
+        # the periodic log_stats path can't carry this.
+        # `source` distinguishes disk vs distributed vs tensor vs ipc.
+        t0 = time.perf_counter()
+        try:
+            yield
+        finally:
+            if self.metrics_collector is not None:
+                self.metrics_collector.observe_weight_load(
+                    time.perf_counter() - t0, source
+                )
+
+    def flush_cache_after_weight_update(self, recv_req) -> None:
+        if recv_req.flush_cache:
+            flush_cache_success = self.flush_cache(
+                empty_cache=recv_req.torch_empty_cache
+            )
+            assert flush_cache_success, "Cache flush failed after updating weights"
+
+    def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
+        """In-place update of the weights from disk."""
+        with self._observe_weight_load("disk"):
+            success, message = self.tp_worker.update_weights_from_disk(recv_req)
+            tp_success = success
+            if success and self.draft_worker is not None:
+                success, message = self.draft_worker.update_weights_from_disk(recv_req)
+            if tp_success:
+                self.flush_cache_after_weight_update(recv_req)
+            if not success:
+                logger.error(message)
+            return UpdateWeightFromDiskReqOutput(
+                success=success, message=message, num_paused_requests=0
+            )
+
+    def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
+        """Initialize the online model parameter update group."""
+        success, message = self.tp_worker.init_weights_update_group(recv_req)
+        return InitWeightsUpdateGroupReqOutput(success=success, message=message)
+
+    def destroy_weights_update_group(
+        self,
+        recv_req: DestroyWeightsUpdateGroupReqInput,
+    ):
+        """Destroy the online model parameter update group."""
+        success, message = self.tp_worker.destroy_weights_update_group(recv_req)
+        return DestroyWeightsUpdateGroupReqOutput(success=success, message=message)
+
+    def update_weights_from_distributed(
+        self,
+        recv_req: UpdateWeightsFromDistributedReqInput,
+    ) -> Tuple[bool, str]:
+        """Update the online model parameter."""
+        with self._observe_weight_load("distributed"):
+            success, message = self.tp_worker.update_weights_from_distributed(recv_req)
+            if success:
+                self.flush_cache_after_weight_update(recv_req)
+            else:
+                logger.error(message)
+            return UpdateWeightsFromDistributedReqOutput(
+                success=success, message=message
+            )
+
+    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+        """Update the online model parameter from tensors."""
+        with self._observe_weight_load("tensor"):
+            if recv_req.disable_draft_model:
+                worker = self.tp_worker
+            else:
+                worker = self.draft_worker or self.tp_worker
+            success, message = worker.update_weights_from_tensor(recv_req)
+            if success:
+                self.flush_cache_after_weight_update(recv_req)
+            else:
+                logger.error(message)
+            torch.distributed.barrier(group=self.tp_cpu_group)
+            return UpdateWeightsFromTensorReqOutput(success=success, message=message)
+
+    def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
+        """Update the online model parameter from IPC for checkpoint-engine integration."""
+        with self._observe_weight_load("ipc"):
+            success, message = self.tp_worker.update_weights_from_ipc(recv_req)
+            tp_success = success
+            if success and self.draft_worker is not None:
+                success, message = self.draft_worker.update_weights_from_ipc(recv_req)
+            if tp_success:
+                self.flush_cache_after_weight_update(recv_req)
+            if not success:
+                logger.error(message)
+            torch.distributed.barrier(group=self.tp_cpu_group)
+            return UpdateWeightsFromIPCReqOutput(success=success, message=message)
+
+    def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
+        parameter = self.tp_worker.get_weights_by_name(recv_req)
+        return GetWeightsByNameReqOutput(parameter=parameter)
+
+    def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
+        assert (
+            self.is_fully_idle()
+        ), "release_memory_occupation should be called only when server is idle."
+
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = GPU_MEMORY_ALL_TYPES
+
+        for tag in tags:
+            self.offload_tags.add(tag)
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            scheduler = self.scheduler
+            if scheduler is not None:
+                if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
+                    for queue_name in (
+                        "disagg_decode_transfer_queue",
+                        "disagg_decode_prealloc_queue",
+                    ):
+                        queue = getattr(scheduler, queue_name, None)
+                        if queue is not None:
+                            queue.release_memory_occupation()
+                elif scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
+                    queue = getattr(scheduler, "disagg_prefill_bootstrap_queue", None)
+                    if queue is not None:
+                        queue.release_memory_occupation()
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
+            self.flush_cache()
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.stashed_model_static_state = _export_static_state(
+                self.tp_worker.model_runner.model
+            )
+            torch.distributed.barrier(self.tp_cpu_group)
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
+
+        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH)
+
+        torch.get_device_module().synchronize()
+
+        return ReleaseMemoryOccupationReqOutput()
+
+    def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = GPU_MEMORY_ALL_TYPES
+
+        for tag in tags:
+            self.offload_tags.remove(tag)
+
+        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH)
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
+            torch.distributed.barrier(self.tp_cpu_group)
+            _import_static_state(
+                self.tp_worker.model_runner.model,
+                self.stashed_model_static_state,
+            )
+            del self.stashed_model_static_state
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
+            scheduler = self.scheduler
+            if scheduler is not None:
+                if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
+                    for queue_name in (
+                        "disagg_decode_transfer_queue",
+                        "disagg_decode_prealloc_queue",
+                    ):
+                        queue = getattr(scheduler, queue_name, None)
+                        if queue is not None:
+                            queue.resume_memory_occupation()
+                elif scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
+                    queue = getattr(scheduler, "disagg_prefill_bootstrap_queue", None)
+                    if queue is not None:
+                        queue.resume_memory_occupation()
+
+        return ResumeMemoryOccupationReqOutput()
+
+    def check_weights(self, recv_req: CheckWeightsReqInput):
+        try:
+            payload = self.tp_worker.model_runner.check_weights(
+                action=recv_req.action, allow_quant_error=recv_req.allow_quant_error
+            )
+
+            if self.draft_worker is not None:
+                draft_runner = _get_draft_model_runner(self.draft_worker)
+                if draft_runner is not None:
+                    draft_payload = draft_runner.check_weights(
+                        action=recv_req.action,
+                        allow_quant_error=recv_req.allow_quant_error,
+                    )
+                    if payload is not None and draft_payload is not None:
+                        payload = _merge_checksum_payloads(payload, draft_payload)
+
+            tp_size = torch.distributed.get_world_size(group=self.tp_cpu_group)
+            if tp_size > 1 and payload is not None:
+                all_payloads = [None] * tp_size
+                torch.distributed.all_gather_object(
+                    all_payloads, payload, group=self.tp_cpu_group
+                )
+                payload = all_payloads
+            return CheckWeightsReqOutput(
+                success=True, message="Success.", payload=payload
+            )
+        except Exception as e:
+            logger.warning(f"check_weights see error: {e}")
+            traceback.print_exc()
+            return CheckWeightsReqOutput(success=False, message=f"{e}")
+
+    def save_remote_model(self, params):
+        url = params["url"]
+
+        self.tp_worker.model_runner.save_remote_model(url)
+
+        if self.draft_worker is not None:
+            draft_url = params.get("draft_url", None)
+            assert (
+                draft_url is not None
+            ), "draft_url must be provided when draft model is enabled"
+            self.draft_worker.model_runner.save_remote_model(draft_url)
+
+    def save_sharded_model(self, params):
+        self.tp_worker.model_runner.save_sharded_model(
+            path=params["path"],
+            pattern=params["pattern"],
+            max_size=params["max_size"],
+        )
+
+
+def _export_static_state(model):
+    return dict(
+        buffers=[
+            (name, buffer.detach().clone()) for name, buffer in model.named_buffers()
+        ]
+    )
+
+
+def _import_static_state(model, static_params):
+    with torch.inference_mode():
+        self_named_buffers = dict(model.named_buffers())
+        for name, tensor in static_params["buffers"]:
+            self_named_buffers[name][...] = tensor

python/sglang/srt/managers/scheduler_input_blocker.py

@@ -0,0 +1,106 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import logging
+from contextlib import contextmanager
+from enum import Enum, auto
+from typing import Any, Callable, List, Optional
+
+from sglang.srt.managers.io_struct import BlockReqInput, BlockReqType
+from sglang.srt.utils.poll_based_barrier import PollBasedBarrier
+
+logger = logging.getLogger(__name__)
+
+
+class SchedulerInputBlocker:
+    def __init__(self, noop: bool):
+        self._state = _State.UNBLOCKED
+        self._pending_reqs = []
+        self._noop = noop
+        self._global_unblock_barrier = PollBasedBarrier(noop=noop)
+
+    def handle(self, recv_reqs: Optional[List[Any]]):
+        assert (recv_reqs is None) == self._noop
+
+        if not self._noop:
+            output_reqs = []
+            for recv_req in recv_reqs:
+                output_reqs += self._handle_recv_req(recv_req)
+
+        global_arrived_unblock_barrier = (
+            self._global_unblock_barrier.poll_global_arrived()
+        )
+        if (
+            self._state == _State.GLOBAL_UNBLOCK_BARRIER
+            and global_arrived_unblock_barrier
+        ):
+            output_reqs += self._handle_arrive_unblock_barrier()
+
+        if not self._noop:
+            return output_reqs
+
+    def _handle_recv_req(self, recv_req):
+        if isinstance(recv_req, BlockReqInput):
+            if recv_req.req_type == BlockReqType.BLOCK:
+                self._execute_block_req()
+                return []
+            elif recv_req.req_type == BlockReqType.UNBLOCK:
+                self._execute_unblock_req()
+                return []
+            else:
+                raise NotImplementedError(f"{recv_req=}")
+        else:
+            if self._state == _State.UNBLOCKED:
+                return [recv_req]
+            else:
+                self._pending_reqs.append(recv_req)
+                return []
+
+    def _execute_block_req(self):
+        logger.info("Handle block req")
+        self._change_state(original=_State.UNBLOCKED, target=_State.BLOCKED)
+
+    def _execute_unblock_req(self):
+        logger.info("Handle unblock req")
+        self._change_state(
+            original=_State.BLOCKED, target=_State.GLOBAL_UNBLOCK_BARRIER
+        )
+        self._global_unblock_barrier.local_arrive()
+
+    def _handle_arrive_unblock_barrier(self):
+        logger.info(f"Arrived at unblock barrier ({len(self._pending_reqs)=})")
+        self._change_state(
+            original=_State.GLOBAL_UNBLOCK_BARRIER, target=_State.UNBLOCKED
+        )
+        output_reqs = [*self._pending_reqs]
+        self._pending_reqs.clear()
+        return output_reqs
+
+    def _change_state(self, original: "_State", target: "_State"):
+        assert self._state == original, f"{self._state=} {original=} {target=}"
+        self._state = target
+
+
+class _State(Enum):
+    UNBLOCKED = auto()
+    BLOCKED = auto()
+    GLOBAL_UNBLOCK_BARRIER = auto()
+
+
+@contextmanager
+def input_blocker_guard_region(dispatch_to_scheduler: Callable[[BlockReqInput], None]):
+    dispatch_to_scheduler(BlockReqInput(req_type=BlockReqType.BLOCK))
+    try:
+        yield
+    finally:
+        dispatch_to_scheduler(BlockReqInput(req_type=BlockReqType.UNBLOCK))

python/sglang/srt/managers/scheduler_recv_skipper.py

@@ -0,0 +1,38 @@
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.server_args import ServerArgs
+
+
+class SchedulerRecvSkipper:
+    @staticmethod
+    def maybe_create(server_args: ServerArgs):
+        if server_args.scheduler_recv_interval <= 1:
+            return None
+        return SchedulerRecvSkipper(server_args)
+
+    def __init__(self, server_args: ServerArgs):
+        # Can be supported if needed, but may need e.g. `global_forward_mode`
+        assert not server_args.enable_dp_attention
+        self._counter = 0
+        self._threshold = server_args.scheduler_recv_interval
+        # All can be tuned if needed
+        self._default_weight = envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DEFAULT.get()
+        self._weight_of_forward_mode = {
+            ForwardMode.DECODE: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DECODE.get(),
+            ForwardMode.TARGET_VERIFY: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_TARGET_VERIFY.get(),
+            None: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE.get(),
+        }
+
+    def handle(self, last_forward_mode: ForwardMode):
+        should_recv = False
+
+        last_weight = self._weight_of_forward_mode.get(
+            last_forward_mode, self._default_weight
+        )
+        self._counter += last_weight
+
+        if self._counter >= self._threshold:
+            self._counter = 0
+            should_recv = True
+
+        return should_recv

python/sglang/srt/managers/tokenizer_control_mixin.py

@@ -0,0 +1,876 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+import time
+import uuid
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import fastapi
+
+from sglang.srt.managers.communicator import FanOutCommunicator
+from sglang.srt.managers.io_struct import (
+    AddExternalCorpusReqInput,
+    AddExternalCorpusReqOutput,
+    AttachHiCacheStorageReqInput,
+    AttachHiCacheStorageReqOutput,
+    CheckWeightsReqInput,
+    CheckWeightsReqOutput,
+    ClearHiCacheReqInput,
+    ClearHiCacheReqOutput,
+    CloseSessionReqInput,
+    DestroyWeightsUpdateGroupReqInput,
+    DestroyWeightsUpdateGroupReqOutput,
+    DetachHiCacheStorageReqInput,
+    DetachHiCacheStorageReqOutput,
+    DumperControlReqInput,
+    DumperControlReqOutput,
+    ExpertDistributionReq,
+    ExpertDistributionReqOutput,
+    ExpertDistributionReqType,
+    FlushCacheReqInput,
+    FlushCacheReqOutput,
+    GetInternalStateReq,
+    GetInternalStateReqOutput,
+    GetWeightsByNameReqInput,
+    GetWeightsByNameReqOutput,
+    InitWeightsSendGroupForRemoteInstanceReqInput,
+    InitWeightsSendGroupForRemoteInstanceReqOutput,
+    InitWeightsUpdateGroupReqInput,
+    InitWeightsUpdateGroupReqOutput,
+    ListExternalCorporaReqInput,
+    ListExternalCorporaReqOutput,
+    LoadLoRAAdapterFromTensorsReqInput,
+    LoadLoRAAdapterFromTensorsReqOutput,
+    LoadLoRAAdapterReqInput,
+    LoadLoRAAdapterReqOutput,
+    LoRAUpdateOutput,
+    OpenSessionReqInput,
+    ProfileReq,
+    ProfileReqOutput,
+    ProfileReqType,
+    ReleaseMemoryOccupationReqInput,
+    ReleaseMemoryOccupationReqOutput,
+    RemoveExternalCorpusReqInput,
+    RemoveExternalCorpusReqOutput,
+    ResumeMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqOutput,
+    SendWeightsToRemoteInstanceReqInput,
+    SendWeightsToRemoteInstanceReqOutput,
+    SetInternalStateReq,
+    SetInternalStateReqOutput,
+    SlowDownReqInput,
+    SlowDownReqOutput,
+    UnloadLoRAAdapterReqInput,
+    UnloadLoRAAdapterReqOutput,
+    UpdateWeightsFromDistributedReqInput,
+    UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromIPCReqOutput,
+    UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput,
+)
+from sglang.srt.managers.load_snapshot import LoadSnapshot
+from sglang.srt.server_args import LoRARef, ServerArgs
+from sglang.srt.utils import (
+    get_bool_env_var,
+    normalize_serialized_named_tensor_payloads,
+)
+from sglang.utils import TypeBasedDispatcher
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tokenizer_manager import TokenizerManager
+
+logger = logging.getLogger(__name__)
+
+# Declarative spec: (attr_name_prefix, response_type[, mode])
+# Each entry creates self.{prefix}_communicator and registers
+# response_type -> communicator.handle_recv in the dispatch table.
+_COMMUNICATOR_SPECS = [
+    ("init_weights_update_group", InitWeightsUpdateGroupReqOutput),
+    ("destroy_weights_update_group", DestroyWeightsUpdateGroupReqOutput),
+    ("update_weights_from_distributed", UpdateWeightsFromDistributedReqOutput),
+    (
+        "init_weights_send_group_for_remote_instance",
+        InitWeightsSendGroupForRemoteInstanceReqOutput,
+    ),
+    ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput),
+    ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput),
+    ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput),
+    ("get_weights_by_name", GetWeightsByNameReqOutput),
+    ("release_memory_occupation", ReleaseMemoryOccupationReqOutput),
+    ("resume_memory_occupation", ResumeMemoryOccupationReqOutput),
+    ("check_weights", CheckWeightsReqOutput),
+    ("slow_down", SlowDownReqOutput),
+    ("flush_cache", FlushCacheReqOutput),
+    ("add_external_corpus", AddExternalCorpusReqOutput),
+    ("remove_external_corpus", RemoveExternalCorpusReqOutput),
+    ("list_external_corpora", ListExternalCorporaReqOutput),
+    ("clear_hicache_storage", ClearHiCacheReqOutput),
+    ("attach_hicache_storage", AttachHiCacheStorageReqOutput),
+    ("detach_hicache_storage", DetachHiCacheStorageReqOutput),
+    ("profile", ProfileReqOutput),
+    ("get_internal_state", GetInternalStateReqOutput),
+    ("set_internal_state", SetInternalStateReqOutput),
+    ("expert_distribution", ExpertDistributionReqOutput),
+    ("update_lora_adapter", LoRAUpdateOutput),
+    ("dumper_control", DumperControlReqOutput),
+]
+
+
+class TokenizerControlMixin:
+    """Mixin for TokenizerManager's control-plane operations (weights, cache, lora,
+    profile, internal state, etc.) -- everything that talks to the scheduler via
+    FanOutCommunicator, as opposed to data-plane inference requests multiplexed by rid.
+    """
+
+    def init_communicators(self: TokenizerManager, server_args: ServerArgs):
+        dispatch_pairs = []
+        for spec in _COMMUNICATOR_SPECS:
+            name, resp_type = spec[0], spec[1]
+            mode = spec[2] if len(spec) > 2 else "queueing"
+            comm = FanOutCommunicator(
+                self._dispatch_to_scheduler,
+                server_args.dp_size,
+                mode,
+            )
+            setattr(self, f"{name}_communicator", comm)
+            dispatch_pairs.append((resp_type, comm.handle_recv))
+        self._result_dispatcher += TypeBasedDispatcher(dispatch_pairs)
+
+    async def add_external_corpus(
+        self: TokenizerManager, obj: AddExternalCorpusReqInput
+    ) -> AddExternalCorpusReqOutput:
+        self.auto_create_handle_loop()
+        if self.server_args.speculative_algorithm != "NGRAM":
+            return AddExternalCorpusReqOutput(
+                success=False,
+                message="Ngram speculative decoding is not enabled.",
+            )
+        truncated = False
+        try:
+            if not obj.corpus_id:
+                import uuid
+
+                obj.corpus_id = uuid.uuid4().hex
+            if obj.file_path is not None:
+                from sglang.srt.speculative.cpp_ngram.external_corpus import (
+                    iter_external_corpus_chunks,
+                )
+
+                max_tokens = (
+                    self.server_args.speculative_ngram_external_corpus_max_tokens
+                )
+                obj.token_chunks = list(
+                    iter_external_corpus_chunks(
+                        obj.file_path, self.tokenizer, max_tokens
+                    )
+                )
+            elif obj.documents is not None:
+                from sglang.srt.speculative.cpp_ngram.external_corpus import (
+                    SEPARATOR_TOKEN,
+                )
+
+                max_tokens = (
+                    self.server_args.speculative_ngram_external_corpus_max_tokens
+                )
+                token_chunks = []
+                total_tokens = 0
+                has_prev = False
+                for doc in obj.documents:
+                    if not doc:
+                        continue
+                    token_ids = list(
+                        self.tokenizer.encode(doc, add_special_tokens=False)
+                    )
+                    if not token_ids:
+                        continue
+                    if has_prev:
+                        token_ids = [SEPARATOR_TOKEN] + token_ids
+                    if total_tokens + len(token_ids) > max_tokens:
+                        truncated = True
+                        break
+                    token_chunks.append(token_ids)
+                    total_tokens += len(token_ids)
+                    has_prev = True
+                obj.token_chunks = token_chunks
+            else:
+                return AddExternalCorpusReqOutput(
+                    success=False,
+                    message="Either file_path or documents must be provided.",
+                )
+            obj.file_path = None
+            obj.documents = None
+            results = await self.add_external_corpus_communicator(obj)
+            all_success, all_message = FanOutCommunicator.merge_results(results)
+            if truncated and all_success:
+                all_message += f" (truncated: exceeded {max_tokens} token limit)"
+            return AddExternalCorpusReqOutput(
+                success=all_success,
+                corpus_id=results[0].corpus_id if all_success else "",
+                message=all_message,
+                loaded_token_count=results[0].loaded_token_count if all_success else 0,
+            )
+        except Exception as e:
+            return AddExternalCorpusReqOutput(success=False, message=str(e))
+
+    async def remove_external_corpus(
+        self: TokenizerManager, corpus_id: str
+    ) -> RemoveExternalCorpusReqOutput:
+        self.auto_create_handle_loop()
+        if self.server_args.speculative_algorithm != "NGRAM":
+            return RemoveExternalCorpusReqOutput(
+                success=False,
+                message="Ngram speculative decoding is not enabled.",
+            )
+        results = await self.remove_external_corpus_communicator(
+            RemoveExternalCorpusReqInput(corpus_id=corpus_id)
+        )
+        all_success, all_message = FanOutCommunicator.merge_results(results)
+        return RemoveExternalCorpusReqOutput(success=all_success, message=all_message)
+
+    async def list_external_corpora(
+        self: TokenizerManager,
+    ) -> ListExternalCorporaReqOutput:
+        self.auto_create_handle_loop()
+        if self.server_args.speculative_algorithm != "NGRAM":
+            return ListExternalCorporaReqOutput(
+                success=False,
+                message="Ngram speculative decoding is not enabled.",
+            )
+        results = await self.list_external_corpora_communicator(
+            ListExternalCorporaReqInput()
+        )
+        all_success, all_message = FanOutCommunicator.merge_results(results)
+        # Merge corpus token counts from all DP ranks (each rank loads the same set).
+        corpus_token_counts = results[0].corpus_token_counts if all_success else {}
+        return ListExternalCorporaReqOutput(
+            success=all_success,
+            corpus_token_counts=corpus_token_counts,
+            message=all_message,
+        )
+
+    async def flush_cache(
+        self: TokenizerManager, timeout_s: Optional[float] = None
+    ) -> FlushCacheReqOutput:
+        self.auto_create_handle_loop()
+        return (
+            await self.flush_cache_communicator(FlushCacheReqInput(timeout_s=timeout_s))
+        )[0]
+
+    async def clear_hicache_storage(self: TokenizerManager) -> ClearHiCacheReqOutput:
+        """Clear the hierarchical cache storage."""
+        self.auto_create_handle_loop()
+        # Delegate to the scheduler to handle HiCacheStorage clearing
+        return (await self.clear_hicache_storage_communicator(ClearHiCacheReqInput()))[
+            0
+        ]
+
+    async def attach_hicache_storage(
+        self: TokenizerManager,
+        hicache_storage_backend: str,
+        hicache_storage_backend_extra_config_json: Optional[str] = None,
+        hicache_storage_prefetch_policy: Optional[str] = None,
+        hicache_write_policy: Optional[str] = None,
+    ) -> AttachHiCacheStorageReqOutput:
+        """Attach (enable) HiCache storage backend at runtime."""
+        self.auto_create_handle_loop()
+        results = await self.attach_hicache_storage_communicator(
+            AttachHiCacheStorageReqInput(
+                hicache_storage_backend=hicache_storage_backend,
+                hicache_storage_backend_extra_config_json=hicache_storage_backend_extra_config_json,
+                hicache_storage_prefetch_policy=hicache_storage_prefetch_policy,
+                hicache_write_policy=hicache_write_policy,
+            )
+        )
+
+        all_success, all_message = FanOutCommunicator.merge_results(results)
+        out = AttachHiCacheStorageReqOutput(success=all_success, message=all_message)
+        # TODO: partial rollback if failed
+        if all_success:
+            # Keep tokenizer side server_info consistent with scheduler side.
+            hicache_fields = {"hicache_storage_backend": hicache_storage_backend}
+            if hicache_storage_backend_extra_config_json is not None:
+                hicache_fields["hicache_storage_backend_extra_config"] = (
+                    hicache_storage_backend_extra_config_json
+                )
+            if hicache_storage_prefetch_policy is not None:
+                hicache_fields["hicache_storage_prefetch_policy"] = (
+                    hicache_storage_prefetch_policy
+                )
+            if hicache_write_policy is not None:
+                hicache_fields["hicache_write_policy"] = hicache_write_policy
+            self.server_args.override("tokenizer.attach_hicache", **hicache_fields)
+        return out
+
+    async def detach_hicache_storage(
+        self: TokenizerManager,
+    ) -> DetachHiCacheStorageReqOutput:
+        """Detach (disable) HiCache storage backend at runtime."""
+        self.auto_create_handle_loop()
+        results = await self.detach_hicache_storage_communicator(
+            DetachHiCacheStorageReqInput()
+        )
+
+        all_success, all_message = FanOutCommunicator.merge_results(results)
+        out = DetachHiCacheStorageReqOutput(success=all_success, message=all_message)
+        # TODO: partial rollback if failed
+        if all_success:
+            self.server_args.override(
+                "tokenizer.detach_hicache",
+                hicache_storage_backend=None,
+                hicache_storage_backend_extra_config=None,
+            )
+        return out
+
+    async def start_profile(
+        self: TokenizerManager,
+        req: Optional[ProfileReq] = None,
+    ):
+        self.auto_create_handle_loop()
+        req = req or ProfileReq()
+        req.req_type = ProfileReqType.START_PROFILE
+        env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
+        req.with_stack = (
+            False if req.with_stack is False or env_with_stack is False else True
+        )
+        env_record_shapes: bool = get_bool_env_var(
+            "SGLANG_PROFILE_RECORD_SHAPES", "true"
+        )
+        req.record_shapes = (req.record_shapes is not False) and env_record_shapes
+        req.profile_id = req.profile_id or str(time.time())
+        return await self._execute_profile(req)
+
+    async def stop_profile(self: TokenizerManager):
+        self.auto_create_handle_loop()
+        req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
+        return await self._execute_profile(req)
+
+    async def _execute_profile(self: TokenizerManager, req: ProfileReq):
+        result = (await self.profile_communicator(req))[0]
+        if not result.success:
+            raise RuntimeError(result.message)
+        return result
+
+    async def start_expert_distribution_record(self: TokenizerManager):
+        self.auto_create_handle_loop()
+        req = ExpertDistributionReq(action=ExpertDistributionReqType.START_RECORD)
+        await self.expert_distribution_communicator(req)
+
+    async def stop_expert_distribution_record(self: TokenizerManager):
+        self.auto_create_handle_loop()
+        req = ExpertDistributionReq(action=ExpertDistributionReqType.STOP_RECORD)
+        await self.expert_distribution_communicator(req)
+
+    async def dump_expert_distribution_record(self: TokenizerManager):
+        self.auto_create_handle_loop()
+        req = ExpertDistributionReq(action=ExpertDistributionReqType.DUMP_RECORD)
+        await self.expert_distribution_communicator(req)
+
+    async def init_weights_update_group(
+        self: TokenizerManager,
+        obj: InitWeightsUpdateGroupReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        assert (
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
+
+        results = await self.init_weights_update_group_communicator(obj)
+        return FanOutCommunicator.merge_results(results)
+
+    async def destroy_weights_update_group(
+        self: TokenizerManager,
+        obj: DestroyWeightsUpdateGroupReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        assert (
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group"
+
+        results = await self.destroy_weights_update_group_communicator(obj)
+        return FanOutCommunicator.merge_results(results)
+
+    async def update_weights_from_distributed(
+        self: TokenizerManager,
+        obj: UpdateWeightsFromDistributedReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        assert (
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
+
+        if obj.abort_all_requests:
+            self.abort_request(abort_all=True)
+
+        # Hold is_pause_cond while updating to prevent unpause from racing.
+        async with self.is_pause_cond:
+            is_paused = self.is_pause
+            if is_paused:
+                results = await self.update_weights_from_distributed_communicator(obj)
+
+        if not is_paused:
+            async with self.model_update_lock.writer_lock:
+                results = await self.update_weights_from_distributed_communicator(obj)
+
+        success, message = FanOutCommunicator.merge_results(results)
+        if success and obj.weight_version is not None:
+            self._update_weight_version_if_provided(obj.weight_version)
+            message += f" Weight version updated to {obj.weight_version}."
+
+        return success, message
+
+    async def init_weights_send_group_for_remote_instance(
+        self: TokenizerManager,
+        obj: InitWeightsSendGroupForRemoteInstanceReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        # TODO: support DP
+        assert (
+            self.server_args.dp_size == 1
+        ), "dp_size must be 1 for init_weights_send_group_for_remote_instance"
+        result = (
+            await self.init_weights_send_group_for_remote_instance_communicator(obj)
+        )[0]
+        return result.success, result.message
+
+    async def send_weights_to_remote_instance(
+        self: TokenizerManager,
+        obj: SendWeightsToRemoteInstanceReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        # TODO: support DP
+        assert (
+            self.server_args.dp_size == 1
+        ), "dp_size must be 1 for send_weights_to_remote_instance"
+        result = (await self.send_weights_to_remote_instance_communicator(obj))[0]
+        return result.success, result.message
+
+    async def update_weights_from_tensor(
+        self: TokenizerManager,
+        obj: UpdateWeightsFromTensorReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        assert (
+            self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+        ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
+
+        if obj.abort_all_requests:
+            self.abort_request(abort_all=True)
+
+        obj.serialized_named_tensors = normalize_serialized_named_tensor_payloads(
+            obj.serialized_named_tensors
+        )
+
+        async with self.is_pause_cond:
+            is_paused = self.is_pause
+            if is_paused:
+                results = await self.update_weights_from_tensor_communicator(obj)
+
+        if not is_paused:
+            async with self.model_update_lock.writer_lock:
+                results = await self.update_weights_from_tensor_communicator(obj)
+
+        success, message = FanOutCommunicator.merge_results(results)
+        if success and obj.weight_version is not None:
+            self._update_weight_version_if_provided(obj.weight_version)
+            message += f" Weight version updated to {obj.weight_version}."
+
+        return success, message
+
+    async def update_weights_from_ipc(
+        self: TokenizerManager,
+        obj: UpdateWeightsFromIPCReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        """Update weights via IPC for checkpoint-engine integration."""
+        self.auto_create_handle_loop()
+        try:
+            # For now, we only support single data parallel instance
+            assert (
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+            ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
+            logger.info("Starting IPC weight update")
+
+            async with self.is_pause_cond:
+                is_paused = self.is_pause
+                if is_paused:
+                    result = (await self.update_weights_from_ipc_communicator(obj))[0]
+                    success, message = result.success, result.message
+
+            if not is_paused:
+                async with self.model_update_lock.writer_lock:
+                    result = (await self.update_weights_from_ipc_communicator(obj))[0]
+                    success, message = result.success, result.message
+        except Exception as e:
+            error_msg = f"IPC weight update failed: {str(e)}"
+            logger.error(error_msg)
+            success, message = False, error_msg
+
+        if success and obj.weight_version is not None:
+            self._update_weight_version_if_provided(obj.weight_version)
+            message += f" Weight version updated to {obj.weight_version}."
+
+        return success, message
+
+    async def _unload_lora_adapter_locked(
+        self: TokenizerManager,
+        obj: UnloadLoRAAdapterReqInput,
+    ) -> UnloadLoRAAdapterReqOutput:
+        assert (
+            self.lora_update_lock.locked()
+        ), "self.lora_update_lock must be locked in order for self._unload_lora_adapter_locked() to be called"
+
+        # Unregister the LoRA adapter from the registry to stop new requests for this adapter
+        # from being started.
+        lora_id = await self.lora_registry.unregister(obj.lora_name)
+        obj.lora_id = lora_id
+
+        # Initiate the actual unloading operation at the backend processes only after all
+        # ongoing requests using this LoRA adapter are finished.
+        await self.lora_registry.wait_for_unload(lora_id)
+        result = (await self.update_lora_adapter_communicator(obj))[0]
+
+        return result
+
+    async def load_lora_adapter(
+        self: TokenizerManager,
+        obj: LoadLoRAAdapterReqInput,
+        _: Optional[fastapi.Request] = None,
+    ) -> LoadLoRAAdapterReqOutput:
+        self.auto_create_handle_loop()
+
+        try:
+            if not self.server_args.enable_lora:
+                raise ValueError(
+                    "LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
+                )
+
+            # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
+            # with dp_size > 1.
+            assert (
+                self.server_args.dp_size == 1
+            ), "dp_size must be 1 for dynamic lora loading"
+            logger.info(
+                "Start load Lora adapter. Lora name=%s, path=%s",
+                obj.lora_name,
+                obj.lora_path,
+            )
+
+            async with self.lora_update_lock:
+                # Generate new uniquely identifiable LoRARef object.
+                new_adapter = LoRARef(
+                    lora_name=obj.lora_name,
+                    lora_path=obj.lora_path,
+                    pinned=obj.pinned,
+                )
+
+                # Trigger the actual loading operation at the backend processes.
+                obj.lora_id = new_adapter.lora_id
+                result = (await self.update_lora_adapter_communicator(obj))[0]
+
+                # Register the LoRA adapter only after loading is successful.
+                if result.success:
+                    await self.lora_registry.register(new_adapter)
+                    self.lora_ref_cache[obj.lora_name] = new_adapter
+
+                if self.server_args.max_loaded_loras is not None:
+                    while (
+                        self.lora_registry.num_registered_loras
+                        > self.server_args.max_loaded_loras
+                    ):
+                        lru_lora_name = await self.lora_registry.lru_lora_name(
+                            exclude_pinned=True
+                        )
+                        if lru_lora_name is None:
+                            raise ValueError(
+                                "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. "
+                                f"LoRA registry is: {self.lora_registry._registry}"
+                            )
+
+                        logger.info(
+                            f"Unloading least recently used LoRA adapter '{lru_lora_name}' "
+                            f"(current number of adapters: {self.lora_registry.num_registered_loras}, "
+                            f"max allowed: {self.server_args.max_loaded_loras})"
+                        )
+
+                        unload_result = await self._unload_lora_adapter_locked(
+                            UnloadLoRAAdapterReqInput(lora_name=lru_lora_name)
+                        )
+                        if not unload_result.success:
+                            raise ValueError(
+                                f"Error while unloading LRU LoRA adapter '{lru_lora_name}': "
+                                f"{unload_result.error_message}"
+                            )
+                        del result.loaded_adapters[lru_lora_name]
+
+                return result
+        except ValueError as e:
+            return LoadLoRAAdapterReqOutput(
+                success=False,
+                error_message=str(e),
+            )
+
+    async def load_lora_adapter_from_tensors(
+        self: TokenizerManager,
+        obj: LoadLoRAAdapterFromTensorsReqInput,
+        _: Optional[fastapi.Request] = None,
+    ) -> LoadLoRAAdapterFromTensorsReqOutput:
+        self.auto_create_handle_loop()
+
+        try:
+            if not self.server_args.enable_lora:
+                raise ValueError(
+                    "LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
+                )
+
+            assert (
+                self.server_args.dp_size == 1
+            ), "dp_size must be 1 for dynamic lora loading"
+            logger.info(
+                "Start load Lora adapter from tensors. Lora name=%s",
+                obj.lora_name,
+            )
+
+            async with self.lora_update_lock:
+                new_adapter = LoRARef(
+                    lora_name=obj.lora_name,
+                    lora_path="__tensor__",
+                    pinned=obj.pinned,
+                )
+                obj.lora_id = new_adapter.lora_id
+                result = (await self.update_lora_adapter_communicator(obj))[0]
+
+                if result.success:
+                    await self.lora_registry.register(new_adapter)
+                    self.lora_ref_cache[obj.lora_name] = new_adapter
+                if self.server_args.max_loaded_loras is not None:
+                    while (
+                        self.lora_registry.num_registered_loras
+                        > self.server_args.max_loaded_loras
+                    ):
+                        lru_lora_name = await self.lora_registry.lru_lora_name(
+                            exclude_pinned=True
+                        )
+                        if lru_lora_name is None:
+                            raise ValueError(
+                                "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. "
+                                f"LoRA registry is: {self.lora_registry._registry}"
+                            )
+
+                        logger.info(
+                            f"Unloading least recently used LoRA adapter '{lru_lora_name}' "
+                            f"(current number of adapters: {self.lora_registry.num_registered_loras}, "
+                            f"max allowed: {self.server_args.max_loaded_loras})"
+                        )
+
+                        unload_result = await self._unload_lora_adapter_locked(
+                            UnloadLoRAAdapterReqInput(lora_name=lru_lora_name)
+                        )
+                        if not unload_result.success:
+                            raise ValueError(
+                                f"Error while unloading LRU LoRA adapter '{lru_lora_name}': "
+                                f"{unload_result.error_message}"
+                            )
+                        del result.loaded_adapters[lru_lora_name]
+
+                return result
+        except ValueError as e:
+            return LoadLoRAAdapterFromTensorsReqOutput(
+                success=False,
+                error_message=str(e),
+            )
+
+    async def unload_lora_adapter(
+        self: TokenizerManager,
+        obj: UnloadLoRAAdapterReqInput,
+        _: Optional[fastapi.Request] = None,
+    ) -> UnloadLoRAAdapterReqOutput:
+        self.auto_create_handle_loop()
+
+        try:
+            if not self.server_args.enable_lora:
+                raise ValueError(
+                    "LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
+                )
+
+            assert (
+                obj.lora_name is not None
+            ), "lora_name must be provided to unload LoRA adapter"
+
+            # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
+            # with dp_size > 1.
+            assert (
+                self.server_args.dp_size == 1
+            ), "dp_size must be 1 for dynamic lora loading"
+            logger.info(
+                "Start unload Lora adapter. Lora name=%s",
+                obj.lora_name,
+            )
+
+            async with self.lora_update_lock:
+                return await self._unload_lora_adapter_locked(obj)
+        except ValueError as e:
+            return UnloadLoRAAdapterReqOutput(success=False, error_message=str(e))
+
+    async def get_weights_by_name(
+        self: TokenizerManager,
+        obj: GetWeightsByNameReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        results = await self.get_weights_by_name_communicator(obj)
+        all_parameters = [r.parameter for r in results]
+        if self.server_args.dp_size == 1:
+            return all_parameters[0]
+        else:
+            return all_parameters
+
+    async def release_memory_occupation(
+        self: TokenizerManager,
+        obj: ReleaseMemoryOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.release_memory_occupation_communicator(obj)
+
+    async def resume_memory_occupation(
+        self: TokenizerManager,
+        obj: ResumeMemoryOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.resume_memory_occupation_communicator(obj)
+
+    async def check_weights(
+        self: TokenizerManager,
+        obj: CheckWeightsReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str, Optional[List[Dict]], Optional[str]]:
+        self.auto_create_handle_loop()
+        results = await self.check_weights_communicator(obj)
+        success, message = FanOutCommunicator.merge_results(results)
+        ranks: Optional[List[Dict]] = None
+        per_engine_checksum: Optional[str] = None
+        if any(r.payload is not None for r in results):
+            ranks = []
+            for r in results:
+                if isinstance(r.payload, list):
+                    ranks.extend(r.payload)
+                else:
+                    ranks.append(r.payload)
+            h = hashlib.sha256()
+            for rank in ranks:
+                h.update(rank["per_gpu_checksum"].encode())
+            per_engine_checksum = h.hexdigest()
+        return success, message, ranks, per_engine_checksum
+
+    async def slow_down(
+        self: TokenizerManager,
+        obj: SlowDownReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.slow_down_communicator(obj)
+
+    async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]:
+        self.auto_create_handle_loop()
+        req = GetInternalStateReq()
+        responses: List[GetInternalStateReqOutput] = (
+            await self.get_internal_state_communicator(req)
+        )
+        # Many DP ranks
+        return [res.internal_state for res in responses]
+
+    async def set_internal_state(
+        self: TokenizerManager, obj: SetInternalStateReq
+    ) -> List[bool]:
+        self.auto_create_handle_loop()
+        responses: List[SetInternalStateReqOutput] = (
+            await self.set_internal_state_communicator(obj)
+        )
+        return [res.updated for res in responses]
+
+    async def dumper_control(
+        self: TokenizerManager, obj: DumperControlReqInput
+    ) -> List[DumperControlReqOutput]:
+        self.auto_create_handle_loop()
+        return await self.dumper_control_communicator(obj)
+
+    async def get_loads(
+        self: TokenizerManager,
+        include: Optional[List[str]] = None,
+        dp_rank: Optional[int] = None,
+    ) -> List[LoadSnapshot]:
+        """
+        Get load snapshots for /v1/loads endpoint.
+
+        Args:
+            include: List of sections to include. Options: core, memory, spec, lora, disagg, queues, all
+            dp_rank: Optional filter for specific DP rank
+
+        Returns:
+            List of LoadSnapshot, one per scheduler (filtered by dp_rank if specified)
+        """
+        self.auto_create_handle_loop()
+        if dp_rank is not None and (dp_rank < 0 or dp_rank >= self.server_args.dp_size):
+            return []
+
+        reader = self.load_snapshot_reader
+        if dp_rank is not None:
+            load = reader.read(dp_rank)
+            results = [load] if load is not None else []
+        else:
+            results = reader.read_all()
+
+        return results
+
+    async def open_session(
+        self: TokenizerManager,
+        obj: OpenSessionReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        if obj.streaming:
+            if not self.server_args.enable_streaming_session:
+                raise ValueError(
+                    "Streaming sessions are disabled. "
+                    "Please relaunch with --enable-streaming-session."
+                )
+
+        if obj.session_id is None:
+            obj.session_id = uuid.uuid4().hex
+        elif obj.session_id in self.session_futures:
+            return None
+
+        future = asyncio.Future()
+        self.session_futures[obj.session_id] = future
+        self._dispatch_to_scheduler(obj)
+
+        try:
+            return await future
+        finally:
+            self.session_futures.pop(obj.session_id, None)
+
+    async def close_session(
+        self: TokenizerManager,
+        obj: CloseSessionReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        await self._async_dispatch_to_scheduler(obj)
+
+    def _update_weight_version_if_provided(
+        self: TokenizerManager, weight_version: Optional[str]
+    ) -> None:
+        """Update weight version if provided."""
+        if weight_version is not None:
+            self.server_args.override(
+                "tokenizer.weight_version", weight_version=weight_version
+            )

python/sglang/srt/managers/tokenizer_manager_score_mixin.py

@@ -0,0 +1,780 @@
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from sglang.srt.configs.model_config import is_cross_encoding_pooler_model
+from sglang.srt.managers.embed_types import PositionalEmbeds
+from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
+from sglang.srt.server_args import MIS_DELIMITER_TOKEN_ID
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ScoreResult:
+    scores: List[List[float]]
+    prompt_tokens: int = 0
+    # Per-item pooled hidden states (pre-head transformer output).
+    # CPU tensors when return_pooled_hidden_states=True; kept as tensors so
+    # in-process consumers (gRPC, engine API) avoid a .tolist() round-trip.
+    # The HTTP path converts to lists in serving_score.py before JSON serialization.
+    # Same layout as scores: one tensor per item (not a single packed 2D tensor).
+    pooled_hidden_states: Optional[List[Optional[torch.Tensor]]] = None
+
+
+class TokenizerManagerScoreMixin:
+    async def score_prompts(
+        self,
+        prompts: Union[str, List[str], List[List[int]]],
+        label_token_ids: List[int],
+        apply_softmax: bool = False,
+        request: Optional[Any] = None,
+    ) -> ScoreResult:
+        """
+        Score probabilities of specified token IDs after each *full prompt*.
+
+        This is a thin wrapper over `score_request` that treats `prompts` as
+        already-composed inputs (i.e., no query/item concatenation needed).
+
+        Args:
+            prompts: A single prompt string, a list of prompt strings, or a list of
+                pre-tokenized prompt token ID sequences.
+            label_token_ids: Token IDs to compute probabilities for.
+            apply_softmax: Whether to normalize probabilities using softmax.
+            request: Optional FastAPI request object.
+
+        Returns:
+            ScoreResult with:
+                scores: List of score lists, one for each prompt, each in the order of label_token_ids.
+                prompt_tokens: The number of prompt tokens processed.
+        """
+        # Text prompts
+        if isinstance(prompts, str) or (
+            isinstance(prompts, list) and (not prompts or isinstance(prompts[0], str))
+        ):
+            return await self.score_request(
+                query="",
+                items=prompts,  # type: ignore[arg-type]
+                label_token_ids=label_token_ids,
+                apply_softmax=apply_softmax,
+                item_first=False,
+                request=request,
+            )
+
+        # Tokenized prompts
+        if isinstance(prompts, list) and (not prompts or isinstance(prompts[0], list)):
+            return await self.score_request(
+                query=[],
+                items=prompts,
+                label_token_ids=label_token_ids,
+                apply_softmax=apply_softmax,
+                item_first=False,
+                request=request,
+            )
+
+        raise ValueError("Invalid prompts type for score_prompts.")
+
+    def _build_multi_item_token_sequence(
+        self, query: List[int], items: List[List[int]], delimiter_token_id: int
+    ) -> Tuple[List[int], List[int]]:
+        """
+        Build a single token sequence for multi-item scoring.
+        Format: query<delimiter>item1<delimiter>item2<delimiter>item3<delimiter>
+
+        Args:
+            query: Query token IDs
+            items: List of item token ID sequences
+            delimiter_token_id: Token ID to use as delimiter
+
+        Returns:
+            Tuple of (combined token sequence, delimiter indices)
+        """
+        combined_sequence = query[:]  # Start with query
+        delimiter_indices = []
+
+        for item in items:
+            delimiter_indices.append(len(combined_sequence))
+            combined_sequence.append(delimiter_token_id)  # Add delimiter
+            combined_sequence.extend(item)  # Add item tokens
+
+        # Add final delimiter after the last item for logprob extraction
+        delimiter_indices.append(len(combined_sequence))
+        combined_sequence.append(delimiter_token_id)
+
+        return combined_sequence, delimiter_indices
+
+    def _batch_tokenize_query_and_items(
+        self,
+        query: Optional[Union[str, List[int]]],
+        items: Optional[Union[str, List[str], List[List[int]]]],
+    ) -> Tuple[List[int], List[List[int]]]:
+        """
+        Tokenize query and items into token IDs.
+
+        Args:
+            query: The query text (str) or pre-tokenized token IDs (List[int]).
+            items: Item texts or pre-tokenized token IDs.
+
+        Returns:
+            (query_ids, items_ids): query token IDs and list of per-item token IDs.
+        """
+        if isinstance(query, str):
+            query_ids = self.tokenizer.encode(query)
+        else:
+            query_ids = list(query)
+
+        items_list = [items] if isinstance(items, str) else items
+
+        items_ids = []
+        for item in items_list:
+            if isinstance(item, str):
+                items_ids.append(self.tokenizer.encode(item))
+            else:
+                items_ids.append(list(item))
+
+        return query_ids, items_ids
+
+    def _process_multi_item_scoring_results(
+        self,
+        results: Any,
+        items: List,
+        label_token_ids: Optional[List[int]],
+        apply_softmax: bool,
+        batch_request=None,
+        return_pooled_hidden_states: bool = False,
+    ) -> ScoreResult:
+        """
+        Process results from multi-item scoring request.
+
+        Extracts per-delimiter scores from whichever field the scheduler
+        populated (input_token_ids_logprobs for generation models,
+        embedding for classification models), then uniformly validates,
+        skips the query-boundary delimiter, and normalizes.
+
+        Args:
+            results: Results from generate_request
+            items: List of items being scored
+            label_token_ids: Token IDs to extract scores for
+            apply_softmax: Whether to apply softmax normalization
+            batch_request: The original batch request containing input sequence
+            return_pooled_hidden_states: Whether to extract pooled hidden states
+                from the result and include them in the ScoreResult.
+
+        Returns:
+            ScoreResult with per-item scores, prompt token count, and optional
+            pooled_hidden_states (when return_pooled_hidden_states=True and the
+            model populated the field).
+        """
+        single_result = results[0] if isinstance(results, list) else results
+        meta_info = single_result.get("meta_info", {})
+        num_items = len(items) if isinstance(items, list) else 1
+        expected_count = num_items + 1
+        request_id = meta_info.get("id", "<unknown>")
+        prompt_tokens = meta_info.get("prompt_tokens", 0)
+
+        # Extract per-delimiter scores from whichever field has them
+        input_logprobs = meta_info.get("input_token_ids_logprobs", [])
+        embedding = single_result.get("embedding")
+
+        if input_logprobs:
+            # Generation model: extract label-token logprobs at each delimiter
+            per_delimiter_scores = []
+            for logprobs_data in input_logprobs:
+                logprobs = self._extract_logprobs_for_tokens(
+                    logprobs_data, label_token_ids
+                )
+                score_list = self._convert_logprobs_to_scores(
+                    logprobs, label_token_ids, apply_softmax
+                )
+                per_delimiter_scores.append(score_list)
+        elif embedding is not None:
+            # Classification model: scores are directly in 2D embedding.
+            if apply_softmax:
+                scores_tensor = (
+                    torch.tensor(embedding)
+                    if isinstance(embedding, list)
+                    else embedding
+                )
+                scores_tensor = torch.nn.functional.softmax(scores_tensor, dim=-1)
+                per_delimiter_scores = scores_tensor.tolist()
+            else:
+                per_delimiter_scores = (
+                    embedding if isinstance(embedding, list) else embedding.tolist()
+                )
+        else:
+            raise RuntimeError(
+                f"No scoring data found for multi-item scoring request {request_id}. "
+                "Expected either input_token_ids_logprobs or embedding."
+            )
+
+        # Validate delimiter count
+        if len(per_delimiter_scores) != expected_count:
+            raise RuntimeError(
+                f"Expected {expected_count} delimiter entries for multi-item scoring "
+                f"with {num_items} items, but got {len(per_delimiter_scores)}. "
+                f"Request ID: {request_id}"
+            )
+
+        # Skip the first delimiter (query-item boundary)
+        scores = per_delimiter_scores[1:]
+
+        phs_list = None
+        if return_pooled_hidden_states:
+            raw_phs = single_result.get("pooled_hidden_state")
+            if raw_phs is not None and len(raw_phs) == expected_count:
+                phs_list = raw_phs[1:]
+
+        return ScoreResult(
+            scores=scores,
+            prompt_tokens=prompt_tokens,
+            pooled_hidden_states=phs_list,
+        )
+
+    def _process_single_item_scoring_results(
+        self,
+        results: Any,
+        label_token_ids: Optional[List[int]],
+        apply_softmax: bool,
+        return_pooled_hidden_states: bool = False,
+    ) -> ScoreResult:
+        """
+        Process results from single-item scoring request.
+
+        For generation (CausalLM) models: reads output_token_ids_logprobs.
+        For non-generation (SequenceClassification) models: reads the embedding field
+        which contains pooled class logits from the classification head.
+
+        Args:
+            results: Results from generate_request
+            label_token_ids: Token IDs to extract scores for (generation models only)
+            apply_softmax: Whether to apply softmax normalization
+            return_pooled_hidden_states: Whether to extract pooled hidden states
+
+        Returns:
+            ScoreResult with per-item scores, prompt token count, and optional pooled_hidden_states.
+        """
+        scores = []
+        phs_list = []
+        has_phs = False
+        prompt_tokens = 0
+
+        is_generation = self.is_generation
+        if is_generation:
+            for result in results:
+                # For single-item scoring, logprobs are in output_token_ids_logprobs
+                output_logprobs = result["meta_info"].get(
+                    "output_token_ids_logprobs", []
+                )
+                prompt_tokens += result["meta_info"].get("prompt_tokens", 0)
+
+                if not output_logprobs or len(output_logprobs) == 0:
+                    raise RuntimeError(
+                        f"output_logprobs is empty for request "
+                        f"{result['meta_info'].get('id', '<unknown>')}."
+                    )
+
+                # Extract logprobs for the first (and only) position
+                logprobs = self._extract_logprobs_for_tokens(
+                    output_logprobs[0], label_token_ids
+                )
+                score_list = self._convert_logprobs_to_scores(
+                    logprobs, label_token_ids, apply_softmax
+                )
+                scores.append(score_list)
+        else:
+            for result in results:
+                embedding = result.get("embedding", None)
+                if embedding is None:
+                    raise ValueError("Embedding not found in the result.")
+
+                prompt_tokens += result.get("meta_info", {}).get("prompt_tokens", 0)
+
+                if apply_softmax:
+                    embedding = torch.softmax(
+                        torch.as_tensor(embedding), dim=-1
+                    ).tolist()
+
+                # The classification head produces per-token logits, which the pooler reduces
+                # into a single vector per input. That vector is returned in the `.embeddings`
+                # field — not as semantic embeddings, but as pooled classification logits.
+                # The field name is reused for compatibility with the existing
+                # EmbeddingPoolerOutput API.
+                scores.append(embedding)
+
+                if return_pooled_hidden_states:
+                    phs = result.get("pooled_hidden_state")
+                    phs_list.append(phs)
+                    if phs is not None:
+                        has_phs = True
+
+        return ScoreResult(
+            scores=scores,
+            prompt_tokens=prompt_tokens,
+            pooled_hidden_states=phs_list if has_phs else None,
+        )
+
+    # ------------------------------------------------------------------
+    # Embed override position resolution
+    # ------------------------------------------------------------------
+
+    def _resolve_overrides_for_sequence(
+        self,
+        token_ids: List[int],
+        embeds: Optional[List[torch.Tensor]],
+        embed_override_token_id: int,
+        position_offset: int = 0,
+        label: str = "input",
+    ) -> Tuple[List[torch.Tensor], List[int]]:
+        """Scan token_ids for placeholder occurrences and pair with embeddings.
+
+        Args:
+            token_ids: The token sequence to scan.
+            embeds: Embedding tensors to place at placeholder positions (None = skip).
+            embed_override_token_id: The placeholder token ID.
+            position_offset: Added to each found position (for absolute coordinates).
+            label: Label for error messages (e.g. "query", "items[2]").
+
+        Returns:
+            (embeds, positions) lists. Empty lists if embeds is None.
+        """
+        if embeds is None:
+            return [], []
+        positions = [
+            idx + position_offset
+            for idx, tok in enumerate(token_ids)
+            if tok == embed_override_token_id
+        ]
+        if len(positions) != len(embeds):
+            raise ValueError(
+                f"{label} contains {len(positions)} occurrences of "
+                f"embed_override_token_id={embed_override_token_id}, "
+                f"but {len(embeds)} override embeddings were provided."
+            )
+        return embeds, positions
+
+    def _resolve_embed_overrides_for_request(
+        self,
+        query: List[int],
+        item: List[int],
+        embed_override_token_id: int,
+        query_embed_overrides: Optional[List[torch.Tensor]],
+        item_embeds: Optional[List[torch.Tensor]],
+        item_position_offset: int,
+        item_label: str,
+    ) -> Optional[PositionalEmbeds]:
+        """Resolve embed overrides for a single query+item pair.
+
+        Returns PositionalEmbeds if any overrides exist, None otherwise.
+        """
+        q_embeds, q_positions = self._resolve_overrides_for_sequence(
+            query,
+            query_embed_overrides,
+            embed_override_token_id,
+            position_offset=0,
+            label="query",
+        )
+        i_embeds, i_positions = self._resolve_overrides_for_sequence(
+            item,
+            item_embeds,
+            embed_override_token_id,
+            position_offset=item_position_offset,
+            label=item_label,
+        )
+        all_embeds = q_embeds + i_embeds
+        all_positions = q_positions + i_positions
+        if not all_embeds:
+            return None
+        return PositionalEmbeds(embeds=all_embeds, positions=all_positions)
+
+    # ------------------------------------------------------------------
+    # Input preparation (tokenization + input_ids construction)
+    # ------------------------------------------------------------------
+
+    def _build_token_id_inputs(
+        self,
+        query: List[int],
+        items: List[List[int]],
+        item_first: bool,
+        use_multi_item_scoring: bool,
+        embed_override_token_id: Optional[int],
+        query_embed_overrides: Optional[List[torch.Tensor]],
+        item_embed_overrides: Optional[List[Optional[List[torch.Tensor]]]],
+    ) -> Tuple[None, List[List[int]], Optional[list], Optional[List[int]]]:
+        """Build input_ids and resolve embed overrides for token-ID inputs.
+
+        Works identically for multi-item-scoring and single-item modes — the only difference is
+        how input_ids are assembled and what position offset each item gets.
+
+        Returns:
+            (text_prompts, input_ids, positional_embed_overrides, delimiter_indices)
+        """
+        # Both query and items are token IDs
+        has_embeds = (
+            query_embed_overrides is not None or item_embed_overrides is not None
+        )
+
+        # Query placeholder positions are invariant across items — resolve once.
+        # (No-op returning ([], []) if has_embeds is False or query_embed_overrides is None.)
+        q_embeds, q_positions = self._resolve_overrides_for_sequence(
+            query,
+            query_embed_overrides,
+            embed_override_token_id,
+            position_offset=0,
+            label="query",
+        )
+
+        if use_multi_item_scoring:
+            # Multi-item scoring: concatenate with placeholder delimiter token.
+            # Positions are derived from item lengths (delimiter_indices), not
+            # by scanning for this token — it exists only for FlashInfer compat.
+            delimiter_token_id = MIS_DELIMITER_TOKEN_ID
+            combined_input_ids, delimiter_indices = (
+                self._build_multi_item_token_sequence(query, items, delimiter_token_id)
+            )
+            input_ids = [combined_input_ids]
+
+            if not has_embeds:
+                return None, input_ids, None, delimiter_indices
+
+            # Resolve embed overrides across the combined multi-item-scoring sequence.
+            all_embeds: List[torch.Tensor] = list(q_embeds)
+            all_positions: List[int] = list(q_positions)
+            current_offset = len(query) + 1  # +1 for first delimiter
+            for i, item in enumerate(items):
+                item_embs = item_embed_overrides[i] if item_embed_overrides else None
+                i_embeds, i_positions = self._resolve_overrides_for_sequence(
+                    item,
+                    item_embs,
+                    embed_override_token_id,
+                    position_offset=current_offset,
+                    label=f"items[{i}]",
+                )
+                all_embeds.extend(i_embeds)
+                all_positions.extend(i_positions)
+                current_offset += len(item) + 1  # +1 for delimiter
+
+            if all_embeds:
+                # PositionalEmbeds.__post_init__ does the single torch.cat stack.
+                positional_embed_overrides = [
+                    PositionalEmbeds(embeds=all_embeds, positions=all_positions)
+                ]
+            else:
+                positional_embed_overrides = None
+            return None, input_ids, positional_embed_overrides, delimiter_indices
+
+        else:
+            # Single-item scoring: process each item separately
+            if item_first:
+                input_ids = [item + query for item in items]
+            else:
+                input_ids = [query + item for item in items]
+
+            if not has_embeds:
+                return None, input_ids, None, None
+
+            positional_embed_overrides = []
+            any_overrides = False
+            for i, item in enumerate(items):
+                item_embs = item_embed_overrides[i] if item_embed_overrides else None
+                i_embeds, i_positions = self._resolve_overrides_for_sequence(
+                    item,
+                    item_embs,
+                    embed_override_token_id,
+                    position_offset=len(query),
+                    label=f"items[{i}]",
+                )
+                combined_embeds = q_embeds + i_embeds
+                if combined_embeds:
+                    positional_embed_overrides.append(
+                        PositionalEmbeds(
+                            embeds=combined_embeds,
+                            positions=q_positions + i_positions,
+                        )
+                    )
+                    any_overrides = True
+                else:
+                    positional_embed_overrides.append(None)
+
+            return (
+                None,
+                input_ids,
+                positional_embed_overrides if any_overrides else None,
+                None,
+            )
+
+    # ------------------------------------------------------------------
+    # Main entry point
+    # ------------------------------------------------------------------
+
+    async def score_request(
+        self,
+        query: Optional[Union[str, List[int]]] = None,
+        items: Optional[Union[str, List[str], List[List[int]]]] = None,
+        label_token_ids: Optional[List[int]] = None,
+        apply_softmax: bool = False,
+        item_first: bool = False,
+        embed_override_token_id: Optional[int] = None,
+        query_embed_overrides: Optional[List[torch.Tensor]] = None,
+        item_embed_overrides: Optional[List[Optional[List[torch.Tensor]]]] = None,
+        request: Optional[Any] = None,
+        return_pooled_hidden_states: bool = False,
+    ) -> ScoreResult:
+        """
+        Score the probability of specified token IDs appearing after the given (query + item) pair.
+
+        This method supports two scoring approaches:
+        1. Single-Item scoring (default): Process each query+item pair independently
+        2. Multi-Item scoring: When --enable-mis is set, combine query and
+           multiple items into a single sequence using delimiter for efficient processing.
+           Note: item_first parameter is ignored in multi-item scoring mode since it uses
+           a fixed format: query<delimiter>item1<delimiter>item2<delimiter>item3<delimiter>
+
+           Multi-item scoring works with both text and pre-tokenized inputs:
+           - Text: query<delimiter_text>item1<delimiter_text>item2<delimiter_text>item3<delimiter_text>
+           - Tokens: query<delimiter_token_id>item1<delimiter_token_id>item2<delimiter_token_id>item3<delimiter_token_id>
+
+        Supports two model types:
+        - Generation (CausalLM): Requires label_token_ids; returns logprob-based scores.
+        - SequenceClassification: label_token_ids is optional; returns pooled class logits.
+
+        Args:
+            query: The query text or pre-tokenized query token IDs
+            items: The item text(s) or pre-tokenized item token IDs
+            label_token_ids: List of token IDs to compute probabilities for
+            apply_softmax: Whether to normalize probabilities using softmax
+            item_first: If True, prepend items to query. Ignored for multi-item scoring.
+            embed_override_token_id: Placeholder token ID for embedding override positions.
+            query_embed_overrides: Embedding vectors replacing placeholder tokens in query.
+            item_embed_overrides: Per-item embedding vectors replacing placeholder tokens in items.
+            request: Optional FastAPI request object
+            return_pooled_hidden_states: Whether to include the raw pooled transformer
+                hidden states (before the task-specific head) in the result. Only
+                supported for non-generation models (SequenceClassification,
+                RewardModel). Raises ValueError for CausalLM models.
+
+        Returns:
+            ScoreResult with:
+                scores: List of score lists, one per item.
+                prompt_tokens: The number of prompt tokens processed.
+                pooled_hidden_states: Per-item CPU tensors when
+                    return_pooled_hidden_states=True and the model supports it;
+                    None otherwise.
+        """
+        is_generation = self.is_generation
+
+        if is_generation and label_token_ids is None:
+            raise ValueError(
+                "label_token_ids is required for generation (CausalLM) models."
+            )
+        if items is None:
+            raise ValueError("items must be provided")
+        if not items:
+            return ScoreResult(scores=[], prompt_tokens=0)
+
+        has_embeds = (
+            query_embed_overrides is not None or item_embed_overrides is not None
+        )
+        if has_embeds and embed_override_token_id is None:
+            raise ValueError(
+                "embed_override_token_id is required when query_embed_overrides "
+                "or item_embed_overrides are supplied."
+            )
+        if item_first and has_embeds:
+            raise ValueError("item_first is not supported when embeddings are supplied")
+        if item_embed_overrides is not None and len(item_embed_overrides) != len(items):
+            raise ValueError(
+                f"item_embed_overrides length ({len(item_embed_overrides)}) "
+                f"must match items length ({len(items)})."
+            )
+        if self.tokenizer is not None and label_token_ids is not None:
+            vocab_size = self.tokenizer.vocab_size
+            for token_id in label_token_ids:
+                if token_id >= vocab_size:
+                    raise ValueError(
+                        f"Token ID {token_id} is out of vocabulary (vocab size: {vocab_size})"
+                    )
+
+        # Check if multi-item scoring is enabled
+        use_multi_item_scoring = self.server_args.enable_mis
+
+        input_ids = None
+        text_prompts = None
+        positional_embed_overrides = None
+        delimiter_indices = None
+
+        use_text_prompts = isinstance(query, str) and not has_embeds
+
+        if use_text_prompts:
+            # Both query and items are text
+            items_list = [items] if isinstance(items, str) else items
+            if use_multi_item_scoring:
+                # Tokenize separately, then combine at token level with placeholder
+                # delimiter. Positions come from item lengths (delimiter_indices),
+                # not from scanning for this token — it's for FlashInfer compat only.
+                delimiter_token_id = MIS_DELIMITER_TOKEN_ID
+                query_ids, items_ids = self._batch_tokenize_query_and_items(
+                    query, items_list
+                )
+                combined_input_ids, delimiter_indices = (
+                    self._build_multi_item_token_sequence(
+                        query_ids, items_ids, delimiter_token_id
+                    )
+                )
+                input_ids = [combined_input_ids]
+            else:
+                # Single-item scoring: create separate prompts for each item
+                if item_first:
+                    text_prompts = [f"{item}{query}" for item in items_list]
+                else:
+                    text_prompts = [f"{query}{item}" for item in items_list]
+
+        elif (
+            isinstance(query, list)
+            and isinstance(items, list)
+            and items
+            and isinstance(items[0], list)
+        ):
+            # Both query and items are token IDs — tokenize text inputs if needed for embed overrides
+            query_ids, items_ids = query, items
+            _, input_ids, positional_embed_overrides, delimiter_indices = (
+                self._build_token_id_inputs(
+                    query_ids,
+                    items_ids,
+                    item_first,
+                    use_multi_item_scoring,
+                    embed_override_token_id,
+                    query_embed_overrides,
+                    item_embed_overrides,
+                )
+            )
+        elif has_embeds:
+            # Text inputs with embed overrides — need to tokenize first to resolve positions
+            query_ids, items_ids = self._batch_tokenize_query_and_items(query, items)
+            _, input_ids, positional_embed_overrides, delimiter_indices = (
+                self._build_token_id_inputs(
+                    query_ids,
+                    items_ids,
+                    item_first,
+                    use_multi_item_scoring,
+                    embed_override_token_id,
+                    query_embed_overrides,
+                    item_embed_overrides,
+                )
+            )
+        else:
+            raise ValueError(
+                "Invalid combination of query/items types for score_request."
+            )
+
+        if return_pooled_hidden_states:
+            if is_generation:
+                raise ValueError(
+                    "return_pooled_hidden_states is not supported for CausalLM models. "
+                    "It requires a model with a task-specific head "
+                    "(e.g. SequenceClassification or RewardModel)."
+                )
+            model_config = self.model_config
+            if model_config is not None:
+                archs = getattr(model_config.hf_config, "architectures", []) or []
+                if is_cross_encoding_pooler_model(archs):
+                    raise ValueError(
+                        f"return_pooled_hidden_states is not supported for "
+                        f"{archs[0]}. This model uses CrossEncodingPooler which "
+                        f"does not expose pre-head hidden states."
+                    )
+
+        # Create the appropriate request type
+        mis_delimiter_indices = [delimiter_indices] if use_multi_item_scoring else None
+        if is_generation:
+            batch_request = GenerateReqInput(
+                text=text_prompts,
+                input_ids=input_ids,
+                token_ids_logprob=label_token_ids,
+                return_logprob=True,
+                # Set logprob_start_len=0 for multi-item scoring since we want logprobs at all delimiter positions
+                logprob_start_len=0 if use_multi_item_scoring else -1,
+                stream=False,
+                sampling_params={"max_new_tokens": 0},
+                positional_embed_overrides=positional_embed_overrides,
+                multi_item_delimiter_indices=mis_delimiter_indices,
+            )
+        else:
+            batch_request = EmbeddingReqInput(
+                text=text_prompts,
+                input_ids=input_ids,
+                positional_embed_overrides=positional_embed_overrides,
+                return_pooled_hidden_states=return_pooled_hidden_states,
+                multi_item_delimiter_indices=mis_delimiter_indices,
+            )
+
+        results = await self.generate_request(batch_request, request).__anext__()
+
+        if use_multi_item_scoring:
+            # Multi-item scoring: extract scores from input_token_ids_logprobs or embedding
+            return self._process_multi_item_scoring_results(
+                results,
+                items,
+                label_token_ids,
+                apply_softmax,
+                batch_request,
+                return_pooled_hidden_states,
+            )
+        else:
+            # Single-item scoring: process each result separately
+            return self._process_single_item_scoring_results(
+                results, label_token_ids, apply_softmax, return_pooled_hidden_states
+            )
+
+    def _convert_logprobs_to_scores(
+        self,
+        logprobs: Dict[int, float],
+        label_token_ids: List[int],
+        apply_softmax: bool,
+    ) -> List[float]:
+        """
+        Convert logprobs dictionary to ordered score list.
+
+        Args:
+            logprobs: Dictionary mapping token_id to logprob
+            label_token_ids: Token IDs in desired order
+            apply_softmax: Whether to apply softmax normalization
+
+        Returns:
+            List of scores in the same order as label_token_ids
+        """
+        score_list = [
+            logprobs.get(token_id, float("-inf")) for token_id in label_token_ids
+        ]
+
+        if apply_softmax:
+            score_list = torch.softmax(torch.tensor(score_list), dim=0).tolist()
+        else:
+            # Convert logprobs to probabilities if not using softmax
+            score_list = [
+                math.exp(x) if x != float("-inf") else 0.0 for x in score_list
+            ]
+
+        return score_list
+
+    def _extract_logprobs_for_tokens(
+        self, logprobs_data: List, label_token_ids: List[int]
+    ) -> Dict[int, float]:
+        """
+        Extract logprobs for specified token IDs from logprobs data.
+
+        Args:
+            logprobs_data: List of (logprob, token_id, text) tuples
+            label_token_ids: Token IDs to extract logprobs for
+
+        Returns:
+            Dictionary mapping token_id to logprob
+        """
+        logprobs = {}
+        if logprobs_data:
+            for logprob, token_id, _ in logprobs_data:
+                if token_id in label_token_ids:
+                    logprobs[token_id] = logprob
+        return logprobs

python/sglang/srt/managers/tp_worker.py

@@ -0,0 +1,615 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A tensor parallel worker."""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.distributed import get_pp_group, get_world_group
+from sglang.srt.managers.io_struct import (
+    DestroyWeightsUpdateGroupReqInput,
+    GetWeightsByNameReqInput,
+    InitWeightsSendGroupForRemoteInstanceReqInput,
+    InitWeightsUpdateGroupReqInput,
+    LoadLoRAAdapterFromTensorsReqInput,
+    LoadLoRAAdapterReqInput,
+    SendWeightsToRemoteInstanceReqInput,
+    UnloadLoRAAdapterReqInput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightsFromDistributedReqInput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromTensorReqInput,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler import GenerationBatchResult
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
+from sglang.srt.utils.hf_transformers_utils import (
+    get_processor,
+    get_tokenizer,
+    get_tokenizer_from_processor,
+)
+from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
+from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.cache_controller import LayerDoneCounter
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTpWorker(ABC):
+    @abstractmethod
+    def forward_batch_generation(self, forward_batch: ForwardBatch):
+        pass
+
+    @property
+    @abstractmethod
+    def model_runner(self) -> ModelRunner:
+        pass
+
+    @property
+    def war_fastpath_runner(self):
+        # The runner that runs the step's LAST shared-buffer-reading phase --
+        # it owns the read-done event the scheduler's WAR barrier waits on.
+        # For a plain worker that's its own runner.
+        return self.model_runner
+
+    @property
+    def sliding_window_size(self) -> Optional[int]:
+        return self.model_runner.sliding_window_size
+
+    @property
+    def is_hybrid_swa(self) -> bool:
+        return self.model_runner.is_hybrid_swa
+
+    def get_tokens_per_layer_info(self):
+        return (
+            self.model_runner.full_max_total_num_tokens,
+            self.model_runner.swa_max_total_num_tokens,
+        )
+
+    def get_pad_input_ids_func(self):
+        return getattr(self.model_runner.model, "pad_input_ids", None)
+
+    def get_memory_pool(self) -> Tuple[ReqToTokenPool, BaseTokenToKVPoolAllocator]:
+        return (
+            self.model_runner.req_to_token_pool,
+            self.model_runner.token_to_kv_pool_allocator,
+        )
+
+    def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
+        success, message = self.model_runner.update_weights_from_disk(
+            recv_req.model_path,
+            recv_req.load_format,
+            recapture_cuda_graph=recv_req.recapture_cuda_graph,
+        )
+        return success, message
+
+    def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
+        success, message = self.model_runner.init_weights_update_group(
+            recv_req.master_address,
+            recv_req.master_port,
+            recv_req.rank_offset,
+            recv_req.world_size,
+            recv_req.group_name,
+            recv_req.backend,
+        )
+        return success, message
+
+    def destroy_weights_update_group(self, recv_req: DestroyWeightsUpdateGroupReqInput):
+        success, message = self.model_runner.destroy_weights_update_group(
+            recv_req.group_name,
+        )
+        return success, message
+
+    def init_weights_send_group_for_remote_instance(
+        self, recv_req: InitWeightsSendGroupForRemoteInstanceReqInput
+    ):
+        success, message = (
+            self.model_runner.init_weights_send_group_for_remote_instance(
+                recv_req.master_address,
+                recv_req.ports,
+                recv_req.group_rank,
+                recv_req.world_size,
+                recv_req.group_name,
+                recv_req.backend,
+            )
+        )
+        return success, message
+
+    def send_weights_to_remote_instance(
+        self, recv_req: SendWeightsToRemoteInstanceReqInput
+    ):
+        success, message = self.model_runner.send_weights_to_remote_instance(
+            recv_req.master_address,
+            recv_req.ports,
+            recv_req.group_name,
+        )
+        return success, message
+
+    def update_weights_from_distributed(
+        self, recv_req: UpdateWeightsFromDistributedReqInput
+    ):
+        success, message = self.model_runner.update_weights_from_distributed(
+            recv_req.names,
+            recv_req.dtypes,
+            recv_req.shapes,
+            recv_req.group_name,
+            recv_req.load_format,
+        )
+        return success, message
+
+    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+
+        monkey_patch_torch_reductions()
+        success, message = self.model_runner.update_weights_from_tensor(
+            named_tensors=MultiprocessingSerializer.deserialize(
+                recv_req.serialized_named_tensors[self.tp_rank]
+            ),
+            load_format=recv_req.load_format,
+        )
+        return success, message
+
+    def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
+        """Update weights from IPC for checkpoint-engine integration."""
+        success, message = self.model_runner.update_weights_from_ipc(recv_req)
+        return success, message
+
+    def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
+        parameter = self.model_runner.get_weights_by_name(
+            recv_req.name, recv_req.truncate_size
+        )
+        return parameter
+
+    def load_lora_adapter(self, recv_req: LoadLoRAAdapterReqInput):
+        result = self.model_runner.load_lora_adapter(recv_req.to_ref())
+        return result
+
+    def unload_lora_adapter(self, recv_req: UnloadLoRAAdapterReqInput):
+        result = self.model_runner.unload_lora_adapter(recv_req.to_ref())
+        return result
+
+    def load_lora_adapter_from_tensors(
+        self, recv_req: LoadLoRAAdapterFromTensorsReqInput
+    ):
+        # The LoRA code handles TP sharding internally using slice_lora_a_weights
+        # and slice_lora_b_weights methods (see lora/layers.py:46-49, mem_pool.py:437-440).
+        if recv_req.load_format == "flattened_bucket":
+            flattened_data = MultiprocessingSerializer.deserialize(
+                recv_req.serialized_tensors
+            )
+            bucket = FlattenedTensorBucket(
+                flattened_tensor=flattened_data["flattened_tensor"],
+                metadata=flattened_data["metadata"],
+            )
+            tensors = dict(bucket.reconstruct_tensors())
+        else:
+            tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors)
+        result = self.model_runner.load_lora_adapter_from_tensors(
+            recv_req.to_ref(),
+            tensors,
+            recv_req.config_dict,
+            recv_req.added_tokens_config,
+        )
+        return result
+
+    def forward_batch_embedding(self, batch: ScheduleBatch):
+        forward_batch = ForwardBatch.init_new(batch, self.model_runner)
+        output = self.model_runner.forward(forward_batch).logits_output
+        return output  # Returns EmbeddingPoolerOutput
+
+
+class TpModelWorker(BaseTpWorker):
+    """A tensor parallel model worker."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        moe_ep_rank: int,
+        pp_rank: int,
+        attn_cp_rank: int,
+        moe_dp_rank: int,
+        dp_rank: Optional[int],
+        nccl_port: int,
+        is_draft_worker: bool = False,
+        req_to_token_pool: Optional[ReqToTokenPool] = None,
+        token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
+        memory_pool_config: Optional[MemoryPoolConfig] = None,
+        is_multi_layer_eagle: bool = False,
+        context_length: Optional[int] = None,
+    ):
+        # Parse args
+        self.server_args = server_args
+        self.tp_size = server_args.tp_size
+        self.ep_size = server_args.ep_size
+        self.pp_size = server_args.pp_size
+        self.tp_rank = tp_rank
+        self.moe_ep_rank = moe_ep_rank
+        self.pp_rank = pp_rank
+        self.dp_rank = dp_rank
+        self.gpu_id = gpu_id
+        self.nccl_port = nccl_port
+        self.is_draft_worker = is_draft_worker
+        self.is_multi_layer_eagle = is_multi_layer_eagle
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.attn_cp_rank = attn_cp_rank
+        self.moe_dp_rank = moe_dp_rank
+        # Draft worker: target's resolved MemoryPoolConfig (forwarded to ModelRunner).
+        self.memory_pool_config = memory_pool_config
+        # Draft worker: target's effective context length; the draft runs at
+        # absolute target positions. None keeps server_args.context_length.
+        self.context_length = context_length
+
+        # MTP model runners
+        self.model_runner_list: List[ModelRunner] = []
+
+        self._init_model_config()
+        self._init_model_runner()
+
+        if is_multi_layer_eagle:
+            self._init_multi_layer_eagle_model_runners()
+
+        self._init_dllm_algorithm()
+
+        if server_args.skip_tokenizer_init or self.is_draft_worker:
+            # A draft worker's tokenizer would only duplicate the target's:
+            # tokenizer_path always points at the target model.
+            self.tokenizer = self.processor = None
+        else:
+            if self.model_config.is_multimodal:
+                self.processor = get_processor(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
+                    tokenizer_backend=server_args.tokenizer_backend,
+                    model_name=server_args.model_path,
+                )
+                self.tokenizer = get_tokenizer_from_processor(self.processor)
+            else:
+                self.tokenizer = get_tokenizer(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
+                    tokenizer_backend=server_args.tokenizer_backend,
+                )
+        self.device = self.model_runner.device
+
+        # Init nccl groups
+        self.pp_group = get_pp_group()
+        self.world_group = get_world_group()
+
+        # Sync random seed across TP workers
+        self.random_seed = broadcast_pyobj(
+            [server_args.random_seed],
+            self.tp_size * self.pp_rank + tp_rank,
+            self.world_group.cpu_group,
+            src=self.world_group.ranks[0],
+        )[0]
+        set_random_seed(self.random_seed)
+
+        self.enable_overlap = not server_args.disable_overlap_schedule
+        self.enable_spec = server_args.speculative_algorithm is not None
+        self.hicache_layer_transfer_counter = None
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config: Optional[MemoryPoolConfig] = None,
+        req_to_token_pool: Optional[ReqToTokenPool] = None,
+        token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
+    ):
+        """Allocate KV cache pools only (no backends or cuda graphs)."""
+        if req_to_token_pool is not None:
+            self.req_to_token_pool = req_to_token_pool
+            self.model_runner.req_to_token_pool = req_to_token_pool
+        if token_to_kv_pool_allocator is not None:
+            self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+            self.model_runner.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.model_runner.alloc_memory_pool(memory_pool_config)
+        for mr in self.model_runner_list[1:]:
+            mr.req_to_token_pool = self.req_to_token_pool
+            mr.token_to_kv_pool_allocator = self.token_to_kv_pool_allocator
+            mr.alloc_memory_pool(memory_pool_config)
+
+        # Validation
+        assert self.model_runner.max_running_requests > 0, "max_running_request is zero"
+        max_req_len = min(
+            self.model_config.context_len - 1,
+            self.model_runner.max_token_pool_size - 1,
+        )
+        assert max_req_len > 0, "Memory pool size is too small"
+
+    def init_attention_backends(self):
+        """Initialize attention backends for all model runners."""
+        self.model_runner.init_attention_backends()
+        for mr in self.model_runner_list[1:]:
+            mr.init_attention_backends()
+
+    def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True):
+        """Capture cuda graphs for all model runners."""
+        self.model_runner.init_cuda_graphs(
+            capture_decode_cuda_graph=capture_decode_cuda_graph
+        )
+        for mr in self.model_runner_list[1:]:
+            mr.init_cuda_graphs(capture_decode_cuda_graph=capture_decode_cuda_graph)
+
+    def _init_model_config(self):
+        from sglang.srt.configs.model_config import ModelConfig
+
+        self.model_config = ModelConfig.from_server_args(
+            self.server_args,
+            model_path=(
+                self.server_args.model_path
+                if not self.is_draft_worker
+                else self.server_args.speculative_draft_model_path
+            ),
+            model_revision=(
+                self.server_args.revision
+                if not self.is_draft_worker
+                else self.server_args.speculative_draft_model_revision
+            ),
+            is_draft_model=self.is_draft_worker,
+            context_length=self.context_length,
+        )
+
+    def _init_model_runner(self):
+        from sglang.srt.model_executor.model_runner import ModelRunner
+
+        self._model_runner = ModelRunner(
+            model_config=self.model_config,
+            mem_fraction_static=self.server_args.mem_fraction_static,
+            gpu_id=self.gpu_id,
+            tp_rank=self.tp_rank,
+            tp_size=self.tp_size,
+            moe_ep_rank=self.moe_ep_rank,
+            moe_ep_size=self.ep_size,
+            pp_rank=self.pp_rank,
+            pp_size=self.pp_size,
+            nccl_port=self.nccl_port,
+            dp_rank=self.dp_rank,
+            server_args=self.server_args,
+            is_draft_worker=self.is_draft_worker,
+            req_to_token_pool=self.req_to_token_pool,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            memory_pool_config=self.memory_pool_config,
+            draft_model_idx=0 if self.is_multi_layer_eagle else None,
+        )
+
+    def _init_multi_layer_eagle_model_runners(self):
+        from sglang.srt.model_executor.model_runner import ModelRunner
+
+        self.model_runner_list.append(self.model_runner)
+        for i in range(1, self.server_args.speculative_num_steps):
+            self.model_runner_list.append(
+                ModelRunner(
+                    model_config=self.model_config,
+                    mem_fraction_static=self.server_args.mem_fraction_static,
+                    gpu_id=self.gpu_id,
+                    tp_rank=self.tp_rank,
+                    tp_size=self.tp_size,
+                    moe_ep_rank=self.moe_ep_rank,
+                    moe_ep_size=self.ep_size,
+                    pp_rank=self.pp_rank,
+                    pp_size=self.pp_size,
+                    nccl_port=self.nccl_port,
+                    dp_rank=self.dp_rank,
+                    server_args=self.server_args,
+                    is_draft_worker=self.is_draft_worker,
+                    req_to_token_pool=self.req_to_token_pool,
+                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                    memory_pool_config=self.memory_pool_config,
+                    draft_model_idx=i,
+                )
+            )
+
+    def _init_dllm_algorithm(self):
+        from sglang.srt.dllm.algorithm.base import DllmAlgorithm
+
+        if self.server_args.dllm_algorithm is not None:
+            self.dllm_algorithm = DllmAlgorithm.from_server_args(self.server_args)
+        else:
+            self.dllm_algorithm = None
+
+    @property
+    def model_runner(self) -> ModelRunner:
+        return self._model_runner
+
+    def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter):
+        self.hicache_layer_transfer_counter = counter
+
+    def set_hicache_consumer(self, consumer_index: int):
+        if self.hicache_layer_transfer_counter is not None:
+            self.hicache_layer_transfer_counter.set_consumer(consumer_index)
+
+    def register_hisparse_coordinator(self, coordinator):
+        self.model_runner.hisparse_coordinator = coordinator
+
+    def get_worker_info(self):
+        max_req_len = min(
+            self.model_config.context_len - 1,
+            self.model_runner.max_token_pool_size - 1,
+        )
+        return (
+            self.model_runner.max_total_num_tokens,
+            self.server_args.max_prefill_tokens,
+            self.model_runner.max_running_requests,
+            self.server_args.max_queued_requests,
+            max_req_len,
+            max_req_len - 5,
+            self.random_seed,
+            self.device,
+            self.model_runner.forward_stream,
+            self.model_runner.req_to_token_pool.size,
+            self.model_runner.req_to_token_pool.max_context_len,
+            self.model_runner.token_to_kv_pool.size,
+        )
+
+    def is_dllm(self):
+        return self.dllm_algorithm is not None
+
+    def _forward_batch_generation_dllm(
+        self,
+        forward_batch: ForwardBatch,
+        batch: Optional[ScheduleBatch] = None,
+    ) -> GenerationBatchResult:
+        algo_states = None
+        if self.dllm_algorithm.fdfo and batch is not None:
+            algo_states = [req.dllm_algo_state for req in batch.reqs]
+
+        (
+            logits_output,
+            next_token_ids,
+            accept_length_per_req_cpu,
+            dllm_algo_state,
+            can_run_cuda_graph,
+        ) = self.dllm_algorithm.run(self.model_runner, forward_batch, algo_states)
+
+        return GenerationBatchResult(
+            logits_output=logits_output,
+            next_token_ids=next_token_ids,
+            accept_length_per_req_cpu=accept_length_per_req_cpu,
+            dllm_algo_state=dllm_algo_state,
+            can_run_cuda_graph=can_run_cuda_graph,
+        )
+
+    def forward_batch_generation(
+        self,
+        batch: Optional[ScheduleBatch],
+        forward_batch: Optional[ForwardBatch] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+        is_verify: bool = False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
+    ) -> GenerationBatchResult:
+        # Get forward batch from schedule batch
+        if batch is not None:
+            # update the consumer index of hicache to the running batch
+            self.set_hicache_consumer(batch.hicache_consumer_index)
+
+            forward_batch = ForwardBatch.init_new(batch, self.model_runner)
+        else:
+            # FIXME(lsyin): unify the interface of forward_batch
+            assert forward_batch is not None
+
+        # Deprecated kwarg: pre-planners mark the batch themselves now.
+        forward_batch.apply_deprecated_skip_attn_backend_init(skip_attn_backend_init)
+
+        if self.is_dllm():
+            return self._forward_batch_generation_dllm(forward_batch, batch)
+
+        if self.pp_group.is_last_rank:
+            out = self.model_runner.forward(
+                forward_batch,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
+            logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+            batch_result = GenerationBatchResult(
+                logits_output=logits_output,
+                can_run_cuda_graph=can_run_cuda_graph,
+                expert_distribution_metrics=out.expert_distribution_metrics,
+                routed_experts_output=out.routed_experts_output,
+                indexer_topk_output=out.indexer_topk_output,
+            )
+
+            if is_verify:
+                # Skip sampling; spec_v2 worker fires its own publish post-verify.
+                return batch_result
+
+            if (
+                self.enable_overlap
+                and not self.enable_spec
+                and forward_batch.sampling_info.grammars is not None
+            ):
+
+                def sample_batch_func():
+                    batch_result.next_token_ids = self.model_runner.sample(
+                        logits_output, forward_batch
+                    )
+                    return batch_result
+
+                batch_result.delay_sample_func = sample_batch_func
+                return batch_result
+
+            if not forward_batch.is_prefill_only:
+                # For normal requests, sample the next token ids.
+                batch_result.next_token_ids = self.model_runner.sample(
+                    logits_output, forward_batch
+                )
+            else:
+                # For prefill-only requests, create dummy token IDs on CPU
+                # The size should match the batch size (number of sequences), not total tokens
+                batch_result.next_token_ids = torch.zeros(
+                    len(forward_batch.seq_lens),
+                    dtype=torch.long,
+                    device=forward_batch.input_ids.device,
+                )
+                if (
+                    forward_batch.return_logprob
+                    and logits_output.next_token_logits is not None
+                ):
+                    # NOTE: Compute logprobs without full sampling
+                    self.model_runner.compute_logprobs_only(
+                        logits_output, forward_batch
+                    )
+
+            return batch_result
+        else:
+            out = self.model_runner.forward(
+                forward_batch,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
+            pp_proxy_tensors, can_run_cuda_graph = out.logits_output, out.can_run_graph
+            return GenerationBatchResult(
+                pp_hidden_states_proxy_tensors=pp_proxy_tensors,
+                can_run_cuda_graph=can_run_cuda_graph,
+                expert_distribution_metrics=out.expert_distribution_metrics,
+            )
+
+    def forward_batch_split_prefill(self, batch: ScheduleBatch):
+        if batch.split_index == 0:
+            forward_batch = ForwardBatch.init_new(batch, self.model_runner)
+            batch.split_forward_batch = forward_batch
+
+        out = self.model_runner.forward(
+            batch.split_forward_batch, split_forward_count=batch.split_forward_count
+        )
+        logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+        if logits_output:
+            next_token_ids = self.model_runner.sample(
+                logits_output, batch.split_forward_batch
+            )
+        else:
+            next_token_ids = None
+        batch_result = GenerationBatchResult(
+            logits_output=logits_output,
+            can_run_cuda_graph=can_run_cuda_graph,
+            expert_distribution_metrics=out.expert_distribution_metrics,
+        )
+        batch_result.next_token_ids = next_token_ids
+        return batch_result

python/sglang/srt/managers/utils.py

@@ -0,0 +1,312 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional, Union
+
+import torch
+
+from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
+from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.state_capturer.base import TopkCaptureOutput
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import GenerationBatchResult
+    from sglang.srt.speculative.eagle_info import EagleDraftInput
+
+
+logger = logging.getLogger(__name__)
+
+
+def _async_d2h(t: torch.Tensor) -> torch.Tensor:
+    """Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H
+    to pageable host memory blocks the caller until done) and record_stream keeps
+    the source alive until the copy stream drains, so the caching allocator can't
+    recycle it early. Non-CUDA falls back to a plain copy."""
+    if not t.is_cuda:
+        return t.to("cpu", non_blocking=True)
+    cpu_t = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
+    cpu_t.copy_(t, non_blocking=True)
+    t.record_stream(torch.cuda.current_stream(t.device))
+    return cpu_t
+
+
+@dataclasses.dataclass
+class GenerationBatchResult:
+    logits_output: Optional[LogitsProcessorOutput] = None
+    pp_hidden_states_proxy_tensors: Optional[PPProxyTensors] = None
+    next_token_ids: Optional[
+        Union[torch.Tensor, List[torch.Tensor], List[List[int]]]
+    ] = None
+    num_correct_drafts: int = 0  # no bonus included
+    num_correct_drafts_per_req_cpu: Optional[List[int]] = None
+    num_block_accept_tokens: int = 0
+    num_cap_tokens: int = 0
+    # FDFO dLLM batching: per-request accepted block length and carried algo state.
+    accept_length_per_req_cpu: Optional[List[int]] = None
+    dllm_algo_state: Optional[List[Any]] = None
+    can_run_cuda_graph: bool = False
+
+    # PP skip output comm: True when output send/recv was skipped and
+    # next_token_ids are placeholder zeros. Used by process_batch_result_prefill
+    # to validate that skipped output is never consumed.
+    skipped_output_comm: bool = False
+
+    # For output processing
+    extend_input_len_per_req: Optional[List[int]] = None
+    extend_logprob_start_len_per_req: Optional[List[int]] = None
+
+    # For overlap scheduling
+    copy_done: Optional[torch.cuda.Event] = None
+    delay_sample_func: Optional[callable] = None
+    future_indices: Optional[torch.Tensor] = None
+    speculative_num_draft_tokens: Optional[int] = None
+
+    # FIXME(lsyin): maybe move to a better place?
+    # sync path: forward stream -> output processor
+    accept_lens: Optional[torch.Tensor] = None
+
+    block_accept_lens: Optional[torch.Tensor] = None
+
+    cap_lens: Optional[torch.Tensor] = None
+
+    # Next-iter seq_lens; published via on_publish.
+    new_seq_lens: Optional[torch.Tensor] = None
+
+    # relay path: forward stream -> next step forward
+    next_draft_input: Optional[EagleDraftInput] = None
+
+    # Refs the worker wants scheduler to keep alive for the same 2-iter window
+    # as batch_record_buf. Used for cross-stream tensor lifetime (e.g. a spec
+    # V2 verify ForwardBatch whose tensors must outlive mid-iter SB rebinds).
+    extra_keep_alive_refs: Optional[List[Any]] = None
+
+    # Routed experts: pending async D2H for overlap scheduling
+    routed_experts_output: Optional[TopkCaptureOutput] = None
+    indexer_topk_output: Optional[TopkCaptureOutput] = None
+
+    # metrics
+    expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None
+
+    # Forward pass metrics (FPM) — GPU-accurate timing via CUDA events
+    fpm_start_event: Optional[torch.cuda.Event] = None
+    fpm_end_event: Optional[torch.cuda.Event] = None
+
+    @property
+    def has_sampled_token_ids(self) -> bool:
+        """True when this iter sampled token ids; False when none were produced
+        this rank/split (a non-last PP rank or a non-final prefill split)."""
+        return isinstance(self.next_token_ids, torch.Tensor)
+
+    @torch.profiler.record_function("copy_result_to_cpu")
+    def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True):
+        """Copy tensors to CPU in overlap scheduling.
+        Only the tensors which are needed for processing results are copied,
+        e.g., next_token_ids, logits outputs
+        """
+        if return_logprob:
+            if self.logits_output.next_token_logprobs is not None:
+                self.logits_output.next_token_logprobs = _async_d2h(
+                    self.logits_output.next_token_logprobs
+                )
+            if self.logits_output.input_token_logprobs is not None:
+                self.logits_output.input_token_logprobs = _async_d2h(
+                    self.logits_output.input_token_logprobs
+                )
+            if self.logits_output.next_token_top_logprobs_val is not None:
+                self.logits_output.next_token_top_logprobs_val = [
+                    _async_d2h(v) if torch.is_tensor(v) else v
+                    for v in self.logits_output.next_token_top_logprobs_val
+                ]
+            if self.logits_output.next_token_top_logprobs_idx is not None:
+                self.logits_output.next_token_top_logprobs_idx = [
+                    _async_d2h(x) if torch.is_tensor(x) else x
+                    for x in self.logits_output.next_token_top_logprobs_idx
+                ]
+            if self.logits_output.next_token_token_ids_logprobs_val is not None:
+                self.logits_output.next_token_token_ids_logprobs_val = [
+                    _async_d2h(v) if torch.is_tensor(v) else v
+                    for v in self.logits_output.next_token_token_ids_logprobs_val
+                ]
+        if return_hidden_states and self.logits_output.hidden_states is not None:
+            self.logits_output.hidden_states = _async_d2h(
+                self.logits_output.hidden_states
+            )
+        self.next_token_ids = _async_d2h(self.next_token_ids)
+
+        if self.accept_lens is not None:
+            self.accept_lens = _async_d2h(self.accept_lens)
+
+        if self.block_accept_lens is not None:
+            self.block_accept_lens = _async_d2h(self.block_accept_lens)
+
+        if self.cap_lens is not None:
+            self.cap_lens = _async_d2h(self.cap_lens)
+
+        # Sub-objects only declare their device fields; the single copy+safety
+        # primitive (_async_d2h: pinned D2H + record_stream) is injected here so
+        # all device->host copying and lifetime safety lives in one place.
+        for holder in (
+            self.routed_experts_output,
+            self.indexer_topk_output,
+            self.expert_distribution_metrics,
+        ):
+            if holder is not None:
+                holder.map_device_tensors(_async_d2h)
+
+        self.copy_done.record()
+
+    @classmethod
+    def from_pp_proxy(
+        cls, logits_output, next_pp_outputs: PPProxyTensors, can_run_cuda_graph
+    ):
+        # TODO(lsyin): refactor PP and avoid using dict
+        proxy_dict = next_pp_outputs.tensors
+        return cls(
+            logits_output=logits_output,
+            pp_hidden_states_proxy_tensors=None,
+            next_token_ids=next_pp_outputs["next_token_ids"],
+            extend_input_len_per_req=proxy_dict.get("extend_input_len_per_req", None),
+            extend_logprob_start_len_per_req=proxy_dict.get(
+                "extend_logprob_start_len_per_req", None
+            ),
+            can_run_cuda_graph=can_run_cuda_graph,
+        )
+
+
+def validate_input_length(
+    req: Req, max_req_input_len: int, allow_auto_truncate: bool
+) -> Optional[str]:
+    """Validate and potentially truncate input length.
+
+    Args:
+        req: The request containing input_ids to validate
+        max_req_input_len: Maximum allowed input length
+        allow_auto_truncate: Whether to truncate long inputs
+
+    Returns:
+        Error message if validation fails, None if successful
+    """
+    if len(req.origin_input_ids) >= max_req_input_len:
+        if allow_auto_truncate:
+            logger.warning(
+                "Request length is longer than the KV cache pool size or "
+                "the max context length. Truncated. "
+                f"{len(req.origin_input_ids)=}, {max_req_input_len=}."
+            )
+            req.origin_input_ids = req.origin_input_ids[:max_req_input_len]
+            return None
+        else:
+            error_msg = (
+                f"Input length ({len(req.origin_input_ids)} tokens) exceeds "
+                f"the maximum allowed length ({max_req_input_len} tokens). "
+                f"Use a shorter input or enable --allow-auto-truncate."
+            )
+            return error_msg
+
+    return None
+
+
+def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
+
+    logits_output = result.logits_output
+    assert logits_output is not None
+
+    return {
+        "extend_input_len_per_req": result.extend_input_len_per_req,
+        "extend_logprob_start_len_per_req": result.extend_logprob_start_len_per_req,
+        "next_token_logprobs": result.logits_output.next_token_logprobs,
+        "next_token_top_logprobs_val": result.logits_output.next_token_top_logprobs_val,
+        "next_token_top_logprobs_idx": result.logits_output.next_token_top_logprobs_idx,
+        "next_token_token_ids_logprobs_val": result.logits_output.next_token_token_ids_logprobs_val,
+        "next_token_token_ids_logprobs_idx": result.logits_output.next_token_token_ids_logprobs_idx,
+        "input_token_logprobs": result.logits_output.input_token_logprobs,
+        "input_top_logprobs_val": result.logits_output.input_top_logprobs_val,
+        "input_top_logprobs_idx": result.logits_output.input_top_logprobs_idx,
+        "input_token_ids_logprobs_val": result.logits_output.input_token_ids_logprobs_val,
+        "input_token_ids_logprobs_idx": result.logits_output.input_token_ids_logprobs_idx,
+    }
+
+
+def get_logprob_from_pp_outputs(
+    next_pp_outputs: PPProxyTensors,
+) -> tuple[LogitsProcessorOutput, list[int], list[int]]:
+    logits_output = LogitsProcessorOutput(
+        # Do not send logits and hidden states because they are large
+        next_token_logits=None,
+        hidden_states=None,
+        next_token_logprobs=next_pp_outputs["next_token_logprobs"],
+        next_token_top_logprobs_val=next_pp_outputs["next_token_top_logprobs_val"],
+        next_token_top_logprobs_idx=next_pp_outputs["next_token_top_logprobs_idx"],
+        next_token_token_ids_logprobs_val=next_pp_outputs[
+            "next_token_token_ids_logprobs_val"
+        ],
+        next_token_token_ids_logprobs_idx=next_pp_outputs[
+            "next_token_token_ids_logprobs_idx"
+        ],
+        input_token_logprobs=next_pp_outputs["input_token_logprobs"],
+        input_top_logprobs_val=next_pp_outputs["input_top_logprobs_val"],
+        input_top_logprobs_idx=next_pp_outputs["input_top_logprobs_idx"],
+        input_token_ids_logprobs_val=next_pp_outputs["input_token_ids_logprobs_val"],
+        input_token_ids_logprobs_idx=next_pp_outputs["input_token_ids_logprobs_idx"],
+    )
+    extend_input_len_per_req = next_pp_outputs["extend_input_len_per_req"]
+    extend_logprob_start_len_per_req = next_pp_outputs[
+        "extend_logprob_start_len_per_req"
+    ]
+
+    return logits_output, extend_input_len_per_req, extend_logprob_start_len_per_req
+
+
+@dataclass
+class EmbeddingBatchResult:
+    """Result from an embedding/classification forward pass.
+
+    Attributes:
+        embeddings: Model output — pooled embeddings or classification logits.
+        pooled_hidden_states: Raw hidden states before the task head.  Present
+            only when the batch contained ``return_pooled_hidden_states=True``
+            requests.  Tensor (uniform shapes) or list of tensors (MIS).
+        copy_done: CUDA event recorded after the async CPU copy completes.
+    """
+
+    embeddings: torch.Tensor
+    pooled_hidden_states: Optional[torch.Tensor] = None
+    copy_done: Optional[torch.cuda.Event] = None
+
+    @property
+    def can_run_cuda_graph(self) -> bool:
+        return False
+
+    @torch.profiler.record_function("copy_embedding_to_cpu")
+    def copy_to_cpu(self):
+        """Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
+        if isinstance(self.embeddings, torch.Tensor):
+            self.copy_done = torch.get_device_module(self.embeddings.device).Event()
+            self.embeddings = _async_d2h(self.embeddings)
+        else:
+            assert isinstance(self.embeddings, list)
+            if len(self.embeddings) == 0:
+                return
+
+            self.copy_done = torch.get_device_module(self.embeddings[0].device).Event()
+            self.embeddings = [_async_d2h(emb) for emb in self.embeddings]
+
+        if self.pooled_hidden_states is not None:
+            if isinstance(self.pooled_hidden_states, list):
+                self.pooled_hidden_states = [
+                    _async_d2h(t) for t in self.pooled_hidden_states
+                ]
+            else:
+                self.pooled_hidden_states = _async_d2h(self.pooled_hidden_states)
+
+        self.copy_done.record()
+
+
+def is_health_check_generate_req(recv_req):
+    rid = getattr(recv_req, "rid", None)
+    return rid is not None and rid.startswith(HEALTH_CHECK_RID_PREFIX)