# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""TokenizerManager is a process that tokenizes the text."""

from __future__ import annotations

import asyncio
import copy
import dataclasses
import json
import logging
import os
import pickle
import signal
import socket
import sys
import threading
import time
from array import array
from collections import deque
from contextlib import nullcontext
from datetime import datetime
from enum import Enum
from functools import lru_cache
from http import HTTPStatus
from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union

import fastapi
import numpy as np
import pybase64
import torch
import uvloop
import zmq
import zmq.asyncio
from fastapi import BackgroundTasks

from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
from sglang.srt.disaggregation.encode_receiver import create_mm_receiver
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry
from sglang.srt.managers.async_dynamic_batch_tokenizer import AsyncDynamicbatchTokenizer
from sglang.srt.managers.disagg_service import start_disagg_service
from sglang.srt.managers.embed_types import PositionalEmbeds
from sglang.srt.managers.io_struct import (
    AbortReq,
    ActiveRanksOutput,
    BaseBatchReq,
    BaseReq,
    BatchEmbeddingOutput,
    BatchStrOutput,
    BatchTokenIDOutput,
    BatchTokenizedEmbeddingReqInput,
    BatchTokenizedGenerateReqInput,
    ConfigureLoggingReq,
    ContinueGenerationReqInput,
    ElasticScaleUpdateReq,
    EmbeddingReqInput,
    FreezeGCReq,
    GenerateReqInput,
    HealthCheckOutput,
    LoadLoRAAdapterReqInput,
    OpenSessionReqOutput,
    PauseGenerationReqInput,
    ScaleElasticEPReqInput,
    ScaleElasticEPReqOutput,
    SessionParams,
    ShutdownReq,
    TokenizedEmbeddingReqInput,
    TokenizedGenerateReqInput,
    UpdateWeightFromDiskReqInput,
    UpdateWeightFromDiskReqOutput,
    async_sock_recv,
    async_sock_send,
    build_flat_input_top_logprobs_arrays,
    sock_send,
    unwrap_from_pickle,
)
from sglang.srt.managers.load_snapshot import create_load_snapshot_reader
from sglang.srt.managers.mm_utils import TensorTransportMode, wrap_shm_features
from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
from sglang.srt.managers.schedule_batch import (
    MultimodalDataItem,
    get_request_return_hidden_states_mode,
)
from sglang.srt.managers.scheduler_input_blocker import input_blocker_guard_region
from sglang.srt.managers.tokenizer_control_mixin import TokenizerControlMixin
from sglang.srt.managers.tokenizer_manager_score_mixin import TokenizerManagerScoreMixin
from sglang.srt.managers.utils import (
    compute_num_reserved_tokens,
    is_health_check_generate_req,
)
from sglang.srt.model_executor.forward_batch_info import (
    get_server_return_hidden_states_mode,
)
from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
from sglang.srt.observability.metrics_collector import (
    STAT_LOGGER_ROLE_TOKENIZER,
    TokenizerMetricsCollector,
    resolve_collector_class,
)
from sglang.srt.observability.req_time_stats import (
    APIServerReqTimeStats,
    convert_time_to_realtime,
    real_time,
    set_time_batch,
)
from sglang.srt.observability.request_metrics_exporter import (
    RequestMetricsExporterManager,
)
from sglang.srt.observability.trace import SpanAttributes, extract_trace_headers
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import (
    PortArgs,
    ServerArgs,
    set_global_server_args_for_tokenizer,
)
from sglang.srt.utils import (
    configure_gc_warning,
    freeze_gc,
    get_bool_env_var,
    get_or_create_event_loop,
    kill_process_tree,
)
from sglang.srt.utils.aio_rwlock import RWLock
from sglang.srt.utils.cuda_vmm_transport_utils import CudaVmmFeatureTransport
from sglang.srt.utils.cudacore_pyspy_dump_utils import (
    collect_scheduler_processes,
    pyspy_dump_schedulers,
    trigger_cuda_user_coredump,
)
from sglang.srt.utils.hf_transformers_utils import (
    get_processor,
    get_tokenizer,
    get_tokenizer_from_processor,
)
from sglang.srt.utils.network import get_zmq_socket
from sglang.srt.utils.request_logger import RequestLogger
from sglang.srt.utils.watchdog import Watchdog
from sglang.utils import TypeBasedDispatcher, get_exception_traceback

asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())

_REQUEST_STATE_WAIT_TIMEOUT = envs.SGLANG_REQUEST_STATE_WAIT_TIMEOUT.get()

logger = logging.getLogger(__name__)


def _reject_missing_dispatched_encoder_embedding(server_args, request_obj, mm_inputs):
    """Do not silently turn a failed EPD request into local vision work."""
    if (
        mm_inputs is None
        and server_args.language_only
        and server_args.encoder_transfer_backend == "zmq_to_tokenizer"
        and request_obj.need_wait_for_mm_inputs
    ):
        raise fastapi.HTTPException(
            status_code=HTTPStatus.SERVICE_UNAVAILABLE,
            detail=(
                "The encoder did not return multimodal embeddings. "
                "The request was not run locally in language-only mode."
            ),
        )


@lru_cache(maxsize=1)
def _ragged_verify_cap_accept() -> bool:
    # The mode env is fixed at server launch; cache to keep it off the
    # per-request metrics path.
    from sglang.srt.speculative.ragged_verify import (
        RaggedVerifyMode,
        read_ragged_verify_mode,
    )

    return read_ragged_verify_mode() is RaggedVerifyMode.CAP_ACCEPT


_INCREMENTAL_STREAMING_META_INFO_KEYS = (
    "output_token_logprobs",
    "output_top_logprobs",
    "output_token_ids_logprobs",
    "output_token_sampling_mask",
    "output_token_sampling_logprobs",
)


class RequestAbortedError(ValueError):
    status_code = 499


@dataclasses.dataclass
class ReqState:
    """Store the state a request."""

    out_list: List[Dict[Any, Any]]
    finished: bool
    event: asyncio.Event
    obj: Union[GenerateReqInput, EmbeddingReqInput]

    # For performance metrics
    time_stats: APIServerReqTimeStats
    abort_requested: bool = False
    lifecycle_id: object = dataclasses.field(default_factory=object)
    dispatched: bool = False
    last_completion_tokens: int = 1
    ttft_observed: bool = False

    # For streaming output
    last_output_offset: int = 0

    # Accumulate text lazily so incremental streaming can emit the incoming
    # delta directly without rebuilding the full output prefix.
    text: str = ""
    text_chunks: List[str] = dataclasses.field(default_factory=list)

    def append_text(self, chunk: str):
        if chunk:
            self.text_chunks.append(chunk)

    def get_text(self) -> str:
        if self.text_chunks:
            self.text += "".join(self.text_chunks)
            self.text_chunks.clear()
        return self.text

    def get_crash_dump_output(self) -> Dict[Any, Any]:
        out = {}
        if self.text or self.text_chunks:
            out["text"] = self.get_text()
        if self.output_ids:
            out["output_ids"] = self.output_ids.copy()
        return out

    # For incremental state update.
    # TODO(lianmin): do not initialize some lists if not needed.
    output_ids: List[int] = dataclasses.field(default_factory=list)
    input_token_logprobs_val: List[float] = dataclasses.field(default_factory=list)
    input_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list)
    output_token_logprobs_val: List[float] = dataclasses.field(default_factory=list)
    output_token_logprobs_idx: List[int] = dataclasses.field(default_factory=list)
    input_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list)
    input_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list)
    output_top_logprobs_val: List[List[float]] = dataclasses.field(default_factory=list)
    output_top_logprobs_idx: List[List[int]] = dataclasses.field(default_factory=list)
    input_token_ids_logprobs_val: List = dataclasses.field(default_factory=list)
    input_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list)
    output_token_ids_logprobs_val: List = dataclasses.field(default_factory=list)
    output_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list)
    output_token_sampling_mask: List = dataclasses.field(default_factory=list)
    output_token_sampling_logprobs: List = dataclasses.field(default_factory=list)

    # Cached flat-format prompt top logprob fields; rebuilt only when more
    # prefill chunks arrive, so streaming decode chunks reuse the payload.
    input_top_logprobs_flat_fields: Optional[Dict[str, Any]] = None
    input_top_logprobs_flat_num_rows: int = -1
    # Scheduler-assembled flat arrays (val float32 [rows, k], idx int32
    # [rows, k], null_prefix), sent once at prefill completion. When present,
    # the nested input_top_logprobs_val/idx above stay empty.
    input_top_logprobs_scheduler_flat: Optional[Tuple[np.ndarray, np.ndarray, int]] = (
        None
    )

    # For detokenized logprobs
    input_token_logprobs: List[Any] = dataclasses.field(default_factory=list)
    output_token_logprobs: List[Any] = dataclasses.field(default_factory=list)
    input_top_logprobs: List[Any] = dataclasses.field(default_factory=list)
    output_top_logprobs: List[Any] = dataclasses.field(default_factory=list)
    input_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list)
    output_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list)
    customized_info_accumulated: Dict[str, List[Any]] = dataclasses.field(
        default_factory=dict
    )

    # For return_prompt_token_ids: stores prompt token IDs captured after tokenization
    prompt_token_ids: Optional[List[int]] = None


def _slice_streaming_output_meta_info(
    meta_info: Dict[Any, Any],
    last_output_offset: int,
    customized_info_keys: Optional[Iterable[str]] = None,
) -> None:
    """Align output-side metadata with the current incremental streaming chunk."""
    streaming_meta_info_keys = set(_INCREMENTAL_STREAMING_META_INFO_KEYS)
    if customized_info_keys is not None:
        streaming_meta_info_keys.update(customized_info_keys)
    for key in meta_info.keys() & streaming_meta_info_keys:
        meta_info[key] = meta_info[key][last_output_offset:]


def _build_flat_input_top_logprobs_fields(
    input_top_logprobs_val: List[Optional[List[float]]],
    input_top_logprobs_idx: List[Optional[List[int]]],
    top_logprobs_num: int,
    return_b64: bool = False,
) -> Dict[str, Any]:
    """Build the flat raw prompt top logprob response fields.

    `input_top_logprobs_shape` is the literal [rows, k] shape of the flat
    arrays. The leading null positions (counted by
    `input_top_logprobs_null_prefix`) precede the arrays, so covered position
    i, entry j lives at flat[(i - null_prefix) * k + j] and the covered range
    spans null_prefix + rows positions. With ``return_b64``, the arrays are
    base64 contiguous little-endian binary; the dtype marker fields let the
    widths change later without a wire break.
    """
    val_arr, idx_arr, null_prefix = build_flat_input_top_logprobs_arrays(
        input_top_logprobs_val, input_top_logprobs_idx, top_logprobs_num
    )
    if return_b64:
        return _build_flat_input_top_logprobs_fields_from_arrays(
            val_arr, idx_arr, null_prefix, return_b64=True
        )
    # Flatten the original python rows so the JSON numbers keep their full
    # (float64) precision, matching the pre-scheduler-flat output.
    return {
        "input_top_logprobs_val_flat": [
            v for row in input_top_logprobs_val[null_prefix:] for v in row
        ],
        "input_top_logprobs_idx_flat": [
            i for row in input_top_logprobs_idx[null_prefix:] for i in row
        ],
        "input_top_logprobs_shape": [val_arr.shape[0], val_arr.shape[1]],
        "input_top_logprobs_null_prefix": null_prefix,
    }


def _build_flat_input_top_logprobs_fields_from_arrays(
    val_arr: np.ndarray,
    idx_arr: np.ndarray,
    null_prefix: int,
    return_b64: bool = False,
) -> Dict[str, Any]:
    """Build the flat response fields from scheduler-assembled [rows, k]
    arrays (see `_build_flat_input_top_logprobs_fields` for the field
    semantics)."""
    fields: Dict[str, Any] = {}
    if return_b64:
        fields["input_top_logprobs_val_flat_b64"] = pybase64.b64encode(
            val_arr.tobytes()
        ).decode("utf-8")
        fields["input_top_logprobs_idx_flat_b64"] = pybase64.b64encode(
            idx_arr.tobytes()
        ).decode("utf-8")
        fields["input_top_logprobs_val_flat_b64_dtype"] = "float32"
        fields["input_top_logprobs_idx_flat_b64_dtype"] = "int32"
    else:
        fields["input_top_logprobs_val_flat"] = val_arr.reshape(-1).tolist()
        fields["input_top_logprobs_idx_flat"] = idx_arr.reshape(-1).tolist()
    fields["input_top_logprobs_shape"] = [val_arr.shape[0], val_arr.shape[1]]
    fields["input_top_logprobs_null_prefix"] = null_prefix
    return fields


class InputFormat(Enum):
    """Input format types for tokenization handling."""

    SINGLE_STRING = 1  # Regular single text like "Hello world"
    BATCH_STRINGS = 2  # Regular batch like ["Hello", "World"]
    CROSS_ENCODER_PAIRS = 3  # Cross-encoder pairs like [["query", "document"]]


_SERVER_ARGS_FIELDS = frozenset(f.name for f in dataclasses.fields(ServerArgs))

_MANAGER_OWNED_FIELDS = ("model_path", "served_model_name")


