lightseekorg--tokenspeed/python/tokenspeed/runtime/engine/generation_output_processor.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

import time
from typing import TYPE_CHECKING

from tokenspeed.runtime.engine.io_struct import BatchTokenIDOut
from tokenspeed.runtime.engine.request_stats import (
    NOOP_STATS,
    RequestStats,
    RequestStatsTracker,
)
from tokenspeed.runtime.engine.request_types import (
    ABORT_CODE,
    FINISH_ABORT,
    FINISH_LENGTH,
    FINISH_MATCHED_STR,
    FINISH_MATCHED_TOKEN,
    INIT_INCREMENTAL_DETOKENIZATION_OFFSET,
    BaseFinishReason,
)
from tokenspeed.runtime.engine.scheduler_utils import (
    make_abort_event,
    make_extend_result_event,
    make_finish_event,
    make_update_reserve_tokens_event,
)
from tokenspeed.runtime.sampling.sampling_params import SamplingParams

if TYPE_CHECKING:
    from tokenspeed.runtime.engine.io_struct import TokenizedGenerateReqInput
    from tokenspeed.runtime.execution.types import ModelExecutionResult
    from tokenspeed.runtime.metrics.collector import EngineMetrics
    from tokenspeed.runtime.grammar.base_grammar_backend import (
        BaseGrammarObject,
    )

from tokenspeed.runtime.utils import get_colorful_logger
from tokenspeed.runtime.utils.nvtx import nvtx_range

logger = get_colorful_logger(__name__)

DEFAULT_FORCE_STREAM_INTERVAL = 50


