sgl-project--sglang/python/sglang/srt/managers/scheduler_components/output_streamer.py

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