from __future__ import annotations

from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.runtime_context import (
    get_disagg,
    get_exec,
    get_parallel,
    get_schedule,
    get_serving,
    get_spec,
    mamba_cache_chunk_size,
    mamba_checkpoint_grid,
    mamba_track_grid,
)
from sglang.srt.utils.common import (
    ceil_align,
    flatten_arrays_to_pinned_cpu,
    is_hip,
    is_pin_memory_available,
)
from sglang.srt.utils.weight_versions import (
    WeightVersionEvent,
    truncate_weight_version_events,
)

# 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.
# ==============================================================================
"""
Store information about requests and batches.

The following is the flow of data structures for a batch:

ScheduleBatch -> ForwardBatch

- ScheduleBatch is managed by `scheduler.py::Scheduler`.
  It contains high-level scheduling data. Most of the data is on the CPU.
- ForwardBatch is managed by `model_runner.py::ModelRunner`.
  It contains low-level tensor data. Most of the data consists of GPU tensors.
  It is constructed directly from a ScheduleBatch by `ForwardBatch.init_new`.
"""

import copy
import dataclasses
import logging
import re
import sys
from array import array
from concurrent.futures import Future
from enum import Enum, auto
from functools import lru_cache
from http import HTTPStatus
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Literal,
    NamedTuple,
    Optional,
    Set,
    Tuple,
    TypeAlias,
    Union,
)

import msgspec
import numpy as np
import torch

from sglang.srt.beam_search.batch_tail import (
    BeamTail,
    append_beam_tail,
    beam_retraction_order,
    num_beam_member_rows,
    strip_beam_tail,
)
from sglang.srt.beam_search.fork import free_member_rows
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.disaggregation.base import BaseKVSender
from sglang.srt.disaggregation.decode_schedule_batch_mixin import (
    ScheduleBatchDisaggregationDecodeMixin,
)
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode
from sglang.srt.dllm.mixin.req import ReqDllmMixin
from sglang.srt.environ import envs
from sglang.srt.managers.embed_types import PositionalEmbeds
from sglang.srt.managers.kv_hints import KvHintsEnvelope
from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import (
    NewTokenRatioTracker,
)
from sglang.srt.mem_cache.allocation import alloc_for_decode, alloc_for_extend
from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
    BasePrefixCache,
    CacheRequestHandle,
    TreeLock,
)
from sglang.srt.mem_cache.common import (
    RetractionBackup,
    backup_kv_cache,
    evict_from_tree_cache,
    match_kv_cache,
    release_kv_cache,
)
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
from sglang.srt.model_executor.forward_batch_info import (
    CaptureHiddenMode,
    ForwardBatch,
    ForwardMode,
)
from sglang.srt.multimodal.transport.cuda_ipc import (
    CUDA_IPC_FEATURE_COPY_EVENT_KEY,
    DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
    RETAINED_CUDA_IPC_FEATURE_PROXY_KEY,
    CudaIpcTensorTransportProxy,
)
from sglang.srt.observability.metrics_collector import (
    DPBalanceStats,
    DPCooperationInfo,
    SchedulerMetricsCollector,
)
from sglang.srt.observability.req_time_stats import (
    APIServerReqTimeStats,
    DPControllerReqTimeStats,
    SchedulerReqTimeStats,
)
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_mask import SamplingMaskRows
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.utils import flatten_nested_list
from sglang.srt.utils.token_sequence_matcher import TokenSequenceMatcher

if TYPE_CHECKING:
    from typing import Any, Dict

    from sglang.srt.configs.model_config import ModelConfig
    from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
    from sglang.srt.managers.scheduler_components.metrics_reporter import PrefillStats
    from sglang.srt.mem_cache.storage_prefetch import StagedPrefetchPlan
    from sglang.srt.session.session_controller import Session
    from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm

INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5

# Constant used as the base offset for MM (multimodal) pad values.
# This ensures pad_values don't overlap with valid text token IDs.
MM_PAD_SHIFT_VALUE = 1_000_000
_MM_HASH_MASK = (1 << 64) - 1

logger = logging.getLogger(__name__)

_is_hip = is_hip()


ReturnHiddenStatesMode = Union[bool, Literal["last"]]
SamplingLogprobsMode = Literal["selected", "support"]


def get_return_hidden_states_mode(
    return_hidden_states: ReturnHiddenStatesMode,
) -> CaptureHiddenMode:
    if return_hidden_states is True:
        return CaptureHiddenMode.FULL
    if return_hidden_states == "last":
        return CaptureHiddenMode.LAST
    if return_hidden_states is False:
        return CaptureHiddenMode.NULL
    raise ValueError(
        "return_hidden_states must be a boolean or the string literal 'last'."
    )


def get_request_return_hidden_states_mode(
    return_hidden_states: Union[List[ReturnHiddenStatesMode], ReturnHiddenStatesMode],
) -> CaptureHiddenMode:
    if isinstance(return_hidden_states, list):
        return max(
            (get_return_hidden_states_mode(mode) for mode in return_hidden_states),
            default=CaptureHiddenMode.NULL,
        )
    return get_return_hidden_states_mode(return_hidden_states)


def get_batch_return_hidden_states_mode(reqs: List[Req]) -> CaptureHiddenMode:
    mode = CaptureHiddenMode.NULL
    for req in reqs:
        mode = max(mode, req.return_hidden_states_mode)
    return mode


def need_return_hidden_states(
    return_hidden_states: Union[List[ReturnHiddenStatesMode], ReturnHiddenStatesMode],
) -> bool:
    return get_request_return_hidden_states_mode(return_hidden_states).need_capture()


@lru_cache(maxsize=1)
def sanity_check_mm_pad_shift_value(vocab_size: int) -> None:
    if vocab_size > MM_PAD_SHIFT_VALUE:
        raise ValueError(
            f"Model vocab_size ({vocab_size}) exceeds MM_PAD_SHIFT_VALUE ({MM_PAD_SHIFT_VALUE}). "
            f"MM pad_values may overlap with valid token IDs. "
            f"Please increase MM_PAD_SHIFT_VALUE in schedule_batch.py."
        )


def split_cached_prefix_by_tier(
    prefix_len: int,
    host_hit_len: int,
    storage_hit_len: int,
    storage_hit_start: Optional[int] = None,
    host_hit_is_storage: bool = False,
) -> tuple[int, int, int]:
    """Split a request's cached prefix into (device, host, storage) tokens.

    ``host_hit_len`` is the materialized host hit. An exact L3-loaded span is
    preserved after H2D promotion, while an evicted tail is clipped by
    ``prefix_len``. In buffer mode host memory is only L3 staging.
    """
    host_hit_len = min(prefix_len, host_hit_len)
    host_start = prefix_len - host_hit_len
    if storage_hit_start is None:
        if host_hit_is_storage:
            storage = host_hit_len
            host = 0
        else:
            storage = min(host_hit_len, storage_hit_len)
            host = host_hit_len - storage
    else:
        storage_end = storage_hit_start + storage_hit_len
        storage = max(
            0,
            min(prefix_len, storage_end) - max(0, storage_hit_start),
        )
        storage_in_host = max(
            0,
            min(prefix_len, storage_end) - max(host_start, storage_hit_start),
        )
        host = 0 if host_hit_is_storage else host_hit_len - storage_in_host
    device = prefix_len - host - storage
    return device, host, storage


def _compute_pad_value(hash: int) -> int:
    """Compute pad value from hash."""
    return MM_PAD_SHIFT_VALUE + (hash % (1 << 30))


class BaseFinishReason:
    def to_json(self):
        raise NotImplementedError()


class FINISH_MATCHED_TOKEN(BaseFinishReason):
    def __init__(self, matched: Union[int, List[int]]):
        super().__init__()
        self.matched = matched

    def to_json(self):
        return {
            "type": "stop",  # to match OpenAI API's return value
            "matched": self.matched,
        }


class FINISH_MATCHED_STR(BaseFinishReason):
    def __init__(self, matched: str):
        super().__init__()
        self.matched = matched

    def to_json(self):
        return {
            "type": "stop",  # to match OpenAI API's return value
            "matched": self.matched,
        }


class FINISHED_MATCHED_REGEX(BaseFinishReason):
    def __init__(self, matched: str):
        super().__init__()
        self.matched = matched

    def to_json(self):
        return {
            "type": "stop",  # to match OpenAI API's return value
            "matched": self.matched,
        }


class FINISH_LENGTH(BaseFinishReason):
    def __init__(self, length: int):
        super().__init__()
        self.length = length

    def to_json(self):
        return {
            "type": "length",  # to match OpenAI API's return value
            "length": self.length,
        }


class FINISH_ABORT(BaseFinishReason):
    def __init__(self, message=None, status_code=None, err_type=None):
        super().__init__()
        self.message = message or "Aborted"
        self.status_code = status_code
        self.err_type = err_type

    def to_json(self):
        return {
            "type": "abort",
            "message": self.message,
            "status_code": self.status_code,
            "err_type": self.err_type,
        }


class Modality(Enum):
    IMAGE = auto()
    VIDEO = auto()
    AUDIO = auto()

    @staticmethod
    def from_str(modality_str: str):
        try:
            return Modality[modality_str.upper()]
        except KeyError:
            raise ValueError(
                f"Invalid modality string: {modality_str}. Valid modalities are: {[m.name for m in Modality]}"
            )

    @staticmethod
    def all():
        return [Modality.IMAGE, Modality.VIDEO, Modality.AUDIO]


class MultimodalInputFormat(Enum):
    NORMAL = auto()
    PROCESSOR_OUTPUT = auto()
    PRECOMPUTED_EMBEDDING = auto()


# Msgpack-native containers and Ext-decoded tensor/transport leaves. Tuple
# containers intentionally decode as lists, matching msgpack's native model.
MultimodalDataValue: TypeAlias = object


class MultimodalDataItem(msgspec.Struct, kw_only=True, dict=True, array_like=True):
    """
    One MultimodalDataItem represents a single multimodal input (one image, one video, or one audio).
    For example, if there are 3 images and 1 audio, there will be 4 MultimodalDataItems.

    Each item has its own hash and pad_value, enabling per-image RadixAttention caching.

    We put the common fields first and the model-specific fields in model_specific_data.
    """

    modality: Modality
    hash: Optional[int] = None
    pad_value: Optional[int] = None
    offsets: Optional[List[Tuple[int, int]]] = None

    format: MultimodalInputFormat = MultimodalInputFormat.NORMAL

    # the raw features returned by processor, e.g. pixel_values or audio_features
    feature: Optional[MultimodalDataValue] = None
    # the precomputed embeddings, passed as final encoder embeddings
    # One and only one of the feature and precomputed_embeddings will be empty
    precomputed_embeddings: Optional[MultimodalDataValue] = None
    # Keep precomputed_embeddings on GPU after use (EPD pool/GPU receive path)
    keep_device_embedding: bool = False

    # Processor-owned tensors/arrays/scalars/transports. msgspec rejects a
    # precise union with multiple custom types, but accepts Ext-decoded values
    # under object.
    model_specific_data: Dict[str, MultimodalDataValue] = msgspec.field(
        default_factory=dict
    )

    def __post_init__(self) -> None:
        if self.hash is not None:
            msgspec.Struct.__setattr__(self, "hash", self.hash & _MM_HASH_MASK)

    def __getattr__(self, name: str) -> MultimodalDataValue:
        if name in self.model_specific_data:
            return self.model_specific_data[name]

        raise AttributeError(
            f"'{self.__class__.__name__}' object has no attribute '{name}'"
        )

    def __setattr__(self, name: str, value: MultimodalDataValue) -> None:
        if name in self.__struct_fields__:
            if name == "hash" and isinstance(value, int):
                value &= _MM_HASH_MASK
            msgspec.Struct.__setattr__(self, name, value)
        else:
            self.model_specific_data[name] = value

    def __setitem__(self, key: str, value: MultimodalDataValue) -> None:
        setattr(self, key, value)

    def set(self, key: str, value: MultimodalDataValue) -> None:
        self.__setitem__(key, value)

    def set_hash(self, hash_value: int) -> None:
        self.hash = hash_value
        self.pad_value = _compute_pad_value(hash_value)

    @staticmethod
    def is_empty_list(l):
        if l is None:
            return True
        return len([item for item in flatten_nested_list(l) if item is not None]) == 0

    def set_pad_value(self):
        if self.pad_value is not None:
            return

        from sglang.srt.multimodal.cache import resolve_multimodal_item_hash

        self.hash = resolve_multimodal_item_hash(
            existing_hash=self.hash,
            feature=self.feature,
            precomputed_embeddings=self.precomputed_embeddings,
        )
        self.pad_value = _compute_pad_value(self.hash)

    def is_modality(self, modality: Modality) -> bool:
        return self.modality == modality

    def is_audio(self):
        return self.modality == Modality.AUDIO

    def is_image(self):
        return self.modality == Modality.IMAGE

    def is_video(self):
        return self.modality == Modality.VIDEO

    def is_valid(self) -> bool:
        return self.is_image() or self.is_video() or self.is_audio()

    def validate(self):
        ...
        # TODO

    def is_precomputed_embedding(self):
        return self.format == MultimodalInputFormat.PRECOMPUTED_EMBEDDING

    @staticmethod
    def from_dict(obj: dict):
        kwargs = dict(obj)
        modality = kwargs.pop("modality")
        if isinstance(modality, str):
            modality = Modality[modality]
        ret = MultimodalDataItem(modality=modality, **kwargs)
        ret.validate()
        return ret

    def has_cuda_ipc_proxy(self):
        return (
            isinstance(self.feature, CudaIpcTensorTransportProxy)
            or isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy)
            or any(
                isinstance(value, CudaIpcTensorTransportProxy)
                for value in self.model_specific_data.values()
            )
        )

    def reconstruct(self, target_device: int, ipc_consumer_count: int = 1):
        """materialize cuda ipc proxy tensors in-place on target_device"""
        if isinstance(self.feature, CudaIpcTensorTransportProxy):
            consumer_count = self._resolve_transport_consumer_count(
                self.feature, ipc_consumer_count
            )
            if consumer_count == 1:
                self.feature = self.feature.reconstruct_on_target_device(target_device)
            else:
                self.feature = self.feature.reconstruct_on_target_device(
                    target_device, consumer_count=consumer_count
                )
        if isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy):
            self.precomputed_embeddings = (
                self.precomputed_embeddings.reconstruct_on_target_device(target_device)
            )
        for extra_key in self.model_specific_data:
            if extra_key == RETAINED_CUDA_IPC_FEATURE_PROXY_KEY:
                continue
            if isinstance(
                self.model_specific_data[extra_key], CudaIpcTensorTransportProxy
            ):
                extra_data = self.model_specific_data[
                    extra_key
                ].reconstruct_on_target_device(target_device)
                self.model_specific_data[extra_key] = extra_data

    def can_defer_cuda_ipc_feature_reconstruction(self) -> bool:
        """Whether a DP-aware model will materialize this feature lazily.

        Hashing and pad-value generation must already have completed on the
        tokenizer worker.  Any additional IPC proxy would still need eager
        reconstruction, so keep the narrow fast path feature-only.
        """
        return (
            self.model_specific_data.get(
                DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, False
            )
            and self.hash is not None
            and self.pad_value is not None
            and isinstance(self.feature, CudaIpcTensorTransportProxy)
            and not isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy)
            and not any(
                isinstance(value, CudaIpcTensorTransportProxy)
                for value in self.model_specific_data.values()
            )
        )

    def acknowledge_deferred_cuda_ipc_feature(self, consumer_count: int = 1):
        """Release a lazy IPC feature when an embedding-cache hit skips ViT."""
        if isinstance(self.feature, CudaIpcTensorTransportProxy):
            consumer_count = self._resolve_transport_consumer_count(
                self.feature, consumer_count
            )
            self.feature.acknowledge_consumption(consumer_count)

    def release_transport_proxies(self, consumer_count: int = 1) -> None:
        """Best-effort release of proxies left by an abandoned request."""
        retained_proxy = self.model_specific_data.pop(
            RETAINED_CUDA_IPC_FEATURE_PROXY_KEY, None
        )
        self.model_specific_data.pop(CUDA_IPC_FEATURE_COPY_EVENT_KEY, None)
        values = [self.feature, self.precomputed_embeddings, retained_proxy]
        values.extend(self.model_specific_data.values())
        for value in values:
            if not isinstance(value, CudaIpcTensorTransportProxy):
                continue
            count = self._resolve_transport_consumer_count(value, consumer_count)
            try:
                value.release_without_reconstruction(count)
            except Exception:
                logger.warning(
                    "Failed to release an abandoned multimodal transport proxy",
                    exc_info=True,
                )

    @staticmethod
    def _resolve_transport_consumer_count(proxy, requested_count: int) -> int:
        """Clamp a group acknowledgement to the proxy's actual consumer set."""
        proxy_count = getattr(
            proxy,
            "total_consumer_count",
            getattr(proxy, "consumer_count", requested_count),
        )
        return min(requested_count, proxy_count)