class RequestState:
    """Per-request state needed for incremental streaming output.

    Extracts only the fields required by process_output from the incoming
    request. Does not hold a reference to Req or any scheduler object.
    """

    def __init__(
        self,
        prompt_input_ids: list[int],
        sampling_params: SamplingParams,
        stream: bool,
        tokenizer,
        eos_token_ids: list[int] = None,
        return_logprob: bool = False,
        top_logprobs_num: int = 0,
        token_ids_logprob: list[int] | None = None,
        multimodal_inputs=None,
        prompt_input_ids_unpadded: list[int] | None = None,
        created_time: float = 0.0,
    ) -> None:
        # --- Extracted from recv_req (immutable) ---
        self.prompt_input_ids: list[int] = prompt_input_ids
        self.prompt_input_ids_unpadded: list[int] = (
            prompt_input_ids_unpadded
            if prompt_input_ids_unpadded is not None
            else prompt_input_ids
        )
        self.multimodal_inputs = multimodal_inputs
        self.sampling_params = sampling_params
        self.stream = stream
        self.eos_token_ids = eos_token_ids
        self.tokenizer = tokenizer
        self.computed_length = 0
        self.return_logprob = return_logprob
        self.top_logprobs_num = top_logprobs_num
        self.token_ids_logprob = token_ids_logprob

        # --- generation state (updated with forward step) ---
        self.output_ids: list[int] = []
        self.finished_reason: BaseFinishReason | None = None
        self.cached_tokens: int = 0
        self.prefix_len: int = 0
        self.spec_verify_ct: int = 0
        self.accept_draft_tokens: float | None = None

        # request stats (host-side); tracker attached only with --enable-log-request-stats
        self.created_time: float = created_time
        self.stats: RequestStatsTracker = NOOP_STATS
        # Sampled-token logprobs, accumulated per generated token.
        # None when return_logprob is False.
        self.output_token_logprobs_val: list[float] | None = (
            [] if return_logprob else None
        )
        self.output_token_logprobs_idx: list[int] | None = (
            [] if return_logprob else None
        )

        # --- Streaming bookkeeping (internal) ---
        self._surr_offset: int | None = None
        self._read_offset: int | None = None
        self.decoded_text: str = ""
        self.send_token_offset: int = 0
        self.send_decode_id_offset: int = 0
        self.finished_output: bool = False

        # abort related
        self.to_abort = False
        self.to_abort_message = None
        # Client-initiated aborts skip streaming a finish (the TM already tore
        # down its state). Pause-initiated aborts set this so the passive client
        # still receives a terminating finish.
        self.abort_notify_client = False

        # cached tokenizer ids
        self._eos_token_id_cached = None
        self._additional_stop_token_ids_cached = None

        # Constrained-decoding state.
        self.grammar: BaseGrammarObject | None = None
        self.grammar_key: tuple[str, str] | None = None
        self.grammar_queued_ts: float = 0.0

    def set_finish_with_abort(self, message: str, notify_client: bool = False) -> None:
        """Mark this request as aborted with ``message``; finished_reason is
        materialized immediately so callers don't need a check_finished() pass.

        ``notify_client`` streams a terminating finish to the client (used for
        pause-initiated aborts, where the client did not tear down its state).
        """
        self.to_abort = True
        self.to_abort_message = message
        self.abort_notify_client = notify_client
        self.finished_reason = FINISH_ABORT(message=message)

    @classmethod
    def from_recv_req(
        cls,
        recv_req: TokenizedGenerateReqInput,
        tokenizer,
        eos_token_ids: list[int],
    ) -> RequestState:
        return cls(
            prompt_input_ids=recv_req.input_ids,
            sampling_params=recv_req.sampling_params,
            stream=recv_req.stream,
            tokenizer=tokenizer,
            eos_token_ids=eos_token_ids,
            return_logprob=getattr(recv_req, "return_logprob", False),
            top_logprobs_num=getattr(recv_req, "top_logprobs_num", 0),
            token_ids_logprob=getattr(recv_req, "token_ids_logprob", None),
            multimodal_inputs=getattr(recv_req, "multimodal_inputs", None),
            prompt_input_ids_unpadded=getattr(recv_req, "input_ids_unpadded", None),
            created_time=getattr(recv_req, "created_time", 0.0),
        )

    @property
    def finished(self) -> bool:
        return self.finished_reason is not None

    @property
    def input_length(self) -> int:
        return len(self.prompt_input_ids)

    @property
    def output_length(self) -> int:
        return len(self.output_ids)

    @property
    def prefill_finished(self):
        return self.computed_length >= self.input_length

    def add_computed_length(self, incr: int):
        self.computed_length += incr

    def maybe_extend_multimodal_mrope_positions(self) -> None:
        mm = self.multimodal_inputs
        if mm is None or mm.mrope_positions is None:
            return

        target_len = self.input_length + self.output_length
        current_len = mm.mrope_positions.shape[-1]
        if current_len >= target_len:
            return

        from tokenspeed.runtime.multimodal.mrope import (
            extend_mrope_positions_for_retracted_request,
        )

        mm.mrope_positions = extend_mrope_positions_for_retracted_request(
            mm.mrope_positions, target_len - current_len
        )

    def release_pending_multimodal_features(self) -> None:
        mm = self.multimodal_inputs
        if mm is not None and hasattr(mm, "release_shm_features"):
            mm.release_shm_features()

    def init_incremental_detokenize(self):
        """Return (all_ids_from_surr_offset, read_offset_relative_to_surr)."""
        if self._surr_offset is None or self._read_offset is None:
            self._read_offset = len(self.prompt_input_ids_unpadded)
            self._surr_offset = max(
                self._read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
            )
        all_ids = self.prompt_input_ids_unpadded + self.output_ids
        return (
            all_ids[self._surr_offset :],
            self._read_offset - self._surr_offset,
        )

    def check_finished(self, skip_grammar_termination: bool = False):

        if self.finished:
            return

        if self.to_abort:
            self.finished_reason = FINISH_ABORT(
                message=self.to_abort_message,
            )
            return

        # When the capturable-grammar hostfunc is authoritative, the
        # caller identifies the terminating token itself (see
        # post_process_forward_op); firing here would re-trigger on
        # every later token and trim content via trim_matched_stop.
        if not skip_grammar_termination and self.grammar is not None:
            if self.grammar.is_terminated():
                self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1])
                return

        if len(self.output_ids) >= self.sampling_params.max_new_tokens:
            self.finished_reason = FINISH_LENGTH(
                length=self.sampling_params.max_new_tokens
            )
            return

        last_token_id = self.output_ids[-1]

        if not self.sampling_params.ignore_eos:
            matched_eos = False

            # Check stop token ids
            if self.sampling_params.stop_token_ids:
                matched_eos = last_token_id in self.sampling_params.stop_token_ids
            if self.eos_token_ids:
                matched_eos |= last_token_id in self.eos_token_ids
            if self._eos_token_id_cached is None:
                self.set_cached_id()
            if self._eos_token_id_cached is not None:
                matched_eos |= last_token_id == self._eos_token_id_cached
            if self._additional_stop_token_ids_cached:
                matched_eos |= last_token_id in self._additional_stop_token_ids_cached
            if matched_eos:
                self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id)
                return

        # Check stop strings
        if len(self.sampling_params.stop_strs) > 0:
            tail_str = self.tokenizer.decode(
                self.output_ids[-(self.sampling_params.stop_str_max_len + 1) :]
            )

            for stop_str in self.sampling_params.stop_strs:
                if stop_str in tail_str or stop_str in self.decoded_text:
                    self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
                    return

    def set_cached_id(self):
        """Assign tokenizer and cache ids needed by check_finished()."""
        eos_id = getattr(self.tokenizer, "eos_token_id", None)
        self._eos_token_id_cached = int(eos_id) if eos_id is not None else None
        extra = getattr(self.tokenizer, "additional_stop_token_ids", None)
        self._additional_stop_token_ids_cached = (
            set(int(x) for x in extra) if extra else None
        )