class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
    """TokenizerManager is a process that tokenizes the text."""

    @property
    def serving_chat_class(self):
        """Return the serving chat class for OpenAI API.

        Override in subclass to provide custom serving behavior.
        """
        from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat

        return OpenAIServingChat

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        *,
        start_pd_bootstrap_service: bool = True,
    ):
        # Parse args
        self.server_args = server_args
        self.startup_time: Optional[Dict[str, Any]] = None
        self._config_updates: List[Tuple[str, Dict[str, Any]]] = []
        self.elastic_worker_count = server_args.dp_size
        self.elastic_pending_ep_size = None
        self.elastic_scale_phase = "idle"
        self.elastic_last_error = None
        self.enable_metrics = server_args.enable_metrics
        self.incremental_streaming_output = server_args.incremental_streaming_output
        self.enable_lora = server_args.enable_lora
        self.enable_trace = server_args.enable_trace
        self.allow_auto_truncate = server_args.allow_auto_truncate
        self.skip_tokenizer_init = server_args.skip_tokenizer_init
        self.preferred_sampling_params = server_args.preferred_sampling_params
        self.crash_dump_folder = server_args.crash_dump_folder
        set_global_server_args_for_tokenizer(server_args)

        # Init model config
        self.init_model_config()
        self._validate_cuda_vmm_feature_transport_support()

        # Initialize tokenizer and multimodalprocessor
        self.init_tokenizer_and_processor()

        # Init inter-process communication
        self.init_ipc_channels(port_args)

        # Init running status
        self.init_running_status()

        # Init logging and dumping
        self.init_request_logging_and_dumping()

        # Init weight update
        self.init_weight_update()

        # Init LoRA status
        self.init_lora()

        # Init PD disaggregation and encoder disaggregation
        self.init_disaggregation(start_pd_bootstrap_service=start_pd_bootstrap_service)

        # Init metric collector and watchdog
        self.init_metric_collector_watchdog()

        # Init request dispatcher
        self.init_request_dispatcher()

        # Construct this last so later initialization failures cannot orphan
        # the transport's recycler thread.
        self.cuda_vmm_feature_transport = CudaVmmFeatureTransport(
            self.server_args, self.mm_processor
        )

    def init_model_config(self):
        server_args = self.server_args
        model_config_class = getattr(self, "model_config_class", ModelConfig)

        # Read model args
        self.model_path = server_args.model_path
        self.served_model_name = server_args.served_model_name
        self.model_config = model_config_class.from_server_args(server_args)
        self.is_generation = self.model_config.is_generation
        self.context_len = self.model_config.context_len
        self.image_token_id = self.model_config.image_token_id
        self.max_req_input_len = None  # Will be set later in engine.py
        self.enable_priority_scheduling = server_args.enable_priority_scheduling
        self.default_priority_value = server_args.default_priority_value
        self.num_reserved_tokens = compute_num_reserved_tokens(server_args)
        self.validate_total_tokens = True

    def init_tokenizer_and_processor(self):
        server_args = self.server_args

        # Initialize tokenizer and processor
        if self.model_config.is_multimodal:
            import_processors("sglang.srt.multimodal.processors")
            if mm_process_pkg := envs.SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE.get():
                import_processors(mm_process_pkg, overwrite=True)
            _processor = get_processor_wrapper(server_args)
            transport_mode = determine_tensor_transport_mode(self.server_args)

            # We want to parallelize the image pre-processing so we create an executor for it
            # We create mm_processor for any skip_tokenizer_init to make sure we still encode
            # images even with skip_tokenizer_init=False.
            self.mm_processor = get_mm_processor(
                self.model_config.hf_config,
                server_args,
                _processor,
                transport_mode,
                model_config=self.model_config,
            )

            if server_args.skip_tokenizer_init:
                self.tokenizer = self.processor = None
            else:
                self.processor = _processor
                self.tokenizer = get_tokenizer_from_processor(self.processor)
                os.environ["TOKENIZERS_PARALLELISM"] = "false"
        else:
            self.mm_processor = self.processor = None

            if server_args.skip_tokenizer_init:
                self.tokenizer = None
            else:
                self.tokenizer = get_tokenizer(
                    server_args.tokenizer_path,
                    tokenizer_mode=server_args.tokenizer_mode,
                    trust_remote_code=server_args.trust_remote_code,
                    revision=server_args.revision,
                    tokenizer_backend=server_args.tokenizer_backend,
                )

        # Initialize async dynamic batch tokenizer if enabled (common for both multimodal and non-multimodal)
        if (
            server_args.enable_dynamic_batch_tokenizer
            and not server_args.skip_tokenizer_init
        ):
            self.async_dynamic_batch_tokenizer = AsyncDynamicbatchTokenizer(
                self.tokenizer,
                max_batch_size=server_args.dynamic_batch_tokenizer_batch_size,
                batch_wait_timeout_s=server_args.dynamic_batch_tokenizer_batch_timeout,
            )
        else:
            self.async_dynamic_batch_tokenizer = None

    def _validate_cuda_vmm_feature_transport_support(self) -> None:
        if self.server_args.mm_feature_transport != "cuda_vmm":
            return

        from sglang.srt.model_loader.utils import get_model_architecture

        model_class, _ = get_model_architecture(self.model_config)
        if not getattr(model_class, "supports_cuda_vmm_feature_transport", False):
            raise ValueError(
                "--mm-feature-transport=cuda_vmm is not supported by model class "
                f"{model_class.__name__}"
            )

    def init_ipc_channels(self, port_args: PortArgs):
        context = zmq.asyncio.Context(2)
        self.recv_from_detokenizer = get_zmq_socket(
            context, zmq.PULL, port_args.tokenizer_ipc_name, True
        )
        if self.server_args.tokenizer_worker_num == 1:
            self.send_to_scheduler = get_zmq_socket(
                context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
            )
            self.tokenizer_ipc_name = None
        else:
            # Use tokenizer_worker_ipc_name in multi-tokenizer mode
            self.send_to_scheduler = get_zmq_socket(
                context, zmq.PUSH, port_args.tokenizer_worker_ipc_name, False
            )
            self.tokenizer_ipc_name = port_args.tokenizer_ipc_name

        self.load_snapshot_reader = create_load_snapshot_reader(
            self.server_args,
            port_args,
            caller="TokenizerManager",
        )

    def _dispatch_to_scheduler(self, obj: Any) -> None:
        if self.tokenizer_ipc_name is not None:
            stamp_http_worker_ipc(obj, self.tokenizer_ipc_name)
        sock_send(self.send_to_scheduler, obj)

    async def _async_dispatch_to_scheduler(self, obj: Any) -> None:
        if self.tokenizer_ipc_name is not None:
            stamp_http_worker_ipc(obj, self.tokenizer_ipc_name)
        await async_sock_send(self.send_to_scheduler, obj)

    def init_running_status(self):
        # Request states
        self.rid_to_state: Dict[str, ReqState] = {}
        # Parallel sampling keeps one caller-visible logical RID per original
        # prompt while the scheduler operates on separate prefix/sample RIDs.
        self.logical_rid_to_child_rids: Dict[str, set[str]] = {}
        self.child_rid_to_logical_rid: Dict[str, str] = {}
        self.event_loop = None
        self.asyncio_tasks = set()

        # Health check
        self.server_status = ServerStatus.Starting
        self.gracefully_exit = False
        self.last_receive_tstamp = real_time()

        # Session
        self.session_futures = {}  # session_id -> asyncio event

        # Subprocess liveness watchdog — set by Engine or http_server after construction
        self._subprocess_watchdog = None

    def init_request_logging_and_dumping(self):
        # TODO: Refactor and organize the log export code.
        # Request logging
        self.request_logger = RequestLogger(
            log_requests=self.server_args.log_requests,
            log_requests_level=self.server_args.log_requests_level,
            log_requests_format=self.server_args.log_requests_format,
            log_requests_target=self.server_args.log_requests_target,
        )

        # Dumping
        self.dump_requests_folder = ""  # By default do not dump
        self.dump_requests_threshold = 1000
        self.dump_requests_exclude_meta_keys: List[str] = [
            "routed_experts",
            "hidden_states",
        ]
        self.dump_request_list: List[Tuple] = []
        self.crash_dump_request_list: deque[Tuple] = deque()
        self.crash_dump_performed = False  # Flag to ensure dump is only called once

        # Initialize performance metrics loggers with proper skip names
        _, obj_skip_names, out_skip_names = self.request_logger.metadata
        self.request_metrics_exporter_manager = RequestMetricsExporterManager(
            self.server_args, obj_skip_names, out_skip_names
        )

    def init_weight_update(self):
        # Initial weights status
        self.initial_weights_loaded = True
        if self.server_args.checkpoint_engine_wait_weights_before_ready:
            self.initial_weights_loaded = False

        # Weight updates
        # The event to notify the weight sync is finished.
        self.model_update_lock = RWLock()
        self.model_update_result: Optional[Awaitable[UpdateWeightFromDiskReqOutput]] = (
            None
        )
        self.model_update_expected_workers = self.elastic_worker_count
        self.model_update_tmp: List[UpdateWeightFromDiskReqOutput] = []
        self.is_pause = False
        self.is_pause_cond = asyncio.Condition()

    def init_lora(self):
        # LoRA
        # Initialize the `LoRARegistry` with initial LoRA adapter paths provided in `server_args`.
        # The registry dynamically updates as adapters are loaded / unloaded during runtime. It
        # serves as the source of truth for available adapters and maps user-friendly LoRA names
        # to internally used unique LoRA IDs.
        self.lora_registry = LoRARegistry(self.server_args.lora_paths)
        # Lock to serialize LoRA update operations.
        # Please note that, unlike `model_update_lock`, this does not block inference, allowing
        # LoRA updates and inference to overlap.
        self.lora_update_lock = asyncio.Lock()
        # A cache for mapping the lora_name for LoRA adapters that have been loaded at any
        # point to their latest LoRARef objects, so that they can be
        # dynamically loaded if needed for inference
        self.lora_ref_cache: Dict[str, LoRARef] = {}
        if self.server_args.lora_paths is not None:
            for lora_ref in self.server_args.lora_paths:
                self.lora_ref_cache[lora_ref.lora_name] = lora_ref

    def init_disaggregation(self, *, start_pd_bootstrap_service: bool = True):
        # PD Disaggregation
        self.disaggregation_mode = DisaggregationMode(
            self.server_args.disaggregation_mode
        )
        # Keep a reference so the bootstrap server is not garbage-collected.
        self.bootstrap_server = (
            start_disagg_service(self.server_args)
            if start_pd_bootstrap_service
            else None
        )
        # Single-source counter for auto-assigning fake bootstrap_room.
        self.fake_bootstrap_room_counter = 0

        # Encoder Disaggregation
        self.encoder_bootstrap_server = None
        if self.server_args.language_only:
            from sglang.srt.disaggregation.encode_receiver import (
                EncoderBootstrapServer,
            )

            # Shared mutable URL list: the bootstrap server appends / removes
            # entries as encoders register, the receiver reads from the same
            # list.  Pre-populated with static --encoder-urls so the legacy
            # CLI flag still works (alongside dynamic registrations).
            self.encoder_urls: List[str] = list(self.server_args.encoder_urls)
            self.encoder_bootstrap_server = EncoderBootstrapServer(
                host=self.server_args.host,
                port=self.server_args.encoder_bootstrap_port,
                urls=self.encoder_urls,
            )
            self.mm_receiver = create_mm_receiver(
                self.server_args,
                dtype=self.model_config.dtype,
                hf_config=self.model_config.hf_config,
                encode_urls=self.encoder_urls,
            )

    def init_metric_collector_watchdog(self):
        # Metrics
        if self.enable_metrics:
            engine_type = DisaggregationMode.to_engine_type(
                self.server_args.disaggregation_mode
            )

            labels = {
                "model_name": self.server_args.served_model_name,
                "engine_type": engine_type,
            }
            if self.enable_priority_scheduling:
                labels["priority"] = ""
            if self.server_args.tokenizer_metrics_allowed_custom_labels:
                for label in self.server_args.tokenizer_metrics_allowed_custom_labels:
                    labels[label] = ""
            if self.server_args.extra_metric_labels:
                labels.update(self.server_args.extra_metric_labels)
            tokenizer_collector_cls = resolve_collector_class(
                self.server_args,
                STAT_LOGGER_ROLE_TOKENIZER,
                TokenizerMetricsCollector,
            )
            self.metrics_collector = tokenizer_collector_cls(
                server_args=self.server_args,
                labels=labels,
                bucket_time_to_first_token=self.server_args.bucket_time_to_first_token,
                bucket_e2e_request_latency=self.server_args.bucket_e2e_request_latency,
                bucket_inter_token_latency=self.server_args.bucket_inter_token_latency,
            )

            start_cpu_monitor_thread("tokenizer")

        if self.server_args.gc_warning_threshold_secs > 0.0:
            configure_gc_warning(self.server_args.gc_warning_threshold_secs)
        self.soft_watchdog = Watchdog.create(
            debug_name="TokenizerManager",
            watchdog_timeout=self.server_args.soft_watchdog_timeout,
            soft=True,
            test_stuck_time=envs.SGLANG_TEST_STUCK_TOKENIZER.get(),
        )

    def set_startup_time(self, startup_time: Dict[str, Any]) -> None:
        self.startup_time = startup_time
        if self.enable_metrics:
            self.metrics_collector.emit_startup_time(startup_time)

    def init_request_dispatcher(self):
        self._result_dispatcher = TypeBasedDispatcher(
            [
                (AbortReq, self._handle_abort_req),
                (OpenSessionReqOutput, self._handle_open_session_req_output),
                (
                    UpdateWeightFromDiskReqOutput,
                    self._handle_update_weights_from_disk_req_output,
                ),
                (FreezeGCReq, lambda x: None),
                # For handling case when scheduler skips detokenizer and forwards back to the tokenizer manager, we ignore it.
                (HealthCheckOutput, lambda x: None),
                # Same skip-detokenizer forwarding case as above.
                (ConfigureLoggingReq, lambda x: None),
                (ActiveRanksOutput, self.update_active_ranks),
                (ElasticScaleUpdateReq, self.forward_elastic_scale_update),
            ]
        )
        self.init_communicators(self.server_args)

        self.sampling_params_class = SamplingParams
        self.signal_handler_class = SignalHandler

    async def generate_request(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        request: Optional[fastapi.Request] = None,
    ):
        self.auto_create_handle_loop()

        # Normalize the request
        obj.normalize_batch_and_arguments()
        self._set_default_priority(obj)
        if (
            isinstance(obj, GenerateReqInput)
            and obj.max_thinking_tokens is not None
            and not self.server_args.enable_strict_thinking
        ):
            raise ValueError(
                "max_thinking_tokens requires the server to be launched with "
                "--enable-strict-thinking"
            )

        if isinstance(obj, GenerateReqInput) and obj.routed_dp_rank is not None:
            dp_size = self.elastic_worker_count
            if dp_size <= 1 and obj.routed_dp_rank == 0:
                logger.debug(
                    f"routed_dp_rank={obj.routed_dp_rank} is ignored because dp_size={dp_size}"
                )
            elif obj.routed_dp_rank < 0 or obj.routed_dp_rank >= dp_size:
                raise ValueError(
                    f"routed_dp_rank={obj.routed_dp_rank} out of range [0, {dp_size})"
                )

        request_lifecycles = self._init_req_state(obj, request)
        try:
            if self.server_args.language_only:
                self._handle_epd_disaggregation_encode_request(obj)

            # Log the request
            self.request_logger.log_received_request(obj, self.tokenizer, request)

            async with self.is_pause_cond:
                await self.is_pause_cond.wait_for(lambda: not self.is_pause)
            self._raise_if_logical_request_aborted(obj)

            async with self.model_update_lock.reader_lock:
                await self._validate_and_resolve_lora(obj)
                self._raise_if_logical_request_aborted(obj)

                # Tokenize the request and send it to the scheduler
                if obj.is_single:
                    tokenized_obj = await self._tokenize_one_request(obj)
                    self._raise_if_logical_rid_aborted(obj.rid)
                    state = self.rid_to_state[obj.rid]
                    if obj.return_prompt_token_ids:
                        state.prompt_token_ids = list(tokenized_obj.input_ids)
                    self._send_one_request(tokenized_obj)
                    async for response in self._wait_one_response(obj, request):
                        yield response
                else:
                    async for response in self._handle_batch_request(obj, request):
                        yield response
        except BaseException:
            # _init_req_state created a rid_to_state entry per (sub-)request up
            # front. The normal remover is the scheduler-response path
            # (_handle_batch_output), so a failure *before* a request reaches the
            # scheduler -- e.g. input-length validation rejecting an over-context
            # request -- would otherwise leak those entries forever. Drop any that
            # are still pending; entries already removed on the normal completion
            # path are left untouched (pop is a no-op).
            self._discard_pending_req_states(obj, request_lifecycles)
            raise

    def _detect_input_format(
        self, texts: Union[str, List[str]], is_cross_encoder: bool
    ) -> InputFormat:
        """Detect the format of input texts for proper tokenization handling.

        Returns:
            - InputFormat.SINGLE_STRING: Regular single text like "Hello world"
            - InputFormat.BATCH_STRINGS: Regular batch like ["Hello", "World"]
            - InputFormat.CROSS_ENCODER_PAIRS: Cross-encoder pairs like [["query", "document"]]
        """
        if isinstance(texts, str):
            return InputFormat.SINGLE_STRING

        if (
            is_cross_encoder
            and len(texts) > 0
            and isinstance(texts[0], list)
            and len(texts[0]) == 2
        ):
            return InputFormat.CROSS_ENCODER_PAIRS

        return InputFormat.BATCH_STRINGS

    def _prepare_tokenizer_input(
        self, texts: Union[str, List[str]], input_format: InputFormat
    ) -> Union[List[str], List[List[str]]]:
        """Prepare input for the tokenizer based on detected format."""
        if input_format == InputFormat.SINGLE_STRING:
            return [texts]  # Wrap single string for batch processing
        elif input_format == InputFormat.CROSS_ENCODER_PAIRS:
            return texts  # Already in correct format: [["query", "doc"]]
        else:  # BATCH_STRINGS
            return texts  # Already in correct format: ["text1", "text2"]

    def _extract_tokenizer_results(
        self,
        input_ids: List[List[int]],
        token_type_ids: Optional[List[List[int]]],
        input_format: InputFormat,
        original_batch_size: int,
    ) -> Union[
        Tuple[List[int], Optional[List[int]]],
        Tuple[List[List[int]], Optional[List[List[int]]]],
    ]:
        """Extract results from tokenizer output based on input format."""

        # For single inputs (string or single cross-encoder pair), extract first element
        if (
            input_format in [InputFormat.SINGLE_STRING, InputFormat.CROSS_ENCODER_PAIRS]
            and original_batch_size == 1
        ):
            single_input_ids = input_ids[0] if input_ids else []
            single_token_type_ids = token_type_ids[0] if token_type_ids else None
            return single_input_ids, single_token_type_ids

        # For true batches, return as-is
        return input_ids, token_type_ids

    async def _tokenize_texts(
        self, texts: Union[str, List[str]], is_cross_encoder: bool = False
    ) -> Union[
        Tuple[List[int], Optional[List[int]]],
        Tuple[List[List[int]], Optional[List[List[int]]]],
    ]:
        """
        Tokenize text(s) using the appropriate tokenizer strategy.

        This method handles multiple input formats and chooses between async dynamic
        batch tokenizer (for single texts only) and regular tokenizer.

        Args:
            texts: Text input in various formats:

                   Regular cases:
                   - Single string: "How are you?"
                   - Batch of strings: ["Hello", "World", "How are you?"]

                   Cross-encoder cases (sentence pairs for similarity/ranking):
                   - Single pair: [["query text", "document text"]]
                   - Multiple pairs: [["q1", "d1"], ["q2", "d2"], ["q3", "d3"]]

            is_cross_encoder: Whether to return token_type_ids for cross-encoder models.
                             Enables proper handling of sentence pairs with segment IDs.

        Returns:
            Single input cases:
                Tuple[List[int], Optional[List[int]]]: (input_ids, token_type_ids)
                Example: ([101, 2129, 102], [0, 0, 0]) for single text
                Example: ([101, 2129, 102, 4068, 102], [0, 0, 0, 1, 1]) for cross-encoder pair

            Batch input cases:
                Tuple[List[List[int]], Optional[List[List[int]]]]: (batch_input_ids, batch_token_type_ids)
                Example: ([[101, 2129, 102], [101, 4068, 102]], None) for regular batch

            Note: token_type_ids is None unless is_cross_encoder=True.
        """
        if not texts or self.tokenizer is None:
            raise ValueError("texts cannot be empty and tokenizer must be initialized")

        # Step 1: Detect input format and prepare for tokenization
        input_format = self._detect_input_format(texts, is_cross_encoder)
        tokenizer_input = self._prepare_tokenizer_input(texts, input_format)
        original_batch_size = len(texts) if not isinstance(texts, str) else 1

        # Step 2: Set up tokenizer arguments
        tokenizer_kwargs = (
            {"return_token_type_ids": is_cross_encoder} if is_cross_encoder else {}
        )

        # Step 3: Choose tokenization strategy
        use_async_tokenizer = (
            self.async_dynamic_batch_tokenizer is not None
            and input_format == InputFormat.SINGLE_STRING
        )

        if use_async_tokenizer:
            logger.debug("Using async dynamic batch tokenizer for single text")
            result = await self.async_dynamic_batch_tokenizer.encode(
                tokenizer_input[0], **tokenizer_kwargs
            )
            # Convert to batch format for consistency
            input_ids = [result["input_ids"]]
            token_type_ids = (
                [result["token_type_ids"]]
                if is_cross_encoder and result.get("token_type_ids")
                else None
            )
        else:
            logger.debug(f"Using regular tokenizer for {len(tokenizer_input)} inputs")

            if not is_cross_encoder and (not getattr(self.tokenizer, "is_fast", False)):
                input_ids = [self.tokenizer.encode(t) for t in tokenizer_input]
                token_type_ids = None
            else:
                encoded = self.tokenizer(tokenizer_input, **tokenizer_kwargs)
                input_ids = encoded["input_ids"]
                token_type_ids = (
                    encoded.get("token_type_ids") if is_cross_encoder else None
                )

        # vLLM's OpenAI embeddings endpoint includes special tokens for
        # encoder models. EmbeddingGemma's restored Gemma tokenizer adds BOS
        # but, by its checkpoint default, omits EOS. Add EOS explicitly here
        # rather than mutating tokenizer-global post-processing state.
        if (
            self.model_config.is_embedding_gemma
            and self.tokenizer.eos_token_id is not None
        ):
            input_ids = [
                (
                    ids
                    if ids and ids[-1] == self.tokenizer.eos_token_id
                    else [*ids, self.tokenizer.eos_token_id]
                )
                for ids in input_ids
            ]

        # Step 4: Extract results based on input format
        return self._extract_tokenizer_results(
            input_ids, token_type_ids, input_format, original_batch_size
        )

    async def _tokenize_one_request(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
    ):
        """Tokenize one request."""
        # Tokenize
        input_embeds = None
        input_text = obj.text
        token_type_ids = None
        is_cross_encoder_request = (
            isinstance(obj, EmbeddingReqInput) and obj.is_cross_encoder_request
        )
        if obj.input_embeds is not None:
            if not self.server_args.disable_radix_cache:
                raise ValueError(
                    "input_embeds is provided while disable_radix_cache is False. "
                    "Please add `--disable-radix-cache` when you launch the server "
                    "if you want to use input_embeds as inputs."
                )
            input_embeds = obj.input_embeds
            input_ids = obj.input_ids
        elif obj.input_ids is not None:
            input_ids = obj.input_ids
        else:
            if self.tokenizer is None:
                raise ValueError(
                    "The engine initialized with skip_tokenizer_init=True cannot "
                    "accept text prompts. Please provide input_ids or re-initialize "
                    "the engine with skip_tokenizer_init=False."
                )

            # For audio-only requests (e.g., Whisper), text may be empty.
            # The multimodal processor will provide input_ids later.
            if not input_text and self.mm_processor and obj.contains_mm_input():
                # Use empty placeholder - multimodal processor will override
                input_ids = []
            else:
                input_ids, token_type_ids = await self._tokenize_texts(
                    input_text, is_cross_encoder_request
                )

        contains_mm_input = obj.contains_mm_input()
        is_mossvl = (
            "MossVLForConditionalGeneration"
            in self.model_config.hf_config.architectures
        )
        should_run_mm_processor = self.mm_processor is not None and (
            contains_mm_input or is_mossvl
        )

        if should_run_mm_processor:
            if obj.image_data is not None and not isinstance(obj.image_data, list):
                obj.image_data = [obj.image_data]
            if obj.video_data is not None and not isinstance(obj.video_data, list):
                obj.video_data = [obj.video_data]
            if obj.audio_data is not None and not isinstance(obj.audio_data, list):
                obj.audio_data = [obj.audio_data]
            if contains_mm_input:
                self._validate_mm_limits(obj)

            mm_inputs = None
            mm_processor_input = (
                input_ids
                if self.mm_processor.prefer_tokenized_input and input_ids is not None
                else (input_text or input_ids)
            )

            if (
                not self.server_args.language_only
                or self.server_args.encoder_transfer_backend == "zmq_to_tokenizer"
            ):
                if self.server_args.language_only:
                    mm_inputs = await self.mm_receiver.recv_mm_data(
                        request_obj=obj,
                        mm_processor=self.mm_processor,
                        prompt=mm_processor_input,
                        need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
                    )
                    _reject_missing_dispatched_encoder_embedding(
                        self.server_args, obj, mm_inputs
                    )
                if mm_inputs is None:
                    if self.server_args.language_only:
                        logger.warning(
                            "Encoder embedding not available, "
                            "falling back to local mm processing"
                        )
                    mm_inputs = await self.mm_processor.process_mm_data_async(
                        image_data=obj.image_data,
                        audio_data=obj.audio_data,
                        input_text=mm_processor_input,
                        request_obj=obj,
                        max_req_input_len=self.max_req_input_len,
                    )
            elif (
                self.server_args.language_only
                and self.server_args.encoder_transfer_backend
                in ["zmq_to_scheduler", "mooncake"]
                and not obj.need_wait_for_mm_inputs
            ):
                # In language_only mode with zmq_to_scheduler/mooncake, if we didn't dispatch
                # to encoder (e.g., only one image), process locally like non-language_only mode
                mm_inputs = await self.mm_processor.process_mm_data_async(
                    image_data=obj.image_data,
                    audio_data=obj.audio_data,
                    input_text=mm_processor_input,
                    request_obj=obj,
                    max_req_input_len=self.max_req_input_len,
                )

            if mm_inputs and mm_inputs.input_ids is not None:
                input_ids = mm_inputs.input_ids
            if mm_inputs and mm_inputs.token_type_ids is not None:
                token_type_ids = mm_inputs.token_type_ids
                if not isinstance(token_type_ids, list):
                    token_type_ids = token_type_ids.flatten().tolist()
            # Caller-supplied per-image hashes (external KV routers, e.g.
            # routing-aware orchestrators that compute a content-addressed
            # hash before dispatch). Setting MultimodalDataItem.hash here
            # short-circuits the internal hash_feature() recompute inside
            # set_pad_value(), making the derived pad_value deterministic
            # from the caller's hash. That alignment lets the router's
            # routing decision agree with sglang's prefix-cache key for
            # the same image. On any per-item parse error or list-length
            # mismatch we fall back to the internal recompute so a
            # malformed mm_hashes never blocks a request.
            caller_mm_hashes = getattr(obj, "mm_hashes", None)
            if caller_mm_hashes and mm_inputs and mm_inputs.mm_items:
                if len(caller_mm_hashes) != len(mm_inputs.mm_items):
                    logger.warning(
                        "mm_hashes length (%d) != mm_items length (%d); "
                        "ignoring caller hashes for this request.",
                        len(caller_mm_hashes),
                        len(mm_inputs.mm_items),
                    )
                else:
                    for item, hex_hash in zip(mm_inputs.mm_items, caller_mm_hashes):
                        if not isinstance(item, MultimodalDataItem):
                            continue
                        try:
                            item.set_hash(int(hex_hash, 16))
                        except (TypeError, ValueError):
                            logger.warning(
                                "Ignoring malformed mm_hashes entry %r; "
                                "this item will fall back to hash_feature().",
                                hex_hash,
                            )
            if (
                envs.SGLANG_MM_PRECOMPUTE_HASH.get()
                and mm_inputs
                and mm_inputs.mm_items
            ):
                for item in mm_inputs.mm_items:
                    if isinstance(item, MultimodalDataItem):
                        item.set_pad_value()
        else:
            mm_inputs = None

        self._validate_one_request(obj, input_ids)
        return self._create_tokenized_object(
            obj, input_text, input_ids, input_embeds, mm_inputs, token_type_ids
        )

    def _validate_one_request(
        self, obj: Union[GenerateReqInput, EmbeddingReqInput], input_ids: List[int]
    ) -> None:
        """Validates that the input token count and the requested token count doesn't exceed the model's context length."""
        # FIXME: unify the length validation logic with the one in the scheduler.
        _max_req_len = self.context_len
        input_token_num = len(input_ids) if input_ids is not None else 0
        input_token_num += self.num_reserved_tokens

        # Validate input length
        if input_token_num >= self.context_len:
            if self.allow_auto_truncate:
                logger.warning(
                    f"The input ({input_token_num} tokens) is longer than the "
                    f"model's context length ({self.context_len} tokens). "
                    "Truncating the input."
                )
                del input_ids[_max_req_len:]
                input_token_num = len(input_ids)
            else:
                raise ValueError(
                    f"The input ({input_token_num} tokens) is longer than the "
                    f"model's context length ({self.context_len} tokens)."
                )

        # Validate total tokens (input + max_new_tokens)
        max_new_tokens = obj.sampling_params.get("max_new_tokens")
        if (
            self.validate_total_tokens
            and max_new_tokens is not None
            and (max_new_tokens + input_token_num) > _max_req_len
        ):
            if self.allow_auto_truncate:
                logger.warning(
                    f"Requested token count ({input_token_num} input + {max_new_tokens} new) "
                    f"exceeds the model's context length ({self.context_len} tokens). "
                    "Truncating max_new_tokens."
                )
                obj.sampling_params["max_new_tokens"] = max(
                    0, _max_req_len - input_token_num
                )
            else:
                total_tokens = max_new_tokens + input_token_num
                error_msg = (
                    f"Requested token count exceeds the model's maximum context length "
                    f"of {self.context_len} tokens. You requested a total of {total_tokens} "
                    f"tokens: {input_token_num} tokens from the input messages and "
                    f"{max_new_tokens} tokens for the completion. Please reduce the number "
                    f"of tokens in the input messages or the completion to fit within the limit."
                )
                raise ValueError(error_msg)

        # Validate embedding requests
        if isinstance(obj, EmbeddingReqInput) and self.is_generation:
            raise ValueError(
                "This model does not appear to be an embedding model by default. "
                "Please add `--is-embedding` when launching the server or try another model."
            )

        # Validate Matryoshka embeddings
        if isinstance(obj, EmbeddingReqInput):
            self._validate_for_matryoshka_dim(obj)

        # Validate generation-specific fields
        if isinstance(obj, GenerateReqInput):
            self._validate_token_ids_logprob(obj)
            requested_hidden_mode = get_request_return_hidden_states_mode(
                obj.return_hidden_states
            )
            server_hidden_mode = get_server_return_hidden_states_mode(self.server_args)
            if requested_hidden_mode > server_hidden_mode:
                if server_hidden_mode.need_capture():
                    raise ValueError(
                        "The requested return_hidden_states mode exceeds the "
                        f"server maximum `{self.server_args.return_hidden_states_mode}`. "
                        "Please launch with `--return-hidden-states-mode full` "
                        "to allow return_hidden_states=True."
                    )
                raise ValueError(
                    "The server is not configured to return hidden states. "
                    "Please set `--return-hidden-states-mode last`, "
                    "`--return-hidden-states-mode full`, or the legacy "
                    "`--enable-return-hidden-states` flag."
                )
            if (
                obj.custom_logit_processor
                and not self.server_args.enable_custom_logit_processor
            ):
                raise ValueError(
                    "The server is not configured to enable custom logit processor. "
                    "Please set `--enable-custom-logit-processor` to enable this feature."
                )

    def _validate_mm_limits(
        self, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> None:
        if not self.server_args.limit_mm_data_per_request:
            return

        for modality, limit in self.server_args.limit_mm_data_per_request.items():
            data = getattr(obj, f"{modality}_data", None)
            if data:
                count = len(data) if isinstance(data, list) else 1
                if count > limit:
                    raise ValueError(
                        f"{modality.capitalize()} count {count} exceeds limit {limit} per request."
                    )

    def _validate_for_matryoshka_dim(self, obj: EmbeddingReqInput) -> None:
        """Validate the request for Matryoshka dim if it has the field set."""
        if obj.dimensions is None:
            return

        if not self.model_config.is_matryoshka:
            raise ValueError(
                f"Model '{self.model_config.model_path}' does not support matryoshka representation, "
                f"changing output dimensions will lead to poor results."
            )

        if obj.dimensions < 1:
            raise ValueError("Requested dimensions must be greater than 0")

        if (
            self.model_config.matryoshka_dimensions
            and obj.dimensions not in self.model_config.matryoshka_dimensions
        ):
            raise ValueError(
                f"Model '{self.model_config.model_path}' only supports {self.model_config.matryoshka_dimensions} matryoshka dimensions, "
                f"using other output dimensions will lead to poor results."
            )

        if obj.dimensions > self.model_config.hidden_size:
            raise ValueError(
                f"Provided dimensions are greater than max embedding dimension: {self.model_config.hidden_size}"
            )

    def _validate_token_ids_logprob(self, obj: GenerateReqInput) -> None:
        # Batch requests are split into per-request sub-objects before this
        # runs (normalize_batch_and_arguments + __getitem__), so the only
        # legal shape here is the per-request contract of
        # TokenizedGenerateReqInput.token_ids_logprob: a flat list of ints.
        token_ids_logprob = obj.token_ids_logprob
        if not token_ids_logprob:
            return
        if not isinstance(token_ids_logprob, list):
            raise ValueError("token_ids_logprob must be a flat list of integers.")
        vocab_size = self.model_config.vocab_size
        for token_id in token_ids_logprob:
            if not isinstance(token_id, int):
                raise ValueError("token_ids_logprob must be a flat list of integers.")
            if token_id < 0 or token_id >= vocab_size:
                raise ValueError(
                    f"token_ids_logprob contains out-of-vocabulary token id "
                    f"{token_id}; valid range is [0, {vocab_size})."
                )

    def _validate_input_ids_in_vocab(
        self, input_ids: Union[List[int], List[List[int]]], vocab_size: int
    ) -> None:
        # Handle both single sequence and batch of sequences
        if isinstance(input_ids[0], list):
            # Batch of sequences
            for seq in input_ids:
                if any(id >= vocab_size for id in seq):
                    raise ValueError(
                        f"The input_ids {seq} contains values greater than the vocab size ({vocab_size})."
                    )
        else:
            # Single sequence
            if any(id >= vocab_size for id in input_ids):
                raise ValueError(
                    f"The input_ids {input_ids} contains values greater than the vocab size ({vocab_size})."
                )

    def _create_tokenized_object(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        input_text: str,
        input_ids: Optional[List[int]],
        input_embeds: Optional[List[List[float]]] = None,
        mm_inputs=None,
        token_type_ids: Optional[List[int]] = None,
    ) -> Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]:
        """Create a tokenized request object from common parameters."""
        input_ids_arr: Optional[array[int]] = (
            array("q", input_ids) if input_ids is not None else None
        )
        # Parse sampling parameters
        # Note: if there are preferred sampling params, we use them if they are not
        # explicitly passed in sampling_params
        if self.preferred_sampling_params:
            sampling_kwargs = {**self.preferred_sampling_params, **obj.sampling_params}
        else:
            sampling_kwargs = obj.sampling_params
        if isinstance(obj, GenerateReqInput) and obj.max_thinking_tokens is not None:
            sampling_kwargs = dict(sampling_kwargs)
            custom_params = dict(sampling_kwargs.get("custom_params") or {})
            custom_params["thinking_budget"] = obj.max_thinking_tokens
            sampling_kwargs["custom_params"] = custom_params
        sampling_params = self.sampling_params_class(**sampling_kwargs)
        sampling_params.normalize(self.tokenizer)
        sampling_params.verify(self.model_config.vocab_size)

        # Build return object
        if isinstance(obj, GenerateReqInput):
            session_params = (
                SessionParams(**obj.session_params) if obj.session_params else None
            )

            bootstrap_room = obj.bootstrap_room
            if (
                bootstrap_room is None
                and self.server_args.disaggregation_transfer_backend == "fake"
            ):
                bootstrap_room = self.fake_bootstrap_room_counter
                self.fake_bootstrap_room_counter += 1

            tokenized_obj = TokenizedGenerateReqInput(
                input_text=input_text,
                input_ids=input_ids_arr,
                mm_inputs=mm_inputs,
                sampling_params=sampling_params,
                return_logprob=obj.return_logprob,
                logprob_start_len=obj.logprob_start_len,
                top_logprobs_num=obj.top_logprobs_num,
                token_ids_logprob=obj.token_ids_logprob,
                return_sampling_mask=obj.return_sampling_mask,
                return_flat_raw_top_logprobs=obj.return_flat_raw_top_logprobs,
                stream=obj.stream,
                rid=obj.rid,
                http_worker_ipc=obj.http_worker_ipc,
                bootstrap_host=obj.bootstrap_host,
                bootstrap_port=obj.bootstrap_port,
                bootstrap_room=bootstrap_room,
                lora_id=obj.lora_id,
                input_embeds=input_embeds,
                positional_embed_overrides=obj.positional_embed_overrides,
                session_id=obj.session_id,
                session_params=session_params,
                custom_logit_processor=obj.custom_logit_processor,
                require_reasoning=obj.require_reasoning,
                return_hidden_states=obj.return_hidden_states,
                return_routed_experts=obj.return_routed_experts,
                routed_experts_start_len=obj.routed_experts_start_len,
                return_indexer_topk=obj.return_indexer_topk,
                routed_dp_rank=obj.routed_dp_rank,
                disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank,
                priority=obj.priority,
                extra_key=obj.extra_key,
                routing_key=obj.routing_key,
                token_type_ids=token_type_ids,
                need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
                num_items_assigned=obj.num_items_assigned,
                multi_item_delimiter_indices=obj.multi_item_delimiter_indices,
                mm_data_mooncake=obj.mm_data_mooncake,
                encoder_urls=obj.encoder_urls,
            )
        elif isinstance(obj, EmbeddingReqInput):
            # Resolve unresolved embed overrides now that input_ids are available
            positional_embed_overrides = obj.positional_embed_overrides
            if (
                positional_embed_overrides is None
                and obj.embed_overrides is not None
                and obj.embed_override_token_id is not None
            ):
                positional_embed_overrides = self._resolve_embed_overrides(
                    input_ids_arr, obj.embed_override_token_id, obj.embed_overrides
                )

            tokenized_obj = TokenizedEmbeddingReqInput(
                input_text=input_text,
                input_ids=input_ids_arr,
                mm_inputs=mm_inputs,
                token_type_ids=token_type_ids,
                sampling_params=sampling_params,
                positional_embed_overrides=positional_embed_overrides,
                rid=obj.rid,
                priority=obj.priority,
                dimensions=obj.dimensions,
                lora_id=obj.lora_id,
                http_worker_ipc=obj.http_worker_ipc,
                return_pooled_hidden_states=obj.return_pooled_hidden_states,
                multi_item_delimiter_indices=obj.multi_item_delimiter_indices,
            )

        tokenized_obj.time_stats = self.rid_to_state[obj.rid].time_stats
        self.rid_to_state[obj.rid].time_stats.set_tokenize_finish_time()

        return tokenized_obj

    @staticmethod
    def _resolve_embed_overrides(
        input_ids: array[int],
        token_id: int,
        embeds: List[torch.Tensor],
    ) -> PositionalEmbeds:
        """Resolve placeholder positions in input_ids and create PositionalEmbeds.

        Scans input_ids for occurrences of token_id and pairs them with the
        provided embedding tensors.
        """
        positions = [idx for idx, tok in enumerate(input_ids) if tok == token_id]
        if len(positions) != len(embeds):
            raise ValueError(
                f"input contains {len(positions)} occurrences of "
                f"embed_override_token_id={token_id}, "
                f"but embed_overrides has {len(embeds)} entries."
            )
        return PositionalEmbeds(embeds=embeds, positions=positions)

    async def _batch_tokenize_and_process(
        self, batch_size: int, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]]:
        """Handle batch tokenization for text inputs only."""
        logger.debug(f"Starting batch tokenization for {batch_size} text requests")

        # If batch does not have text nothing to tokenize
        # so lets construct the return object
        if not self._batch_has_text(batch_size, obj):
            # All requests already have input_ids, no need to tokenize
            return [await self._tokenize_one_request(obj[i]) for i in range(batch_size)]

        self._validate_batch_tokenization_constraints(batch_size, obj)

        # Collect requests and texts
        requests = [obj[i] for i in range(batch_size)]
        texts = [req.text for req in requests]

        # Check if any request is a cross-encoder request
        is_cross_encoder_request = any(
            isinstance(req, EmbeddingReqInput) and req.is_cross_encoder_request
            for req in requests
        )

        # Batch tokenize all texts using unified method
        input_ids_list, token_type_ids_list = await self._tokenize_texts(
            texts, is_cross_encoder_request
        )

        # Process all requests
        tokenized_objs = []
        for i, req in enumerate(requests):
            self._validate_one_request(obj[i], input_ids_list[i])
            token_type_ids = (
                token_type_ids_list[i] if token_type_ids_list is not None else None
            )
            tokenized_objs.append(
                self._create_tokenized_object(
                    req, req.text, input_ids_list[i], None, None, token_type_ids
                )
            )
        logger.debug(f"Completed batch processing for {batch_size} requests")
        return tokenized_objs

    def _validate_batch_tokenization_constraints(
        self, batch_size: int, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> None:
        """Validate constraints for batch tokenization processing."""
        for i in range(batch_size):
            if self.is_generation and obj[i].contains_mm_input():
                raise ValueError(
                    "For multimodal input processing do not set `enable_tokenizer_batch_encode`."
                )
            if obj[i].input_ids is not None:
                raise ValueError(
                    "Batch tokenization is not needed for pre-tokenized input_ids. Do not set `enable_tokenizer_batch_encode`."
                )
            if obj[i].input_embeds is not None:
                raise ValueError(
                    "Batch tokenization is not needed for input_embeds. Do not set `enable_tokenizer_batch_encode`."
                )

    def _batch_has_text(
        self, batch_size: int, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> bool:
        """Check if any request in the batch contains text input."""
        for i in range(batch_size):
            if obj[i].text:
                return True
            elif self.is_generation and obj[i].contains_mm_input():
                return True

        return False

    def _should_use_batch_tokenization(self, batch_size, requests) -> bool:
        """Return True if we should run the tokenizer in batch mode.

        Current policy:
        - Respect explicit server flag `enable_tokenizer_batch_encode`.
        - Or, if no request has text or multimodal input (all use pre-tokenized input_ids or input_embeds), batch the requests without tokenization.
        - Batch tokenization does not support DP attention yet, and it will make everything goes to the first rank currently
        """
        return batch_size > 0 and (
            self.server_args.enable_tokenizer_batch_encode
            or (
                (not self.server_args.enable_dp_attention)
                and (not self._batch_has_text(batch_size, requests))
            )
        )

    def _send_one_request(
        self,
        tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
    ):
        prepared_mm_items = []
        dispatched = False
        try:
            prepared_mm_items = self.cuda_vmm_feature_transport.prepare_for_dispatch(
                (tokenized_obj.mm_inputs,)
            )
            tokenized_obj.time_stats.set_api_server_dispatch_time()
            tokenized_obj = wrap_shm_features(tokenized_obj)
            time_stats = tokenized_obj.time_stats
            tokenized_obj.wrap_pickle_fields()
            self._dispatch_to_scheduler(tokenized_obj)
            dispatched = True
            state = self.rid_to_state.get(tokenized_obj.rid)
            if state is not None:
                state.dispatched = True
            tokenized_obj.time_stats = time_stats
            tokenized_obj.time_stats.set_api_server_dispatch_finish_time()
        finally:
            if not dispatched:
                self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items)

    def _send_batch_request(
        self,
        tokenized_objs: List[
            Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]
        ],
    ):
        """Send a batch of tokenized requests as a single batched request to the scheduler."""
        prepared_mm_items = []
        dispatched = False
        try:
            prepared_mm_items = self.cuda_vmm_feature_transport.prepare_for_dispatch(
                tokenized_obj.mm_inputs for tokenized_obj in tokenized_objs
            )

            set_time_batch(tokenized_objs, "set_api_server_dispatch_time")
            time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs]
            for tokenized_obj in tokenized_objs:
                tokenized_obj.wrap_pickle_fields()

            if isinstance(tokenized_objs[0], TokenizedGenerateReqInput):
                batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs)
            else:
                batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs)

            self._dispatch_to_scheduler(batch_req)
            dispatched = True
            for tokenized_obj in tokenized_objs:
                state = self.rid_to_state.get(tokenized_obj.rid)
                if state is not None:
                    state.dispatched = True
            for tokenized_obj, time_stat in zip(tokenized_objs, time_stats):
                tokenized_obj.time_stats = time_stat
            set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time")
        finally:
            if not dispatched:
                self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items)

    def _coalesce_streaming_chunks(
        self,
        out_list: list,
        rid: str,
        customized_info_keys: Optional[Iterable[str]] = None,
    ) -> dict:
        """Coalesce multiple incremental streaming chunks into one.

        Both text and output_ids are incremental deltas, so we concatenate them;
        all other fields (meta_info, etc.) are taken from the last chunk.
        """
        if len(out_list) >= 20:
            logger.warning(
                "Streaming backlog: rid=%s, coalescing %d queued chunks into one. "
                "This may inflate P99 ITL for affected requests.",
                rid,
                len(out_list),
            )
        out = dict(out_list[-1])
        if "output_ids" in out:
            out["output_ids"] = [id for chunk in out_list for id in chunk["output_ids"]]
        if "text" in out:
            out["text"] = "".join(chunk["text"] for chunk in out_list)
        if "meta_info" in out:
            meta_info_list = [chunk["meta_info"] for chunk in out_list]
            meta_info = dict(meta_info_list[-1])
            incremental_streaming_keys = set(_INCREMENTAL_STREAMING_META_INFO_KEYS)
            if customized_info_keys is not None:
                incremental_streaming_keys.update(customized_info_keys)
            for key in incremental_streaming_keys:
                if any(key in m for m in meta_info_list):
                    meta_info[key] = [
                        item for m in meta_info_list for item in m.get(key, [])
                    ]
            out["meta_info"] = meta_info
        return out

    async def _handle_abort_finish_reason(
        self,
        out: dict,
        state: ReqState,
        is_stream: bool,
    ) -> Optional[dict]:
        """Handle abort/error finish reasons from the scheduler.

        Returns the output dict if it should be yielded (stream abort), or None
        for normal flow. Raises ValueError or HTTPException for non-stream aborts.
        """
        finish_reason = out["meta_info"]["finish_reason"]

        if (
            finish_reason.get("type") == "abort"
            and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST
        ):
            if not is_stream:
                raise ValueError(finish_reason["message"])
            return out

        if finish_reason.get("type") == "abort" and finish_reason.get(
            "status_code"
        ) in (
            HTTPStatus.SERVICE_UNAVAILABLE,
            HTTPStatus.INTERNAL_SERVER_ERROR,
        ):
            # Delete the key to prevent resending abort request to the scheduler and
            # to ensure aborted request state is cleaned up.
            if state.obj.rid in self.rid_to_state:
                self._remove_req_state(state.obj.rid)

            # Mark ongoing LoRA request as finished.
            if self.enable_lora and state.obj.lora_path:
                await self.lora_registry.release(state.obj.lora_id)
            if not is_stream:
                raise fastapi.HTTPException(
                    status_code=finish_reason["status_code"],
                    detail=finish_reason["message"],
                )
            return out

        return None

    async def _wait_one_response(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        request: Optional[fastapi.Request] = None,
    ):
        """Wait for the response of one request."""
        state = self.rid_to_state[obj.rid]
        # Not all request types have `stream` (e.g., EmbeddingReqInput). Default to non-streaming.
        is_stream = getattr(obj, "stream", False)
        while True:
            try:
                await asyncio.wait_for(
                    state.event.wait(), timeout=_REQUEST_STATE_WAIT_TIMEOUT
                )
            except asyncio.TimeoutError:
                if (
                    request is not None
                    and not obj.background
                    and await request.is_disconnected()
                ):
                    # Abort the request for disconnected requests (non-streaming, waiting queue)
                    self.abort_request(obj.rid)
                    # Use exception to kill the whole call stack and asyncio task
                    raise ValueError(
                        f"Request is disconnected from the client side (type 1). Abort request {obj.rid=}"
                    )
                continue

            # Drain all pending outputs atomically.
            out_list = state.out_list
            state.out_list = []
            finished = state.finished
            state.event.clear()

            # With incremental streaming, each chunk is a delta — coalesce
            # multiple queued chunks to avoid dropping token ids.
            incremental_stream = is_stream and self.incremental_streaming_output
            if incremental_stream and len(out_list) > 1:
                out = self._coalesce_streaming_chunks(
                    out_list,
                    obj.rid,
                    state.customized_info_accumulated.keys(),
                )
            else:
                out = out_list[-1]

            # Resolve deferred text for non-incremental streaming.
            # _handle_batch_output sets "text": None on intermediate chunks
            # to avoid O(n) string rebuild per step (O(n^2) total).
            if (
                is_stream
                and not incremental_stream
                and "text" in out
                and out["text"] is None
            ):
                out["text"] = state.get_text()

            if finished:
                # Record response sent time right before we log finished results and metrics.
                if not state.time_stats.response_sent_to_client_time:
                    state.time_stats.set_response_sent_to_client_time()
                    out["meta_info"][
                        "response_sent_to_client_ts"
                    ] = state.time_stats.get_response_sent_to_client_realtime()
                self.request_logger.log_finished_request(
                    obj,
                    out,
                    request=request,
                )

                if self.request_metrics_exporter_manager.exporter_enabled():
                    asyncio.create_task(
                        self.request_metrics_exporter_manager.write_record(obj, out)
                    )

                # Check if this was an abort/error created by scheduler
                if isinstance(out["meta_info"].get("finish_reason"), dict):
                    abort_out = await self._handle_abort_finish_reason(
                        out, state, is_stream
                    )
                    if abort_out is not None:
                        yield abort_out
                        break

                yield out
                break

            if is_stream:
                # Record response sent time right before we send response.
                if not state.time_stats.response_sent_to_client_time:
                    state.time_stats.set_response_sent_to_client_time()
                    out["meta_info"][
                        "response_sent_to_client_ts"
                    ] = state.time_stats.get_response_sent_to_client_realtime()
                yield out
            else:
                if (
                    request is not None
                    and not obj.background
                    and await request.is_disconnected()
                ):
                    # Abort the request for disconnected requests (non-streaming, running)
                    self.abort_request(obj.rid)
                    # Use exception to kill the whole call stack and asyncio task
                    raise ValueError(
                        f"Request is disconnected from the client side (type 3). Abort request {obj.rid=}"
                    )

    async def _handle_batch_request(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        request: Optional[fastapi.Request] = None,
    ):
        batch_size = obj.batch_size

        generators = []
        rids = []
        if getattr(obj, "parallel_sample_num", 1) == 1:
            if self._should_use_batch_tokenization(batch_size, obj):
                tokenized_objs = await self._batch_tokenize_and_process(batch_size, obj)
                self._raise_if_logical_request_aborted(obj)
                self._send_batch_request(tokenized_objs)

                # Set up generators for each request in the batch
                for i in range(batch_size):
                    tmp_obj = obj[i]
                    state = self.rid_to_state[tmp_obj.rid]
                    if tmp_obj.return_prompt_token_ids:
                        state.prompt_token_ids = list(tokenized_objs[i].input_ids)
                    generators.append(self._wait_one_response(tmp_obj, request))
                    rids.append(tmp_obj.rid)
            else:
                # Sequential tokenization and processing
                with (
                    input_blocker_guard_region(
                        dispatch_to_scheduler=self._dispatch_to_scheduler,
                    )
                    if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN")
                    else nullcontext()
                ):
                    for i in range(batch_size):
                        tmp_obj = obj[i]
                        tokenized_obj = await self._tokenize_one_request(tmp_obj)
                        self._raise_if_logical_rid_aborted(tmp_obj.rid)
                        state = self.rid_to_state[tmp_obj.rid]
                        if tmp_obj.return_prompt_token_ids:
                            state.prompt_token_ids = list(tokenized_obj.input_ids)
                        self._send_one_request(tokenized_obj)
                        generators.append(self._wait_one_response(tmp_obj, request))
                        rids.append(tmp_obj.rid)
        else:
            # FIXME: When using batch and parallel_sample_num together, the perf is not optimal.
            if batch_size > 128:
                logger.warning(
                    "Sending a single large batch with parallel sampling (n > 1) has not been well optimized. "
                    "The performance might be better if you just duplicate the requests n times or use "
                    "many threads to send them one by one with parallel sampling (n > 1)."
                )

            # Tokenize all requests
            objs = [obj[i] for i in range(batch_size)]
            tokenized_objs = await asyncio.gather(
                *(self._tokenize_one_request(obj) for obj in objs)
            )
            self._raise_if_logical_request_aborted(obj)

            # Cache the common prefix for parallel sampling
            for i in range(batch_size):
                logical_rid = objs[i].rid
                self._raise_if_logical_rid_aborted(logical_rid)
                tmp_obj = copy.copy(objs[i])
                tokenized_obj = copy.copy(tokenized_objs[i])
                # Ensure independent mm_items so wrap_shm_features won't mutate the original
                if hasattr(tokenized_obj, "mm_inputs") and tokenized_obj.mm_inputs:
                    tokenized_obj.mm_inputs = copy.copy(tokenized_obj.mm_inputs)
                    tokenized_obj.mm_inputs.mm_items = [
                        copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
                    ]
                tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
                tokenized_obj.sampling_params = copy.copy(tokenized_obj.sampling_params)
                tokenized_obj.sampling_params.max_new_tokens = 0
                tokenized_obj.stream = False
                self._init_child_req_state(logical_rid, tmp_obj)
                self._send_one_request(tokenized_obj)
                await self._wait_one_response(tmp_obj, request).__anext__()
                self._raise_if_logical_rid_aborted(logical_rid)

            # Expand requests, assign new rids for them, and send them
            for i in range(batch_size):
                logical_rid = objs[i].rid
                for _ in range(obj.parallel_sample_num):
                    self._raise_if_logical_rid_aborted(logical_rid)
                    tmp_obj = copy.copy(objs[i])
                    tokenized_obj = copy.copy(tokenized_objs[i])
                    # Ensure independent mm_items so wrap_shm_features won't mutate the original
                    if hasattr(tokenized_obj, "mm_inputs") and tokenized_obj.mm_inputs:
                        tokenized_obj.mm_inputs = copy.copy(tokenized_obj.mm_inputs)
                        tokenized_obj.mm_inputs.mm_items = [
                            copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
                        ]
                    tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
                    self._init_child_req_state(logical_rid, tmp_obj)
                    state = self.rid_to_state[tmp_obj.rid]
                    tokenized_obj.time_stats = state.time_stats
                    if tmp_obj.return_prompt_token_ids:
                        state.prompt_token_ids = list(tokenized_objs[i].input_ids)
                    self._send_one_request(tokenized_obj)
                    generators.append(self._wait_one_response(tmp_obj, request))
                    rids.append(tmp_obj.rid)

                parent_state = self.rid_to_state.get(logical_rid)
                if parent_state is not None:
                    parent_state.time_stats.set_finished_time()
                self._remove_req_state(logical_rid)

        # Wait for all requests
        is_stream = hasattr(obj, "stream") and obj.stream
        if not is_stream:
            outputs = await self._collect_batch_responses(generators)
            yield outputs
        else:
            async for response in self._stream_batch_responses(generators, rids):
                yield response

    async def _collect_batch_responses(self, generators):
        tasks = [asyncio.create_task(gen.__anext__()) for gen in generators]
        try:
            return await asyncio.gather(*tasks)
        finally:
            for task in tasks:
                if not task.done():
                    task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            await asyncio.gather(
                *(gen.aclose() for gen in generators),
                return_exceptions=True,
            )

    async def _stream_batch_responses(self, generators, rids):
        rid_to_index = {rid: i for i, rid in enumerate(rids)}
        task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
        try:
            while task_map:
                done, _ = await asyncio.wait(
                    task_map.keys(), return_when=asyncio.FIRST_COMPLETED
                )

                for task in done:
                    gen = task_map.pop(task)
                    try:
                        result = task.result()
                        result["index"] = rid_to_index[result["meta_info"]["id"]]
                        yield result
                        new_task = asyncio.create_task(gen.__anext__())
                        task_map[new_task] = gen
                    except StopAsyncIteration:
                        pass
        finally:
            pending_tasks = list(task_map)
            for task in pending_tasks:
                task.cancel()
            if pending_tasks:
                await asyncio.gather(*pending_tasks, return_exceptions=True)
            await asyncio.gather(
                *(gen.aclose() for gen in generators),
                return_exceptions=True,
            )

    def abort_request(self, rid: str = "", abort_all: bool = False):
        # Empty rid would startswith-match every request on the scheduler.
        if not abort_all and not rid:
            logger.warning("Ignore abort_request with empty rid and abort_all=False")
            return
        if abort_all:
            for state_rid, state in self.rid_to_state.items():
                if state_rid not in self.child_rid_to_logical_rid:
                    state.abort_requested = True
            target_rids = (rid,)
        elif rid in self.child_rid_to_logical_rid:
            # Preserve direct child aborts for internal callers.
            target_rids = (rid,)
        elif rid in self.rid_to_state:
            state = self.rid_to_state[rid]
            state.abort_requested = True
            parallel_sample_num = getattr(state.obj, "parallel_sample_num", None)
            if parallel_sample_num is None:
                sampling_params = getattr(state.obj, "sampling_params", None)
                parallel_sample_num = (
                    sampling_params.get("n", 1)
                    if isinstance(sampling_params, dict)
                    else 1
                )
            if parallel_sample_num > 1:
                # Snapshot because scheduler abort echoes remove child ownership.
                target_rids = tuple(sorted(self.logical_rid_to_child_rids.get(rid, ())))
            else:
                target_rids = (rid,)
        elif child_rids := self.logical_rid_to_child_rids.get(rid):
            target_rids = tuple(sorted(child_rids))
        elif self.server_args.tokenizer_worker_num == 1:
            return
        else:
            target_rids = (rid,)

        for target_rid in target_rids:
            self._dispatch_to_scheduler(AbortReq(rid=target_rid, abort_all=abort_all))
        if self.enable_metrics:
            # TODO: also use custom_labels from the request
            self.metrics_collector.observe_one_aborted_request(
                self.metrics_collector.labels
            )

    async def pause_generation(self, obj: PauseGenerationReqInput):
        async with self.is_pause_cond:
            self.is_pause = True
            if obj.mode != "abort":
                await self._async_dispatch_to_scheduler(obj)
            else:
                # we are using the model_update_lock to check if there is still on-going requests.
                while True:
                    # TODO: maybe make it async instead of fire-and-forget
                    self.abort_request(abort_all=True)
                    is_locked = await self.model_update_lock.is_locked()
                    if not is_locked:
                        break
                    await asyncio.sleep(1.0)

    async def continue_generation(self, obj: ContinueGenerationReqInput):
        async with self.is_pause_cond:
            self.is_pause = False
            await self._async_dispatch_to_scheduler(obj)
            self.is_pause_cond.notify_all()

    async def update_weights_from_disk(
        self,
        obj: UpdateWeightFromDiskReqInput,
        request: Optional[fastapi.Request] = None,
    ) -> Tuple[bool, str]:
        self.auto_create_handle_loop()

        if obj.load_format is None:
            obj.load_format = self.config_value("load_format")
        logger.info("Start update_weights. Load format=%s", obj.load_format)

        if obj.abort_all_requests:
            self.abort_request(abort_all=True)

        # Immediately update the weights if the engine is in paused state
        async with self.is_pause_cond:
            is_paused = self.is_pause

        lock_context = (
            self.model_update_lock.writer_lock if not is_paused else nullcontext()
        )
        async with lock_context:
            success, message, num_paused_requests = (
                await self._wait_for_model_update_from_disk(obj)
            )

        if success and obj.weight_version is not None:
            self._update_weight_version_if_provided(obj.weight_version)
            message += f" Weight version updated to {obj.weight_version}."

        return success, message, num_paused_requests

    def record_config_updates(self, source: str, **fields) -> None:
        """Record a control-plane config change for this engine.

        Per-engine state: several ``Engine``s can share one tokenizer process.
        The readback endpoints overlay these onto the startup config. The
        process-global sibling is ``RuntimeContext.override`` /
        ``resolved_server_args_dict``, which writes the config bags every
        process shares.
        """
        unknown = sorted(f for f in fields if f not in _SERVER_ARGS_FIELDS)
        if unknown:
            raise ValueError(
                f"{unknown} are not ServerArgs fields; the readback endpoints "
                "overlay these onto a serialized ServerArgs, so an unknown key "
                "would surface as a phantom config entry."
            )
        self._config_updates.append((source, dict(fields)))

    def config_value(self, name: str):
        """The value in effect for one config field, control-plane updates first."""
        if name in _MANAGER_OWNED_FIELDS:
            return getattr(self, name)
        for _source, fields in reversed(self._config_updates):
            if name in fields:
                return fields[name]
        return getattr(self.server_args, name)

    def _dump_config_snapshot(self) -> Optional[Dict[str, Any]]:
        """The config in effect, or None when it cannot be serialized.

        A dump is worth having even when the config is not: request data is the
        part that cannot be reconstructed afterwards.
        """
        try:
            return self.resolved_config_dict(dataclasses.asdict(self.server_args))
        except Exception as e:
            logger.error(f"Failed to snapshot the resolved config for the dump: {e!r}")
            return None

    def resolved_config_dict(self, base: Dict[str, Any]) -> Dict[str, Any]:
        """``base`` (a serialized ``ServerArgs``) with the control-plane updates on top."""
        resolved = dict(base)
        for _source, fields in self._config_updates:
            resolved.update(fields)
        for name in _MANAGER_OWNED_FIELDS:
            resolved[name] = getattr(self, name)
        return resolved

    def _update_model_path_info(self, model_path: str, load_format: str):
        self.served_model_name = model_path
        self.record_config_updates("tokenizer.update_weights", load_format=load_format)
        self.model_path = model_path

    async def _wait_for_model_update_from_disk(
        self, obj: UpdateWeightFromDiskReqInput
    ) -> Tuple[bool, str]:
        expected_workers = self.elastic_worker_count
        self.model_update_expected_workers = expected_workers
        self.model_update_tmp = []
        self.model_update_result = asyncio.Future()
        self._dispatch_to_scheduler(obj)
        if expected_workers == 1:
            result = await self.model_update_result
            if result.success:
                self._update_model_path_info(obj.model_path, obj.load_format)
            return result.success, result.message, result.num_paused_requests
        else:
            result = await self.model_update_result

            all_success = all([r.success for r in result])
            if all_success is True:
                self._update_model_path_info(obj.model_path, obj.load_format)
            all_message = [r.message for r in result]
            all_message = " | ".join(all_message)
            all_paused_requests = [r.num_paused_requests for r in result]
            return all_success, all_message, all_paused_requests

    def configure_logging(self, obj: ConfigureLoggingReq):
        self.request_logger.configure(
            log_requests=obj.log_requests,
            log_requests_level=obj.log_requests_level,
            log_requests_format=obj.log_requests_format,
        )
        if obj.dump_requests_folder is not None:
            self.dump_requests_folder = obj.dump_requests_folder
        if obj.dump_requests_threshold is not None:
            self.dump_requests_threshold = obj.dump_requests_threshold
        if obj.dump_requests_exclude_meta_keys is not None:
            self.dump_requests_exclude_meta_keys = list(
                obj.dump_requests_exclude_meta_keys
            )
        if obj.crash_dump_folder is not None:
            self.crash_dump_folder = obj.crash_dump_folder
        if obj.log_level is not None:
            # setLevel() may raise exception if obj.log_level is illegal string.
            # Let the exception propagate to the caller.
            # Only legal requests will be sent to scheduler.
            logging.getLogger().setLevel(obj.log_level.upper())
            self._dispatch_to_scheduler(obj)
        logging.info(f"Config logging: {obj=}")

    async def freeze_gc(self):
        """Send a freeze_gc message to the scheduler first, then freeze locally."""
        self._dispatch_to_scheduler(FreezeGCReq())
        freeze_gc("Tokenizer Manager")
        return None

    def create_abort_task(self, obj: GenerateReqInput):
        # Abort the request if the client is disconnected.
        async def abort_request():
            await asyncio.sleep(2)
            if obj.is_single:
                self.abort_request(obj.rid)
            else:
                for rid in obj.rid:
                    self.abort_request(rid)

        background_tasks = BackgroundTasks()
        background_tasks.add_task(abort_request)
        return background_tasks

    def auto_create_handle_loop(self):
        if self.event_loop is not None:
            return

        # Create and start the handle_loop task
        loop = get_or_create_event_loop()
        self.asyncio_tasks.add(
            loop.create_task(print_exception_wrapper(self.handle_loop))
        )
        self.event_loop = loop

        # We only add signal handler when the tokenizer manager is in the main thread
        # due to the CPython limitation.
        if threading.current_thread() is threading.main_thread():
            signal_handler = self.signal_handler_class(self)
            loop.add_signal_handler(signal.SIGTERM, signal_handler.sigterm_handler)
            # Update the signal handler for the process. It overrides the sigquit handler in the launch phase.
            loop.add_signal_handler(
                signal.SIGQUIT, signal_handler.running_phase_sigquit_handler
            )

        self.asyncio_tasks.add(
            loop.create_task(print_exception_wrapper(self.sigterm_watchdog))
        )

    async def handle_loop(self):
        """The event loop that handles requests"""
        while True:
            with self.soft_watchdog.disable():
                recv_obj = await async_sock_recv(self.recv_from_detokenizer)
            if isinstance(
                recv_obj,
                (BatchStrOutput, BatchEmbeddingOutput, BatchTokenIDOutput),
            ):
                await self._handle_batch_output(recv_obj)
            else:
                self._result_dispatcher(recv_obj)
            self.last_receive_tstamp = real_time()
            self.soft_watchdog.feed()

    async def _handle_batch_output(
        self,
        recv_obj: Union[
            BatchStrOutput,
            BatchEmbeddingOutput,
            BatchTokenIDOutput,
        ],
    ):
        recv_obj.time_stats = unwrap_from_pickle(recv_obj.time_stats)
        if isinstance(recv_obj, (BatchStrOutput, BatchTokenIDOutput)):
            customized_info = unwrap_from_pickle(recv_obj.customized_info)
        else:
            customized_info = None
        pending_notify: dict[str, ReqState] = {}
        batch_notify_size = self.server_args.batch_notify_size
        for i, rid in enumerate(recv_obj.rids):
            state = self.rid_to_state.get(rid, None)
            if state is None:
                # Known race: /health_generate pops its rid as soon as ANY message bumps last_receive_tstamp.
                if rid.startswith(HEALTH_CHECK_RID_PREFIX):
                    continue
                logger.error(
                    f"Received output for {rid=} but the state was deleted in TokenizerManager."
                )
                continue

            # Build meta_info and return value
            meta_info = {
                "id": rid,
                "finish_reason": recv_obj.finished_reasons[i],
                "prompt_tokens": recv_obj.prompt_tokens[i],
                "weight_version": self.config_value("weight_version"),
                "num_retractions": recv_obj.retraction_counts[i],
            }

            if self.enable_metrics:
                if recv_obj.time_stats is not None:
                    scheduler_time_stats = recv_obj.time_stats[i]
                    meta_info.update(scheduler_time_stats.convert_to_output_meta_info())

            if getattr(state.obj, "return_logprob", False):
                self.convert_logprob_style(
                    meta_info,
                    state,
                    state.obj.top_logprobs_num,
                    state.obj.token_ids_logprob,
                    state.obj.return_text_in_logprobs and not self.skip_tokenizer_init,
                    recv_obj,
                    i,
                )
            if (
                isinstance(state.obj, GenerateReqInput)
                and state.obj.return_sampling_mask
            ):
                output_sampling_mask = recv_obj.output_token_sampling_mask
                if output_sampling_mask is not None:
                    state.output_token_sampling_mask.extend(output_sampling_mask[i])
                    output_sampling_logprobs = recv_obj.output_token_sampling_logprobs
                    if output_sampling_logprobs is not None:
                        state.output_token_sampling_logprobs.extend(
                            output_sampling_logprobs[i]
                        )
                    meta_info["output_token_sampling_mask"] = (
                        state.output_token_sampling_mask
                    )
                    meta_info["output_token_sampling_logprobs"] = (
                        state.output_token_sampling_logprobs
                    )
                    meta_info["output_token_sampling_mask_length"] = len(
                        state.output_token_sampling_mask
                    )

            if not isinstance(recv_obj, BatchEmbeddingOutput):
                meta_info.update(
                    {
                        "reasoning_tokens": recv_obj.reasoning_tokens[i],
                        "completion_tokens": recv_obj.completion_tokens[i],
                        "cached_tokens": recv_obj.cached_tokens[i],
                    }
                )
                # Add detailed cache breakdown if available
                if (
                    hasattr(recv_obj, "cached_tokens_details")
                    and recv_obj.cached_tokens_details
                ):
                    meta_info["cached_tokens_details"] = recv_obj.cached_tokens_details[
                        i
                    ]
                if customized_info is not None:
                    for k, v in customized_info.items():
                        if k not in state.customized_info_accumulated:
                            state.customized_info_accumulated[k] = []
                        state.customized_info_accumulated[k].extend(v[i])
                        meta_info[k] = state.customized_info_accumulated[k]

                # Add multimodal prompt token counts only for requests that
                # actually consumed them, so plain-text meta_info stays unchanged.
                image_tokens_list = getattr(recv_obj, "image_tokens", None)
                audio_tokens_list = getattr(recv_obj, "audio_tokens", None)
                video_tokens_list = getattr(recv_obj, "video_tokens", None)
                if image_tokens_list and image_tokens_list[i]:
                    meta_info["image_tokens"] = image_tokens_list[i]
                if audio_tokens_list and audio_tokens_list[i]:
                    meta_info["audio_tokens"] = audio_tokens_list[i]
                if video_tokens_list and video_tokens_list[i]:
                    meta_info["video_tokens"] = video_tokens_list[i]

            if getattr(recv_obj, "output_hidden_states", None):
                hidden_states = recv_obj.output_hidden_states[i]
                if hidden_states is not None:
                    meta_info["hidden_states"] = hidden_states
            if getattr(recv_obj, "routed_experts", None):
                val = recv_obj.routed_experts[i]
                if val is not None:
                    # BatchStrOutput is pre-encoded by the detokenizer;
                    # BatchTokenIDOutput (skip_tokenizer_init) bypasses it.
                    if isinstance(val, torch.Tensor):
                        val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8")
                    meta_info["routed_experts"] = val
            if getattr(recv_obj, "indexer_topk", None):
                val = recv_obj.indexer_topk[i]
                if val is not None:
                    if isinstance(val, torch.Tensor):
                        val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8")
                    meta_info["indexer_topk"] = val
            if getattr(recv_obj, "dp_ranks", None):
                meta_info["dp_rank"] = recv_obj.dp_ranks[i]

            state.finished = recv_obj.finished_reasons[i] is not None
            if isinstance(recv_obj, BatchStrOutput):
                # Not all request types have `stream` (e.g., EmbeddingReqInput). Default to non-streaming.
                is_stream = getattr(state.obj, "stream", False)
                incremental = is_stream and self.incremental_streaming_output
                delta_text = recv_obj.output_strs[i]
                delta_output_ids = list(recv_obj.output_ids[i])
                output_offset = state.last_output_offset
                state.append_text(delta_text)
                state.output_ids.extend(delta_output_ids)

                if is_stream:
                    if incremental:
                        output_token_ids = delta_output_ids
                        _slice_streaming_output_meta_info(
                            meta_info,
                            output_offset,
                            state.customized_info_accumulated.keys(),
                        )
                        state.last_output_offset = len(state.output_ids)
                        out_dict = {
                            "text": delta_text,
                            "output_ids": output_token_ids,
                            "meta_info": meta_info,
                        }
                    elif state.finished:
                        out_dict = {
                            "text": state.get_text(),
                            "output_ids": state.output_ids.copy(),
                            "meta_info": meta_info,
                        }
                    else:
                        # Non-incremental intermediate: pass reference (no
                        # copy) and defer text to _wait_one_response to avoid
                        # O(n) per-step cost that compounds to O(n^2).
                        out_dict = {
                            "text": None,
                            "output_ids": state.output_ids,
                            "meta_info": meta_info,
                        }
                elif state.finished:
                    out_dict = {
                        "text": state.get_text(),
                        "output_ids": state.output_ids.copy(),
                        "meta_info": meta_info,
                    }
                else:
                    out_dict = None
                if out_dict is not None and state.prompt_token_ids is not None:
                    out_dict["prompt_token_ids"] = state.prompt_token_ids
            elif isinstance(recv_obj, BatchTokenIDOutput):
                is_stream = getattr(state.obj, "stream", False)
                incremental = is_stream and self.incremental_streaming_output
                delta_output_ids = list(recv_obj.output_ids[i])
                output_offset = state.last_output_offset
                state.output_ids.extend(delta_output_ids)

                if is_stream:
                    if incremental:
                        output_token_ids = delta_output_ids
                        _slice_streaming_output_meta_info(
                            meta_info,
                            output_offset,
                            state.customized_info_accumulated.keys(),
                        )
                        state.last_output_offset = len(state.output_ids)
                        out_dict = {
                            "output_ids": output_token_ids,
                            "meta_info": meta_info,
                        }
                    elif state.finished:
                        out_dict = {
                            "output_ids": state.output_ids.copy(),
                            "meta_info": meta_info,
                        }
                    else:
                        out_dict = {
                            "output_ids": state.output_ids,
                            "meta_info": meta_info,
                        }
                elif state.finished:
                    out_dict = {
                        "output_ids": state.output_ids.copy(),
                        "meta_info": meta_info,
                    }
                else:
                    out_dict = None
                if out_dict is not None and state.prompt_token_ids is not None:
                    out_dict["prompt_token_ids"] = state.prompt_token_ids
            else:
                assert isinstance(recv_obj, BatchEmbeddingOutput)
                out_dict = {
                    "embedding": recv_obj.embeddings[i],
                    "meta_info": meta_info,
                }
                # Unpack pooled hidden states (PHS).
                # See paired sender logic in output_streamer.py.
                #   Stacked:     len == 1 and N > 1 → unwrap the tensor
                #   Non-stacked: len == N → index directly
                pooled_hidden_states = recv_obj.pooled_hidden_states
                if pooled_hidden_states is not None:
                    if len(pooled_hidden_states) == 1 and len(recv_obj.rids) > 1:
                        pooled_hidden_states = pooled_hidden_states[0]
                    if pooled_hidden_states[i] is not None:
                        out_dict["pooled_hidden_state"] = pooled_hidden_states[i]

            # Set first_token_time on the first output batch.
            # This is the single write point for first_token_time.
            if state.time_stats.first_token_time == 0.0:
                state.time_stats.set_first_token_time()

            if state.finished:
                if state.time_stats.trace_ctx.tracing_enable:
                    state.time_stats.trace_ctx.trace_set_root_attrs(
                        self.convert_to_span_attrs(state, recv_obj, i)
                    )
                state.time_stats.set_finished_time()
                meta_info["e2e_latency"] = state.time_stats.get_e2e_latency()

                if self.server_args.speculative_algorithm:
                    self._calculate_spec_decoding_metrics(meta_info, recv_obj, i)
                if self.enable_metrics:
                    scheduler_time_stats = (
                        recv_obj.time_stats[i]
                        if recv_obj.time_stats is not None
                        else None
                    )
                    completion_tokens = (
                        recv_obj.completion_tokens[i]
                        if not isinstance(recv_obj, BatchEmbeddingOutput)
                        else 0
                    )
                    meta_info.update(
                        state.time_stats.convert_to_output_meta_info(
                            scheduler_time_stats, completion_tokens
                        )
                    )

                self._remove_req_state(rid)

                # Mark ongoing LoRA request as finished.
                if self.enable_lora and state.obj.lora_path:
                    asyncio.create_task(self.lora_registry.release(state.obj.lora_id))

            if out_dict is not None:
                state.out_list.append(out_dict)
                pending_notify[rid] = state

                if len(pending_notify) >= batch_notify_size:
                    for s in pending_notify.values():
                        s.event.set()
                    pending_notify = {}
                    await asyncio.sleep(0)

            if self.enable_metrics and state.obj.log_metrics:
                self.collect_metrics(state, recv_obj, i)
            if self.dump_requests_folder and state.finished and state.obj.log_metrics:
                self.dump_requests(state, out_dict)
            if self.crash_dump_folder and state.finished and state.obj.log_metrics:
                self.record_request_for_crash_dump(state, out_dict)

        # handle_loop awaits next recv immediately
        for s in pending_notify.values():
            s.event.set()

    def add_logprob_to_meta_info(
        self,
        meta_info: dict,
        state: ReqState,
        top_logprobs_num: int,
        token_ids_logprob: List[int],
        return_text_in_logprobs: bool,
    ):
        # 1. Handle regular logprobs
        if len(state.input_token_logprobs_val) > len(state.input_token_logprobs):
            state.input_token_logprobs.extend(
                self.detokenize_logprob_tokens(
                    state.input_token_logprobs_val[len(state.input_token_logprobs) :],
                    state.input_token_logprobs_idx[len(state.input_token_logprobs) :],
                    return_text_in_logprobs,
                )
            )

        if len(state.output_token_logprobs_val) > len(state.output_token_logprobs):
            state.output_token_logprobs.extend(
                self.detokenize_logprob_tokens(
                    state.output_token_logprobs_val[len(state.output_token_logprobs) :],
                    state.output_token_logprobs_idx[len(state.output_token_logprobs) :],
                    return_text_in_logprobs,
                )
            )

        meta_info["input_token_logprobs"] = state.input_token_logprobs
        meta_info["output_token_logprobs"] = state.output_token_logprobs
        meta_info["output_token_logprobs_length"] = len(state.output_token_logprobs)

        # 2. Handle top logprobs
        if top_logprobs_num > 0:
            # Guarded by the caller's return_logprob check, so obj is a
            # GenerateReqInput here.
            use_flat = state.obj.return_flat_raw_top_logprobs
            if use_flat and state.input_top_logprobs_scheduler_flat is not None:
                # The scheduler already assembled the flat arrays (sent once
                # at prefill completion); encode them directly.
                if state.input_top_logprobs_flat_fields is None:
                    val_arr, idx_arr, null_prefix = (
                        state.input_top_logprobs_scheduler_flat
                    )
                    state.input_top_logprobs_flat_fields = (
                        _build_flat_input_top_logprobs_fields_from_arrays(
                            val_arr,
                            idx_arr,
                            null_prefix,
                            return_b64=state.obj.return_flat_raw_top_logprobs_b64,
                        )
                    )
                meta_info.update(state.input_top_logprobs_flat_fields)
            elif use_flat:
                # Flat replaces nested for the input side only.
                if state.input_top_logprobs_flat_num_rows != len(
                    state.input_top_logprobs_val
                ):
                    try:
                        state.input_top_logprobs_flat_fields = (
                            _build_flat_input_top_logprobs_fields(
                                state.input_top_logprobs_val,
                                state.input_top_logprobs_idx,
                                top_logprobs_num,
                                return_b64=state.obj.return_flat_raw_top_logprobs_b64,
                            )
                        )
                    except ValueError as e:
                        # A raise here would disrupt unrelated requests in the
                        # shared batch-output loop; degrade to nested instead.
                        state.input_top_logprobs_flat_fields = None
                        logger.error(
                            "Falling back to nested input top logprobs for "
                            "rid=%s: %s",
                            meta_info.get("id"),
                            e,
                        )
                    state.input_top_logprobs_flat_num_rows = len(
                        state.input_top_logprobs_val
                    )
                if state.input_top_logprobs_flat_fields is not None:
                    meta_info.update(state.input_top_logprobs_flat_fields)
                else:
                    use_flat = False
            if not use_flat:
                if len(state.input_top_logprobs_val) > len(state.input_top_logprobs):
                    state.input_top_logprobs.extend(
                        self.detokenize_top_logprobs_tokens(
                            state.input_top_logprobs_val[
                                len(state.input_top_logprobs) :
                            ],
                            state.input_top_logprobs_idx[
                                len(state.input_top_logprobs) :
                            ],
                            return_text_in_logprobs,
                        )
                    )
                meta_info["input_top_logprobs"] = state.input_top_logprobs
            if len(state.output_top_logprobs_val) > len(state.output_top_logprobs):
                state.output_top_logprobs.extend(
                    self.detokenize_top_logprobs_tokens(
                        state.output_top_logprobs_val[len(state.output_top_logprobs) :],
                        state.output_top_logprobs_idx[len(state.output_top_logprobs) :],
                        return_text_in_logprobs,
                    )
                )
            meta_info["output_top_logprobs"] = state.output_top_logprobs

        # 3. Handle token_ids_logprob
        if token_ids_logprob is not None:
            if len(state.input_token_ids_logprobs_val) > len(
                state.input_token_ids_logprobs
            ):
                state.input_token_ids_logprobs.extend(
                    self.detokenize_top_logprobs_tokens(
                        state.input_token_ids_logprobs_val[
                            len(state.input_token_ids_logprobs) :
                        ],
                        state.input_token_ids_logprobs_idx[
                            len(state.input_token_ids_logprobs) :
                        ],
                        return_text_in_logprobs,
                    )
                )
            if len(state.output_token_ids_logprobs_val) > len(
                state.output_token_ids_logprobs
            ):
                state.output_token_ids_logprobs.extend(
                    self.detokenize_top_logprobs_tokens(
                        state.output_token_ids_logprobs_val[
                            len(state.output_token_ids_logprobs) :
                        ],
                        state.output_token_ids_logprobs_idx[
                            len(state.output_token_ids_logprobs) :
                        ],
                        return_text_in_logprobs,
                    )
                )

            meta_info["input_token_ids_logprobs"] = state.input_token_ids_logprobs
            meta_info["output_token_ids_logprobs"] = state.output_token_ids_logprobs

    def convert_logprob_style(
        self,
        meta_info: dict,
        state: ReqState,
        top_logprobs_num: int,
        token_ids_logprob: List[int],
        return_text_in_logprobs: bool,
        recv_obj: BatchStrOutput,
        recv_obj_index: int,
    ):
        if recv_obj.input_token_logprobs_val is None:
            return

        if (
            len(recv_obj.input_token_logprobs_val) > 0
            and recv_obj.input_token_logprobs_val[recv_obj_index] is not None
        ):
            state.input_token_logprobs_val.extend(
                recv_obj.input_token_logprobs_val[recv_obj_index]
            )
            state.input_token_logprobs_idx.extend(
                recv_obj.input_token_logprobs_idx[recv_obj_index]
            )
        state.output_token_logprobs_val.extend(
            recv_obj.output_token_logprobs_val[recv_obj_index]
        )
        state.output_token_logprobs_idx.extend(
            recv_obj.output_token_logprobs_idx[recv_obj_index]
        )

        if top_logprobs_num > 0:
            if len(recv_obj.input_top_logprobs_val) > 0:
                state.input_top_logprobs_val.extend(
                    recv_obj.input_top_logprobs_val[recv_obj_index]
                )
                state.input_top_logprobs_idx.extend(
                    recv_obj.input_top_logprobs_idx[recv_obj_index]
                )
            if (
                recv_obj.input_top_logprobs_val_flat is not None
                and recv_obj.input_top_logprobs_val_flat[recv_obj_index] is not None
            ):
                state.input_top_logprobs_scheduler_flat = (
                    recv_obj.input_top_logprobs_val_flat[recv_obj_index],
                    recv_obj.input_top_logprobs_idx_flat[recv_obj_index],
                    recv_obj.input_top_logprobs_flat_null_prefix[recv_obj_index],
                )
            state.output_top_logprobs_val.extend(
                recv_obj.output_top_logprobs_val[recv_obj_index]
            )
            state.output_top_logprobs_idx.extend(
                recv_obj.output_top_logprobs_idx[recv_obj_index]
            )

        if token_ids_logprob is not None:
            if len(recv_obj.input_token_ids_logprobs_val) > 0:
                state.input_token_ids_logprobs_val.extend(
                    recv_obj.input_token_ids_logprobs_val[recv_obj_index]
                )
                state.input_token_ids_logprobs_idx.extend(
                    recv_obj.input_token_ids_logprobs_idx[recv_obj_index]
                )
            state.output_token_ids_logprobs_val.extend(
                recv_obj.output_token_ids_logprobs_val[recv_obj_index]
            )
            state.output_token_ids_logprobs_idx.extend(
                recv_obj.output_token_ids_logprobs_idx[recv_obj_index]
            )

        self.add_logprob_to_meta_info(
            meta_info,
            state,
            state.obj.top_logprobs_num,
            state.obj.token_ids_logprob,
            return_text_in_logprobs,
        )

    def detokenize_logprob_tokens(
        self,
        token_logprobs_val: List[float],
        token_logprobs_idx: List[int],
        decode_to_text: bool,
    ):
        if not decode_to_text:
            return [
                (logprob, token_id, None)
                for logprob, token_id in zip(token_logprobs_val, token_logprobs_idx)
            ]
        else:
            assert self.tokenizer is not None
            # In transformers v5, batch_decode([1, 2, 3]) concatenates all tokens
            # into one string. Wrap each ID in its own list so they decode separately.
            token_texts = self.tokenizer.batch_decode(
                [[idx] for idx in token_logprobs_idx]
            )
            return list(zip(token_logprobs_val, token_logprobs_idx, token_texts))

    def detokenize_top_logprobs_tokens(
        self,
        token_logprobs_val: List[float],
        token_logprobs_idx: List[int],
        decode_to_text: bool,
    ):
        # TODO: The current implementation only batches the detokenization for top-k tokens per single position.
        # We should batch all top-k tokens in all positions.
        ret = []
        for i in range(len(token_logprobs_val)):
            if token_logprobs_val[i]:
                ret.append(
                    self.detokenize_logprob_tokens(
                        token_logprobs_val[i], token_logprobs_idx[i], decode_to_text
                    )
                )
            else:
                ret.append(None)
        return ret

    def _calculate_spec_decoding_metrics(
        self,
        meta_info: Dict[str, Any],
        recv_obj: Union[
            BatchStrOutput,
            BatchEmbeddingOutput,
            BatchTokenIDOutput,
        ],
        i: int,
    ) -> None:
        """Calculate speculative decoding metrics, such as acceptance rate and acceptance length metrics."""
        if (
            hasattr(recv_obj, "spec_verify_ct")
            and recv_obj.spec_verify_ct[i] > 0
            and hasattr(recv_obj, "spec_num_correct_drafts")
            and len(recv_obj.spec_num_correct_drafts) > i
        ):
            # Total number of proposed draft tokens per request.
            num_proposed_drafts = recv_obj.spec_verify_ct[i] * (
                self.server_args.speculative_num_draft_tokens - 1
            )
            num_correct_drafts = recv_obj.spec_num_correct_drafts[i]

            # Calculate per-request acceptance rate and average acceptance length.
            if num_proposed_drafts > 0:
                # accept_rate: num_correct_drafts / num_proposed_drafts (strict count, no bonus).
                meta_info["spec_accept_rate"] = num_correct_drafts / num_proposed_drafts
                # accept_length: completion_tokens / verify_ct (includes bonus token).
                meta_info["spec_accept_length"] = (
                    recv_obj.completion_tokens[i] / recv_obj.spec_verify_ct[i]
                )

                meta_info["spec_num_correct_drafts"] = num_correct_drafts
                meta_info["spec_num_proposed_drafts"] = num_proposed_drafts
                meta_info["spec_verify_ct"] = recv_obj.spec_verify_ct[i]

                if (
                    getattr(recv_obj, "spec_num_cap_tokens", None) is not None
                    and len(recv_obj.spec_num_cap_tokens) > i
                    and recv_obj.spec_num_cap_tokens[i] > 0
                ):
                    meta_info["spec_cap_length"] = (
                        recv_obj.spec_num_cap_tokens[i] / recv_obj.spec_verify_ct[i]
                    )
                if (
                    _ragged_verify_cap_accept()
                    and getattr(recv_obj, "spec_num_block_accept_tokens", None)
                    is not None
                    and len(recv_obj.spec_num_block_accept_tokens) > i
                ):
                    meta_info["spec_block_accept_length"] = (
                        recv_obj.spec_num_block_accept_tokens[i]
                        / recv_obj.spec_verify_ct[i]
                    )

                # FIXME: backward-compat aliases, remove in next release.
                meta_info["spec_accepted_drafts"] = num_correct_drafts
                meta_info["spec_proposed_drafts"] = num_proposed_drafts

            # Acceptance histogram: tracks how many decoding steps accepted a certain number of draft tokens.
            if (
                recv_obj.spec_correct_drafts_histogram
                and len(recv_obj.spec_correct_drafts_histogram) > i
                and recv_obj.spec_correct_drafts_histogram[i]
            ):
                meta_info["spec_correct_drafts_histogram"] = (
                    recv_obj.spec_correct_drafts_histogram[i]
                )
                # FIXME: backward-compat alias, remove in next release.
                meta_info["spec_accept_histogram"] = (
                    recv_obj.spec_correct_drafts_histogram[i]
                )
            if (
                getattr(recv_obj, "spec_cap_lens_histogram", None)
                and len(recv_obj.spec_cap_lens_histogram) > i
                and recv_obj.spec_cap_lens_histogram[i]
            ):
                meta_info["spec_cap_lens_histogram"] = recv_obj.spec_cap_lens_histogram[
                    i
                ]

    def _request_has_grammar(self, obj: GenerateReqInput) -> bool:
        return (
            obj.sampling_params.get("json_schema", None)
            or obj.sampling_params.get("regex", None)
            or obj.sampling_params.get("ebnf", None)
            or obj.sampling_params.get("structural_tag", None)
        )

    def collect_metrics(self, state: ReqState, recv_obj: BatchStrOutput, i: int):
        completion_tokens = (
            recv_obj.completion_tokens[i]
            if getattr(recv_obj, "completion_tokens", None)
            else 0
        )

        custom_labels = getattr(state.obj, "custom_labels", None)
        labels = dict(self.metrics_collector.labels)
        if custom_labels:
            labels.update(custom_labels)
        if self.enable_priority_scheduling:
            priority = getattr(state.obj, "priority", None)
            if priority is not None:
                labels["priority"] = str(priority)
        if (
            not state.ttft_observed
            and self.disaggregation_mode != DisaggregationMode.PREFILL
        ):
            state.ttft_observed = True
            state.last_completion_tokens = completion_tokens
            self.metrics_collector.observe_time_to_first_token(
                labels,
                state.time_stats.get_first_token_latency(),
                stream=getattr(state.obj, "stream", False),
            )
        else:
            num_new_tokens = completion_tokens - state.last_completion_tokens
            if num_new_tokens:
                self.metrics_collector.observe_inter_token_latency(
                    labels,
                    state.time_stats.get_interval(),
                    num_new_tokens,
                )
                state.time_stats.set_last_time()
                state.last_completion_tokens = completion_tokens

        if state.finished:
            # Get detailed cache breakdown if available
            cached_tokens_details = None
            if (
                hasattr(recv_obj, "cached_tokens_details")
                and recv_obj.cached_tokens_details
            ):
                cached_tokens_details = recv_obj.cached_tokens_details[i]

            spec_verify_ct = (
                recv_obj.spec_verify_ct[i]
                if hasattr(recv_obj, "spec_verify_ct")
                and recv_obj.spec_verify_ct
                and len(recv_obj.spec_verify_ct) > i
                else 0
            )

            self.metrics_collector.observe_one_finished_request(
                labels,
                recv_obj.prompt_tokens[i],
                completion_tokens,
                recv_obj.cached_tokens[i],
                state.time_stats.get_e2e_latency(),
                self._request_has_grammar(state.obj),
                cached_tokens_details,
                spec_verify_ct=spec_verify_ct,
                is_streaming=getattr(state.obj, "stream", False),
            )

    def dump_requests(self, state: ReqState, out_dict: dict):
        if self.dump_requests_exclude_meta_keys and isinstance(
            out_dict.get("meta_info"), dict
        ):
            exclude = self.dump_requests_exclude_meta_keys
            if any(k in out_dict["meta_info"] for k in exclude):
                filtered_meta = {
                    k: v for k, v in out_dict["meta_info"].items() if k not in exclude
                }
                out_dict = {**out_dict, "meta_info": filtered_meta}

        self.dump_request_list.append(
            (
                state.obj,
                out_dict,
                convert_time_to_realtime(state.time_stats.created_time),
                convert_time_to_realtime(state.time_stats.finished_time),
            )
        )

        if len(self.dump_request_list) >= self.dump_requests_threshold:
            filename = os.path.join(
                self.dump_requests_folder,
                datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ".pkl",
            )
            self._dump_data_to_file(
                data_list=self.dump_request_list,
                filename=filename,
                log_message=f"Dump {len(self.dump_request_list)} requests to {filename}",
            )
            self.dump_request_list = []

    def record_request_for_crash_dump(self, state: ReqState, out_dict: dict):
        current_time = real_time()
        self.crash_dump_request_list.append(
            (
                state.obj,
                out_dict,
                convert_time_to_realtime(state.time_stats.created_time),
                current_time,
            )
        )
        # Remove requests older than 5 minutes based on finish time
        while (
            self.crash_dump_request_list
            and current_time - self.crash_dump_request_list[0][3] >= 300
        ):
            self.crash_dump_request_list.popleft()

    def _dump_data_to_file(
        self, data_list: List[Tuple], filename: str, log_message: str
    ):
        logger.info(log_message)
        to_dump_with_server_args = {
            "server_args": self.server_args,
            "config_updates": list(self._config_updates),
            "resolved_config": self._dump_config_snapshot(),
            "requests": data_list.copy(),
        }

        def background_task():
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            with open(filename, "wb") as f:
                try:
                    pickle.dump(to_dump_with_server_args, f)
                except Exception as e:
                    # When the server is launched with --trust-remote-code,
                    # server_args sometimes fails to pickle. Retry without
                    # server_args so the request data still gets persisted.
                    logger.error(
                        f"Failed to pickle dump with server_args: {e!r}; "
                        "retrying without server_args"
                    )
                    f.seek(0)
                    f.truncate()
                    to_dump_with_server_args["server_args"] = None
                    # The snapshot copies the same object field by field.
                    to_dump_with_server_args["resolved_config"] = None
                    pickle.dump(to_dump_with_server_args, f)

        asyncio.create_task(asyncio.to_thread(background_task))

    def dump_requests_before_crash(
        self, hostname: str = os.getenv("HOSTNAME", socket.gethostname())
    ):
        should_dump_pyspy = envs.SGLANG_PYSPY_DUMP_BEFORE_CRASH.get()
        should_dump_cuda_coredump = envs.SGLANG_CUDA_COREDUMP_BEFORE_CRASH.get()
        should_dump_diagnostics = should_dump_pyspy or should_dump_cuda_coredump
        if not self.crash_dump_folder and not should_dump_diagnostics:
            return

        if self.crash_dump_performed:
            logger.info(
                "SIGTERM/SIGQUIT/Exception triggered, but crash dump already performed, skipping."
            )
            return
        else:
            self.crash_dump_performed = True

        # Dump requests info
        if self.crash_dump_folder:
            logger.error(f"Dumping requests before crash. {self.crash_dump_folder=}")

            # Add finished requests from crash_dump_request_list
            data_to_dump = []
            if self.crash_dump_request_list:
                data_to_dump.extend(self.crash_dump_request_list)

            # Add unfinished requests from rid_to_state
            unfinished_requests = []
            for rid, state in self.rid_to_state.items():
                if not state.finished:
                    state.time_stats.set_finished_time()
                    unfinished_requests.append(
                        (
                            state.obj,
                            (
                                state.out_list[-1]
                                if state.out_list
                                else state.get_crash_dump_output()
                            ),
                            convert_time_to_realtime(state.time_stats.created_time),
                            convert_time_to_realtime(state.time_stats.finished_time),
                        )
                    )
            if unfinished_requests:
                data_to_dump.extend(unfinished_requests)

            if data_to_dump:
                # Create a file
                filename = os.path.join(
                    self.crash_dump_folder,
                    hostname,
                    f'crash_dump_{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.pkl',
                )
                os.makedirs(os.path.dirname(filename), exist_ok=True)

                # Write the data to the file
                data_to_dump_with_server_args = {
                    "server_args": self.server_args,
                    "config_updates": list(self._config_updates),
                    "resolved_config": self._dump_config_snapshot(),
                    "requests": data_to_dump,
                    "launch_command": " ".join(sys.argv),
                }
                with open(filename, "wb") as f:
                    try:
                        pickle.dump(data_to_dump_with_server_args, f)
                    except Exception as e:
                        # When the server is launched with --trust-remote-code,
                        # server_args sometimes fails to pickle. Retry without
                        # server_args so the request data still gets persisted.
                        logger.error(
                            f"Failed to pickle dump with server_args: {e!r}; "
                            "retrying without server_args"
                        )
                        f.seek(0)
                        f.truncate()
                        data_to_dump_with_server_args["server_args"] = None
                        # The snapshot copies the same object field by field.
                        data_to_dump_with_server_args["resolved_config"] = None
                        pickle.dump(data_to_dump_with_server_args, f)
                logger.error(
                    f"Dumped {len(self.crash_dump_request_list)} finished and {len(unfinished_requests)} unfinished requests before crash to {filename}"
                )

        # Dump pyspy and cuda coredump
        if should_dump_diagnostics:
            logger.info(
                "Sleeping 5 seconds before crash diagnostics to let GPU activity settle."
            )
            time.sleep(5)

            scheduler_procs = collect_scheduler_processes()
            if scheduler_procs:
                if should_dump_pyspy:
                    pyspy_dump_schedulers(scheduler_only=True)

                if should_dump_cuda_coredump:
                    trigger_cuda_user_coredump(scheduler_only=True)
                    cuda_coredump_wait_secs = (
                        envs.SGLANG_CUDA_COREDUMP_BEFORE_CRASH_WAIT_SECS.get()
                    )
                    if cuda_coredump_wait_secs > 0:
                        logger.info(
                            "Waiting %.1f seconds for CUDA coredumps before exiting.",
                            cuda_coredump_wait_secs,
                        )
                        time.sleep(cuda_coredump_wait_secs)
            else:
                logger.error(
                    "No live scheduler processes found; skipping py-spy and CUDA coredump."
                )

    async def sigterm_watchdog(self):
        while not self.gracefully_exit:
            await asyncio.sleep(5)

        # Drain requests
        while True:
            remain_num_req = len(self.rid_to_state)
            remaining_rids = list(self.rid_to_state.keys())

            if self.server_status == ServerStatus.UnHealthy:
                # if health check failed, we should exit immediately
                logger.error(
                    "Signal SIGTERM received while health check failed. Force exiting."
                )
                self.dump_requests_before_crash()
                self.force_exit_handler()
                break

            elif get_bool_env_var("SGL_FORCE_SHUTDOWN"):
                # if force shutdown flag set, exit immediately
                logger.error(
                    "Signal SIGTERM received while force shutdown flag set. Force exiting."
                )
                self.force_exit_handler()
                break

            logger.info(
                f"Gracefully exiting... Remaining number of requests {remain_num_req}. Remaining requests {remaining_rids=}."
            )
            if remain_num_req > 0:
                await asyncio.sleep(5)
            else:
                break

        # Stop the watchdog: child exits are expected during shutdown, not crashes.
        if self._subprocess_watchdog is not None:
            self._subprocess_watchdog.stop()
        # Ask schedulers to release resources in userspace and exit (see
        # ShutdownReq), then wait for them before hard-killing the rest.
        self._dispatch_to_scheduler(ShutdownReq())
        deadline = time.monotonic() + 15
        while time.monotonic() < deadline and collect_scheduler_processes():
            time.sleep(0.1)
        kill_process_tree(os.getpid(), include_parent=True)
        sys.exit(0)

    def force_exit_handler(self):
        """Put some custom force exit logic here."""
        pass

    def _handle_abort_req(self, recv_obj: AbortReq):
        if is_health_check_generate_req(recv_obj):
            return
        # Two scheduler messages can race in handle_loop for the same rid: a
        # batch output that finishes it normally (deletes rid_to_state[rid])
        # and this abort echo. If the finish wins, the rid is already gone and
        # there is nothing left to abort. Common under mass client
        # disconnects, amplified by prefix / abort_all fan-out.
        state = self.rid_to_state.get(recv_obj.rid)
        if state is None:
            logger.info(
                "Abort request for rid=%s not found in rid_to_state; "
                "likely already finished/removed.",
                recv_obj.rid,
            )
            return
        state.finished = True
        state.time_stats.set_finished_time()

        abort_message = recv_obj.abort_message or "Abort in waiting queue"
        finish_reason = {
            "type": "abort",
            "message": abort_message,
        }
        if recv_obj.finished_reason:
            finish_reason = recv_obj.finished_reason
        meta_info = {
            "id": recv_obj.rid,
            "finish_reason": finish_reason,
            "weight_version": self.config_value("weight_version"),
            "e2e_latency": state.time_stats.get_e2e_latency(),
        }
        is_stream = getattr(state.obj, "stream", False)
        if getattr(state.obj, "return_logprob", False):
            self.add_logprob_to_meta_info(
                meta_info,
                state,
                state.obj.top_logprobs_num,
                state.obj.token_ids_logprob,
                state.obj.return_text_in_logprobs
                and not self.server_args.skip_tokenizer_init,
            )

        output_ids = state.output_ids
        meta_info["completion_tokens"] = len(output_ids)
        if is_stream:
            output_ids = [output_ids[-1]] if len(output_ids) > 0 else []
        out = {
            "text": state.get_text(),
            "output_ids": output_ids,
            "meta_info": meta_info,
        }
        self._remove_req_state(recv_obj.rid)

        state.out_list.append(out)
        state.event.set()

    def update_active_ranks(self, ranks: ActiveRanksOutput):
        self._dispatch_to_scheduler(ranks)

    def forward_elastic_scale_update(self, msg: ElasticScaleUpdateReq):
        if not msg.success:
            self.elastic_pending_ep_size = None
            self.elastic_scale_phase = "failed"
            self.elastic_last_error = msg.error
            return

        self._dispatch_to_scheduler(msg)
        self.elastic_worker_count = msg.effective_ep_size
        self.elastic_pending_ep_size = None
        self.elastic_scale_phase = "serving_expanded"
        self.elastic_last_error = None
        self.update_control_communicator_fan_out(msg.effective_ep_size)

    def get_elastic_ep_state(self):
        return {
            "is_scaling_elastic_ep": self.elastic_pending_ep_size is not None,
            "effective_ep_size": self.elastic_worker_count,
            "pending_ep_size": self.elastic_pending_ep_size,
            "scale_phase": self.elastic_scale_phase,
            "last_error": self.elastic_last_error,
        }

    async def scale_elastic_ep(
        self, obj: ScaleElasticEPReqInput
    ) -> ScaleElasticEPReqOutput:
        """Send a scale request to every DP scheduler."""
        if self.elastic_pending_ep_size is not None:
            return ScaleElasticEPReqOutput(
                success=False,
                message=(
                    "A previous scale operation has not completed yet. Wait until "
                    "all pending ranks have joined before issuing another scale."
                ),
                old_ep_size=self.elastic_worker_count,
                new_ep_size=obj.new_ep_size,
                pending_ep_size=self.elastic_pending_ep_size,
                scale_phase=self.elastic_scale_phase,
            )
        self.auto_create_handle_loop()
        responses: List[ScaleElasticEPReqOutput] = (
            await self.scale_elastic_ep_communicator(obj)
        )
        for res in responses:
            if not res.success:
                self.elastic_scale_phase = res.scale_phase
                self.elastic_pending_ep_size = res.pending_ep_size
                self.elastic_last_error = res.message
                return res
        self.elastic_pending_ep_size = responses[0].pending_ep_size
        self.elastic_scale_phase = responses[0].scale_phase
        self.elastic_last_error = None
        return responses[0]

    def _handle_open_session_req_output(self, recv_obj):
        future = self.session_futures.get(recv_obj.session_id)
        if future is None:
            logger.warning(
                "Open session response arrived after waiter cleanup: %s",
                recv_obj.session_id,
            )
            return
        if not future.done():
            future.set_result(recv_obj.session_id if recv_obj.success else None)

    def _handle_update_weights_from_disk_req_output(self, recv_obj):
        if self.model_update_expected_workers == 1:
            self.model_update_result.set_result(recv_obj)
        else:
            self.model_update_tmp.append(recv_obj)
            if len(self.model_update_tmp) == self.model_update_expected_workers:
                self.model_update_result.set_result(self.model_update_tmp)

    async def _validate_and_resolve_lora(
        self, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> None:
        if not obj.lora_path:
            return

        if not self.enable_lora:
            first_adapter = (
                obj.lora_path
                if isinstance(obj.lora_path, str)
                else next((a for a in obj.lora_path if a), None)
            )

            raise ValueError(
                f"LoRA adapter '{first_adapter}' was requested, but LoRA is not enabled. "
                "Please launch the server with --enable-lora flag and preload adapters "
                "using --lora-paths or /load_lora_adapter endpoint."
            )

        await self._resolve_lora_path(obj)

    async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInput]):
        if isinstance(obj.lora_path, str):
            unique_lora_paths = set([obj.lora_path])
        else:
            unique_lora_paths = set(obj.lora_path)

        if (
            self.server_args.max_loaded_loras is not None
            and len(unique_lora_paths) > self.server_args.max_loaded_loras
        ):
            raise ValueError(
                f"Received request with {len(unique_lora_paths)} unique loras requested "
                f"but max loaded loras is {self.server_args.max_loaded_loras}"
            )

        # Reload all existing LoRA adapters that have been dynamically unloaded
        unregistered_loras = await self.lora_registry.get_unregistered_loras(
            unique_lora_paths
        )
        for lora_path in unregistered_loras:
            if lora_path is None:
                continue

            if lora_path not in self.lora_ref_cache:
                raise ValueError(
                    f"Got LoRA adapter that has never been loaded: {lora_path}\n"
                    f"All loaded adapters: {self.lora_ref_cache.keys()}."
                )

            logger.info(f"Reloading evicted adapter: {lora_path}")
            new_lora_ref = self.lora_ref_cache[lora_path]
            load_result = await self.load_lora_adapter(
                LoadLoRAAdapterReqInput(
                    lora_name=new_lora_ref.lora_name,
                    lora_path=new_lora_ref.lora_path,
                    pinned=new_lora_ref.pinned,
                )
            )
            if (
                not load_result.success
                and "already loaded" not in load_result.error_message
            ):
                raise ValueError(
                    f"Failed to implicitly load LoRA adapter {lora_path}: {load_result.error_message}"
                )

        # Look up the LoRA ID from the registry and start tracking ongoing LoRA requests.
        obj.lora_id = await self.lora_registry.acquire(obj.lora_path)
        # Propagate lora_id to any sub-objects already cached by __getitem__.
        for i, sub_obj in obj.__dict__.get("_sub_obj_cache", {}).items():
            sub_obj.lora_id = (
                obj.lora_id[i] if isinstance(obj.lora_id, list) else obj.lora_id
            )

    @staticmethod
    def _logical_rids(obj) -> List[str]:
        if not hasattr(obj, "is_single") or obj.is_single:
            return [obj.rid]
        return list(obj.rid)

    def _register_child_rid(self, logical_rid: str, child_rid: str) -> None:
        if child_rid == logical_rid:
            raise ValueError(
                "Parallel-sampling child RID must differ from its logical RID"
            )
        owner = self.child_rid_to_logical_rid.get(child_rid)
        if owner is not None and owner != logical_rid:
            raise ValueError(
                f"Request ID {child_rid} is already owned by logical request {owner}"
            )
        self.child_rid_to_logical_rid[child_rid] = logical_rid
        self.logical_rid_to_child_rids.setdefault(logical_rid, set()).add(child_rid)

    def _init_child_req_state(
        self,
        logical_rid: str,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        request: Optional[fastapi.Request] = None,
    ) -> None:
        self._raise_if_logical_rid_aborted(logical_rid)
        logical_state = self.rid_to_state[logical_rid]
        self._init_req_state(
            obj,
            request,
            lifecycle_id=logical_state.lifecycle_id,
        )
        try:
            self._register_child_rid(logical_rid, obj.rid)
        except BaseException:
            self._remove_req_state(obj.rid)
            raise

    def _remove_req_state(
        self,
        rid: str,
        lifecycle_id: Optional[object] = None,
    ) -> Optional[ReqState]:
        """Remove a request state and its parallel-sampling ownership."""
        state = self.rid_to_state.get(rid)
        if state is None or (
            lifecycle_id is not None and state.lifecycle_id is not lifecycle_id
        ):
            return None
        self.rid_to_state.pop(rid)
        logical_rid = self.child_rid_to_logical_rid.pop(rid, None)
        if logical_rid is not None:
            children = self.logical_rid_to_child_rids.get(logical_rid)
            if children is not None:
                children.discard(rid)
                if not children:
                    self.logical_rid_to_child_rids.pop(logical_rid, None)
        return state

    def _raise_if_logical_rid_aborted(self, logical_rid: str) -> None:
        state = self.rid_to_state.get(logical_rid)
        if state is None or state.abort_requested:
            raise RequestAbortedError(f"Request {logical_rid} was aborted")

    def _raise_if_logical_request_aborted(self, obj) -> None:
        for logical_rid in self._logical_rids(obj):
            self._raise_if_logical_rid_aborted(logical_rid)

    def _init_req_state(
        self,
        obj: Union[GenerateReqInput, EmbeddingReqInput],
        request: Optional[fastapi.Request] = None,
        lifecycle_id: Optional[object] = None,
    ) -> Dict[str, object]:
        created_time = obj.received_time

        external_trace_header = None
        if self.enable_trace:
            if obj.external_trace_header:
                # When the request comes from the rust grpc server or Engine there isn't a
                # real request object but we still need to propagate the trace context from
                # the trace context that is explicitly passed in
                external_trace_header = obj.external_trace_header
            elif request:
                external_trace_header = extract_trace_headers(request.headers)
                obj.external_trace_header = external_trace_header

        # Normalize single/batch into a uniform list of (rid, sub_obj, bootstrap_room)
        if not hasattr(obj, "is_single") or obj.is_single:
            items = [(obj.rid, obj, getattr(obj, "bootstrap_room", None))]
        else:
            items = [
                (
                    obj.rid[i],
                    obj[i],
                    (
                        obj.bootstrap_room[i]
                        if hasattr(obj, "bootstrap_room") and obj.bootstrap_room
                        else None
                    ),
                )
                for i in range(len(obj.rid))
            ]

        rids = [rid for rid, _, _ in items]
        seen_rids = set()
        for rid in rids:
            if rid in seen_rids:
                raise ValueError(f"Duplicate request ID detected: {rid}")
            seen_rids.add(rid)
            if (
                rid in self.rid_to_state
                or rid in self.logical_rid_to_child_rids
                or rid in self.child_rid_to_logical_rid
            ):
                raise ValueError(f"Duplicate request ID detected: {rid}")

        # Mutate only after every RID passes duplicate validation so a rejected
        # batch cannot leave a partial rid_to_state insertion behind.
        lifecycle_ids = {}
        for rid, sub_obj, bootstrap_room in items:
            time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode)
            state = ReqState(
                [],
                False,
                asyncio.Event(),
                sub_obj,
                time_stats,
                lifecycle_id=lifecycle_id if lifecycle_id is not None else object(),
            )
            self.rid_to_state[rid] = state
            lifecycle_ids[rid] = state.lifecycle_id
            if self.enable_trace:
                time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header)
            time_stats.set_created_time(created_time)
        return lifecycle_ids

    def _discard_pending_req_states(
        self,
        obj,
        lifecycle_ids: Optional[Dict[str, object]] = None,
    ):
        """Drop all logical and child state owned by *obj*.

        Safe to call after a partial/failed dispatch: only requests known to have
        reached the scheduler are aborted, all owned state is removed, and a later
        output for a discarded RID is ignored by the scheduler-response path.
        """
        if lifecycle_ids is None:
            lifecycle_ids = {
                logical_rid: state.lifecycle_id
                for logical_rid in self._logical_rids(obj)
                if (state := self.rid_to_state.get(logical_rid)) is not None
            }
        for logical_rid in self._logical_rids(obj):
            lifecycle_id = lifecycle_ids.get(logical_rid)
            if lifecycle_id is None:
                continue
            child_rids = tuple(
                child_rid
                for child_rid in self.logical_rid_to_child_rids.get(logical_rid, ())
                if (
                    (state := self.rid_to_state.get(child_rid)) is not None
                    and state.lifecycle_id is lifecycle_id
                )
            )
            logical_state = self.rid_to_state.get(logical_rid)
            owns_logical_state = (
                logical_state is not None and logical_state.lifecycle_id is lifecycle_id
            )
            target_rids = tuple(
                rid for rid in child_rids if self.rid_to_state[rid].dispatched
            )
            if not child_rids and owns_logical_state and logical_state.dispatched:
                target_rids = (logical_rid,)
            for target_rid in target_rids:
                try:
                    self._dispatch_to_scheduler(
                        AbortReq(rid=target_rid, abort_all=False)
                    )
                except Exception:
                    logger.exception(
                        "Failed to abort request rid=%s",
                        target_rid,
                    )
            for child_rid in child_rids:
                self._remove_req_state(child_rid, lifecycle_id)
            self._remove_req_state(logical_rid, lifecycle_id)

    def _should_dispatch_to_encoder(
        self, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ) -> bool:
        """Check if the request should be dispatched to encoder for processing.

        Returns True if the request should be dispatched to encoder (multiple multimodal items),
        False if it should be processed locally (single multimodal item or no multimodal items).

        Args:
            obj: The request input object

        Returns:
            bool: True if should dispatch to encoder, False otherwise
        """
        if obj.batch_size > 1:
            logger.warning(
                "Batch request (batch_size=%d) is not supported in EPD disaggregation mode; skipping encoder dispatch.",
                obj.batch_size,
            )
            return False
        if not isinstance(obj, GenerateReqInput) or not obj.contains_mm_input():
            return False

        # Count image / video / audio items for dispatch threshold
        def _count_mm_items(data):
            return (
                len(data) if isinstance(data, list) else (1 if data is not None else 0)
            )

        total_mm_items = (
            _count_mm_items(getattr(obj, "image_data", None))
            + _count_mm_items(getattr(obj, "video_data", None))
            + _count_mm_items(getattr(obj, "audio_data", None))
        )
        return total_mm_items >= envs.SGLANG_ENCODER_DISPATCH_MIN_ITEMS.get()

    def _handle_epd_disaggregation_encode_request(
        self, obj: Union[GenerateReqInput, EmbeddingReqInput]
    ):
        """Handle EPD-disaggregation mode encoding request."""
        if isinstance(obj, GenerateReqInput) and obj.contains_mm_input():
            # dispatch to encoder by default
            should_dispatch = True
            if self.server_args.enable_adaptive_dispatch_to_encoder:
                should_dispatch = self._should_dispatch_to_encoder(obj)

            # Set need_wait_for_mm_inputs flag based on whether we dispatch to encoder
            # This flag will be used in _tokenize_one_request to determine processing path
            if should_dispatch:
                obj.need_wait_for_mm_inputs = True
                if self.server_args.encoder_transfer_backend in [
                    "zmq_to_scheduler",
                    "mooncake",
                ]:
                    time_stats_json = None
                    if self.enable_trace:
                        state = self.rid_to_state.get(obj.rid)
                        if state is not None:
                            time_stats_json = state.time_stats.encode_json()

                    self.mm_receiver.send_encode_request(
                        obj, time_stats_json=time_stats_json
                    )
            else:
                obj.need_wait_for_mm_inputs = False

    def convert_to_span_attrs(
        self,
        state: ReqState,
        recv_obj: Union[
            BatchStrOutput,
            BatchEmbeddingOutput,
            BatchTokenIDOutput,
        ],
        i: int,
    ) -> Dict[str, Any]:
        """Convert attributes to span attributes."""
        span_attrs = {}

        if not self.enable_trace:
            return span_attrs

        # Token usage attributes
        if not isinstance(recv_obj, BatchEmbeddingOutput):
            span_attrs[SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS] = (
                recv_obj.completion_tokens[i]
            )
        span_attrs[SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS] = recv_obj.prompt_tokens[
            i
        ]
        span_attrs[SpanAttributes.GEN_AI_USAGE_CACHED_TOKENS] = recv_obj.cached_tokens[
            i
        ]

        # Request identifiers
        span_attrs[SpanAttributes.GEN_AI_REQUEST_ID] = (
            str(state.obj.rid) if state.obj.rid else None
        )

        # Sampling parameters
        sampling_params = state.obj.sampling_params or {}

        if max_new_tokens := sampling_params.get("max_new_tokens"):
            span_attrs[SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS] = max_new_tokens

        if top_p := sampling_params.get("top_p"):
            span_attrs[SpanAttributes.GEN_AI_REQUEST_TOP_P] = top_p

        if temperature := sampling_params.get("temperature"):
            span_attrs[SpanAttributes.GEN_AI_REQUEST_TEMPERATURE] = temperature

        if top_k := sampling_params.get("top_k"):
            span_attrs[SpanAttributes.GEN_AI_REQUEST_TOP_K] = top_k

        if n := sampling_params.get("n"):
            span_attrs[SpanAttributes.GEN_AI_REQUEST_N] = n

        # Response attributes
        span_attrs[SpanAttributes.GEN_AI_RESPONSE_MODEL] = self.served_model_name

        finish_reason = (
            recv_obj.finished_reasons[i].get("type")
            if recv_obj.finished_reasons[i]
            else None
        )
        if finish_reason:
            span_attrs[SpanAttributes.GEN_AI_RESPONSE_FINISH_REASONS] = json.dumps(
                [finish_reason]
            )

        # Latency attributes
        span_attrs.update(state.time_stats.convert_to_gen_ai_span_attrs())

        return span_attrs

    def _set_default_priority(self, obj: Union[GenerateReqInput, EmbeddingReqInput]):
        """Set the default priority value."""
        if (
            self.enable_priority_scheduling
            and obj.priority is None
            and self.default_priority_value is not None
        ):
            obj.priority = self.default_priority_value