class MultimodalProcessorOutput(
    msgspec.Struct, kw_only=True, dict=True, array_like=True, weakref=True
):
    """Raw output from multimodal processors before scheduler-side preparation (pad, hash).

    This is the typed replacement for the dict previously returned by
    ``BaseMultimodalProcessor.process_mm_data_async``.  Preprocessed inputs may
    already carry ``pad_value`` and ``hash`` to avoid hashing the same tensor once
    per scheduler TP rank.
    """

    mm_items: List[MultimodalDataItem]
    input_ids: Optional[List[int]] = None
    padded_input_ids: Optional[List[int]] = None

    # image
    im_token_id: Optional[int] = None
    im_start_id: Optional[int] = None
    im_end_id: Optional[int] = None
    slice_start_id: Optional[int] = None
    slice_end_id: Optional[int] = None

    # video
    video_token_id: Optional[int] = None

    # audio
    audio_token_id: Optional[int] = None
    audio_start_id: Optional[int] = None
    audio_end_id: Optional[int] = None

    # QWen2-VL related
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[torch.Tensor] = None

    # Moss-VL related
    vision_position_ids: Optional[torch.Tensor] = None
    media_nums_per_sample: Optional[List[int]] = None
    visible_frame_counts: Optional[torch.Tensor] = None

    # for transformers-compatibility
    token_type_ids: Optional[torch.Tensor] = None

    @staticmethod
    def from_dict(d: dict) -> MultimodalProcessorOutput:
        return MultimodalProcessorOutput(
            mm_items=d["mm_items"],
            input_ids=d.get("input_ids"),
            padded_input_ids=d.get("padded_input_ids"),
            im_token_id=d.get("im_token_id"),
            im_start_id=d.get("im_start_id"),
            im_end_id=d.get("im_end_id"),
            slice_start_id=d.get("slice_start_id"),
            slice_end_id=d.get("slice_end_id"),
            video_token_id=d.get("video_token_id"),
            audio_token_id=d.get("audio_token_id"),
            audio_start_id=d.get("audio_start_id"),
            audio_end_id=d.get("audio_end_id"),
            mrope_positions=d.get("mrope_positions"),
            mrope_position_delta=d.get("mrope_position_delta"),
            vision_position_ids=d.get("vision_position_ids"),
            media_nums_per_sample=d.get("media_nums_per_sample"),
            visible_frame_counts=d.get("visible_frame_counts"),
        )

    @staticmethod
    def build_padded_input_ids(input_ids, mm_items: List[MultimodalDataItem]):
        """pad the input_ids with mm_items if it's not already padded"""
        if input_ids is None or not mm_items:
            return None

        for item in mm_items:
            if item.pad_value is None or item.offsets is None:
                return None

        if isinstance(input_ids, torch.Tensor):
            padded_input_ids = input_ids.flatten().tolist()
        else:
            padded_input_ids = list(input_ids)

        for item in mm_items:
            for start, end in item.offsets:
                padded_input_ids[start : end + 1] = [item.pad_value] * (end - start + 1)
        return padded_input_ids

    @staticmethod
    def build_token_modalities(
        input_ids, mm_items: List[MultimodalDataItem]
    ) -> Optional[List[int]]:
        """Build the pre-padding token modality map from item offsets."""
        if input_ids is None or not mm_items:
            return None
        if isinstance(input_ids, torch.Tensor):
            num_tokens = input_ids.numel()
        else:
            num_tokens = len(flatten_nested_list(input_ids))
        token_modalities = [0] * num_tokens
        for item in mm_items:
            if not item.offsets:
                continue
            modality = item.modality.value
            for start, end in item.offsets:
                if start < 0 or end < start or end >= num_tokens:
                    raise ValueError(
                        "Invalid multimodal token offsets: "
                        f"offset=({start}, {end}), num_tokens={num_tokens}"
                    )
                if any(token_modalities[index] for index in range(start, end + 1)):
                    raise ValueError(
                        f"Overlapping multimodal token offsets at ({start}, {end})"
                    )
                token_modalities[start : end + 1] = [modality] * (end - start + 1)
        return token_modalities


@dataclasses.dataclass
class MultimodalInputs:
    """The multimodal data related inputs."""

    # items of data
    mm_items: List[MultimodalDataItem]
    padded_input_ids: Optional[List[int]] = None
    image_pad_len: Optional[list] = None
    num_image_tokens: Optional[int] = None
    token_modalities: Optional[List[int]] = None

    # image
    im_token_id: Optional[int] = None
    im_start_id: Optional[int] = None
    im_end_id: Optional[int] = None
    slice_start_id: Optional[int] = None
    slice_end_id: Optional[int] = None

    # video
    video_token_id: Optional[int] = None

    # audio
    audio_token_id: Optional[int] = None
    audio_start_id: Optional[int] = None
    audio_end_id: Optional[int] = None

    # QWen2-VL related
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[torch.Tensor] = None
    mrope_position_delta_repeated_cache: Optional[torch.Tensor] = None

    # Moss-VL related
    vision_position_ids: Optional[torch.Tensor] = None
    media_nums_per_sample: Optional[List[int]] = None
    visible_frame_counts: Optional[torch.Tensor] = None

    def release_features(self):
        """Release feature tensors to free GPU memory."""
        for item in self.mm_items:
            try:
                # Release both deferred features that were never used and
                # borrowed features retained for possible re-prefill.
                item.release_transport_proxies()
            except Exception:
                logger.warning(
                    "Failed to release an unused multimodal feature transport",
                    exc_info=True,
                )
            finally:
                item.feature = None

    @staticmethod
    def from_processor_output(
        obj: MultimodalProcessorOutput, *, requires_mm_token_modalities: bool = False
    ):
        mm_items = obj.mm_items
        assert isinstance(mm_items, list)
        mm_items = [item for item in mm_items if item.is_valid()]

        # try reconstructing from cuda-ipc
        reconstruct_device = None
        try:
            for mm_item in mm_items:
                if (
                    mm_item.has_cuda_ipc_proxy()
                    and not mm_item.can_defer_cuda_ipc_feature_reconstruction()
                ):
                    if reconstruct_device is None:
                        reconstruct_device = torch.cuda.current_device()
                    mm_item.reconstruct(reconstruct_device)
        except BaseException:
            for mm_item in mm_items:
                mm_item.release_transport_proxies()
            raise

        if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0:
            # Multi-modal feature hashing optimization:
            # When SGLANG_MM_BUFFER_SIZE_MB > 0, we temporarily move feature tensors to GPU
            # for faster hash computation, while avoiding OOM issues.
            from sglang.srt.managers.mm_utils import (
                init_feature_buffer,
                is_feature_buffer_initialized,
                reset_buffer_offset,
                try_add_to_buffer,
            )

            device = torch.cuda.current_device() if torch.cuda.is_available() else "cpu"
            if not is_feature_buffer_initialized():
                init_feature_buffer(device)
            reset_buffer_offset()
            for item in mm_items:
                if item.feature is not None:
                    if isinstance(item.feature, torch.Tensor):
                        item.feature = try_add_to_buffer(item.feature)

        token_modalities = (
            MultimodalProcessorOutput.build_token_modalities(obj.input_ids, mm_items)
            if requires_mm_token_modalities
            else None
        )

        for item in mm_items:
            item.set_pad_value()

        if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0:
            for item in mm_items:
                if item.feature is not None:
                    item.feature = item.feature.to("cpu", non_blocking=True)

        mm_inputs = MultimodalInputs(
            mm_items=mm_items,
            padded_input_ids=obj.padded_input_ids,
            token_modalities=token_modalities,
        )
        optional_args = [
            "mrope_positions",
            "mrope_position_delta",
            "im_token_id",
            "im_start_id",
            "im_end_id",
            "video_token_id",
            "slice_start_id",
            "slice_end_id",
            "audio_start_id",
            "audio_end_id",
            "audio_token_id",
            "vision_position_ids",
            "media_nums_per_sample",
            "visible_frame_counts",
        ]
        for arg in optional_args:
            val = getattr(obj, arg, None)
            if val is not None:
                setattr(mm_inputs, arg, val)

        return mm_inputs

    def contains_image_inputs(self) -> bool:
        return any(item.is_image() for item in self.mm_items)

    def contains_video_inputs(self) -> bool:
        return any(item.is_video() for item in self.mm_items)

    def contains_audio_inputs(self) -> bool:
        return any(item.is_audio() for item in self.mm_items)

    def contains_mm_input(self) -> bool:
        return any(True for item in self.mm_items if item.is_valid())

    def compute_mm_token_counts(self) -> Tuple[int, int, int]:
        """Count prompt tokens consumed by each modality (image, audio, video).

        A modality's token count is the total span covered by its items'
        offsets. Returns a (image_tokens, audio_tokens, video_tokens) tuple.
        """
        image_tokens = audio_tokens = video_tokens = 0
        for item in self.mm_items:
            if not item.offsets:
                continue
            num_tokens = sum(end - start + 1 for start, end in item.offsets)
            if item.is_image():
                image_tokens += num_tokens
            elif item.is_audio():
                audio_tokens += num_tokens
            elif item.is_video():
                video_tokens += num_tokens
        return image_tokens, audio_tokens, video_tokens

    def merge(self, other: MultimodalInputs):

        # args needed to be merged
        optional_args = [
            "mm_items",
            "image_pad_len",
        ]
        for arg in optional_args:
            self_arg = getattr(self, arg, None)
            if self_arg is not None:
                setattr(self, arg, self_arg + getattr(other, arg))

        if other.token_modalities is not None:
            if self.token_modalities is None:
                self.token_modalities = list(other.token_modalities)
            else:
                self.token_modalities += other.token_modalities

        mrope_positions = self.mrope_positions
        if mrope_positions is not None:
            if other.mrope_positions is None:
                self.mrope_positions = mrope_positions
            else:
                self.mrope_positions = torch.cat(
                    [self.mrope_positions, other.mrope_positions], dim=1
                )

        mrope_position_delta = self.mrope_position_delta
        if mrope_position_delta is not None:
            if other.mrope_position_delta is None:
                self.mrope_position_delta = mrope_position_delta
            else:
                self.mrope_position_delta = torch.cat(
                    [self.mrope_position_delta, other.mrope_position_delta], dim=0
                )

        for key, val in other.__dict__.items():
            if "_id" in key:
                # set token_ids
                if getattr(self, key, None) is None:
                    setattr(self, key, getattr(other, key, None))
        # other args would be kept intact


@dataclasses.dataclass(slots=True, kw_only=True)
class ReqLogprob:
    top_logprobs_num: int
    token_ids_logprob: Optional[List[int]]
    input_token_logprobs_val: Optional[List[float]] = None
    input_token_logprobs_idx: Optional[List[int]] = None
    input_top_logprobs_val: Optional[List[List[float]]] = None
    input_top_logprobs_idx: Optional[List[List[int]]] = None
    # Flat replacements for the rows above (see
    # build_flat_top_logprobs_arrays); when set, the nested rows are
    # emptied and the arrays ship instead.
    input_top_logprobs_val_flat: Optional[np.ndarray] = None
    input_top_logprobs_idx_flat: Optional[np.ndarray] = None
    input_top_logprobs_flat_null_prefix: Optional[int] = None
    input_token_ids_logprobs_val: Optional[List[List[float]]] = None
    input_token_ids_logprobs_idx: Optional[List[List[int]]] = None
    output_token_logprobs_val: Optional[list] = None
    output_token_logprobs_idx: Optional[list] = None
    output_top_logprobs_val: Optional[list] = None
    output_top_logprobs_idx: Optional[list] = None
    # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
    output_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (
        None
    )
    output_token_ids_logprobs_idx: Optional[list] = None


