chore: import upstream snapshot with attribution

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