class OutputProcesser:
    """Streams generation output to the detokenizer.

    Logprob support is intentionally omitted.
    """

    # Upper bound on how long a pending abort stays buffered waiting for
    # its matching register(). Generous — a client reorder of more than
    # a few seconds is already pathological; 5 min gives plenty of slack
    # while preventing unbounded growth on stray/post-completion aborts.
    _PENDING_ABORT_TTL_S = 300.0

    def __init__(
        self,
        send_to_tokenizer,
        attn_tp_rank: int = 0,
        spec_algorithm=None,
        spec_num_tokens: int | None = None,
        stream_interval: int = 1,
        enable_log_request_stats: bool = False,
        *,
        metrics: EngineMetrics,
    ) -> None:
        # BatchTokenIDOut is pushed directly to
        # ``send_to_tokenizer`` (AsyncLLM's input socket). The
        # inline detokenizer inside AsyncLLM is the only
        # detokenization path.
        self.send_to_tokenizer = send_to_tokenizer
        # Per-request logs fire on each DP replica's TP leader (attn_tp_rank == 0),
        # NOT the global rank 0 — otherwise DP replicas > 0 would log nothing and
        # their requests would be missing from the logs entirely.
        self.attn_tp_rank = attn_tp_rank
        self.spec_algorithm = spec_algorithm
        self.spec_num_tokens = spec_num_tokens
        self.stream_interval = stream_interval
        self.metrics = metrics
        self.enable_log_request_stats = enable_log_request_stats
        # previous forward step ts, for host-side preempt timing
        self._last_step_ts: float = 0.0
        self.log_cnt = 0
        self.rid_to_state: dict[str, RequestState] = {}
        # rid → monotonic ts at which the abort was seen. Covers the
        # "abort arrives before register()" race (pre-arrival reorder),
        # plus grammar-queued aborts that publish_finished_at_admission
        # handles. Entries for rids that never register are swept by TTL
        # to keep this bounded across a long-running server.
        self.pending_aborts: dict[str, float] = {}

    def log_accept_length(self, rid, request_state: RequestState):
        # When --enable-log-request-stats is on, the richer RequestStats line (which
        # already carries acc_len) replaces this one — see _log_request_stats.
        if self.attn_tp_rank == 0 and not self.enable_log_request_stats:
            logger.info(
                "Req: %s Finish! Accept_num_tokens_avg: %s",
                rid,
                request_state.accept_draft_tokens,
            )

    def _log_request_stats(
        self, rid: str, rs: RequestState, finish_time: float
    ) -> None:
        # Single guard for the whole stats path: no tracker (flag off) or non-zero
        # rank => nothing to do. Keeps the forward-loop call sites trivial and the
        # derivation in from_state total (it always sees a tracker).
        if rs.stats is NOOP_STATS or self.attn_tp_rank != 0:
            return
        rs.stats.mark_finish(finish_time)
        stats = RequestStats.from_state(rs, self.spec_algorithm, self.spec_num_tokens)
        # Fused into the scheduler's per-request finish line (supersedes the
        # Accept_num_tokens_avg variant in log_accept_length).
        logger.info("Req: %s Finish! %s", rid, stats)

    def sweep_pending_aborts(self) -> None:
        """Drop TTL-expired entries from ``pending_aborts``.

        Safe to call anytime. pending_aborts is insertion-ordered so we
        can stop at the first non-expired entry. Called both inside
        ``mark_abort`` (so adds are bounded) and periodically from the
        event loop (so entries also age out when aborts stop arriving).
        """
        cutoff = time.monotonic() - self._PENDING_ABORT_TTL_S
        while self.pending_aborts:
            oldest_rid = next(iter(self.pending_aborts))
            if self.pending_aborts[oldest_rid] >= cutoff:
                break
            self.pending_aborts.pop(oldest_rid)

    def mark_abort(self, rid: str, notify_client: bool = False):
        """Mark a request for abort. Safe to call before or after register().

        Routes through ``RequestState.set_finish_with_abort`` so
        ``finished_reason`` is materialized immediately. Without that,
        the gate ``request_state.to_abort and request_state.finished``
        in ``post_process_forward_op`` never fires (``.finished`` is
        ``finished_reason is not None``), so the scheduler keeps
        running the request until natural ``max_tokens``/EOS — the
        cancelled request burns up to ``max_tokens`` forward steps and
        latches a ``--max-num-seqs`` slot in the meantime.

        ``notify_client`` streams a terminating finish to the client (for
        pause-initiated aborts; client-initiated aborts leave it False since
        the tokenizer manager has already cleaned up its own state).
        """
        state = self.rid_to_state.get(rid)
        if state is not None:
            msg = "Aborted by pause" if notify_client else "AbortReq from client"
            state.set_finish_with_abort(msg, notify_client=notify_client)
            return

        self.sweep_pending_aborts()
        self.pending_aborts[rid] = time.monotonic()

    def register(self, rid, state):
        self.rid_to_state[rid] = state
        if self.enable_log_request_stats:
            state.stats = RequestStatsTracker()
        if self.pending_aborts.pop(rid, None) is not None:
            # Same reasoning as ``mark_abort``: drive the abort all the
            # way to ``finished_reason`` so the slot-release gate fires.
            state.set_finish_with_abort("AbortReq from client")

    def publish_finished_at_admission(self, rid: str, state: RequestState) -> None:
        """Stream a finish for a request that was finished before admission.

        Used for grammar-aborted requests (invalid/timed-out compile, missing
        backend) so the client gets a finish_reason without us wasting a
        scheduler slot or a forward step on them.
        """
        self.rid_to_state[rid] = state
        try:
            state.finished_output = False
            self.stream_output([rid], [state])
        finally:
            state.release_pending_multimodal_features()
            self.rid_to_state.pop(rid, None)
            # This path replaces register() for grammar-aborted rids —
            # drop any queued abort marker so pending_aborts doesn't leak
            # and a reused rid isn't instantly re-aborted on next register.
            self.pending_aborts.pop(rid, None)

    def reap_finished_orphan(self, rid: str, state: RequestState) -> None:
        """Resolve a finished request that no future forward op will reap.

        Stream the terminating finish to a passive client (pause-initiated
        aborts still have the client waiting on the stream); client-initiated
        aborts already tore down their own state, so just drop the registered
        state so the rid does not leak.
        """
        if state.abort_notify_client:
            self.publish_finished_at_admission(rid, state)
        else:
            self.rid_to_state.pop(rid, None)

    def _host_advance_matcher(self, completion, model_execution_results):
        """Host-side fallback for the grammar matcher advance.

        Reads already-synced CPU tensors. Fires when no next step arrives to run
        the hostfunc (e.g., last live request finished).
        """
        grammars = completion.grammars or []
        stride = completion.tokens_per_req
        bs = completion.bs
        advance_mask = completion.advance_mask or [True] * bs
        output_tokens = model_execution_results.output_tokens
        accept_lengths = model_execution_results.output_lengths
        terminated_at = [-1] * bs
        for i, grammar in enumerate(grammars):
            if (
                grammar is None
                or grammar.finished
                or grammar.is_terminated()
                or not advance_mask[i]
            ):
                continue
            n_accepted = int(accept_lengths[i].item())
            for j in range(n_accepted):
                tok = int(output_tokens[i * stride + j].item())
                try:
                    grammar.accept_token(tok)
                except Exception:
                    break
                if grammar.is_terminated():
                    terminated_at[i] = j
                    break
        completion.terminated_at = terminated_at

    def add_computed_length(self, rids, input_lengths, extend_prefix_lens):
        for i, rid in enumerate(rids):
            if rid not in self.rid_to_state:
                continue
            if i < len(extend_prefix_lens):
                self.rid_to_state[rid].computed_length = (
                    input_lengths[i] + extend_prefix_lens[i]
                )  # Avoid accumulation here so chunked prefill does not distort the value.
            else:
                self.rid_to_state[rid].add_computed_length(input_lengths[i])

    @staticmethod
    def _aggregate_spec_decode_step(
        *,
        forward_op,
        output_lengths,
        rid_to_state,
    ) -> tuple[int, int]:
        n_ext = forward_op.num_extends()
        accepted = 0
        num_slots = 0
        for i in range(n_ext, len(forward_op.request_ids)):
            rid = forward_op.request_ids[i]
            rs = rid_to_state.get(rid)
            if rs is None or not rs.prefill_finished:
                continue
            out_len = int(output_lengths[i].item())
            accepted += max(0, out_len - 1)
            num_slots += 1
        return num_slots, accepted

    def _emit_spec_decode_metrics(
        self, forward_op, model_execution_results: ModelExecutionResult
    ) -> None:
        if not self.metrics.enabled:
            return
        if forward_op.num_extends() > 0:
            return
        if self.spec_algorithm is None or self.spec_num_tokens is None:
            return
        if model_execution_results.output_lengths is None:
            return
        num_slots, accepted_draft_tokens = self._aggregate_spec_decode_step(
            forward_op=forward_op,
            output_lengths=model_execution_results.output_lengths,
            rid_to_state=self.rid_to_state,
        )
        if num_slots > 0:
            self.metrics.record_spec_decode_step(
                num_decode_slots=num_slots,
                accepted_draft_tokens=accepted_draft_tokens,
                draft_width=self.spec_num_tokens,
            )

    def add_cached_tokens(self, rids: list[str], extend_prefix_lens: list[int]) -> None:
        for rid, prefix_len in zip(rids, extend_prefix_lens):
            if rs := self.rid_to_state.get(rid):
                rs.cached_tokens += max(0, prefix_len - rs.computed_length)

    def post_process_forward_op(
        self,
        forward_op,
        model_execution_results: ModelExecutionResult,
        is_prefill_instance: bool = False,
        on_first_token=None,
    ):
        self.add_cached_tokens(
            forward_op.request_ids,
            forward_op.extend_prefix_lens,
        )
        with nvtx_range("commit:sync", color="red"):
            model_execution_results.sync()

        self._emit_spec_decode_metrics(forward_op, model_execution_results)

        # Wait briefly for the next step's build hostfunc to advance
        # the matcher; if it doesn't come, advance on host. The lock
        # on the completion ensures exactly one path wins.
        grammar_completion = model_execution_results.grammar_completion
        grammar_terminated_at = None
        if grammar_completion is not None:
            if not grammar_completion.event.wait(timeout=0.005):
                with grammar_completion.lock:
                    if not grammar_completion.event.is_set():
                        self._host_advance_matcher(
                            grammar_completion, model_execution_results
                        )
                        grammar_completion.event.set()
            grammar_terminated_at = grammar_completion.terminated_at
        self.log_cnt += 1
        self.add_computed_length(
            forward_op.request_ids,
            forward_op.input_lengths,
            forward_op.extend_prefix_lens,
        )
        num_extends = forward_op.num_extends()

        # per-request stats timing (host-side, only when --enable-log-request-stats)
        stats_now = time.time() if self.enable_log_request_stats else 0.0
        step_dt = 0.0
        prefilling_others = False
        if self.enable_log_request_stats:
            step_dt = (
                stats_now - self._last_step_ts if self._last_step_ts > 0.0 else 0.0
            )
            self._last_step_ts = stats_now
            prefilling_others = num_extends > 0

        request_changes = []
        stream_out_rids = []
        stream_out_states = []
        output_logprobs_list = (
            model_execution_results.output_logprobs.tolist()
            if model_execution_results.output_logprobs is not None
            else None
        )
        # NaN-guard flags, aligned with forward_op.request_ids (None when disabled).
        nan_flags_list = (
            model_execution_results.output_nan_flags.tolist()
            if model_execution_results.output_nan_flags is not None
            else None
        )
        # Per-slot total prefill length as the OP sees it (C++ Request::PrefillSize()).
        # After a flat retract the victim's generated tokens are rebased into the
        # prefill window (RebasePrefill), so this can exceed the original prompt
        # length that RequestState.prefill_finished compares against.
        prefill_lengths = getattr(forward_op, "prefill_lengths", None)
        pt = 0
        for i, rid in enumerate(forward_op.request_ids):
            output_length = model_execution_results.output_lengths[i].item()
            model_output_ids = model_execution_results.output_tokens.tolist()[
                pt : pt + output_length
            ]
            model_output_logprobs = (
                output_logprobs_list[pt : pt + output_length]
                if output_logprobs_list is not None
                else None
            )
            is_decode_slot = i >= num_extends
            if self.spec_num_tokens is not None and is_decode_slot:
                pt += self.spec_num_tokens
            else:
                pt += output_length

            if rid not in self.rid_to_state:
                # means it's delayed token, do not process
                continue

            request_state: RequestState = self.rid_to_state[rid]
            # scheduled_time is stamped pre-forward in the event loop (queue end)

            # Mid-chunk extend slot by the op's own prefill_lengths (rebased after
            # flat retract; C++ owes no result and the sampled token is garbage).
            # Fresh requests: prefill_length == prompt length, same as the gate below.
            if (
                not is_decode_slot
                and prefill_lengths is not None
                and forward_op.extend_prefix_lens[i] + forward_op.input_lengths[i]
                < prefill_lengths[i]
            ):
                continue

            # Do not output chunking result
            if not request_state.prefill_finished:
                continue

            request_state.stats.mark_prefill_done(stats_now)
            if i >= num_extends:
                request_state.stats.record_decode_step(step_dt, prefilling_others)

            nan_detected = nan_flags_list is not None and nan_flags_list[i]
            if nan_detected and not request_state.finished:
                request_state.finished_reason = FINISH_ABORT(
                    message=(
                        "Request terminated: numerical corruption (NaN logits"
                        " or out-of-vocab sample) detected during generation."
                    ),
                    err_type=ABORT_CODE.NumericalError,
                )
                # Keep one sanitized token so accounting matches a mid-step finish.
                model_output_ids = model_output_ids[:1]
                if model_output_logprobs is not None:
                    model_output_logprobs = model_output_logprobs[:1]
                self.metrics.record_nan_abort()
                if self.attn_tp_rank == 0:
                    logger.warning(
                        "Req %s terminated: NaN detected in logits (or an"
                        " out-of-vocab sample escaped the sampler);"
                        " isolating it from the batch.",
                        rid,
                    )

            # Notify caller of first output token (used by prefill node to hand off
            # bootstrap token and speculative candidates to the KV transfer layer).
            # NaN-terminated requests skip the handoff: their KV is suspect.
            if on_first_token is not None and model_output_ids and not nan_detected:
                bootstrap_token = int(model_output_ids[0])
                spec_candidate_ids = None
                if model_execution_results.next_input_ids is not None and i < len(
                    model_execution_results.next_input_ids
                ):
                    spec_candidate_ids = [
                        int(x)
                        for x in model_execution_results.next_input_ids[i].tolist()
                    ]

                on_first_token(
                    rid,
                    forward_op.request_pool_indices[i],
                    bootstrap_token,
                    spec_candidate_ids,
                )

            if is_decode_slot and self.spec_algorithm is not None:
                request_state.spec_verify_ct += 1

            # With the capturable grammar pipeline the matcher is
            # advanced by the hostfunc; here we just read which token
            # (if any) terminated it so FINISH_MATCHED_TOKEN fires on
            # the right token and check_finished skips the now-stale
            # grammar.is_terminated() probe.
            use_hostfunc = grammar_terminated_at is not None
            advance_grammar = not use_hostfunc and request_state.grammar is not None
            term_idx = (
                grammar_terminated_at[i]
                if use_hostfunc and request_state.grammar is not None
                else -1
            )
            new_ids = []
            for j, model_output_id in enumerate(model_output_ids):
                request_state.output_ids.append(model_output_id)
                if advance_grammar:
                    request_state.grammar.accept_token(model_output_id)
                if (
                    request_state.return_logprob
                    and request_state.output_token_logprobs_val is not None
                    and model_output_logprobs is not None
                ):
                    request_state.output_token_logprobs_val.append(
                        model_output_logprobs[j]
                    )
                    request_state.output_token_logprobs_idx.append(model_output_id)
                if term_idx == j:
                    # Grammar termination takes precedence over
                    # length/EOS/stop_str at the same step (matching
                    # check_finished's original order).
                    request_state.finished_reason = FINISH_MATCHED_TOKEN(
                        matched=model_output_id
                    )
                else:
                    request_state.check_finished(skip_grammar_termination=use_hostfunc)
                new_ids.append(model_output_id)
                if request_state.finished:
                    request_state.accept_draft_tokens = (
                        (len(request_state.output_ids) - 1)
                        / request_state.spec_verify_ct
                        if request_state.spec_verify_ct > 0
                        else 0
                    )
                    self.log_accept_length(rid, request_state)
                    break

            # first output token == TTFT anchor
            if request_state.output_ids:
                request_state.stats.mark_first_token(stats_now)

            # For aborted requests, skip output to detokenizer (the tokenizer
            # manager already cleaned up), just notify the scheduler to finish.
            # Exception: pause-initiated aborts (abort_notify_client) leave a
            # passive client that still needs a terminating finish streamed.
            if request_state.to_abort and request_state.finished:
                request_changes.append(make_extend_result_event(rid, new_ids))
                request_changes.append(make_finish_event(rid))
                if request_state.abort_notify_client:
                    stream_out_rids.append(rid)
                    stream_out_states.append(request_state)
                self._log_request_stats(rid, request_state, stats_now)
                request_state.release_pending_multimodal_features()
                self.rid_to_state.pop(rid)
                continue

            request_changes.append(make_extend_result_event(rid, new_ids))
            if is_prefill_instance:
                # Prefill instances: never stream intermediate output to detokenizer.
                # The finish packet is sent exactly once by finish_prefill_request()
                # when SucceededEvent arrives (KV transfer complete).  Sending output
                # here would either give the client partial data or trigger a double-
                # finish on the TM side.
                pass
            elif request_state.finished:
                stream_out_rids.append(rid)
                stream_out_states.append(request_state)
                # Abort (vs Finish) keeps corrupted KV out of the prefix caches.
                request_changes.append(
                    make_abort_event(rid) if nan_detected else make_finish_event(rid)
                )
                self._log_request_stats(rid, request_state, stats_now)
                request_state.release_pending_multimodal_features()
                self.rid_to_state.pop(rid)
            else:
                stream_out_rids.append(rid)
                stream_out_states.append(request_state)
                if is_decode_slot:
                    request_changes.append(
                        make_update_reserve_tokens_event(rid, output_length)
                    )

        self.stream_output(stream_out_rids, stream_out_states)
        return request_changes

    def on_remote_prefill_done(self, req_id: str, bootstrap_token: int) -> None:
        """Record the bootstrap token on a decode-node request (RemotePrefillDoneEvent).

        The bootstrap_token is the first real output token produced by the prefill node.
        It is appended to output_ids so the decode side starts generation from the
        correct position.

        bootstrap_token == -1 means the prefill side did not (or could not) supply a
        token (e.g. it was generated on a rank whose ZMQ message arrived after the
        success barrier had already been satisfied).
        """
        if req_id not in self.rid_to_state:
            return
        if bootstrap_token == -1:
            logger.warning(
                "[on_remote_prefill_done] rid=%s received bootstrap_token=-1, skipping append to output_ids",
                req_id,
            )
            return
        state = self.rid_to_state[req_id]
        state.output_ids.append(bootstrap_token)
        state.check_finished()

    def finish_prefill_request(self, req_id: str) -> list:
        """Finish a prefill-instance request when KV transfer succeeds (SucceededEvent).

        Called by event_loop._process_kv_transfer_events at the correct moment — the
        SucceededEvent itself drives the C++ FSM transition Decoding → Finished,
        so we must NOT emit an additional make_finish_event here.

        We send a finished BatchTokenIDOut to the detokenizer so the Prefill TM
        can resolve its HTTP coroutine and let the HTTP load balancer unblock.
        Without this, the load balancer waits forever for the prefill side's HTTP
        response while the decode side has already finished — client hangs
        indefinitely.
        """
        if req_id not in self.rid_to_state:
            return []
        rs = self.rid_to_state.pop(req_id)
        rs.release_pending_multimodal_features()

        # Ensure a finish reason is set so TokenizerManager marks the request done.
        if not rs.finished:
            rs.finished_reason = FINISH_LENGTH(length=len(rs.output_ids))
        rs.finished_output = False
        # PD prefill node's terminal path (the other finish/abort logging lives in
        # post_process_forward_op). Self-guarded, so a no-op when the flag is off.
        self._log_request_stats(req_id, rs, time.time())
        self.stream_output([req_id], [rs])
        # SucceededEvent already finishes the C++ FSM; no extra FinishEvent needed
        return []

    def stream_output(
        self, stream_out_rids: list[str], output_states: list[RequestState]
    ) -> None:
        """Collect per-step results and forward them to the detokenizer."""
        if len(output_states) == 0:
            return

        rids_to_send = []
        finished_reasons = []
        decoded_texts: list[str] = []
        decode_ids_list = []
        read_offsets: list[int] = []
        output_ids = []
        output_multi_ids = []
        skip_special_tokens: list[bool] = []
        spaces_between_special_tokens: list[bool] = []
        no_stop_trim: list[bool] = []
        prompt_tokens: list[int] = []
        completion_tokens: list[int] = []
        cached_tokens: list[int] = []
        spec_verify_ct: list[int] = []
        batch_accept_draft_tokens: list[float] = []
        output_extra_infos: list[dict] = []
        output_token_logprobs_val: list[list[float]] = []
        output_token_logprobs_idx: list[list[int]] = []

        for i, rs in enumerate(output_states):
            # For finished requests, always output (unless already output)
            if rs.finished:
                if rs.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
                rs.finished_output = True
                should_output = True
            else:
                # For ongoing requests, use stream interval logic
                if rs.stream:
                    stream_interval = getattr(
                        rs.sampling_params, "stream_interval", None
                    )
                    if stream_interval is None:
                        stream_interval = self.stream_interval
                    should_output = (
                        rs.output_length % stream_interval == 1
                        if stream_interval > 1
                        else rs.output_length % stream_interval == 0
                    )
                else:
                    stream_interval = DEFAULT_FORCE_STREAM_INTERVAL
                    should_output = (
                        rs.output_length == 1 or rs.output_length % stream_interval == 0
                    )

            if not should_output:
                continue

            rids_to_send.append(stream_out_rids[i])
            send_token_offset = rs.send_token_offset

            finished_reasons.append(
                rs.finished_reason.to_json() if rs.finished_reason else None
            )
            decoded_texts.append(rs.decoded_text)

            decode_ids, read_offset = rs.init_incremental_detokenize()
            decode_ids_list.append(decode_ids[rs.send_decode_id_offset :])
            rs.send_decode_id_offset = len(decode_ids)

            read_offsets.append(read_offset)
            output_ids.append(rs.output_ids[send_token_offset:])
            rs.send_token_offset = rs.output_length

            output_multi_ids.append([])

            skip_special_tokens.append(rs.sampling_params.skip_special_tokens)
            spaces_between_special_tokens.append(
                rs.sampling_params.spaces_between_special_tokens
            )
            no_stop_trim.append(rs.sampling_params.no_stop_trim)
            prompt_tokens.append(rs.input_length)
            completion_tokens.append(rs.output_length)
            cached_tokens.append(rs.cached_tokens)

            if self.spec_algorithm is not None:
                spec_verify_ct.append(rs.spec_verify_ct)
                batch_accept_draft_tokens.append(rs.accept_draft_tokens)

            output_extra_infos.append({"decode_prefix_len": rs.prefix_len})

            if rs.return_logprob and rs.output_token_logprobs_val is not None:
                # Send only the slice not yet shipped; send_token_offset was
                # just advanced above, so use the logprob list tail.
                n_new = rs.output_length - send_token_offset
                output_token_logprobs_val.append(
                    rs.output_token_logprobs_val[-n_new:] if n_new > 0 else []
                )
                output_token_logprobs_idx.append(
                    rs.output_token_logprobs_idx[-n_new:] if n_new > 0 else []
                )
            else:
                output_token_logprobs_val.append([])
                output_token_logprobs_idx.append([])

        # Don't send empty batch to detokenizer
        if len(rids_to_send) == 0:
            return

        batch_id_out = BatchTokenIDOut(
            rids=rids_to_send,
            finished_reasons=finished_reasons,
            decoded_texts=decoded_texts,
            decode_ids=decode_ids_list,
            read_offsets=read_offsets,
            output_ids=output_ids,
            output_multi_ids=output_multi_ids,
            skip_special_tokens=skip_special_tokens,
            spaces_between_special_tokens=spaces_between_special_tokens,
            no_stop_trim=no_stop_trim,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            cached_tokens=cached_tokens,
            spec_verify_ct=spec_verify_ct,
            input_token_logprobs_val=[],
            input_token_logprobs_idx=[],
            output_token_logprobs_val=output_token_logprobs_val,
            output_token_logprobs_idx=output_token_logprobs_idx,
            input_top_logprobs_val=[],
            input_top_logprobs_idx=[],
            output_top_logprobs_val=[],
            output_top_logprobs_idx=[],
            input_token_ids_logprobs_val=[],
            input_token_ids_logprobs_idx=[],
            output_token_ids_logprobs_val=[],
            output_token_ids_logprobs_idx=[],
            output_hidden_states=[],
            batch_accept_draft_tokens=batch_accept_draft_tokens,
            output_extra_infos=output_extra_infos,
            generated_time=time.time(),
        )

        # Push BatchTokenIDOut directly to AsyncLLM via the shared
        # tokenizer-ipc socket. AsyncLLM runs IncrementalDetokenizer
        # inline — there is no detokenizer subprocess anymore.
        self.send_to_tokenizer.send_pyobj(batch_id_out)