@dataclasses.dataclass(slots=True, kw_only=True)
class ReqKvInfo:
    # Device KV a request holds outside the prefix cache. Always present on the Req;
    # whether any KV is held is `holds_kv` (a row is registered).
    # Match observations and scheduling state stay on the Req itself.
    req_pool_idx: Optional[int] = None  # req_to_token row, the register for the slots

    # The request's own KV is [cache_protected_len, kv_allocated_len).
    cache_protected_len: int = 0  # Tree cache owns [0, here) (matched or inserted)
    # This request already inserted [0, here) into the tree; later inserts count
    # a hit only on the nodes past it, so a request counts each node once.
    cache_inserted_len: int = 0
    kv_committed_len: int = 0  # KV content committed up to here, <= kv_allocated_len
    kv_allocated_len: int = 0

    # SWA slots in [swa_dead_lo(page_size), get_evicted_seqlen(SWA)) are freed.
    swa_evict_floor: int = 0  # [0, here) never window-evicted (prefill-aware SWA)
    component_evicted_seqlens: dict[ComponentType, int] = dataclasses.field(
        default_factory=dict
    )

    # Host-side KV backup the request holds across a retraction (unified cache).
    retraction_backup: Optional[RetractionBackup] = None

    # Mamba state: an independent resource; whether it is held is `holds_mamba`.
    mamba_pool_idx: Optional[torch.Tensor] = None  # shape (1)
    mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None  # shape (2)
    mamba_next_track_idx: Optional[int] = None  # 0 or 1
    mamba_last_track_idx: Optional[int] = None  # 0 or 1
    # Seq len of the last cached mamba state
    mamba_last_track_seqlen: Optional[int] = None
    # Seq len of the other ping-pong slot's state. None means what is in
    # that slot cannot be named (never written, donated, freed).
    mamba_prev_track_seqlen: Optional[int] = None
    # Deferred COW: source mamba pool index from radix cache node (copy on forward stream)
    mamba_cow_src_index: Optional[torch.Tensor] = None
    # Deferred clear: newly allocated mamba slot needs zeroing on forward stream
    mamba_needs_clear: bool = False

    def swa_dead_lo(self, page_size: int) -> int:
        # Lowest SWA position this request may free itself: above the tree-owned
        # prefix and above the eviction shield, page-aligned upward.
        lo = max(self.cache_protected_len, self.swa_evict_floor)
        if page_size > 1 and lo > self.cache_protected_len:
            lo = ceil_align(lo, page_size)
        return lo

    @property
    def holds_kv(self) -> bool:
        return self.req_pool_idx is not None

    @property
    def holds_mamba(self) -> bool:
        return self.mamba_pool_idx is not None

    @property
    def is_kv_released(self) -> bool:
        return self.kv_allocated_len == 0 and self.max_evicted_seqlen == 0

    @property
    def max_evicted_seqlen(self) -> int:
        return max(self.component_evicted_seqlens.values(), default=0)

    def get_evicted_seqlen(self, component_type: ComponentType) -> int:
        return self.component_evicted_seqlens.get(component_type, 0)

    def set_evicted_seqlen(self, component_type: ComponentType, length: int) -> None:
        self.component_evicted_seqlens[component_type] = length

    def clamp_evicted_seqlens(self, length: int) -> None:
        for component, cursor in self.component_evicted_seqlens.items():
            self.component_evicted_seqlens[component] = min(cursor, length)

    def mark_kv_released(self) -> None:
        self.kv_allocated_len = 0
        self.component_evicted_seqlens.clear()


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

    def __init__(
        self,
        rid: str,
        origin_input_text: str,
        origin_input_ids: array[int],
        sampling_params: SamplingParams,
        return_logprob: bool = False,
        top_logprobs_num: int = 0,
        dllm_config: Optional[DllmConfig] = None,
        token_ids_logprob: List[int] = None,
        return_sampling_mask: bool = False,
        sampling_logprobs_mode: SamplingLogprobsMode = "selected",
        return_flat_raw_top_logprobs: bool = False,
        stream: bool = False,
        origin_input_ids_unpadded: Optional[array[int]] = None,
        lora_id: Optional[str] = None,
        input_embeds: Optional[List[List[float]]] = None,
        positional_embed_overrides: Optional[PositionalEmbeds] = None,
        token_type_ids: List[int] = None,
        session: Optional[Session] = None,
        custom_logit_processor: Optional[str] = None,
        require_reasoning: bool = False,
        return_hidden_states: ReturnHiddenStatesMode = False,
        return_routed_experts: bool = False,
        routed_experts_start_len: int = 0,
        return_indexer_topk: bool = False,
        eos_token_ids: Optional[Set[int]] = None,
        bootstrap_host: Optional[str] = None,
        bootstrap_port: Optional[int] = None,
        bootstrap_room: Optional[int] = None,
        disagg_mode: Optional[DisaggregationMode] = None,
        routed_dp_rank: Optional[int] = None,
        disagg_prefill_dp_rank: Optional[int] = None,
        kv_hints: Optional[KvHintsEnvelope] = None,
        vocab_size: Optional[int] = None,
        priority: Optional[int] = None,
        metrics_collector: Optional[SchedulerMetricsCollector] = None,
        extra_key: Optional[str] = None,
        routing_key: Optional[str] = None,
        dimensions: Optional[int] = None,
        http_worker_ipc: Optional[str] = None,
        time_stats: Optional[
            Union[APIServerReqTimeStats, DPControllerReqTimeStats]
        ] = None,
        return_pooled_hidden_states: bool = False,
        multi_item_delimiter_indices: Optional[List[int]] = None,
        token_indices_to_pool: Optional[List[int]] = None,
        decision_layout: Optional[List[int]] = None,
        session_id: Optional[str] = None,
        cache_salt: Optional[str] = None,
    ):
        # Input and output info
        self.rid = rid
        self.cache_request_handle = CacheRequestHandle(rid=rid, attempt_id=0)
        self.origin_input_ids = origin_input_ids
        self.origin_input_ids_unpadded = (
            origin_input_ids_unpadded
            if origin_input_ids_unpadded
            else self.origin_input_ids
        )  # Before image padding
        # Each decode stage's output ids. Append-only by contract:
        # refresh_fill_ids infers how many output tokens are already in
        # full_untruncated_fill_ids from lengths alone, so in-place rewrites
        # that preserve length would silently corrupt fill_ids.
        self.output_ids = array("q")
        # Full untruncated sequence: origin + output (+ DLLM mask block).
        # Kept in sync by refresh_fill_ids; admission only updates
        # extend_end, never mutates this array's length.
        self.full_untruncated_fill_ids = array("q")
        # This round extends fill ids [prefix_len, extend_end).
        self.extend_end: Optional[int] = None
        self.dllm_initialized: bool = False

        self.session = session
        # media added by this turn, excluding the session's borrowed history
        self._session_mm_items: List[MultimodalDataItem] = []
        self.session_id = session_id
        # Used by the session radix cache to reject registration after a close/reopen.
        self.session_generation: Optional[int] = None
        self.input_embeds = input_embeds
        self.positional_embed_overrides = positional_embed_overrides
        self.multi_item_delimiter_indices = multi_item_delimiter_indices
        self.token_indices_to_pool = token_indices_to_pool
        self.decision_layout = decision_layout

        # For req-level memory management
        self.kv = ReqKvInfo()

        # Full-KV-derived boundary whose SWA window should be inserted after
        # the current prefill pass.
        self.swa_branching_seqlen: Optional[int] = None

        # for cross-encoder model
        self.token_type_ids = token_type_ids

        # The index of the extend / decode batch
        self.extend_batch_idx = 0
        self.decode_batch_idx = 0

        # For multi-http worker
        self.http_worker_ipc = http_worker_ipc

        # Require reasoning for the request
        self.require_reasoning = require_reasoning

        # State indicating whether the reasoning phase has finished (only meaningful when require_reasoning is True)
        self._is_reasoning_over = False
        self.reasoning_tokens = 0
        self._think_end_matcher: Optional[TokenSequenceMatcher] = None
        self._think_end_match_len = 0

        # Only custom logit processors need the "__req__" reference cycle.
        if custom_logit_processor is not None and isinstance(
            sampling_params.custom_params, dict
        ):
            sampling_params = copy.copy(sampling_params)
            sampling_params.custom_params = sampling_params.custom_params | {
                "__req__": self
            }
        self.sampling_params = sampling_params
        self.custom_logit_processor = custom_logit_processor
        self.return_hidden_states = return_hidden_states
        self.return_hidden_states_mode = get_return_hidden_states_mode(
            return_hidden_states
        )

        # Extra key for caller-defined request classification.
        if lora_id is not None:
            extra_key = (
                extra_key or ""
            ) + lora_id  # lora_id is concatenated to the extra key

        self.extra_key = extra_key
        self.cache_salt = cache_salt or None
        self.lora_id = lora_id
        self.routing_key = routing_key

        # Lazy extra buffer: skip radix cache insert when prealloc failed at
        # boundary — the forward overwrites the only slot, corrupting the state.
        self.mamba_lazy_checkpoint: bool = True

        # Check finish
        self.tokenizer = None
        self.finished_reason: Optional[BaseFinishReason] = None
        # finished position (in output_ids), used when checking stop conditions with speculative decoding
        self.finished_len = None
        # Whether this request has finished output
        self.finished_output = None
        # If we want to abort the request in the middle of the event loop,
        # set to_finish instead of directly setting finished_reason.
        # Note: We should never set finished_reason in the middle, the req will get filtered and never respond
        self.to_finish: Optional[BaseFinishReason] = None
        self.stream = stream
        self.eos_token_ids = eos_token_ids
        self.vocab_size = vocab_size
        self.priority = priority

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

        # For multimodal inputs
        self.multimodal_inputs: Optional[MultimodalInputs] = None
        # Pre-computed multimodal prompt token counts; populated on the prefill
        # node and transferred to decode via the metadata buffer in disagg (PD) mode.
        self.mm_image_tokens: int = 0
        self.mm_audio_tokens: int = 0
        self.mm_video_tokens: int = 0

        # Prefix info
        # Tokens [0, prefix_len) are already in the row, or, before allocation,
        # on the path to last_node; the next extend starts at prefix_len.
        self.prefix_len: int = 0
        # TODO(ispobock): rename to last_device_node
        self.last_node: Any = None
        self.last_host_node: Any = None
        self.best_match_node: Any = None
        # Per-component host hit lengths split off from host_hit_length:
        self.host_hit_length = 0
        self.swa_host_hit_length = 0
        self.mamba_host_hit_length = 0
        # The branching point seqlen to track mamba state. If set, given by prefix
        # match, it will be the tracked seqlen in the ping pong buffer for the
        # right prefill pass.
        self.mamba_branching_seqlen: Optional[int] = None
        # Total cached prefix length (on-device prefix_len + host_hit_length),
        # capped at the max allowed prefix. Set during prefix matching at schedule
        # time and used to estimate uncached tokens / sort by longest prefix for
        # load reporting.
        self.num_matched_prefix_tokens = 0
        # Tokens loaded from storage backend (L3) during prefetch for this request
        self.storage_hit_length = 0
        self.storage_hit_start: Optional[int] = None
        # FULL host-hit tokens actually spliced to device by init_load_back at
        # admission; less than host_hit_length when the load-back was declined
        # or the staged splice dropped (that shortfall is re-prefilled).
        self.host_loaded_length = 0
        # Buffer-mode host memory is transport staging, not an L2 cache tier.
        self.host_hit_is_storage = False
        self.storage_prefetch_retry_attempts = 0
        self.staged_prefetch_plan: Optional[StagedPrefetchPlan] = None
        # The tree lock this request holds; None while it runs on a session
        # slot's record or holds no lock.
        self.lock: Optional[TreeLock] = None
        # Device/host prefix used to plan the latest L3 lookup. Admission uses
        # it to detect newly exposed storage demand after queue-time eviction.
        self.storage_prefetch_last_match_len: Optional[int] = None
        # Logical-page KV sharding: rotation base of the chain this request
        # extends (owner of position-page P is (base + P) % shard_size).
        # Refreshed at every sharded alloc — read through last_node, or drawn
        # least-full for a new chain — and consumed by the radix insert to
        # stamp new tree nodes. Allocation itself must NOT read it back when
        # a tree node is available (the checkpoint dedup rebind
        # would make it stale); the only allocation-time reader is the
        # radix-disabled fallback, which keeps no nodes past the root and has
        # no rebind.
        self.kv_rotation_base: Optional[int] = None

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

        # For retraction
        self.is_retracted = False
        # Indicates if the req has ever been retracted.
        self.retracted_stain = False

        self.weight_version_events: List[WeightVersionEvent] = []

        # Incremental streamining
        self.send_token_offset: int = 0
        self.send_decode_id_offset: int = 0
        # TODO (Byron): send_output_token_logprobs_offset and send_decode_id_offset can be different in disaggregation mode
        # because the decode server does not have the first output token logprobs
        self.send_output_token_logprobs_offset: int = 0

        # Logprobs (arguments)
        self.return_logprob = return_logprob
        # Start index to compute logprob from.
        self.logprob_start_len = 0
        self.logprob = ReqLogprob(
            top_logprobs_num=top_logprobs_num,
            token_ids_logprob=token_ids_logprob,
        )
        self.temp_scaled_logprobs = False
        self.top_p_normalized_logprobs = False
        self.return_sampling_mask = return_sampling_mask
        self.sampling_logprobs_mode = sampling_logprobs_mode
        self.return_flat_raw_top_logprobs = return_flat_raw_top_logprobs

        # Logprobs (return values)
        # True means the input logprob has been already sent to detokenizer.
        self.input_logprob_sent: bool = False
        # Temporary holder to store input_token_logprobs.
        self.input_token_logprobs: Optional[List[Tuple[int]]] = None
        self.temp_input_top_logprobs_val: Optional[List[torch.Tensor]] = None
        self.temp_input_top_logprobs_idx: Optional[List[int]] = None
        self.temp_input_token_ids_logprobs_val: Optional[List[float]] = None
        self.temp_input_token_ids_logprobs_idx: Optional[List[int]] = None

        if return_logprob:
            # shape: (bs, 1)
            self.logprob.output_token_logprobs_val = []
            self.logprob.output_token_logprobs_idx = []
            # shape: (bs, k)
            self.logprob.output_top_logprobs_val = []
            self.logprob.output_top_logprobs_idx = []
            # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
            self.logprob.output_token_ids_logprobs_val = []
            self.logprob.output_token_ids_logprobs_idx = []
        self.sampling_mask_rows: Optional[SamplingMaskRows] = (
            SamplingMaskRows() if return_sampling_mask else None
        )
        self.hidden_states: List[List[float]] = []
        self.hidden_states_tensor = None  # Note: use tensor instead of list to transfer hidden_states when PD + MTP
        self.output_topk_p = None
        self.output_topk_index = None
        self.output_dsa_topk_indices = None

        # capture routed experts
        self.return_routed_experts = return_routed_experts
        self.routed_experts_start_len = routed_experts_start_len
        self.routed_experts: Optional[torch.Tensor] = (
            None  # cpu tensor: shape (seqlen, topk)
        )

        self.return_indexer_topk = return_indexer_topk
        self.indexer_topk: Optional[torch.Tensor] = (
            None  # cpu tensor: shape (seqlen, num_indexer_layers, index_topk)
        )
        # Customized info
        self.customized_info: Optional[Dict[str, List[Any]]] = None

        # Embedding (return values)
        self.embedding = None

        # Constrained decoding
        self.grammar_key: Optional[Tuple[str, str]] = None
        self.grammar: Optional[Union[BaseGrammarObject, Future[BaseGrammarObject]]] = (
            None
        )
        self.grammar_wait_ct = 0

        # The number of cached tokens that were already cached in the KV cache
        self.cached_tokens = 0
        self.already_computed = 0

        # Detailed breakdown of cached tokens by source (for HiCache)
        self.cached_tokens_device = 0  # Tokens from device cache (GPU)
        self.cached_tokens_host = 0  # Tokens from host cache (CPU memory)
        self.cached_tokens_storage = 0  # Tokens from L3 storage backend
        self._cache_breakdown_computed = (
            False  # Track if breakdown was already computed
        )

        # Per-request count of verification forward passes.
        self.spec_verify_ct = 0

        # Per-request count of accepted draft tokens (excludes the bonus token).
        self.spec_num_correct_drafts = 0

        self.spec_num_block_accept_tokens = 0

        self.spec_num_cap_tokens = 0

        # Acceptance histogram for speculative decoding.
        # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens.
        # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens.
        self.spec_correct_drafts_histogram: List[int] = []

        self.spec_cap_lens_histogram: List[int] = []

        # The number of times this request has been retracted / preempted.
        self.retraction_count = 0
        self.retraction_mb_id = None

        # For observability
        self.metrics_collector = metrics_collector
        if time_stats is not None:
            self.time_stats = SchedulerReqTimeStats.new_from_obj(time_stats)
        else:
            self.time_stats = SchedulerReqTimeStats(disagg_mode=disagg_mode)
        self.time_stats.set_metrics_collector(metrics_collector)
        self.time_stats.set_scheduler_recv_time()
        self.has_log_time_stats: bool = False

        # For disaggregation
        self.bootstrap_host: str = bootstrap_host
        self.bootstrap_port: Optional[int] = bootstrap_port
        self.bootstrap_room: Optional[int] = bootstrap_room
        # Decode-local: the already-emitted boundary token to replay when a
        # retracted request is rebootstrapped. Set in pause_generation(retract)
        # and consumed in the decode transfer commit; never plumbed to prefill.
        self.pd_rebootstrap_forced_output_id: Optional[int] = None
        # Override embeds give KV that the token ids alone do not determine.
        self.skip_radix_cache_insert = (
            bootstrap_host == FAKE_BOOTSTRAP_HOST
            or positional_embed_overrides is not None
        )
        # Set when the request's in-flight forward ran on KV that never landed:
        # that forward's output is dropped and the request finishes with it.
        self.discard_output_reason: Optional[FINISH_ABORT] = None
        self.disagg_kv_sender: Optional[BaseKVSender] = None

        self.routed_dp_rank: Optional[int] = routed_dp_rank
        self.disagg_prefill_dp_rank: Optional[int] = disagg_prefill_dp_rank
        # Orchestrator KV hints, forwarded to the HiCache storage backends.
        self.kv_hints: Optional[KvHintsEnvelope] = kv_hints

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

        # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap
        # This is because kv is not ready in `process_prefill_chunk`.
        # We use `tmp_end_idx` to store the end index of the kv cache to send.
        self.tmp_end_idx: int = -1
        # Decode-side cached-prefix length; base of the staging chunk grid
        # (start_send_idx starts here but advances with every send).
        self.disagg_decode_prefix_len: int = 0
        # At-rest device-resident prefix end, snapshotted on the request's
        # first prefill batch; the cached-prefix early-send never goes past it.
        self.early_send_prefix_end: Optional[int] = None
        self.metadata_buffer_index: int = -1
        self.expected_kv_checksum: int = 0
        # Used in overlap sequence to signal that an optimistic request should
        # abort chunking. Set in create_sender, consumed in process_batch_result.
        self.pending_bootstrap = False
        # Number of optimistic prefill forward passes started. preserved across retracts.
        self.prefill_attempt_count = 0

        # For Matryoshka embeddings
        self.dimensions = dimensions

        # Beam search overlay: leader and internal members share one BeamGroup.
        self.beam_group = None

        # Whether to return pooled hidden states (pre-head transformer output)
        self.return_pooled_hidden_states = return_pooled_hidden_states
        self.pooled_hidden_state = None

        # For diffusion LLM
        self.init_diffusion_llm(dllm_config)

        # For hisparse
        self.hisparse_staging = False

        # Snapshot of the scheduler prefill-token counter taken at waiting_queue entry; used by HRRN aging.
        self.arrival_processed_tokens: int = 0

    def advance_cache_request_handle(self) -> None:
        self.cache_request_handle = dataclasses.replace(
            self.cache_request_handle,
            attempt_id=self.cache_request_handle.attempt_id + 1,
        )

    @property
    def seqlen(self) -> int:
        """Get the current sequence length of the request."""
        return len(self.origin_input_ids) + len(self.output_ids)

    @property
    def is_beam_leader(self) -> bool:
        """The user-visible row of a beam group; the group's other rows are
        scheduler-internal and never streamed to the user."""
        return self.beam_group is not None and self.beam_group.leader is self

    @property
    def is_prefill_only(self) -> bool:
        """Check if this request is prefill-only (no token generation needed)."""
        # NOTE: when spec is enabled, prefill_only optimizations are disabled

        spec_alg = get_spec().speculative_algorithm
        return self.sampling_params.max_new_tokens == 0 and spec_alg is None

    @property
    def output_ids_through_stop(self) -> array[int]:
        """Get the output ids through the stop condition. Stop position is included."""
        if self.finished_len is not None:
            return self.output_ids[: self.finished_len]
        return self.output_ids

    def needs_host_load_back(self) -> bool:
        """Whether any cache layer has a host hit that needs L2 H2D load_back."""
        return (
            self.host_hit_length > 0
            or self.swa_host_hit_length > 0
            or self.mamba_host_hit_length > 0
        )

    def materialized_host_hit_len(self) -> int:
        """Host-hit tokens that actually reached the device: the metrics tier
        split must not credit host/storage for a declined or dropped
        load-back, whose span was re-prefilled and counted as input."""
        return min(self.host_hit_length, self.host_loaded_length)

    def fulfilled_storage_hit_len(self, prefix_len: int) -> int:
        """L3-fetched tokens covered by the admitted cached prefix."""
        if self.storage_hit_start is None:
            return min(prefix_len, self.storage_hit_length)
        return max(
            0,
            min(prefix_len, self.storage_hit_start + self.storage_hit_length)
            - self.storage_hit_start,
        )

    def detach_kv(self) -> ReqKvInfo:
        # Hand the KV record to a new holder; the req keeps a fresh empty one.
        kv, self.kv = self.kv, ReqKvInfo()
        return kv

    def owned_kv_len(self) -> int:
        # Report only the prompt prefix so thinking + answer fall into the
        # overallocated range and are reclaimed by release_kv_cache. #22373.
        if get_serving().strip_thinking_cache and self.reasoning_tokens > 0:
            return min(self.kv.kv_committed_len, len(self.origin_input_ids))
        if self.finished_len is None:
            return self.kv.kv_committed_len
        # A verify step commits a whole accepted chunk, so output_ids can run
        # past the stop; the client never sees those tokens, so nothing matches.
        return min(
            self.kv.kv_committed_len,
            len(self.origin_input_ids) + len(self.output_ids_through_stop),
        )

    def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
        """Record one step accepted draft count (excludes bonus token) into the histogram."""
        if len(self.spec_correct_drafts_histogram) <= num_correct_drafts:
            self.spec_correct_drafts_histogram.extend(
                [0] * (num_correct_drafts - len(self.spec_correct_drafts_histogram) + 1)
            )
        self.spec_correct_drafts_histogram[num_correct_drafts] += 1

    def update_spec_cap_lens_histogram(self, cap_len: int):
        cap_len = int(cap_len)
        if len(self.spec_cap_lens_histogram) <= cap_len:
            self.spec_cap_lens_histogram.extend(
                [0] * (cap_len - len(self.spec_cap_lens_histogram) + 1)
            )
        self.spec_cap_lens_histogram[cap_len] += 1

    def extend_image_inputs(self, image_inputs):
        if self.session is not None:
            self._extend_session_image_inputs(image_inputs)
        elif self.multimodal_inputs is None:
            self.multimodal_inputs = image_inputs
        else:
            self.multimodal_inputs.merge(image_inputs)

    def _extend_session_image_inputs(self, image_inputs):
        """Append media while preserving the saved session and its position history."""
        self._session_mm_items.extend(image_inputs.mm_items)
        # Padding can change token values without changing their count.
        self.full_untruncated_fill_ids = array("q")
        if self.multimodal_inputs is not None:
            # Branches and aborted turns must leave the parent's metadata intact.
            self.multimodal_inputs = dataclasses.replace(self.multimodal_inputs)

        positions = image_inputs.mrope_positions
        if positions is not None:
            prefix_len = len(self.origin_input_ids) - positions.shape[1]
            prefix = (
                self.multimodal_inputs.mrope_positions
                if self.multimodal_inputs is not None
                else None
            )
            if prefix is None:
                prefix = positions.new_empty((3, 0))
            prefix = prefix[:, :prefix_len]
            next_position = prefix.max() + 1 if prefix.numel() else 0
            text_len = prefix_len - prefix.shape[1]
            text_positions = (
                torch.arange(
                    text_len, dtype=positions.dtype, device=positions.device
                ).expand(3, -1)
                + next_position
            )
            # Fill the reply/text gap, then shift the new turn's media coordinates.
            positions = torch.cat(
                [prefix, text_positions, positions + next_position + text_len], dim=1
            )

        if self.multimodal_inputs is None:
            self.multimodal_inputs = image_inputs
        else:
            # Use the full table above, or let the scheduler compute missing positions.
            self.multimodal_inputs.mrope_positions = None
            self.multimodal_inputs.mrope_position_delta = None
            self.multimodal_inputs.merge(image_inputs)

        self.multimodal_inputs.mrope_position_delta_repeated_cache = None
        if positions is not None:
            self.multimodal_inputs.mrope_positions = positions
            self.multimodal_inputs.mrope_position_delta = (
                positions.max() + 1 - positions.shape[1]
            ).reshape(1, 1)

    def release_mm_inputs_on_abort(self) -> None:
        """Release this turn's media without invalidating session history."""
        if self.session is None:
            if self.multimodal_inputs is not None:
                self.multimodal_inputs.release_features()
        else:
            for item in self._session_mm_items:
                item.release_transport_proxies()
                item.feature = None
            self._session_mm_items.clear()
        self.multimodal_inputs = None

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

    def next_output_finishes_by_length(self) -> bool:
        # Whether committing one more output token reaches max_new_tokens
        return len(self.output_ids) + 1 >= self.sampling_params.max_new_tokens

    @property
    def extend_len(self) -> int:
        return self.extend_end - self.prefix_len

    def get_fill_ids(self) -> array:
        return self.full_untruncated_fill_ids[: self.extend_end]

    def refresh_fill_ids(self) -> None:
        """Keep full_untruncated_fill_ids == origin_input_ids + output_ids by
        appending only the new output tokens.

        Falls back to a full rebuild when the in-place append is invalid:
        - aliasing: full_untruncated_fill_ids is origin_input_ids itself, so
          extending in place would write output tokens into the origin;
        - lengths disagree: fresh req (array still empty), retraction
          (output_ids reset to empty), or set_finish_with_abort (origin
          replaced by a 1-token stub).
        """
        n_have_output = len(self.full_untruncated_fill_ids) - len(self.origin_input_ids)
        if (
            self.full_untruncated_fill_ids is not self.origin_input_ids
            and 0 <= n_have_output <= len(self.output_ids)
        ):
            self.full_untruncated_fill_ids.extend(self.output_ids[n_have_output:])
        else:
            self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids

    def init_next_round_input(
        self,
        tree_cache: Optional[BasePrefixCache] = None,
        cow_mamba: Optional[bool] = None,
    ):
        if self.is_dllm():
            self._init_fill_ids_for_dllm()
            self.determine_dllm_phase()
        else:
            self.refresh_fill_ids()

        input_len = len(self.full_untruncated_fill_ids)

        # Streaming sessions reuse committed KV from the session slot, so
        # custom logprob_start_len is not supported — override to -1.
        if (
            self.session is not None
            and self.session.streaming
            and self.return_logprob
            and self.logprob_start_len >= 0
        ):
            logger.warning(
                "logprob_start_len=%d is not supported for streaming sessions "
                "and will be ignored (rid=%s). Only new-token logprobs are returned.",
                self.logprob_start_len,
                self.rid,
            )
            self.logprob_start_len = -1

        # Pass the full array with a raw-token cap (limit) instead of slicing,
        # avoiding an O(context) copy per prefill-batch build.
        token_ids_to_match = self.full_untruncated_fill_ids
        key_limit: Optional[int] = self._compute_max_prefix_len(input_len)

        # Disable prefix caching when embed overrides are present: same token IDs
        # with different override vectors must not share cached KV values.
        if self.positional_embed_overrides is not None:
            token_ids_to_match = array("q")
            key_limit = None

        if tree_cache is not None:
            if cow_mamba is None:
                cow_mamba = tree_cache.supports_mamba()
            match_result = match_kv_cache(
                self,
                tree_cache,
                token_ids_to_match,
                cow_mamba=cow_mamba,
                max_prefix_len=key_limit,
            )
            self.mamba_branching_seqlen = match_result.mamba_branching_seqlen
            if match_result.cache_protected_len is None:
                self.kv.cache_protected_len = self.prefix_len

            if self.is_dllm():
                self._update_block_offset_for_dllm()

        if (
            self.is_retracted
            and self.multimodal_inputs is not None
            and self.multimodal_inputs.mrope_positions is not None
        ):
            from sglang.srt.managers.mm_utils import (
                extend_mrope_positions_for_retracted_request,
            )

            self.multimodal_inputs.mrope_positions = (
                extend_mrope_positions_for_retracted_request(
                    self.multimodal_inputs.mrope_positions, len(self.output_ids)
                )
            )

    def _compute_max_prefix_len(self, input_len: int) -> int:
        # NOTE: the matched length is at most 1 less than the input length to enable logprob computation
        max_prefix_len = input_len - 1
        if (
            self.dllm_config is not None
            and self.dllm_config.requires_separate_context_encoding
        ):
            max_prefix_len = min(max_prefix_len, self.dllm_block_offset)
        if self.return_logprob and self.logprob_start_len >= 0:
            max_prefix_len = min(max_prefix_len, self.logprob_start_len)
        return max(max_prefix_len, 0)

    # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
    def init_incremental_detokenize(self):
        first_iter = self.surr_offset is None or self.read_offset is None

        output_ids = self.output_ids_through_stop

        if first_iter:
            self.read_offset = len(self.origin_input_ids_unpadded)
            self.surr_offset = max(
                self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
            )
            self.surr_and_decode_ids = (
                self.origin_input_ids_unpadded[self.surr_offset :] + output_ids
            )
            self.cur_decode_ids_len = len(output_ids)
        else:
            self.surr_and_decode_ids.extend(output_ids[self.cur_decode_ids_len :])
            self.cur_decode_ids_len = len(output_ids)

        return self.surr_and_decode_ids, self.read_offset - self.surr_offset

    def _stop_match_tail_len(self, new_accepted_len: int) -> int:
        max_len_tail_str = max(
            self.sampling_params.stop_str_max_len + 1,
            self.sampling_params.stop_regex_max_len + 1,
        )
        # Cover all newly accepted tokens so an early stop string is not missed
        # when speculative decoding accepts multiple tokens per step.
        return min(
            max_len_tail_str + max(new_accepted_len - 1, 0), len(self.output_ids)
        )

    def tail_str(self, new_accepted_len: int = 1) -> str:
        # Check stop strings and stop regex patterns together
        if (
            len(self.sampling_params.stop_strs) == 0
            and len(self.sampling_params.stop_regex_strs) == 0
        ):
            return ""

        tail_len = self._stop_match_tail_len(new_accepted_len)
        return self.tokenizer.decode(self.output_ids[-tail_len:])

    def check_match_stop_str_prefix(self) -> bool:
        if not self.sampling_params.stop_strs:
            return False

        tail_str = self.tail_str()

        # Early return if tail_str is empty
        if not tail_str:
            return False

        for stop_str in self.sampling_params.stop_strs:
            if not stop_str:
                continue
            # Check if stop_str is contained in tail_str (fastest check first)
            if stop_str in tail_str:
                return True

            # Check if tail_str suffix matches stop_str prefix
            # Only check if stop_str is not empty, it's for stream output
            min_len = min(len(tail_str), len(stop_str))
            for i in range(1, min_len + 1):
                if tail_str[-i:] == stop_str[:i]:
                    return True

        return False

    def _check_token_based_finish(self, new_accepted_tokens: List[int]) -> bool:
        if self.sampling_params.ignore_eos:
            return False

        # Check stop token ids
        matched_eos = False

        for i, token_id in enumerate(new_accepted_tokens):
            if self.sampling_params.stop_token_ids:
                matched_eos |= token_id in self.sampling_params.stop_token_ids
            if self.eos_token_ids:
                matched_eos |= token_id in self.eos_token_ids
            if self.tokenizer is not None:
                matched_eos |= token_id == self.tokenizer.eos_token_id
                if self.tokenizer.additional_stop_token_ids:
                    matched_eos |= token_id in self.tokenizer.additional_stop_token_ids
            if matched_eos:
                self.finished_reason = FINISH_MATCHED_TOKEN(matched=token_id)
                matched_pos = len(self.output_ids) - len(new_accepted_tokens) + i
                self.finished_len = matched_pos + 1
                return True

        return False

    def _locate_str_stop_finished_len(
        self,
        new_accepted_len: int,
        *,
        stop_str: Optional[str] = None,
        stop_regex: Optional[str] = None,
    ) -> int:
        """Map a matched stop string/regex to output_ids length (stop included)."""

        def matched(text: str) -> bool:
            if stop_str is not None:
                return stop_str in text
            return re.search(stop_regex, text) is not None

        tail_len = self._stop_match_tail_len(new_accepted_len)
        start = len(self.output_ids) - tail_len
        token_window = self.output_ids[start:]

        # Old prefixes were checked in the previous step.
        for token_count in range(
            max(1, len(token_window) - new_accepted_len + 1), len(token_window)
        ):
            if matched(self.tokenizer.decode(token_window[:token_count])):
                return start + token_count

        # The full tail window is already known to match by the caller.
        return len(self.output_ids)

    def _check_str_based_finish(self, new_accepted_len: int = 1):
        if (
            len(self.sampling_params.stop_strs) > 0
            or len(self.sampling_params.stop_regex_strs) > 0
        ):
            tail_str = self.tail_str(new_accepted_len)

            # Check stop strings
            if len(self.sampling_params.stop_strs) > 0:
                for stop_str in self.sampling_params.stop_strs:
                    stop_str_in_tail = stop_str in tail_str
                    if stop_str_in_tail or stop_str in self.decoded_text:
                        self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
                        if stop_str_in_tail:
                            self.finished_len = self._locate_str_stop_finished_len(
                                new_accepted_len, stop_str=stop_str
                            )
                        return True

            # Check stop regex
            if len(self.sampling_params.stop_regex_strs) > 0:
                for stop_regex_str in self.sampling_params.stop_regex_strs:
                    # Seatbelt, not validation: patterns are checked at ingress
                    # (Python's `normalize`, or the rust server's stricter
                    # `stop_regex_bound`). This runs per decode step on the hot
                    # path, so an `re.error` escaping here would take the whole
                    # scheduler down over one malformed request. Fail that request
                    # instead.
                    try:
                        matched = re.search(stop_regex_str, tail_str)
                    except (re.error, RecursionError) as e:
                        logger.warning(
                            "req %s: invalid stop_regex %r (%s); aborting the request",
                            self.rid,
                            stop_regex_str,
                            e,
                        )
                        self.finished_reason = FINISH_ABORT(
                            f"invalid stop_regex {stop_regex_str!r}: {e}",
                            HTTPStatus.BAD_REQUEST,
                            "BadRequestError",
                        )
                        break
                    if matched:
                        self.finished_reason = FINISHED_MATCHED_REGEX(
                            matched=stop_regex_str
                        )
                        self.finished_len = self._locate_str_stop_finished_len(
                            new_accepted_len, stop_regex=stop_regex_str
                        )
                        return True

        return False

    def _check_vocab_boundary_finish(self, new_accepted_tokens: List[int] = None):
        for i, token_id in enumerate(new_accepted_tokens):
            if token_id < 0 or (
                self.vocab_size is not None and token_id >= self.vocab_size
            ):
                offset = len(self.output_ids) - len(new_accepted_tokens) + i
                if self.sampling_params.stop_token_ids:
                    self.output_ids[offset] = next(
                        iter(self.sampling_params.stop_token_ids)
                    )
                if self.eos_token_ids:
                    self.output_ids[offset] = next(iter(self.eos_token_ids))
                self.finished_reason = FINISH_MATCHED_STR(matched="NaN happened")
                self.finished_len = offset + 1
                return True

        return False

    def _cap_finished_len_at_max_new_tokens(self) -> None:
        """Demote a stop matched beyond the length budget to a length finish.

        Speculative decoding can accept a run that both crosses
        ``max_new_tokens`` and contains a stop; a stop located past the cap
        must not extend the emitted output beyond the cap.
        """
        max_new_tokens = self.sampling_params.max_new_tokens
        if self.finished_len is not None and self.finished_len > max_new_tokens:
            self.finished_reason = FINISH_LENGTH(length=max_new_tokens)
            self.finished_len = max_new_tokens

    def update_finish_state(self, new_accepted_len: int = 1):
        if self.finished():
            return

        if self.to_finish:
            self.finished_reason = self.to_finish
            self.to_finish = None
            return

        new_accepted_tokens = self.output_ids[-new_accepted_len:]

        # Sanitize out-of-range / NaN token ids before any decode.
        if self._check_vocab_boundary_finish(new_accepted_tokens):
            self._cap_finished_len_at_max_new_tokens()
            return

        # Stop string beats EOS/stop-token matched in the same step (speculative
        # decoding can accept >1 token): token-based would trim only the last
        # token and leak the stop string.
        if self._check_str_based_finish(new_accepted_len):
            self._cap_finished_len_at_max_new_tokens()
            return

        # Stop token/EOS beats the length cap for the same reason: a spec accept
        # run can cross max_new_tokens in the very step the EOS lands, and a
        # length-first finish would keep the over-accepted tokens after the EOS
        # (up to the cap) in the emitted output.
        if self._check_token_based_finish(new_accepted_tokens):
            self._cap_finished_len_at_max_new_tokens()
            return

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

        if self.grammar is not None and self.grammar.is_terminated():
            self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1])
            return

    def reset_for_retract(self):
        # Increment retraction count before resetting other state. We should not reset this
        # since we are tracking the total number of retractions for each request.
        self.retraction_count += 1

        self.prefix_len = 0
        self.routed_experts = None
        self.indexer_topk = None
        self.last_node = None
        self.kv.cache_protected_len = 0
        self.kv.cache_inserted_len = 0
        self.kv_rotation_base = None
        self.num_matched_prefix_tokens = 0
        self.lock = None
        self.swa_branching_seqlen = None
        self.extend_end = None
        self.dllm_initialized = False
        self.is_retracted = True
        self.retracted_stain = True
        self.input_token_logprobs = None
        self.temp_input_top_logprobs_val = None
        self.temp_input_top_logprobs_idx = None
        self.temp_input_token_ids_logprobs_val = None
        self.temp_input_token_ids_logprobs_idx = None
        self.inflight_middle_chunks = 0
        self.kv.mamba_pool_idx = None
        self.kv.mamba_ping_pong_track_buffer = None
        self.kv.mamba_next_track_idx = None
        self.kv.mamba_last_track_idx = None
        self.kv.mamba_last_track_seqlen = None
        self.kv.mamba_prev_track_seqlen = None
        self.mamba_branching_seqlen = None
        self.kv.mamba_cow_src_index = None
        self.kv.mamba_needs_clear = False
        self.already_computed = 0
        assert not self.kv.holds_kv, "expect it is already released"
        self.kv.kv_committed_len = 0
        self.extend_batch_idx = 0
        self.decode_batch_idx = 0

        # When using input_embeds, we cannot easily mix the original input embeddings
        # with the newly generated output token IDs during re-prefill of retracted request.
        # output_ids will have no use, but will lead to wrong size cache indexes.
        # Therefore, we discard the generated output_ids and restart prefill and generation
        # to ensure shape consistency in KV cache.
        if self.input_embeds is not None:
            self.output_ids = array("q")
            self.weight_version_events = truncate_weight_version_events(
                self.weight_version_events, num_kept_tokens=self.send_token_offset
            )

    def _mamba_pool_needing_backup(self, req_to_token_pool, allocator):
        if allocator.get_kvcache().cpu_copy_carries_mamba:
            return None
        if not isinstance(req_to_token_pool, HybridReqToTokenPool):
            return None
        return req_to_token_pool.mamba_pool

    def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
        token_indices = req_to_token_pool.req_to_token[
            self.kv.req_pool_idx, : self.seqlen - 1
        ]
        # Copies over both the kv cache and mamba state if available
        mamba_pool = self._mamba_pool_needing_backup(
            req_to_token_pool, token_to_kv_pool_allocator
        )
        self.kv.retraction_backup = RetractionBackup(
            cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy(
                token_indices,
                mamba_indices=self.kv.mamba_pool_idx,
                req_pool_index=self.kv.req_pool_idx,
            ),
            mamba_cpu=(
                mamba_pool.get_cpu_copy(self.kv.mamba_pool_idx.unsqueeze(0))
                if mamba_pool is not None and self.kv.holds_mamba
                else None
            ),
        )

    def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
        assert self.kv.retraction_backup is not None
        token_indices = req_to_token_pool.req_to_token[
            self.kv.req_pool_idx, : self.seqlen - 1
        ]
        # Loads both the kv cache and mamba state if exists
        mamba_cpu = self.kv.retraction_backup.mamba_cpu
        if mamba_cpu is not None and self.kv.holds_mamba:
            req_to_token_pool.mamba_pool.load_cpu_copy(
                mamba_cpu, self.kv.mamba_pool_idx.unsqueeze(0)
            )
        token_to_kv_pool_allocator.load_cpu_copy(
            self.kv.retraction_backup.cpu_tensors,
            token_indices,
            mamba_indices=self.kv.mamba_pool_idx,
            req_pool_index=self.kv.req_pool_idx,
        )
        self.kv.retraction_backup = None

    def build_rebootstrap_payload(self) -> dict:
        """Build the prefill ``/generate`` payload that asks the original prefill
        worker to recompute this request's prefix KV under the current weights
        (PD true-retraction rebootstrap).

        ``input_ids`` are coerced to plain ``int`` so the payload is always
        JSON-serializable even when ``origin_input_ids``/``output_ids`` hold
        numpy scalars. The sampling-param allow-list forces ``max_new_tokens=1``
        and drops stop/grammar/min_new_tokens so the recompute only re-derives
        the prefix KV and samples a single handoff token. The already-emitted
        boundary token is replayed on the *decode* side (the transfer commit
        overrides the sampled handoff with it), so it is intentionally not sent
        to the prefill here.
        """
        # TODO: multi-modal requests are not supported here. The payload only
        # carries token ``input_ids`` and drops any image/audio/video inputs, so
        # the rebootstrap recompute would not reproduce the original prefix KV
        # for multi-modal requests. Add multi-modal support before enabling it.
        sp = self.sampling_params
        return {
            "input_ids": [int(x) for x in self.origin_input_ids]
            + [int(x) for x in self.output_ids],
            "sampling_params": {
                "max_new_tokens": 1,
                "temperature": sp.temperature,
                "top_p": sp.top_p,
                "top_k": sp.top_k,
                "min_p": sp.min_p,
                "frequency_penalty": sp.frequency_penalty,
                "presence_penalty": sp.presence_penalty,
                "repetition_penalty": sp.repetition_penalty,
                "ignore_eos": sp.ignore_eos,
                "skip_special_tokens": sp.skip_special_tokens,
                "spaces_between_special_tokens": sp.spaces_between_special_tokens,
                "no_stop_trim": sp.no_stop_trim,
            },
            "return_logprob": False,
            "stream": False,
            "rid": self.rid,
            "bootstrap_host": self.bootstrap_host,
            "bootstrap_port": self.bootstrap_port,
            "bootstrap_room": self.bootstrap_room,
            "priority": self.priority,
            "extra_key": self.extra_key,
            "cache_salt": self.cache_salt,
            "routing_key": self.routing_key,
            "routed_dp_rank": self.disagg_prefill_dp_rank,
            "disagg_prefill_dp_rank": self.disagg_prefill_dp_rank,
        }

    def log_time_stats(self):
        # If overlap schedule, we schedule one decode batch ahead so this gets called twice.
        if self.has_log_time_stats:
            return

        bootstrap_info = (
            f", bootstrap_room={self.bootstrap_room}"
            if self.bootstrap_room is not None
            else ""
        )
        prefix = (
            f"ReqTimeStats("
            f"rid={self.rid}{bootstrap_info}, "
            f"input_len={len(self.origin_input_ids)}, "
            f"cached_input_len={self.cached_tokens}, "
            f"output_len={len(self.output_ids)}, "
            f"attempts={self.prefill_attempt_count}, "
            f"type={self.time_stats.disagg_mode_str()})"
        )
        logger.info(f"{prefix}: {self.time_stats.convert_to_duration()}")
        self.has_log_time_stats = True

    def set_finish_with_abort(
        self,
        error_msg: str,
        status_code: int = HTTPStatus.BAD_REQUEST,
        err_type: str = "BadRequestError",
    ):
        if get_parallel().tp_rank == 0:
            logger.error(f"{error_msg}, {self.rid=}")
        # Session requests share historical multimodal inputs with their prior
        # request. The session owns and releases those features when it closes.
        self.release_mm_inputs_on_abort()
        self.grammar = None
        self.origin_input_ids = array(
            "q", [0]
        )  # set it to one token to skip the long prefill
        self.return_logprob = False
        self.logprob_start_len = -1
        self.to_finish = FINISH_ABORT(error_msg, status_code, err_type)

    def update_reasoning_tokens(self, token_id, think_end_ids):
        if self._is_reasoning_over:
            return

        if not isinstance(token_id, list):
            token_id = [token_id]

        if self._think_end_matcher is None:
            self._think_end_matcher = TokenSequenceMatcher(think_end_ids)

        matched = self._think_end_match_len
        for position, token in enumerate(token_id):
            matched = self._think_end_matcher.advance(matched, token)
            if matched == len(self._think_end_matcher):
                self.reasoning_tokens += position + 1
                self._is_reasoning_over = True
                return

        self._think_end_match_len = matched
        self.reasoning_tokens += len(token_id)

    def __repr__(self):
        return (
            f"Req(rid={self.rid}, "
            f"input_ids={self.origin_input_ids}, output_ids={self.output_ids}, "
            f"{self.grammar=}, "
            f"{self.sampling_params=})"
        )