class ServerStatus(Enum):
    Up = "Up"
    Starting = "Starting"
    UnHealthy = "UnHealthy"


async def print_exception_wrapper(func):
    """
    Sometimes an asyncio function does not print exception.
    We do another wrapper to handle the exception.
    """
    try:
        await func()
    except Exception:
        traceback = get_exception_traceback()
        logger.error(f"TokenizerManager hit an exception: {traceback}")
        if hasattr(func, "__self__") and isinstance(func.__self__, TokenizerManager):
            func.__self__.dump_requests_before_crash()
        kill_process_tree(os.getpid(), include_parent=True)
        sys.exit(1)


def get_processor_wrapper(server_args):
    try:
        processor = get_processor(
            server_args.tokenizer_path,
            tokenizer_mode=server_args.tokenizer_mode,
            trust_remote_code=server_args.trust_remote_code,
            revision=server_args.revision,
            use_fast=not server_args.disable_fast_image_processor,
            tokenizer_backend=server_args.tokenizer_backend,
            model_name=server_args.model_path,
        )
    except ValueError as e:
        error_message = str(e)
        if "does not have a slow version" in error_message:
            logger.info(
                f"Processor {server_args.tokenizer_path} does not have a slow version. Automatically use fast version"
            )
            processor = get_processor(
                server_args.tokenizer_path,
                tokenizer_mode=server_args.tokenizer_mode,
                trust_remote_code=server_args.trust_remote_code,
                revision=server_args.revision,
                use_fast=True,
                tokenizer_backend=server_args.tokenizer_backend,
                model_name=server_args.model_path,
            )
        else:
            raise e
    return processor


def determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode:
    is_cross_node = server_args.dist_init_addr

    if is_cross_node:
        # Fallback to default CPU transport for multi-node
        return "default"
    else:
        return "cuda_ipc"


class SignalHandler:
    def __init__(self, tokenizer_manager: TokenizerManager):
        self.tokenizer_manager = tokenizer_manager

    def sigterm_handler(self, signum=None, frame=None):
        logger.warning(
            f"SIGTERM received. {signum=} {frame=}. Draining requests and shutting down..."
        )
        self.tokenizer_manager.gracefully_exit = True

    def running_phase_sigquit_handler(self, signum=None, frame=None):
        logger.error(
            f"SIGQUIT received. {signum=}, {frame=}. It usually means one child failed."
        )
        # Stop subprocess watchdog before killing processes to prevent false-positive
        # crash detection during normal shutdown
        if self.tokenizer_manager._subprocess_watchdog is not None:
            self.tokenizer_manager._subprocess_watchdog.stop()
        self.tokenizer_manager.dump_requests_before_crash()
        kill_process_tree(os.getpid())


# Note: request abort handling logic
# We should handle all of the following cases correctly.
#
# | entrypoint | is_streaming | status          | abort engine    | cancel asyncio task   | rid_to_state                |
# | ---------- | ------------ | --------------- | --------------- | --------------------- | --------------------------- |
# | http       | yes          | validation      | background task | fast api              | del in _handle_abort_req    |
# | http       | yes          | waiting queue   | background task | fast api              | del in _handle_abort_req    |
# | http       | yes          | running         | background task | fast api              | del in _handle_batch_output |
# | http       | no           | validation      | http exception  | http exception        | del in _handle_abort_req    |
# | http       | no           | waiting queue   | type 1          | type 1 exception      | del in _handle_abort_req    |
# | http       | no           | running         | type 3          | type 3 exception      | del in _handle_batch_output |
#


def stamp_http_worker_ipc(obj: Any, ipc_name: str) -> None:
    if isinstance(obj, BaseReq):
        obj.http_worker_ipc = ipc_name
    elif isinstance(
        obj, (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput)
    ):
        for req in obj:
            req.http_worker_ipc = ipc_name
    elif isinstance(obj, BaseBatchReq):
        obj.http_worker_ipcs = [ipc_name] * len(obj.rids)
