lightseekorg--tokenspeed/python/tokenspeed/runtime/engine/request.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.

import copy
import time
from typing import Any

import torch

from tokenspeed.runtime.cache.req_to_token_pool import (
    ReqToTokenPoolInfo,
)
from tokenspeed.runtime.engine.request_types import (  # noqa: F401
    ABORT_CODE,
    FINISH_ABORT,
    FINISH_LENGTH,
    FINISH_MATCHED_STR,
    FINISH_MATCHED_TOKEN,
    INIT_INCREMENTAL_DETOKENIZATION_OFFSET,
    BaseFinishReason,
)
from tokenspeed.runtime.grammar.base_grammar_backend import BaseGrammarObject
from tokenspeed.runtime.metrics.collector import TimeStats
from tokenspeed.runtime.sampling.sampling_params import SamplingParams
from tokenspeed.runtime.utils import get_colorful_logger

logger = get_colorful_logger(__name__)


class Req:
    """The input and output status of a request."""

    def __init__(
        self,
        rid: str,
        origin_input_text: str,
        origin_input_ids: tuple[int],
        sampling_params: SamplingParams,
        return_logprob: bool = False,
        top_logprobs_num: int = 0,
        token_ids_logprob: list[int] = None,
        stream: bool = False,
        origin_input_ids_unpadded: tuple[int] | None = None,
        input_embeds: list[list[float]] | None = None,
        input_extra_infos: list[dict] | None = None,
        session_id: str | None = None,
        custom_logit_processor: str | None = None,
        return_hidden_states: bool = False,
        eos_token_ids: set[int] | None = None,
        bootstrap_host: str | None = None,
        bootstrap_port: int | None = None,
        bootstrap_room: int | None = None,
        origin_input_multi_ids: list[list[int]] | None = None,
        created_time: float | None = None,
    ):
        # Input and output info
        self.rid = rid
        self.origin_input_text = origin_input_text
        self.origin_input_ids_unpadded = (
            origin_input_ids_unpadded
            if origin_input_ids_unpadded
            else origin_input_ids  # Before image padding
        )
        self.origin_input_ids = origin_input_ids
        self.origin_input_multi_ids = origin_input_multi_ids
        # Each decode stage's output ids
        self.output_ids = []
        self.output_multi_ids = []
        # fill_ids = origin_input_ids + output_ids. Updated if chunked.
        self.fill_ids = None
        self.fill_multi_ids = None
        self.fill_input_embeds = None
        # For Eagle and chunked prefill, remove first token when chunked prefill
        self.draft_fill_ids = None
        self.session_id = session_id
        self.input_embeds = input_embeds
        self.input_extra_infos = input_extra_infos

        # Sampling info
        if isinstance(sampling_params.custom_params, dict):
            sampling_params = copy.copy(sampling_params)
            sampling_params.custom_params = sampling_params.custom_params | {
                "__req__": self
            }
        self.sampling_params = sampling_params

        self.custom_logit_processor = custom_logit_processor
        self.return_hidden_states = return_hidden_states

        # Memory pool info
        self.req_pool_idx: int | None = None
        self.req_to_token_pool_info: ReqToTokenPoolInfo | None = None
        # substitute for prefix_indices
        self.prefix_page_ids = []
        self.prefix_len = 0

        # Check finish
        self.tokenizer = None
        # Cached tokenizer-related ids to avoid repeated HF attribute lookups in check_finished().
        self._eos_token_id_cached: int | None = None
        self._additional_stop_token_ids_cached: set[int] | None = None
        self.finished_reason = None
        # Whether this request has finished output
        self.finished_output = None
        # If we want to abort the request in the middle of the event loop, set this to true
        # Note: We should never set finished_reason in the middle, the req will get filtered and never respond
        self.to_abort = False
        # This carries the error message for `.to_abort` and will be attached to the finished_reason at the end of the event loop
        self.to_abort_message: str = "Unknown error"
        self.stream = stream
        self.eos_token_ids = eos_token_ids

        # For incremental decoding
        # ----- | --------- read_ids -------|
        # ----- |   surr_ids  |
        # xxxxx | xxxxxxxxxxx | xxxxxxxxxxx |
        # ----- ^ ----------- ^ ----------- ^
        # ----- 1 ----------- 2 ----------- 3
        # 1: surr_offset
        # 2: read_offset
        # 3: last token
        self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
        self.read_offset = None
        self.decoded_text = ""

        # Prefix info
        # The indices to kv cache for the shared prefix.
        self.prefix_indices = []
        # Number of tokens to run prefill.
        self.extend_input_len = 0
        # The relative logprob_start_len in an extend batch
        self.extend_logprob_start_len = 0
        self.last_node = None

        # Whether or not if it is chunked. It increments whenever
        # it is chunked, and decrement whenever chunked request is
        # processed.
        self.is_chunked = 0

        # For retraction
        self.is_retracted = False

        # Incremental streamining
        self.send_token_offset: int = 0
        self.send_decode_id_offset: int = 0
        # because the decode server does not have the first output token logprobs
        self.send_output_token_logprobs_offset: int = 0

        # Logprobs (arguments)
        self.return_logprob = return_logprob
        # Start index to compute logprob from.
        self.logprob_start_len = 0
        self.top_logprobs_num = top_logprobs_num
        self.token_ids_logprob = token_ids_logprob

        # Logprobs (return values)
        self.input_logprob_sent: bool = False
        self.input_token_logprobs_val: list[float] | None = None
        self.input_token_logprobs_idx: list[int] | None = None
        self.input_top_logprobs_val: list[float] | None = None
        self.input_top_logprobs_idx: list[int] | None = None
        self.input_token_ids_logprobs_val: list[float] | None = None
        self.input_token_ids_logprobs_idx: list[int] | None = None
        # Temporary holder to store input_token_logprobs.
        self.input_token_logprobs: list[tuple[int]] | None = None
        self.temp_input_top_logprobs_val: list[torch.Tensor] | None = None
        self.temp_input_top_logprobs_idx: list[int] | None = None
        self.temp_input_token_ids_logprobs_val: list[float] | None = None
        self.temp_input_token_ids_logprobs_idx: list[int] | None = None

        if return_logprob:
            self.output_token_logprobs_val = []
            self.output_token_logprobs_idx = []
            self.output_top_logprobs_val = []
            self.output_top_logprobs_idx = []
            self.output_token_ids_logprobs_val = []
            self.output_token_ids_logprobs_idx = []
        else:
            self.output_token_logprobs_val = self.output_token_logprobs_idx = (
                self.output_top_logprobs_val
            ) = self.output_top_logprobs_idx = self.output_token_ids_logprobs_val = (
                self.output_token_ids_logprobs_idx
            ) = None
        self.hidden_states = []

        # Embedding (return values)
        self.embedding = None

        # Constrained decoding
        self.grammar: BaseGrammarObject | None = None

        # The number of cached tokens that were already cached in the KV cache
        self.cached_tokens = 0
        self.already_computed = 0
        self.last_host_node: Any = None
        self.host_hit_length = 0

        # The number of verification forward passes in the speculative decoding.
        # This is used to compute the average acceptance length per request.
        self.spec_verify_ct = 0

        # Time of obj created
        # Use the created_time from tokenizer if provided, otherwise use current time
        if created_time is not None:
            self.created_time = created_time
        else:
            self.created_time = time.time()
        # Calculate the time from receiving the request at TokenizerManager to reaching process_input_requests in the scheduling process
        self.tokenizer_to_scheduler_latency = time.time() - self.created_time
        # For metrics
        self.time_stats: TimeStats = TimeStats()
        self.has_log_time_stats: bool = False
        self.queue_time_start = None
        self.queue_time_end = None
        self.last_tic = time.monotonic()
        self.first_latency_recorded = (
            False  # Flag to track if first latency has been recorded
        )
        self.prefill_waiting_recorded = False
        self.first_chunk_forward_start_time = None

        self.reserve_num_tokens = 0
        # For disaggregation
        self.bootstrap_host: str = bootstrap_host
        self.bootstrap_port: int | None = bootstrap_port
        self.bootstrap_room: int | None = bootstrap_room

        # the start index of the sent kv cache
        # We want to send it chunk by chunk for chunked prefill.
        # After every chunk forward, we do the following:
        # kv_send(req.input_ids[req.start_send_idx:len(req.fill_ids)])
        # start_send_idx = len(req.fill_ids)
        self.start_send_idx: int = 0

        # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap
        # This is because kv is not ready in `process_prefill_chunk`.
        # We use `tmp_end_idx` to store the end index of the kv cache to send.
        self.tmp_end_idx: int = -1
        self.metadata_buffer_index: int = -1
        # Only meaningful in speculative reasoning.
        self.accept_draft_tokens: float | None = None

        self.output_extra_info: dict[str, Any] = {}

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

    @property
    def seqlen(self):
        return len(self.origin_input_ids) + len(self.output_ids)

    def finished(self) -> bool:
        # Whether request reached finished condition
        return self.finished_reason is not None

    def init_incremental_detokenize(self):
        first_iter = self.surr_offset is None or self.read_offset is None

        if first_iter:
            self.read_offset = len(self.origin_input_ids_unpadded)
            self.surr_offset = max(
                self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
            )
            # self.surr_offset = self.read_offset

        all_ids = self.origin_input_ids_unpadded + self.output_ids
        return all_ids[self.surr_offset :], self.read_offset - self.surr_offset

    def check_finished(self):
        if self.finished():
            return

        if self.to_abort:
            self.finished_reason = FINISH_ABORT(
                message=self.to_abort_message,
            )
            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

        if self.grammar is not None:
            if self.grammar.is_terminated():
                self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1])
                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.tokenizer is not None and self._eos_token_id_cached is None:
                self.set_tokenizer(self.tokenizer)
            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 __repr__(self):
        return (
            f"Req(rid={self.rid}, "
            f"input_ids={len(self.origin_input_ids)}, output_ids={len(self.output_ids)})"
        )