class _MambaRadixCacheV2TrackEntry(NamedTuple):
    track_mask: bool
    track_index: int
    track_seqlen: int


def mamba_lazy_spec_in_window(
    req, mamba_track_interval: int, max_draft_tokens: int
) -> bool:
    """Whether a track-interval crossing is reachable by an in-flight verify.

    kv_committed_len lags device seq_lens by up to one verify under overlap;
    the 2x window absorbs it.
    """
    seq_len = req.kv.kv_committed_len
    window = 2 * max_draft_tokens
    return seq_len // mamba_track_interval != (seq_len + window) // mamba_track_interval


def set_mamba_track_indices_from_reqs(
    batch, track_positions: Optional[List[int]] = None
):
    """Build mamba_track_indices from req objects (authoritative source).

    track_positions: optional per-req ping-pong position override (the lazy
    spec track plan, see mamba_lazy_spec_prepare).
    """
    req_to_token_pool = batch.req_to_token_pool
    all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[
        batch.req_pool_indices
    ]  # (bs, ping_pong_size), int64, on device
    if track_positions is None:
        # Guard: mamba_next_track_idx may be None for requests that haven't
        # gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
        # Keep the gather in bounds for those requests. Freed rows are
        # invalidated below so they cannot scatter into a recycled slot.
        track_positions = [
            (
                req.kv.mamba_next_track_idx
                if req.kv.mamba_next_track_idx is not None
                else 0
            )
            for req in batch.reqs
        ]
    freed_rows = [
        i for i, req in enumerate(batch.reqs) if req.kv.mamba_next_track_idx is None
    ]
    batch.mamba_track_buffer_indices = list(track_positions)
    idx = (
        torch.tensor(
            track_positions,
            dtype=torch.int64,
            pin_memory=True,
        )
        .unsqueeze(1)
        .to(device=all_buffers.device, non_blocking=True)
    )
    batch.mamba_track_indices = (
        torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64)
    )
    if freed_rows:
        # Overlap can leave a request in TARGET_VERIFY after its Mamba state has
        # been freed and reused by another request. The downstream scatter
        # treats a negative destination as a no-op, while the old fallback to
        # position 0 could overwrite the reused live slot.
        batch.mamba_track_indices[freed_rows] = -1


def release_req(
    *,
    req: Req,
    remaing_req_count: int,
    req_to_token_pool: ReqToTokenPool,
    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
    tree_cache: BasePrefixCache,
    hisparse_coordinator: Optional[HiSparseCoordinator],
    offload_kv: bool = True,
) -> bool:
    """Returns False when the KV backup failed and the request cannot be resumed."""
    if hisparse_coordinator is not None and not req.finished():
        hisparse_coordinator.retract_req(req)

    # In decode disaggregation the retracted KV is offloaded to host so it can be
    # restored later without recompute (see resume_retracted_reqs/load_kv_cache).
    # Callers that will recompute the KV instead (PD true-retraction rebootstrap)
    # pass offload_kv=False to skip the wasteful device->host copy.
    backup_saved = True
    # The config bag reflects role flips; server_args keeps the launch role.
    if get_disagg().disaggregation_mode == "decode" and offload_kv:
        backup_saved = backup_kv_cache(
            req,
            tree_cache,
            req_to_token_pool,
            token_to_kv_pool_allocator,
            get_disagg().disaggregation_decode_retraction_backup,
        )
    # TODO (csy): for preempted requests, we may want to insert into the tree
    release_kv_cache(req, tree_cache, checkpoint=False)
    # NOTE(lsyin): we should use the newly evictable memory instantly.
    num_tokens = remaing_req_count * envs.SGLANG_RETRACT_DECODE_STEPS.get()
    evict_from_tree_cache(tree_cache, num_tokens)

    req.reset_for_retract()
    return backup_saved


def retract_all(
    *,
    reqs: List[Req],
    req_to_token_pool: ReqToTokenPool,
    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
    tree_cache: BasePrefixCache,
    hisparse_coordinator: Optional[HiSparseCoordinator],
    offload_kv: bool = True,
) -> None:
    for idx in range(len(reqs)):
        release_req(
            req=reqs[idx],
            remaing_req_count=len(reqs) - idx,
            req_to_token_pool=req_to_token_pool,
            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
            tree_cache=tree_cache,
            hisparse_coordinator=hisparse_coordinator,
            offload_kv=offload_kv,
        )


def compute_extend_logprob_start_len(
    *,
    logprob_start_len: int,
    prefix_len: int,
    extend_len: int,
    full_untruncated_fill_len: int,
) -> int:
    # Key variables:
    # - logprob_start_len: Absolute position in full sequence where logprob computation begins
    # - extend_logprob_start_len: Relative position within current extend batch where logprob computation begins
    # - extend_input_len: Number of tokens that need to be processed in this extend batch
    if logprob_start_len == -1:
        resolved_start = full_untruncated_fill_len
    else:
        # logprob_start_len should be at least the length of the prefix indices
        resolved_start = max(logprob_start_len, prefix_len)
    return min(resolved_start - prefix_len, extend_len)


def _compute_chunked_req_next_prompt_token(
    chunked_req: Optional[Req],
    vocab_size: int,
) -> Optional[int]:
    """Return the next real prompt token after the fill boundary, skipping
    multimodal placeholder (hash) tokens that lie outside the model vocab."""
    if chunked_req is None:
        return None
    fill_len = chunked_req.extend_end
    origin_ids = chunked_req.origin_input_ids
    if fill_len >= len(origin_ids):
        return None
    if origin_ids[fill_len] < vocab_size:
        return int(origin_ids[fill_len])
    return None


@dataclasses.dataclass
class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
    """Store all information of a batch on the scheduler."""

    # === Core: request list (ForwardBatch derives lora_ids / rids / grammars / positions from it) ===
    reqs: List[Req]

    # === Global config and shared resources (engine-lifetime; identical across batches) ===
    # Memory pool and cache
    req_to_token_pool: ReqToTokenPool = None
    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator = None
    tree_cache: BasePrefixCache = None

    # Batch configs
    model_config: ModelConfig = None
    enable_overlap: bool = False

    # Device
    device: str = "cuda"

    # HiSparse (engine-level coordinator ref, same across batches)
    hisparse_coordinator: Optional[HiSparseCoordinator] = None

    # === Batch-variant scheduler state (per-batch; not read by ForwardBatch) ===
    # Tell whether the current running batch is full so that we can skip
    # the check of whether to prefill new requests.
    # This is an optimization to reduce the overhead of the prefill check.
    batch_is_full: bool = False

    # For chunked prefill in PP
    chunked_req: Optional[Req] = None
    chunked_req_next_prompt_token: Optional[int] = None
    contains_last_prefill_chunk: bool = True

    # Tracks process_prefill_chunk() on the original scheduler batch only.
    prefill_chunk_processed: bool = False

    # For DP attention
    inner_idle_batch: Optional[ScheduleBatch] = None
    # Decode requests carried alongside a chunked-prefill batch
    decoding_reqs: List[Req] = None

    # For split prefill
    split_index: int = 0
    split_prefill_finished: bool = False
    split_forward_count: int = 1
    split_forward_batch: ForwardBatch = None
    # A full prefill has one result but can span several run_batch calls.
    split_prefill_start: Optional[Tuple[int, float]] = None

    req_pool_indices_cpu: torch.Tensor = None  # shape: [b], int64

    # Forward-pass metrics
    fpm_start_time: float = 0.0

    # hicache pointer for synchronizing data loading from CPU to GPU
    hicache_consumer_index: int = -1

    # Metrics
    dp_cooperation_info: Optional[DPCooperationInfo] = None
    dp_balance_stats: Optional[DPBalanceStats] = None
    prefill_stats: Optional[PrefillStats] = None
    forward_iter: Optional[int] = None
    launch_ts: Optional[float] = None
    after_idle_gap: bool = False

    # === GPU tensors crossing to ForwardBatch (clone targets for stream isolation) ===
    # Batched arguments to model runner
    input_ids: torch.Tensor = None  # shape: [b], int64
    # Staging consumed by resolve_forward_inputs (prefill H2D / mixed gather).
    prefill_input_ids_cpu: Optional[torch.Tensor] = None
    mix_running_indices: Optional[torch.Tensor] = None
    # CPU twin of mix_running_indices; lets the overlap tail resolve gather
    # pinned mirrors without a device sync.
    mix_running_indices_cpu: Optional[torch.Tensor] = None
    input_embeds: torch.Tensor = None  # shape: [b, hidden_size], float32

    # Token replacement embeddings and absolute positions (optional).
    replace_embeds: Optional[torch.Tensor] = None
    replace_positions: Optional[torch.Tensor] = None

    # Read by ForwardBatch ngram embedding init
    ne_token_table: torch.Tensor = None
    # Mask marking chunked (not-yet-finished) prefill requests whose sampled
    # pseudo next-token must NOT be written into the ngram token table.
    ne_skip_token_table_update: torch.Tensor = None
    # DeepSeek-V4.1 engram, extend batches only: [bs, n - 1] int32 predecessors
    # of each request's first extend token (NgramEmbeddingManager).
    engram_history: Optional[torch.Tensor] = None
    encoder_swa_reset: Optional[List[bool]] = None

    req_pool_indices: torch.Tensor = None  # shape: [b], int64
    seq_lens: torch.Tensor = None  # shape: [b], int64

    # The original sequence lengths, Qwen-1M related
    orig_seq_lens: torch.Tensor = None  # shape: [b], int32

    # The output locations of the KV cache
    out_cache_loc: torch.Tensor = None  # shape: [b], int64
    # DSV4-NPU: KV-only per-pool slot bundle from
    # DSV4NPUTokenToKVPoolAllocator (None elsewhere).
    out_cache_loc_dsv4: Optional[Any] = None

    # For hybrid GDN prefix cache
    mamba_track_indices: torch.Tensor = None  # shape: [b], int64
    # Per-batch snapshot of the logical ping-pong positions selected for this
    # forward (normally req.kv.mamba_next_track_idx; spec may override it). Result
    # processing uses it to update req.kv.mamba_last_track_idx, since both req-level
    # indices may advance under overlap.
    mamba_track_buffer_indices: Optional[List[int]] = None  # shape: [b], 0 or 1
    mamba_track_mask: torch.Tensor = None  # shape: [b], bool
    mamba_track_seqlens: torch.Tensor = None  # shape: [b], int64
    # TBO rejects Mamba tracking; enabling it must also slice these CPU lists.
    mamba_track_seqlens_cpu: Optional[List[int]] = None
    mamba_prefill_track_mask_cpu: Optional[List[bool]] = None
    mamba_track_mask_cpu: Optional[List[bool]] = None  # shape: [b]
    mamba_track_mask_next_cpu: Optional[List[bool]] = None  # shape: [b]
    mamba_decode_batch_idx_cpu: Optional[List[int]] = None  # shape: [b]
    # Lazy + spec: this iteration's per-req scatter positions
    # (see mamba_lazy_spec_prepare).
    mamba_lazy_spec_track_positions_cpu: Optional[List[int]] = None  # shape: [b]
    # Deferred mamba init ops: COW pairs and clear indices (performed on forward stream)
    mamba_cow_src_indices: torch.Tensor = None
    mamba_cow_dst_indices: torch.Tensor = None
    mamba_clear_indices: torch.Tensor = None

    # Encoder-decoder device tensors (host fields in the host metadata group)
    encoder_lens: Optional[torch.Tensor] = None
    encoder_out_cache_loc: Optional[torch.Tensor] = None

    # It comes empty list if logprob is not required.
    extend_input_logprob_token_ids: Optional[torch.Tensor] = None

    # === Config / flags crossing to ForwardBatch (by-value) ===
    forward_mode: ForwardMode = None
    global_forward_mode: Optional[ForwardMode] = None

    # Full-DP metadata from the existing scheduler gather.
    dp_spec_prefill_coordination_metadata: Optional[tuple] = None
    dp_spec_prefill_coordination_applied: bool = False

    # For DP attention
    is_extend_in_batch: bool = False
    can_run_decode_cuda_graph: bool = False
    # Rank-consistent EAGLE draft replay gate. Keep it separate so missing
    # draft-only state does not disable target verification or draft extend.
    can_run_dp_draft_cuda_graph: bool = False
    can_run_dp_prefill_cuda_graph: bool = False
    dp_prefill_cuda_graph_max_prefix_len: int = 0
    tbo_split_seq_index: Optional[int] = None
    # Rank-consistent forward mode for the recv skipper, derived from the MLP
    # sync all-gather (the TBO-only `global_forward_mode` is None without TBO).
    recv_skipper_forward_mode: Optional[ForwardMode] = None
    spec_verify_tier_num_tokens: int = -1

    # For processing logprobs
    return_logprob: bool = False

    # Whether this batch is prefill-only (no token generation needed)
    is_prefill_only: bool = False

    # Speculative decoding
    spec_algorithm: SpeculativeAlgorithm = None

    # Whether to return hidden states
    return_hidden_states: bool = False
    return_hidden_states_mode: CaptureHiddenMode = CaptureHiddenMode.NULL

    # Has grammar
    has_grammar: bool = False

    # The sum of all sequence lengths
    seq_lens_sum: int = None
    extend_num_tokens: Optional[int] = None

    # Diffusion LLM
    dllm_config: Optional[DllmConfig] = None

    # === Host metadata crossing to ForwardBatch (CPU lists / mirrors) ===
    seq_lens_cpu: torch.Tensor = None  # shape: [b], int64

    # For multimodal inputs
    multimodal_inputs: Optional[List] = None

    # For processing logprobs
    top_logprobs_nums: Optional[List[int]] = None
    token_ids_logprobs: Optional[List[List[int]]] = None

    # For encoder-decoder architectures
    encoder_cached: Optional[List[bool]] = None
    encoder_lens_cpu: Optional[List[int]] = None

    # For extend and mixed chunekd prefill
    prefix_lens: List[int] = None
    extend_lens: List[int] = None
    extend_logprob_start_lens: List[int] = None

    # For DP attention
    global_num_tokens: Optional[List[int]] = None
    global_num_tokens_for_logprob: Optional[List[int]] = None
    # The draft model can use a different MoE A2A backend than the target.
    draft_global_num_tokens: Optional[List[int]] = None
    draft_global_num_tokens_for_logprob: Optional[List[int]] = None
    # Full DP token vector retained for Aiter MegaMoE even when the normal MLP
    # TP gather path stores only this rank's token count.
    global_spec_verify_tier_num_tokens: Optional[List[int]] = None

    # Member rows riding one forward; None whenever reqs and rows are 1:1.
    beam_tail: Optional[BeamTail] = None

    # === Compound crossing to ForwardBatch (carry their own device tensors) ===
    # Sampling info
    sampling_info: SamplingBatchInfo = None

    # Speculative decoding
    # spec_info: Optional[SpecInput] = None
    spec_info: Optional[SpecInput] = None

    @classmethod
    def init_new(
        cls,
        reqs: List[Req],
        req_to_token_pool: ReqToTokenPool,
        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
        tree_cache: BasePrefixCache,
        model_config: ModelConfig,
        enable_overlap: bool,
        spec_algorithm: SpeculativeAlgorithm,
        chunked_req: Optional[Req] = None,
        dllm_config: Optional[DllmConfig] = None,
    ):
        return_logprob = any(req.return_logprob for req in reqs)

        return_hidden_states_mode = get_batch_return_hidden_states_mode(reqs)

        batch = cls(
            reqs=reqs,
            req_to_token_pool=req_to_token_pool,
            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
            tree_cache=tree_cache,
            model_config=model_config,
            enable_overlap=enable_overlap,
            return_logprob=return_logprob,
            has_grammar=any(req.grammar for req in reqs),
            device=req_to_token_pool.device,
            spec_algorithm=spec_algorithm,
            return_hidden_states=return_hidden_states_mode.need_capture(),
            return_hidden_states_mode=return_hidden_states_mode,
            is_prefill_only=all(req.is_prefill_only for req in reqs),
            chunked_req=chunked_req,
            chunked_req_next_prompt_token=_compute_chunked_req_next_prompt_token(
                chunked_req,
                model_config.vocab_size,
            ),
            dllm_config=dllm_config,
        )
        return batch

    def batch_size(self):
        return len(self.reqs)

    def is_empty(self):
        return len(self.reqs) == 0

    def is_dllm(self):
        return self.dllm_config is not None

    def grammar_needs_sync(self) -> bool:
        """Whether grammar forces this batch onto the synchronous path, i.e. the
        previous batch's result is resolved before this forward."""
        return self.has_grammar and not self.spec_algorithm.supports_grammar_overlap()

    def prepare_encoder_info_extend(
        self, input_ids: List[array[int]], seq_lens: List[int]
    ):
        _pin = is_pin_memory_available(self.device)
        encoder_lens_cpu = []
        encoder_cached = []

        for req in self.reqs:
            im = req.multimodal_inputs
            if im is None or im.num_image_tokens is None:
                # No image input
                encoder_lens_cpu.append(0)
                encoder_cached.append(True)
            else:
                encoder_lens_cpu.append(im.num_image_tokens)
                encoder_cached.append(
                    self.forward_mode.is_decode()
                    or req.prefix_len >= im.num_image_tokens
                )
        self.encoder_lens_cpu = encoder_lens_cpu
        self.encoder_cached = encoder_cached

        self.encoder_lens = torch.tensor(
            self.encoder_lens_cpu, dtype=torch.int64, pin_memory=_pin
        ).to(self.device, non_blocking=True)

        # Strip encoder infos
        pt = 0
        decoder_out_cache_loc = []
        encoder_out_cache_loc = []
        extend_lens = self.extend_lens[:]
        prefix_lens = self.prefix_lens[:]
        for i, req in enumerate(self.reqs):
            encoder_len = self.encoder_lens_cpu[i]
            seq_lens[i] -= encoder_len

            if req.prefix_len < encoder_len:
                # NOTE: the encoder part should be considered as a whole
                assert req.prefix_len == 0
                input_ids[i] = input_ids[i][encoder_len:]
                encoder_out_cache_loc.append(self.out_cache_loc[pt : pt + encoder_len])
                decoder_out_cache_loc.append(
                    self.out_cache_loc[pt + encoder_len : pt + req.extend_len]
                )
                extend_lens[i] -= encoder_len
                self.extend_num_tokens = self.extend_num_tokens - encoder_len
            else:
                decoder_out_cache_loc.append(
                    self.out_cache_loc[pt : pt + req.extend_len]
                )
                prefix_lens[i] -= encoder_len

            pt += req.extend_len
        self.extend_lens = extend_lens
        self.prefix_lens = prefix_lens

        # Reassign: ED stripping rebuilds prefill_input_ids_cpu (CPU pinned);
        # resolve_forward_inputs will H2D this on forward stream. self.input_ids
        # stays None.
        self.prefill_input_ids_cpu = flatten_arrays_to_pinned_cpu(input_ids, _pin)
        self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
            self.device, non_blocking=True
        )
        self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)

        if not decoder_out_cache_loc:
            self.out_cache_loc = torch.zeros(0, dtype=torch.int64).to(
                self.device, non_blocking=True
            )
        else:
            self.out_cache_loc = torch.cat(decoder_out_cache_loc)

        if not encoder_out_cache_loc:
            self.encoder_out_cache_loc = torch.zeros(0, dtype=torch.int64).to(
                self.device, non_blocking=True
            )
        else:
            self.encoder_out_cache_loc = torch.cat(encoder_out_cache_loc)

        assert len(self.out_cache_loc) == self.extend_num_tokens, (
            f"Expected {len(self.out_cache_loc)}, got {self.extend_num_tokens}"
        )

        if self.extend_input_logprob_token_ids is not None:
            new_token_ids_parts = []
            offset = 0
            extend_logprob_start_lens = self.extend_logprob_start_lens[:]
            for i, req in enumerate(self.reqs):
                encoder_len = self.encoder_lens_cpu[i]
                old_start_len = extend_logprob_start_lens[i]
                old_contribution = req.extend_len - old_start_len

                if req.prefix_len < encoder_len:
                    tokens_to_strip = max(0, encoder_len - old_start_len)
                    new_token_ids_parts.append(
                        self.extend_input_logprob_token_ids[
                            offset + tokens_to_strip : offset + old_contribution
                        ]
                    )
                    extend_logprob_start_lens[i] = max(0, old_start_len - encoder_len)
                else:
                    new_token_ids_parts.append(
                        self.extend_input_logprob_token_ids[
                            offset : offset + old_contribution
                        ]
                    )

                offset += old_contribution
            self.extend_logprob_start_lens = extend_logprob_start_lens

            if new_token_ids_parts:
                self.extend_input_logprob_token_ids = torch.cat(new_token_ids_parts)
            else:
                self.extend_input_logprob_token_ids = None

        for i, req in enumerate(self.reqs):
            encoder_len = self.encoder_lens_cpu[i]
            if encoder_len == 0:
                continue
            req.logprob_start_len = max(req.logprob_start_len, encoder_len)

    def prepare_for_extend(self):
        self.forward_mode = ForwardMode.EXTEND

        if self.is_dllm():
            # For DLLM, we use a separate forward mode
            self.forward_mode = ForwardMode.DLLM_EXTEND

        # Init tensors
        reqs = self.reqs
        input_ids = [r.get_fill_ids()[r.prefix_len :] for r in reqs]
        extend_num_tokens = sum(len(ids) for ids in input_ids)
        seq_lens = [r.extend_end for r in reqs]
        orig_seq_lens = [max(r.extend_end, len(r.origin_input_ids)) for r in reqs]
        prefix_lens = [r.prefix_len for r in reqs]
        extend_lens = [r.extend_len for r in reqs]
        extend_logprob_start_lens = [
            compute_extend_logprob_start_len(
                logprob_start_len=r.logprob_start_len,
                prefix_len=prefix_lens[i],
                extend_len=extend_lens[i],
                full_untruncated_fill_len=len(r.full_untruncated_fill_ids),
            )
            for i, r in enumerate(reqs)
        ]

        _pin = is_pin_memory_available(self.device)
        # Stay on pinned CPU; H2D is deferred to forward stream via
        # resolve_forward_inputs.
        pinned_input_ids = flatten_arrays_to_pinned_cpu(input_ids, _pin)
        seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
            self.device, non_blocking=True
        )
        seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
        orig_seq_lens_tensor = torch.tensor(
            orig_seq_lens, dtype=torch.int32, pin_memory=_pin
        ).to(self.device, non_blocking=True)

        # Set batch fields needed by alloc_for_extend
        self.prefix_lens = prefix_lens
        self.extend_lens = extend_lens
        self.seq_lens = seq_lens_tensor
        self.seq_lens_cpu = seq_lens_cpu
        self.extend_num_tokens = extend_num_tokens

        if get_exec().features.enable_encoder_swa_bounded_replay:
            for req in reqs:
                if (
                    req.multimodal_inputs is not None
                    or req.input_embeds is not None
                    or req.positional_embed_overrides is not None
                ):
                    raise ValueError(
                        "encoder SWA replay currently supports token-only text requests"
                    )
                if req.return_logprob and req.logprob_start_len not in (
                    -1,
                    len(req.origin_input_ids),
                ):
                    raise ValueError(
                        "encoder SWA replay cannot return cached prompt logprobs"
                    )
            self.encoder_swa_reset = [
                r.kv.req_pool_idx is None or r.is_retracted for r in reqs
            ]
        # Allocate memory
        out_cache_loc, req_pool_indices_tensor, req_pool_indices_cpu = alloc_for_extend(
            self
        )

        # Set fields
        input_embeds = []
        all_replace_embeds: List[torch.Tensor] = []
        all_replace_positions: List[int] = []
        has_replace_embeds = False
        input_id_pointer = 0
        input_id_lens = [len(input_id) for input_id in input_ids]
        extend_input_logprob_token_ids = []
        multimodal_inputs = []
        mamba_track_mask_cpu = []
        mamba_track_indices_cpu = []
        mamba_track_seqlens_cpu = []

        for i, (req, seq_len, pre_len) in enumerate(zip(reqs, seq_lens, prefix_lens)):
            req.extend_batch_idx += 1

            # If input_embeds are available, store them
            if req.input_embeds is not None:
                # Slice to match extend_input_len — PrefillAdder truncates
                # fill_len/extend_input_len on chunk overflow but not input_embeds.
                input_embeds.extend(
                    req.input_embeds[pre_len : pre_len + req.extend_len]
                )

            if req.positional_embed_overrides is not None:
                # Override positions are absolute in the full sequence.
                # Convert to extend-tensor coordinates by subtracting pre_len,
                # then skip any that fall within the cached prefix.
                embeds_to_add = []
                for embed_idx, pos in enumerate(
                    req.positional_embed_overrides.positions
                ):
                    extend_pos = pos - pre_len
                    if extend_pos < 0 or extend_pos >= req.extend_len:
                        continue  # Outside current extend chunk, skip
                    embeds_to_add.append((embed_idx, input_id_pointer + extend_pos))
                if embeds_to_add:
                    has_replace_embeds = True
                    indices, positions = zip(*embeds_to_add)
                    all_replace_embeds.append(
                        req.positional_embed_overrides.embeds[list(indices)]
                    )
                    all_replace_positions.extend(positions)
            input_id_pointer += input_id_lens[i]

            multimodal_inputs.append(req.multimodal_inputs)

            # Only calculate cached_tokens once. Once retracted, the 'retracted_stain'
            # flag will always True
            if not req.retracted_stain:
                new_cached = pre_len - req.already_computed
                req.cached_tokens += new_cached

                # Calculate detailed breakdown of cached tokens by source (for HiCache)
                # Only compute once on FIRST chunk - subsequent chunks in chunked prefill
                # would incorrectly count previously computed tokens as cache hits.
                if not req._cache_breakdown_computed:
                    # storage_hit_length is set after pop_prefetch_loaded_span()
                    # returns a completed prefetch.
                    (
                        req.cached_tokens_device,
                        req.cached_tokens_host,
                        req.cached_tokens_storage,
                    ) = split_cached_prefix_by_tier(
                        prefix_len=req.prefix_len,
                        host_hit_len=req.materialized_host_hit_len(),
                        storage_hit_len=req.storage_hit_length,
                        storage_hit_start=req.storage_hit_start,
                        host_hit_is_storage=req.host_hit_is_storage,
                    )
                    req._cache_breakdown_computed = True

                req.already_computed = seq_len
            req.is_retracted = False

            if get_exec().mamba.enable_mamba_extra_buffer:
                track_entry = self._mamba_radix_cache_v2_req_prepare_for_extend(req)
                mamba_track_mask_cpu.append(track_entry.track_mask)
                mamba_track_indices_cpu.append(track_entry.track_index)
                mamba_track_seqlens_cpu.append(track_entry.track_seqlen)

            if self.return_logprob:
                # Find input logprob token ids.
                # First, find a global index within origin_input_ids and slide it by 1
                # to compute input logprobs. It is because you need the next token
                # to compute input logprobs. E.g., (chunk size 2)
                #
                # input_logprobs = [1, 2, 3, 4]
                # get_fill_ids() = [1, 2]
                # extend_input_logprob_token_id = [2, 3]
                #
                # Note that it can also overflow. In this case, we pad it with 0.
                # input_logprobs = [1, 2, 3, 4]
                # get_fill_ids() = [3, 4]
                # extend_input_logprob_token_id = [4, 0]
                global_start_idx, global_end_idx = (
                    req.prefix_len,
                    req.extend_end,
                )
                if req.logprob_start_len == -1:
                    logprob_start_len = len(req.origin_input_ids)
                else:
                    logprob_start_len = req.logprob_start_len
                # Apply logprob_start_len
                if global_start_idx < logprob_start_len:
                    global_start_idx = logprob_start_len

                logprob_token_ids = req.origin_input_ids[
                    global_start_idx + 1 : global_end_idx + 1
                ]
                extend_input_logprob_token_ids.extend(logprob_token_ids)

                # We will need req.extend_len - extend_logprob_start_lens[i] number of
                # tokens, and logprob_token_ids is for input logprob, so pad the rest of them by 0.
                extend_input_logprob_token_ids.extend(
                    [0]
                    * (
                        req.extend_len
                        - extend_logprob_start_lens[i]
                        - len(logprob_token_ids)
                    )
                )

        if self.return_logprob:
            extend_input_logprob_token_ids = torch.tensor(
                extend_input_logprob_token_ids
            )
            # Clamp placeholder or out-of-range token IDs (e.g., multimodal hashes)
            # so they stay within the vocab boundary before being sent to GPU.
            extend_input_logprob_token_ids.clamp_(0, self.model_config.vocab_size - 1)
        else:
            extend_input_logprob_token_ids = None

        if has_replace_embeds:
            replace_embeds_tensor = torch.cat(all_replace_embeds, dim=0).to(
                self.device, non_blocking=True
            )
            replace_positions_tensor = torch.tensor(
                all_replace_positions, dtype=torch.long, device=self.device
            )
        else:
            replace_embeds_tensor = None
            replace_positions_tensor = None

        self.input_ids = None
        self.prefill_input_ids_cpu = pinned_input_ids
        self.req_pool_indices = req_pool_indices_tensor
        self.req_pool_indices_cpu = req_pool_indices_cpu
        self.orig_seq_lens = orig_seq_lens_tensor
        self.out_cache_loc = out_cache_loc
        self.input_embeds = (
            torch.tensor(input_embeds, pin_memory=_pin).to(
                self.device, non_blocking=True
            )
            if input_embeds
            else None
        )
        self.replace_embeds = replace_embeds_tensor
        self.replace_positions = replace_positions_tensor
        for mm_input in multimodal_inputs:
            if mm_input is None:
                continue
            if isinstance(mm_input.vision_position_ids, torch.Tensor):
                mm_input.vision_position_ids = mm_input.vision_position_ids.to(
                    self.device, non_blocking=True
                )
            if isinstance(mm_input.visible_frame_counts, torch.Tensor):
                mm_input.visible_frame_counts = mm_input.visible_frame_counts.to(
                    self.device, non_blocking=True
                )
        self.multimodal_inputs = multimodal_inputs
        self.seq_lens_sum = sum(seq_lens)

        if self.return_logprob:
            self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
            self.token_ids_logprobs = [r.logprob.token_ids_logprob for r in reqs]

        self.extend_logprob_start_lens = extend_logprob_start_lens
        self.extend_input_logprob_token_ids = extend_input_logprob_token_ids

        if get_exec().mamba.enable_mamba_extra_buffer:
            self.mamba_prefill_track_mask_cpu = mamba_track_mask_cpu
            self.mamba_track_seqlens_cpu = mamba_track_seqlens_cpu
            self.mamba_track_indices = torch.tensor(
                mamba_track_indices_cpu,
                dtype=torch.int64,
                device=self.device,
            )
            self.mamba_track_mask = torch.tensor(
                mamba_track_mask_cpu,
                dtype=torch.bool,
                device=self.device,
            )
            self.mamba_track_seqlens = torch.tensor(
                mamba_track_seqlens_cpu,
                dtype=torch.int64,
                device=self.device,
            )

        # Collect mamba init info for deferred ops on forward stream
        if any(req.kv.holds_mamba for req in reqs):
            self._collect_deferred_mamba_cow_and_clear(reqs)

        if self.model_config.is_encoder_decoder:
            self.prepare_encoder_info_extend(input_ids, seq_lens)

        # Build sampling info
        self.sampling_info = SamplingBatchInfo.from_schedule_batch(
            self,
            self.model_config.vocab_size,
        )

    def _mamba_radix_cache_v2_req_prepare_for_extend(
        self,
        req: Req,
    ) -> _MambaRadixCacheV2TrackEntry:
        cache_chunk_size = mamba_cache_chunk_size()
        state_chunk_size = getattr(
            self.model_config.hf_text_config, "mamba_chunk_size", 64
        )
        # Donated depth must land on the actual DCP-widened radix page, while
        # kernel snapshots stay on the cache chunk grid.
        checkpoint_grid = mamba_checkpoint_grid(self.tree_cache.page_size)

        def _force_track_h(i: int) -> int:
            # h is indexed relative to the extend start, so check that offset.
            assert (i - req.prefix_len) % cache_chunk_size == 0, (
                f"The force track calculation only handles last-position or "
                f"unaligned seqlens, so it needs a chunk-aligned offset to "
                f"start from. But i={i} prefix_len={req.prefix_len} "
                f"chunk_size={cache_chunk_size} checkpoint_grid={checkpoint_grid}"
            )
            # There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
            # 1) aligned with cache_chunk_size-> retrieve from last_recurrent_state
            #    a) is the last position -> retrieve from last_recurrent_state
            #    b) is NOT the last position -> retrieve from h
            # 2) unaligned with cache_chunk_size -> retrieve from h
            # Currently, the math calculation only supports case 1a and 2. So for 1b, we need to add 1
            # to force the math calculation to retrieve the correct mamba state from h.
            return i + 1

        prefix_len = req.prefix_len
        seq_end = req.extend_end
        if get_parallel().dcp_enabled:
            # DCP widens radix pages beyond scheduler chunk boundaries. Pick an
            # absolute page depth only when the kernel produced an h snapshot.
            mamba_track_seqlen_aligned = (seq_end // checkpoint_grid) * checkpoint_grid
            mask = (
                mamba_track_seqlen_aligned > prefix_len
                and (mamba_track_seqlen_aligned - prefix_len) % cache_chunk_size == 0
            )
        else:
            # Chunked prefill can leave an active request off the absolute
            # checkpoint grid. Without DCP, keep tracking snapshots relative to
            # that request prefix so later chunks can continue donating states.
            mask = req.extend_len >= checkpoint_grid
            mamba_track_seqlen_aligned = (
                prefix_len + (req.extend_len // checkpoint_grid) * checkpoint_grid
            )
        track_index = req.kv.mamba_ping_pong_track_buffer[
            req.kv.mamba_next_track_idx
        ].item()
        mamba_track_seqlen = -1
        if mask:
            # mamba_track_seqlen is used to calculate the indices to track in
            # hybrid_linear_attn_backend's _init_track_ssm_indices. Due to the
            # fact that the ssm state between aligned and non-aligned are retrieved differently,
            # if 1) last pos and 2) is aligned, then retrieved from the last_recurrent_state,
            # otherwise retrieved from h (i.e. unaligned).
            # We need to pass the non-aligned seqlen to the calculation. Even though
            # we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen.
            mamba_track_seqlen = seq_end

            # A coarser checkpoint grid may not be a model-state boundary, so
            # force retrieval from the intermediate h state in that case.
            mamba_track_fla_chunk_aligned = (
                req.prefix_len + (req.extend_len // state_chunk_size) * state_chunk_size
            )
            if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
                # We want to track mamba_track_seqlen_aligned, and it's not the last position,
                # so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
                mamba_track_seqlen = _force_track_h(mamba_track_seqlen_aligned)

            # In lazy mode, skip the swap — the second ping-pong slot is not
            # allocated yet; it will be allocated on demand at the track boundary
            # in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
            req.kv.mamba_last_track_idx = req.kv.mamba_next_track_idx
            if not get_exec().mamba.enable_mamba_extra_buffer_lazy:
                req.kv.mamba_next_track_idx = (
                    self.req_to_token_pool.get_mamba_ping_pong_other_idx(
                        req.kv.mamba_next_track_idx
                    )
                )
            if req.mamba_branching_seqlen is not None:
                # track branching point in this forward if the branching point
                # is within the current extend batch.
                branching_seqlen_aligned_mask = (
                    req.mamba_branching_seqlen - req.prefix_len
                ) % cache_chunk_size == 0
                if (
                    req.mamba_branching_seqlen > req.prefix_len
                    and req.mamba_branching_seqlen < mamba_track_seqlen
                    and branching_seqlen_aligned_mask
                ):
                    # We want to track mamba_track_seqlen_aligned, and it's not the last position,
                    # so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
                    # See _force_track_h() for more details.
                    mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen)
                    mamba_track_seqlen_aligned = req.mamba_branching_seqlen
            req.kv.mamba_last_track_seqlen = mamba_track_seqlen_aligned

        return _MambaRadixCacheV2TrackEntry(
            track_mask=mask,
            track_index=track_index,
            track_seqlen=mamba_track_seqlen,
        )

    def _collect_deferred_mamba_cow_and_clear(self, reqs):
        """Collect deferred COW/clear info from requests."""
        cow_src_tensors = []
        cow_dst_tensors = []
        clear_tensors = []
        for req in reqs:
            if req.kv.mamba_cow_src_index is not None:
                cow_src_tensors.append(req.kv.mamba_cow_src_index)
                cow_dst_tensors.append(req.kv.mamba_pool_idx.unsqueeze(0))
                req.kv.mamba_cow_src_index = None
                req.kv.mamba_needs_clear = False
            elif req.kv.mamba_needs_clear:
                clear_tensors.append(req.kv.mamba_pool_idx.unsqueeze(0))
                req.kv.mamba_needs_clear = False
        self.mamba_cow_src_indices = (
            torch.cat(cow_src_tensors) if cow_src_tensors else None
        )
        self.mamba_cow_dst_indices = (
            torch.cat(cow_dst_tensors) if cow_dst_tensors else None
        )
        self.mamba_clear_indices = torch.cat(clear_tensors) if clear_tensors else None

    def prepare_for_split_prefill(self):
        self.prepare_for_extend()
        # For split prefill, we need to set the forward mode to SPLIT_PREFILL
        self.forward_mode = ForwardMode.SPLIT_PREFILL

    def mix_with_running(self, running_batch: ScheduleBatch):
        self.forward_mode = ForwardMode.MIXED
        running_bs = running_batch.batch_size()

        # Same invariant as convert_decode_to_extend: the caller ran
        # prepare_for_decode, so a tail's prefix is its row length - 1.
        if self.spec_algorithm.is_none():
            running_prefix_lens = [s - 1 for s in running_batch.seq_lens_cpu.tolist()]
        else:
            # Spec rows sit at the committed base; seq_lens is rebuilt below.
            running_prefix_lens = [r.seqlen - 1 for r in running_batch.reqs]
        for req, prefix_len in zip(
            running_batch.reqs, running_prefix_lens, strict=True
        ):
            req.refresh_fill_ids()
            # req.prefix_len still ends at the last prefill, so req.extend_len is
            # not 1 here; this row's prefix and extend length live on the batch.
            req.extend_end = prefix_len + 1

        # Decode tokens of the running portion live in future_map.output_tokens_buf.
        self.input_ids = None
        self.mix_running_indices = running_batch.req_pool_indices
        self.mix_running_indices_cpu = running_batch.req_pool_indices_cpu
        if not self.spec_algorithm.is_none():
            # Spec keeps no per-step out_cache_loc on the running batch; gather
            # each tail's bonus slot at the committed length (rebound under overlap).
            tail_base = torch.tensor(
                [r.seqlen - 1 for r in running_batch.reqs],
                dtype=torch.int64,
                device=self.seq_lens.device,
            )
            running_out_cache_loc = self.req_to_token_pool.req_to_token[
                running_batch.req_pool_indices.long(),
                tail_base,
            ].to(self.out_cache_loc.dtype)
            # The spec relay is unresolved at schedule time, so merge_batch
            # would null seq_lens_cpu; rebuild the tails from request state.
            running_seq_lens_cpu = torch.tensor(
                [int(r.seqlen) for r in running_batch.reqs], dtype=torch.int64
            )
            if self.seq_lens_cpu is None:
                merged_seq_lens_cpu = running_seq_lens_cpu
            else:
                merged_seq_lens_cpu = torch.cat(
                    [self.seq_lens_cpu, running_seq_lens_cpu]
                )
        else:
            # Non-spec: the running batch carries prepared seq_lens_cpu
            # (r.seqlen lags it under overlap); merge_batch concats it.
            tail_base = None
            running_out_cache_loc = running_batch.out_cache_loc
            merged_seq_lens_cpu = None
        out_cache_loc = torch.cat([self.out_cache_loc, running_out_cache_loc])

        self.merge_batch(running_batch)
        self.out_cache_loc = out_cache_loc
        if merged_seq_lens_cpu is not None:
            self.seq_lens_cpu = merged_seq_lens_cpu
        if tail_base is not None:
            # Spec seq_lens sit at the committed base (bonus token pending);
            # this step commits it, so tails carry base + 1 or attention drops the row.
            merged = self.seq_lens.clone()
            merged[-running_bs:] = tail_base + 1
            self.seq_lens = merged

        # NOTE: prefix_len is what has been cached, but we don't cache each decode step
        self.prefix_lens = self.prefix_lens + running_prefix_lens
        self.extend_lens = self.extend_lens + [1] * running_bs
        self.extend_num_tokens = self.extend_num_tokens + running_bs
        # TODO (lianmin): Revisit this. It should be seq_len - 1
        self.extend_logprob_start_lens = (
            self.extend_logprob_start_lens + [0] * running_bs
        )
        self.is_prefill_only = False

    def convert_decode_to_extend(self):
        """View every decode request as a 1-token extend with its context as
        prefix, making this a plain extend (prefill) batch. Called after the
        DP mlp-sync when a peer rank runs extend, so this rank replays the
        extend graphs; requires prepare_for_decode(); token count unchanged.
        """
        bs = self.batch_size()
        self.forward_mode = ForwardMode.EXTEND
        # Stale residue from this object's life as a prefill batch; the
        # extend-merge path would exclude (silently drop) that req otherwise.
        self.chunked_req = None
        # Also stale residue; None keeps the prefill result path from
        # re-reporting old prefill stats for what is decode work.
        self.prefill_stats = None
        # A 1-token extend's position is arange(prefix, prefix + 1), so prefix
        # must be seq_len - 1; output_ids trails the row by one or zero and
        # cannot stand in for it. Rows past bs are a beam tail, not requests.
        seq_lens = self.seq_lens_cpu[:bs].tolist()
        for req, seq_len in zip(self.reqs, seq_lens, strict=True):
            req.refresh_fill_ids()
            # end runs one past full_untruncated_fill_ids while output_ids
            # trails; safe only while decoding_reqs suppresses the checkpoint insert.
            # As in mix_with_running, read prefix and extend length off the batch.
            req.extend_end = seq_len

        self.prefix_lens = [seq_len - 1 for seq_len in seq_lens]
        self.extend_lens = [1] * bs
        self.extend_num_tokens = bs
        self.extend_logprob_start_lens = [0] * bs
        self.decoding_reqs = self.reqs
        self.is_prefill_only = False

    def new_tokens_required_next_decode(
        self, selected_indices: Optional[List[int]] = None
    ):
        page_size = self.token_to_kv_pool_allocator.page_size
        requests = (
            self.reqs
            if selected_indices is None
            else [self.reqs[i] for i in selected_indices]
        )

        if self.spec_algorithm.is_none():
            new_pages = sum(
                1 for r in requests if r.kv.kv_committed_len % page_size == 0
            )
            return new_pages * page_size + num_beam_member_rows(requests)

        return self._new_tokens_required_next_decode_spec_v2(requests, page_size)

    def _new_tokens_required_next_decode_spec_v2(self, requests, page_size):
        """Tight estimate matching eagle_utils.eagle_prepare_for_decode allocation."""
        reserve = get_alloc_reserve_per_decode()
        total = 0
        for r in requests:
            x = max(0, r.kv.kv_committed_len + reserve - r.kv.kv_allocated_len)
            cur = r.kv.kv_allocated_len
            nxt = cur + x
            total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
        return total

    def check_decode_mem(self, selected_indices: Optional[List[int]] = None):
        """Whether the next decode step fits in the KV pool. The ALLOCATOR owns
        the capacity gate (eviction + any per-step reservations of its own) —
        the retract loop converges on this same check, so allocator-side
        shortfalls retract gracefully instead of tripping fail-loud alloc
        errors."""
        num_tokens = self.new_tokens_required_next_decode(selected_indices)
        requests = (
            self.reqs
            if selected_indices is None
            else [self.reqs[i] for i in selected_indices]
        )
        return self.token_to_kv_pool_allocator.check_decode_capacity(
            num_tokens=num_tokens,
            tree_cache=self.tree_cache,
            requests=requests,
            spec_algorithm=self.spec_algorithm,
        )

    def retract_decode(self) -> Tuple[List[Req], float, List[Req]]:
        """Retract the decoding requests when there is not enough memory."""
        sorted_indices = self._get_decode_retraction_order(self.reqs)
        sorted_indices = beam_retraction_order(sorted_indices, self.reqs)

        retracted_reqs = []
        reqs_to_abort: List[Req] = []
        first_iter = True
        while first_iter or (
            not self.check_decode_mem(selected_indices=sorted_indices)
        ):
            if len(sorted_indices) == 1:
                # Always keep at least one request
                break

            first_iter = False
            idx = sorted_indices.pop()
            req = self.reqs[idx]
            if req.beam_group is not None:
                # Free member rows before release_req: they need the leader's kv
                # info, which still carries the lockstep allocated length.
                req.to_finish = FINISH_ABORT(
                    "Beam search group aborted: KV cache pool is full. Beam "
                    "groups cannot be retracted, so they are aborted instead "
                    "of being requeued.",
                    status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                )
                reqs_to_abort.append(req)
                free_member_rows(
                    req.beam_group,
                    self.req_to_token_pool,
                    self.token_to_kv_pool_allocator,
                )
                # Aborting, so a host backup to resume from would be wasted.
                self.release_req(idx, len(sorted_indices), offload_kv=False)
                continue
            # release memory and don't insert into the tree because we need the space instantly
            if self.release_req(idx, len(sorted_indices)):
                retracted_reqs.append(req)
            else:
                # No backup exists and the device KV is already freed, so the
                # request cannot resume.
                if get_disagg().disaggregation_decode_retraction_backup == "host_pool":
                    message = "Retraction host KV pool exhausted. Aborting the request."
                    status_code = HTTPStatus.INTERNAL_SERVER_ERROR
                else:
                    message = (
                        "Retracted under decode memory pressure without a KV "
                        "backup. Retry later."
                    )
                    status_code = HTTPStatus.SERVICE_UNAVAILABLE
                req.to_finish = FINISH_ABORT(message, status_code=status_code)
                reqs_to_abort.append(req)
                logger.warning(
                    "retract_decode: aborted request %s: %s", req.rid, message
                )

        if len(sorted_indices) <= 1 and not self.check_decode_mem(
            selected_indices=sorted_indices
        ):
            # Even the last remaining request cannot fit in memory.
            # Instead of crashing the scheduler, gracefully abort it.
            last_idx = sorted_indices.pop()
            last_req = self.reqs[last_idx]
            last_req.to_finish = FINISH_ABORT(
                "Out of memory even after retracting all other requests "
                "in the decode batch. Aborting the last request.",
                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
            )
            reqs_to_abort.append(last_req)
            if last_req.beam_group is not None:
                free_member_rows(
                    last_req.beam_group,
                    self.req_to_token_pool,
                    self.token_to_kv_pool_allocator,
                )
            self.release_req(last_idx, 0, offload_kv=False)
            logger.warning(
                "retract_decode: aborted last request %s due to OOM", last_req.rid
            )

        self.filter_batch(keep_indices=sorted_indices)

        # Reqs in batch are filtered
        new_estimate_ratio = (
            NewTokenRatioTracker.estimate_new_token_ratio_after_retract(self.reqs)
        )

        return retracted_reqs, new_estimate_ratio, reqs_to_abort

    @staticmethod
    def _get_decode_retraction_order(reqs: List[Req]) -> List[int]:
        """Return indices ordered from most-preferred to least-preferred to keep.

        The retraction loop pops from the end of this list, so the least-preferred
        request is retracted first.
        """
        sorted_indices = list(range(len(reqs)))

        # TODO(lsyin): improve retraction policy for radix cache

        def length_key(req: Req) -> Tuple[int, int]:
            return (len(req.output_ids), -len(req.origin_input_ids))

        if get_schedule().retraction_policy == "priority":
            priority_sign = (
                1 if get_schedule().schedule_low_priority_values_first else -1
            )

            def retraction_key(req: Req) -> Tuple[int, int, int]:
                priority = req.priority
                if priority is None:
                    priority = (
                        sys.maxsize
                        if get_schedule().schedule_low_priority_values_first
                        else -sys.maxsize - 1
                    )
                return (priority * (-priority_sign), *length_key(req))

            sorted_indices.sort(
                key=lambda i: retraction_key(reqs[i]),
                reverse=True,
            )
            return sorted_indices

        sorted_indices.sort(
            key=lambda i: length_key(reqs[i]),
            reverse=True,
        )
        return sorted_indices

    def release_req(
        self,
        idx: int,
        remaing_req_count: int,
        offload_kv: bool = True,
    ) -> bool:
        return release_req(
            req=self.reqs[idx],
            remaing_req_count=remaing_req_count,
            req_to_token_pool=self.req_to_token_pool,
            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
            tree_cache=self.tree_cache,
            hisparse_coordinator=self.hisparse_coordinator,
            offload_kv=offload_kv,
        )

    def prepare_encoder_info_decode(self):
        # Reset the encoder cached status
        self.encoder_cached = [True] * len(self.reqs)

    def prepare_for_idle(self):
        self.forward_mode = ForwardMode.IDLE
        self.input_ids = torch.empty(0, dtype=torch.int64, device=self.device)
        self.seq_lens = torch.empty(0, dtype=torch.int64, device=self.device)
        self.seq_lens_cpu = torch.empty(0, dtype=torch.int64)
        self.orig_seq_lens = torch.empty(0, dtype=torch.int32, device=self.device)
        self.out_cache_loc = torch.empty(0, dtype=torch.int64, device=self.device)
        self.req_pool_indices = torch.empty(0, dtype=torch.int64, device=self.device)
        self.req_pool_indices_cpu = torch.empty(0, dtype=torch.int64)
        self.seq_lens_sum = 0
        self.extend_num_tokens = 0
        self.sampling_info = SamplingBatchInfo.from_schedule_batch(
            self,
            self.model_config.vocab_size,
        )

    def mamba_lazy_prealloc_at_boundary(self, mamba_track_interval: int):
        """Allocate a temporary second ping-pong slot for reqs at a track boundary.

        In lazy mode each request normally holds only 1 ping-pong slot.
        When seq_len hits a track interval boundary, we allocate the
        second slot so the forward pass can write the new tracked state
        there. The old slot is freed after the forward in
        mamba_lazy_post_decode_at_boundary.
        """
        pool = self.req_to_token_pool
        for i, req in enumerate(self.reqs):
            buf = req.kv.mamba_ping_pong_track_buffer
            assert buf is not None
            # Skip reqs not at a track boundary
            if self.seq_lens_cpu[i].item() % mamba_track_interval != 0:
                continue
            other_idx = 1 - req.kv.mamba_next_track_idx
            if buf[other_idx].item() != -1:
                # With overlap the previous forward's post-processing
                # (which frees this slot) hasn't run yet. Skip.
                continue
            if envs.SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL.get():
                new_slot = None
            else:
                # No evict-retry: a transient slot is not worth evicting a
                # cached checkpoint for; on failure tracking degrades in place.
                new_slot = pool.mamba_allocator.alloc(1)
            if new_slot is not None:
                pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0])
                req.kv.mamba_next_track_idx = other_idx
            else:
                # The in-flight write invalidates the keep slot's old checkpoint.
                req.kv.mamba_last_track_seqlen = None

    def mamba_lazy_spec_prepare(self, mamba_track_interval: int, max_draft_tokens: int):
        """Lazy-mode spec counterpart of mamba_lazy_prealloc_at_boundary.

        A crossing is only *possible* at prepare time (accept length is
        unknown), so ensure the pending slot exists for reqs whose next
        boundary is reachable, WITHOUT swapping mamba_next_track_idx; the
        mask-gated commit writes it only on a real crossing, and
        _mamba_lazy_spec_confirm_crossing promotes it afterwards. The per-req
        scatter position is recorded on the batch and rides the result-queue
        copy (forward isolation restores batch fields).
        """
        pool = self.req_to_token_pool
        track_positions: List[int] = []
        for req in self.reqs:
            buf = req.kv.mamba_ping_pong_track_buffer
            assert buf is not None
            if not mamba_lazy_spec_in_window(
                req, mamba_track_interval, max_draft_tokens
            ):
                # No crossing reachable: the scatter mask stays -1, the
                # position is never written.
                track_positions.append(req.kv.mamba_next_track_idx)
                continue
            other_idx = 1 - req.kv.mamba_next_track_idx
            has_pending = buf[other_idx].item() != -1
            if not has_pending:
                if envs.SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL.get():
                    new_slot = None
                else:
                    # No evict-retry: a transient slot is not worth
                    # evicting a cached checkpoint for.
                    new_slot = pool.mamba_allocator.alloc(1)
                if new_slot is not None:
                    pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0])
                    has_pending = True
                else:
                    req.kv.mamba_last_track_seqlen = None
            # On failure the verify scatters in place into the keep slot.
            track_positions.append(
                other_idx if has_pending else req.kv.mamba_next_track_idx
            )
        self.mamba_lazy_spec_track_positions_cpu = track_positions

    def cumulate_penalty_output_tokens(self):
        # Under overlap batch.input_ids is just a placeholder here -- the
        # real token is relayed via future_map and resolved at forward
        # entry. So take the last output token from Req directly
        # (origin_input_ids[-1] on the first decode, before any output).
        last_tokens = [
            req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1]
            for req in self.reqs
        ]
        # Non-blocking H2D so this per-step copy doesn't sync behind the forward.
        # pin_memory (matching the prefill-path tensors) keeps the copy async;
        # is_pin_memory_available falls back to pageable on unsupported devices.
        latest_output_ids = torch.tensor(
            last_tokens,
            dtype=torch.int64,
            pin_memory=is_pin_memory_available(self.device),
        ).to(self.device, non_blocking=True)
        self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
            latest_output_ids
        )

    def prepare_for_decode(self):
        self.forward_mode = ForwardMode.DECODE
        self.mamba_track_seqlens_cpu = None
        self.mamba_prefill_track_mask_cpu = None
        # Decode embeds the last output token via embed_tokens; clear the stale
        # prefill-time tensor so it doesn't leak into ForwardBatch.
        self.input_embeds = None

        self.mamba_cow_src_indices = None
        self.mamba_cow_dst_indices = None
        self.mamba_clear_indices = None

        # Clear context parallel metadata - CP is only for prefill, not decode
        if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None:
            self.attn_cp_metadata = None

        if not self.spec_algorithm.is_none():
            # Spec decoding owns decode preparation (allocation, seq-lens bookkeeping).
            from sglang.srt.speculative.spec_utils import spec_prepare_for_decode

            self.mamba_track_mask_cpu = None
            self.mamba_track_mask_next_cpu = None
            self.mamba_decode_batch_idx_cpu = None
            spec_prepare_for_decode(self)
            return

        # Beam member rows ride this decode batch: append them to the row
        # tensors before allocation so alloc_for_decode covers them too.
        strip_beam_tail(self)
        append_beam_tail(self)

        if self.sampling_info.penalizer_orchestrator.is_required:
            self.cumulate_penalty_output_tokens()

        # input_ids is set at end of previous run_batch (placeholder for
        # overlap; next_token_ids cast for non-overlap).

        if self.model_config.is_encoder_decoder:
            self.prepare_encoder_info_decode()

        # Allocate memory (DSV4-NPU c{4,128}_state alloc lens are computed inside
        # the allocator, triggered from mem_cache/common.py.)
        self.out_cache_loc = alloc_for_decode(self, token_per_req=1)

        for req in self.reqs:
            req.decode_batch_idx += 1

        # New-tensor avoids racing model_worker_batch refs queued for
        # overlap forward.
        self.seq_lens = self.seq_lens + 1
        self.seq_lens_cpu = self.seq_lens_cpu + 1
        self.orig_seq_lens = self.orig_seq_lens + 1
        # Sum is recomputed lazily by ForwardBatch.init_new.
        self.seq_lens_sum = None

        if self.hisparse_coordinator is not None:
            self.hisparse_coordinator.map_last_loc_to_buffer(
                self.seq_lens,
                self.out_cache_loc,
                self.req_pool_indices,
                self.seq_lens_cpu,
                self.req_pool_indices_cpu,
            )

        if get_exec().mamba.enable_mamba_extra_buffer:
            mamba_track_interval = mamba_track_grid(self.tree_cache.page_size)

            if len(self.reqs) == 0:
                self.mamba_track_indices = torch.empty(
                    (0,), dtype=torch.int64, device=self.device
                )
                self.mamba_track_buffer_indices = []
            else:
                if get_exec().mamba.enable_mamba_extra_buffer_lazy:
                    self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
                set_mamba_track_indices_from_reqs(self)

            track_remainders_cpu = self.seq_lens_cpu % mamba_track_interval
            track_mask_cpu = track_remainders_cpu == 0
            self.mamba_track_mask_cpu = track_mask_cpu.tolist()
            self.mamba_track_mask_next_cpu = (
                (track_remainders_cpu == mamba_track_interval - 1).tolist()
                if self.enable_overlap
                else None
            )
            # ScheduleBatch.copy() snapshots the list of requests, but the Req
            # objects remain shared. The next overlapped decode can therefore
            # advance their counters before this batch's result is processed.
            self.mamba_decode_batch_idx_cpu = [
                req.decode_batch_idx for req in self.reqs
            ]
            # async H2D
            self.mamba_track_mask = track_mask_cpu.pin_memory().to(
                device=self.device, non_blocking=True
            )

    def filter_batch(
        self,
        reqs_to_exclude: Optional[Union[Req, List[Req]]] = None,
        keep_indices: Optional[List[int]] = None,
    ):
        strip_beam_tail(self)
        if keep_indices is None:
            if isinstance(reqs_to_exclude, Req):
                reqs_to_exclude = [reqs_to_exclude]
            elif reqs_to_exclude is None:
                reqs_to_exclude = []
            keep_indices = [
                i
                for i in range(len(self.reqs))
                if not self.reqs[i].finished() and self.reqs[i] not in reqs_to_exclude
            ]

        if keep_indices is None or len(keep_indices) == 0:
            # Filter out all requests. Stale tensors are left as-is: is_empty()
            # keys off reqs, so callers drop the batch before a forward reads them.
            self.reqs = []
            self.return_hidden_states = False
            self.return_hidden_states_mode = CaptureHiddenMode.NULL
            return

        if len(keep_indices) == len(self.reqs):
            # No need to filter
            return

        keep_indices_device = torch.tensor(
            keep_indices,
            dtype=torch.int64,
            pin_memory=is_pin_memory_available(self.device),
        ).to(self.device, non_blocking=True)

        if self.model_config.is_encoder_decoder:
            self.encoder_lens = self.encoder_lens[keep_indices_device]
            self.encoder_lens_cpu = [self.encoder_lens_cpu[i] for i in keep_indices]

        self.reqs = [self.reqs[i] for i in keep_indices]
        if self.multimodal_inputs is not None:
            self.multimodal_inputs = [self.multimodal_inputs[i] for i in keep_indices]
        self.req_pool_indices = self.req_pool_indices[keep_indices_device]
        self.req_pool_indices_cpu = self.req_pool_indices_cpu[keep_indices]
        self.seq_lens = self.seq_lens[keep_indices_device]
        self.orig_seq_lens = self.orig_seq_lens[keep_indices_device]
        self.out_cache_loc = None
        # Sum is recomputed lazily by ForwardBatch.init_new.
        self.seq_lens_sum = None

        if self.input_ids is not None:
            self.input_ids = self.input_ids[keep_indices_device]
        # Optional under no-verify-sync; resolve_seq_lens repopulates before forward.
        if self.seq_lens_cpu is not None:
            self.seq_lens_cpu = self.seq_lens_cpu[keep_indices]

        self.mamba_track_indices = None
        self.mamba_track_buffer_indices = None
        self.mamba_track_mask = None
        self.mamba_track_seqlens = None
        self.mamba_track_seqlens_cpu = None
        self.mamba_prefill_track_mask_cpu = None
        self.mamba_track_mask_cpu = None
        self.mamba_track_mask_next_cpu = None
        self.mamba_decode_batch_idx_cpu = None
        self.mamba_lazy_spec_track_positions_cpu = None
        self.mamba_cow_src_indices = None
        self.mamba_cow_dst_indices = None
        self.mamba_clear_indices = None
        self.return_logprob = any(req.return_logprob for req in self.reqs)
        if self.return_logprob:
            self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in keep_indices]
            self.token_ids_logprobs = [self.token_ids_logprobs[i] for i in keep_indices]
        else:
            self.top_logprobs_nums = None
            self.token_ids_logprobs = None

        self.has_grammar = any(req.grammar for req in self.reqs)
        self.return_hidden_states_mode = get_batch_return_hidden_states_mode(self.reqs)
        self.return_hidden_states = self.return_hidden_states_mode.need_capture()

        self.sampling_info.filter_batch(keep_indices, keep_indices_device)
        if self.spec_info:
            self.spec_info.filter_batch(
                new_indices=keep_indices_device,
                new_indices_cpu=keep_indices,
            )

    def merge_batch(self, other: ScheduleBatch):
        strip_beam_tail(self)
        strip_beam_tail(other)
        # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because
        # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it
        # needs to be called with pre-merged Batch.reqs.
        self.sampling_info.merge_batch(other.sampling_info)

        # Encoder-decoder infos
        if self.model_config.is_encoder_decoder:
            self.encoder_lens = torch.cat([self.encoder_lens, other.encoder_lens])
            self.encoder_lens_cpu = self.encoder_lens_cpu + other.encoder_lens_cpu
        self.req_pool_indices = torch.cat(
            [self.req_pool_indices, other.req_pool_indices]
        )
        self.req_pool_indices_cpu = torch.cat(
            [self.req_pool_indices_cpu, other.req_pool_indices_cpu]
        )
        self.seq_lens = torch.cat([self.seq_lens, other.seq_lens])
        self.orig_seq_lens = torch.cat([self.orig_seq_lens, other.orig_seq_lens])
        self.out_cache_loc = None
        # Sum is recomputed lazily by ForwardBatch.init_new.
        self.seq_lens_sum = None
        # Cat only when both sides hold a real token tensor; otherwise drop to
        # None and let resolve_forward_inputs rebuild from the merged
        # req_pool_indices. Mismatch arises e.g. with spec_v1, which keeps its
        # tensor while a relay-staged side is None -- there the worker rebuilds.
        if self.input_ids is not None and other.input_ids is not None:
            self.input_ids = torch.cat([self.input_ids, other.input_ids])
        else:
            self.input_ids = None
        # Optional under no-verify-sync; drop the mirror if either side absent.
        if self.seq_lens_cpu is None or other.seq_lens_cpu is None:
            self.seq_lens_cpu = None
        else:
            self.seq_lens_cpu = torch.cat([self.seq_lens_cpu, other.seq_lens_cpu])
        self.mamba_track_indices = None
        self.mamba_track_buffer_indices = None
        self.mamba_track_mask = None
        self.mamba_track_seqlens = None
        self.mamba_track_seqlens_cpu = None
        self.mamba_prefill_track_mask_cpu = None
        self.mamba_track_mask_cpu = None
        self.mamba_track_mask_next_cpu = None
        self.mamba_decode_batch_idx_cpu = None
        self.mamba_lazy_spec_track_positions_cpu = None
        if self.return_logprob and other.return_logprob:
            self.top_logprobs_nums = self.top_logprobs_nums + other.top_logprobs_nums
            self.token_ids_logprobs = self.token_ids_logprobs + other.token_ids_logprobs
        elif self.return_logprob:
            self.top_logprobs_nums = self.top_logprobs_nums + [0] * len(other.reqs)
            self.token_ids_logprobs = self.token_ids_logprobs + [None] * len(other.reqs)
        elif other.return_logprob:
            self.top_logprobs_nums = [0] * len(self.reqs) + other.top_logprobs_nums
            self.token_ids_logprobs = [None] * len(self.reqs) + other.token_ids_logprobs
        self.reqs = self.reqs + other.reqs
        if self.multimodal_inputs is not None:
            self.multimodal_inputs = self.multimodal_inputs + other.multimodal_inputs

        self.return_logprob = self.return_logprob or other.return_logprob
        self.has_grammar = self.has_grammar or other.has_grammar
        self.return_hidden_states_mode = max(
            self.return_hidden_states_mode, other.return_hidden_states_mode
        )
        self.return_hidden_states = self.return_hidden_states_mode.need_capture()
        self.is_prefill_only = self.is_prefill_only and other.is_prefill_only

        if self.spec_info:
            self.spec_info.merge_batch(other.spec_info)

    def copy(self):
        # Only contain fields that will be used by process_batch_result.
        # Shallow-copy the reqs list as a defensive snapshot. filter_batch and
        # merge_batch historically mutated the list in place; they now rebind
        # new lists, but the slice stays so this snapshot never aliases the
        # original.
        return ScheduleBatch(
            reqs=self.reqs[:],
            extend_lens=self.extend_lens,
            prefix_lens=self.prefix_lens,
            req_to_token_pool=self.req_to_token_pool,
            req_pool_indices=self.req_pool_indices,
            req_pool_indices_cpu=self.req_pool_indices_cpu,
            model_config=self.model_config,
            forward_mode=self.forward_mode,
            out_cache_loc=self.out_cache_loc,
            return_logprob=self.return_logprob,
            has_grammar=self.has_grammar,
            return_hidden_states=self.return_hidden_states,
            return_hidden_states_mode=self.return_hidden_states_mode,
            decoding_reqs=self.decoding_reqs,
            spec_algorithm=self.spec_algorithm,
            spec_info=self.spec_info,
            dp_spec_prefill_coordination_metadata=self.dp_spec_prefill_coordination_metadata,
            dp_spec_prefill_coordination_applied=self.dp_spec_prefill_coordination_applied,
            global_num_tokens=self.global_num_tokens,
            global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
            can_run_decode_cuda_graph=self.can_run_decode_cuda_graph,
            can_run_dp_draft_cuda_graph=self.can_run_dp_draft_cuda_graph,
            can_run_dp_prefill_cuda_graph=self.can_run_dp_prefill_cuda_graph,
            dp_prefill_cuda_graph_max_prefix_len=self.dp_prefill_cuda_graph_max_prefix_len,
            is_extend_in_batch=self.is_extend_in_batch,
            is_prefill_only=self.is_prefill_only,
            seq_lens_cpu=self.seq_lens_cpu,
            enable_overlap=self.enable_overlap,
            mamba_track_indices=self.mamba_track_indices,
            mamba_track_buffer_indices=self.mamba_track_buffer_indices,
            mamba_track_mask=self.mamba_track_mask,
            mamba_track_seqlens=self.mamba_track_seqlens,
            mamba_track_seqlens_cpu=self.mamba_track_seqlens_cpu,
            mamba_prefill_track_mask_cpu=self.mamba_prefill_track_mask_cpu,
            mamba_track_mask_cpu=self.mamba_track_mask_cpu,
            mamba_track_mask_next_cpu=self.mamba_track_mask_next_cpu,
            mamba_decode_batch_idx_cpu=self.mamba_decode_batch_idx_cpu,
            mamba_lazy_spec_track_positions_cpu=self.mamba_lazy_spec_track_positions_cpu,
            dp_cooperation_info=self.dp_cooperation_info,
            dp_balance_stats=self.dp_balance_stats,
            prefill_stats=self.prefill_stats,
            fpm_start_time=self.fpm_start_time,
            forward_iter=self.forward_iter,
            launch_ts=self.launch_ts,
            after_idle_gap=self.after_idle_gap,
            split_prefill_start=self.split_prefill_start,
            extend_num_tokens=self.extend_num_tokens,
        )

    def _swa_eviction_trigger(self, sliding_window_size: int, eviction_interval: int):
        """Pick this forward's per-request SWA eviction trigger.

        Default: evict only once >= eviction_interval tokens have slid out of
        the window, amortizing eviction work while keeping each request's
        overshoot within the interval the pool budget reserves. Gating on
        accumulated tokens rather than an iteration-counter phase cannot
        starve, because seqlen progress is monotonic per KV handle.

        Aiter MegaMoE DSV4 keeps the older forward-interval cadence instead:
        per-request token gating synchronizes SWA pressure across DP ranks and
        triggers a retraction / re-prefill storm at high concurrency.
        """
        if _is_hip and envs.SGLANG_AMD_USE_FLYDSL_MEGA_MOE.get():
            due_this_forward = (self.forward_iter or 0) % eviction_interval == 0
            return lambda req: due_this_forward
        return lambda req: (
            req.seqlen - 1 - sliding_window_size
            >= req.kv.get_evicted_seqlen(ComponentType.SWA) + eviction_interval
        )

    def maybe_evict_swa(self):
        if self.tree_cache.supports_swa():
            sliding_window_size = self.tree_cache.sliding_window_size
            # Auxiliary windows check their own cursors and prefix locks.
            has_auxiliary_swa = self.tree_cache.supports_auxiliary_swa()

            release_leaf_lock = envs.SGLANG_OPT_RELEASE_PREFILL_SWA.get() and hasattr(
                self.tree_cache, "release_swa_prefix_lock"
            )

            eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get())
            swa_evict_due = self._swa_eviction_trigger(
                sliding_window_size, eviction_interval
            )
            self.token_to_kv_pool_allocator.free_group_begin()
            for idx, req in enumerate(self.reqs):
                if self.forward_mode.is_decode():
                    # In overlap scheduler, we cannot evict swa when
                    # req.decode_batch_idx == 0 since the prev extend batch is
                    # still running. `_swa_eviction_trigger` owns the rest.
                    if (
                        req.decode_batch_idx >= 1
                        and req.kv.holds_kv
                        and (has_auxiliary_swa or swa_evict_due(req))
                    ):
                        # The scheduler already gated the primary window,
                        # including backends that use a forward-based cadence.
                        self._evict_swa(
                            req,
                            req.seqlen - 1,
                            eviction_interval=(
                                eviction_interval if has_auxiliary_swa else 1
                            ),
                        )

                    # Past the window the request no longer reads its prefill's SWA;
                    # release that part of the tree lock so SWA LRU can reclaim it.
                    if (
                        release_leaf_lock
                        and req.decode_batch_idx >= sliding_window_size
                    ):
                        self.tree_cache.release_swa_prefix_lock(req)
                elif (
                    self.forward_mode.is_extend()
                    and not self.tree_cache.supports_prefix_sharing()
                ):
                    pre_len = self.prefix_lens[idx]
                    if self.enable_overlap:
                        # In chunked prefill case, when the second extend batch is scheduling, the first extend batch is still running, so we cannot evict swa tokens
                        if req.extend_batch_idx < 2:
                            continue
                        else:
                            pre_len = (
                                pre_len - get_schedule().chunked_prefill_size
                                if get_schedule().chunked_prefill_size > 0
                                else pre_len
                            )
                            self._evict_swa(req, pre_len)
                    else:
                        self._evict_swa(req, pre_len)
            self.token_to_kv_pool_allocator.free_group_end()

    def _evict_swa(self, req: Req, pre_len: int, *, eviction_interval: int = 1):
        assert self.tree_cache.supports_swa(), "prefix cache must support swa"
        self.tree_cache.evict_sliding_windows(
            req, pre_len, eviction_interval=eviction_interval
        )

    def __str__(self):
        return (
            f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "
            f"#req={(len(self.reqs))})"
        )


class NextBatchPlan(msgspec.Struct):
    batch_to_run: Optional[ScheduleBatch]
    running_batch: ScheduleBatch