ThomasX
099bcfb41e
In disaggregated prefill with --enable-single-batch-overlap, batch.copy() was used to defer process_batch_result, but the copy omitted spec_info. This caused req.hidden_states_tensor to be None when using Eagle/spec-v2 speculative decoding on the decode server, silently degrading draft quality for the first decode round.
2525 lines
100 KiB
Python
2525 lines
100 KiB
Python
from __future__ import annotations
|
|
|
|
from sglang.srt.dllm.config import DllmConfig
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|
from sglang.srt.utils.common import ceil_align, is_pin_memory_available
|
|
|
|
# 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 -> ModelWorkerBatch -> ForwardBatch
|
|
|
|
- ScheduleBatch is managed by `scheduler.py::Scheduler`.
|
|
It contains high-level scheduling data. Most of the data is on the CPU.
|
|
- ModelWorkerBatch is managed by `tp_worker.py::TpModelWorker`.
|
|
It is a subset of `ScheduleBatch` that only contains data related to the model forward on GPU.
|
|
It will be transformed from CPU scheduler to GPU model runner.
|
|
- ForwardBatch is managed by `model_runner.py::ModelRunner`.
|
|
It contains low-level tensor data. Most of the data consists of GPU tensors.
|
|
|
|
TODO(lmzheng): ModelWorkerBatch seems a bit redundant and we consider removing it in the future.
|
|
"""
|
|
|
|
import copy
|
|
import dataclasses
|
|
import logging
|
|
import re
|
|
from concurrent.futures import Future
|
|
from enum import Enum, auto
|
|
from functools import lru_cache
|
|
from http import HTTPStatus
|
|
from itertools import chain
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
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 DisaggregationMode
|
|
from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank
|
|
from sglang.srt.dllm.mixin.req import ReqDllmMixin
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchPrefixParams
|
|
from sglang.srt.mem_cache.common import (
|
|
alloc_for_decode,
|
|
alloc_for_extend,
|
|
evict_from_tree_cache,
|
|
release_kv_cache,
|
|
)
|
|
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
|
from sglang.srt.mem_cache.radix_cache import RadixKey
|
|
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardBatch,
|
|
ForwardMode,
|
|
)
|
|
from sglang.srt.observability.metrics_collector import (
|
|
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_params import SamplingParams
|
|
from sglang.srt.server_args import ServerArgs, get_global_server_args
|
|
from sglang.srt.utils import flatten_nested_list
|
|
from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
|
|
|
|
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.session_controller import Session
|
|
from sglang.srt.observability.scheduler_metrics_mixin import PrefillStats
|
|
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
|
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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@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 _compute_pad_value(hash: int) -> int:
|
|
"""Compute pad value from hash."""
|
|
return MM_PAD_SHIFT_VALUE + (hash % (1 << 30))
|
|
|
|
|
|
class BaseFinishReason:
|
|
def __init__(self, is_error: bool = False):
|
|
self.is_error = is_error
|
|
|
|
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__(is_error=True)
|
|
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()
|
|
MULTI_IMAGES = 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()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MultimodalDataItem:
|
|
"""
|
|
One MultimodalDataItem contains all inputs for one modality.
|
|
For example, if there are 3 images and 1 audio inputs, there will be 2 MultimodalDataItem.
|
|
One for images and one for audio.
|
|
|
|
We put the common fields first and the model-specific fields in model_specific_data.
|
|
"""
|
|
|
|
modality: Modality
|
|
hash: int = None
|
|
pad_value: int = None
|
|
offsets: Optional[list] = None
|
|
|
|
format: MultimodalInputFormat = MultimodalInputFormat.NORMAL
|
|
|
|
# the raw features returned by processor, e.g. pixel_values or audio_features
|
|
feature: Union[torch.Tensor, np.ndarray] = 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[Union[torch.Tensor, np.ndarray]] = None
|
|
|
|
# Model-specific data stored in a dictionary
|
|
model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict)
|
|
|
|
def __getattr__(self, name: str):
|
|
if (
|
|
"model_specific_data" in self.__dict__
|
|
and name in self.__dict__["model_specific_data"]
|
|
):
|
|
return self.__dict__["model_specific_data"][name]
|
|
else:
|
|
raise AttributeError(
|
|
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
|
)
|
|
|
|
def __setitem__(self, key: str, value: Any):
|
|
if key in self.__dict__:
|
|
self.__dict__[key] = value
|
|
else:
|
|
self.model_specific_data[key] = value
|
|
|
|
def set(self, key: str, value: Any):
|
|
self.__setitem__(key, 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):
|
|
"""
|
|
Set the pad value after first hashing the data
|
|
"""
|
|
if self.pad_value is not None:
|
|
return
|
|
|
|
from sglang.srt.managers.mm_utils import hash_feature
|
|
|
|
if envs.SGLANG_MM_SKIP_COMPUTE_HASH.get():
|
|
import uuid
|
|
|
|
self.hash = uuid.uuid4().int
|
|
self.pad_value = _compute_pad_value(self.hash)
|
|
return
|
|
if self.hash is None:
|
|
if self.feature is not None:
|
|
hashed_feature = self.feature
|
|
else:
|
|
hashed_feature = self.precomputed_embeddings
|
|
self.hash = hash_feature(hashed_feature)
|
|
assert self.hash is not None
|
|
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 in [Modality.IMAGE, Modality.MULTI_IMAGES]
|
|
|
|
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 merge(self, other):
|
|
self.feature += other.feature
|
|
self.offsets += other.offsets
|
|
self.hash = hash((self.hash, other.hash))
|
|
self.set_pad_value()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MultimodalInputs:
|
|
"""The multimodal data related inputs."""
|
|
|
|
# items of data
|
|
mm_items: List[MultimodalDataItem]
|
|
image_pad_len: Optional[list] = None
|
|
num_image_tokens: Optional[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
|
|
|
|
@staticmethod
|
|
def from_dict(obj: dict):
|
|
# Check if MM splitting is enabled
|
|
if not envs.SGLANG_ENABLE_MM_SPLITTING.get():
|
|
mm_items = obj["mm_items"]
|
|
else:
|
|
from sglang.srt.managers.mm_utils import get_new_expanded_mm_items
|
|
|
|
original_mm_items = obj["mm_items"]
|
|
# Now, `mm_items` contains one item per image.
|
|
mm_items = get_new_expanded_mm_items(original_mm_items)
|
|
|
|
ret = MultimodalInputs(
|
|
mm_items=mm_items,
|
|
)
|
|
|
|
assert isinstance(ret.mm_items, list)
|
|
ret.mm_items = [item for item in ret.mm_items if item.is_valid()]
|
|
|
|
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 ret.mm_items:
|
|
if item.feature is not None:
|
|
if isinstance(item.feature, torch.Tensor):
|
|
item.feature = try_add_to_buffer(item.feature)
|
|
|
|
for item in ret.mm_items:
|
|
item.set_pad_value()
|
|
|
|
if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0:
|
|
for item in ret.mm_items:
|
|
if item.feature is not None:
|
|
item.feature = item.feature.to("cpu", non_blocking=True)
|
|
|
|
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",
|
|
]
|
|
for arg in optional_args:
|
|
if arg in obj:
|
|
setattr(ret, arg, obj[arg])
|
|
|
|
return ret
|
|
|
|
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 merge(self, other: MultimodalInputs):
|
|
"""
|
|
merge image inputs when requests are being merged
|
|
"""
|
|
|
|
# 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))
|
|
|
|
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
|
|
|
|
|
|
class Req(ReqDllmMixin):
|
|
"""The input and output status of a request."""
|
|
|
|
def __init__(
|
|
self,
|
|
rid: str,
|
|
origin_input_text: str,
|
|
origin_input_ids: List[int],
|
|
sampling_params: SamplingParams,
|
|
return_logprob: bool = False,
|
|
top_logprobs_num: int = 0,
|
|
dllm_config: Optional[DllmConfig] = None,
|
|
token_ids_logprob: List[int] = None,
|
|
stream: bool = False,
|
|
origin_input_ids_unpadded: Optional[Tuple[int]] = None,
|
|
lora_id: Optional[str] = None,
|
|
input_embeds: Optional[List[List[float]]] = None,
|
|
token_type_ids: List[int] = None,
|
|
session: Optional[Session] = None,
|
|
custom_logit_processor: Optional[str] = None,
|
|
require_reasoning: bool = False,
|
|
return_hidden_states: bool = False,
|
|
return_routed_experts: 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,
|
|
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,
|
|
):
|
|
# Input and output info
|
|
self.rid = rid
|
|
self.origin_input_text = origin_input_text
|
|
self.origin_input_ids_unpadded = (
|
|
origin_input_ids_unpadded
|
|
if origin_input_ids_unpadded
|
|
else origin_input_ids # Before image padding
|
|
)
|
|
self.origin_input_ids = origin_input_ids
|
|
# Each decode stage's output ids
|
|
self.output_ids = []
|
|
# fill_ids = origin_input_ids + output_ids. Updated if chunked.
|
|
self.fill_ids = []
|
|
self.session = session
|
|
self.input_embeds = input_embeds
|
|
|
|
# For req-level memory management
|
|
self.kv_committed_len = 0
|
|
self.kv_allocated_len = 0
|
|
self.kv_committed_freed = False
|
|
self.kv_overallocated_freed = False
|
|
|
|
# for corss-endoder model
|
|
self.token_type_ids = token_type_ids
|
|
|
|
# The length of KV that have been removed in swa cache.
|
|
# SWA KV cache eviction behavior differs by cache type:
|
|
# - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in
|
|
# `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction.
|
|
# - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`.
|
|
self.swa_evicted_seqlen = 0
|
|
|
|
# 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 (hybrid reasoning model only)
|
|
self.require_reasoning = require_reasoning
|
|
|
|
# Sampling info
|
|
if isinstance(sampling_params.custom_params, dict):
|
|
sampling_params = copy.copy(sampling_params)
|
|
sampling_params.custom_params = sampling_params.custom_params | {
|
|
"__req__": self
|
|
}
|
|
self.sampling_params = sampling_params
|
|
self.custom_logit_processor = custom_logit_processor
|
|
self.return_hidden_states = return_hidden_states
|
|
|
|
# extra key for classifying the request (e.g. cache_salt)
|
|
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.lora_id = lora_id
|
|
self.routing_key = routing_key
|
|
|
|
# Memory pool info
|
|
self.req_pool_idx: Optional[int] = None
|
|
self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1)
|
|
self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2)
|
|
self.mamba_next_track_idx: Optional[int] = None # 0 or 1
|
|
self.mamba_last_track_seqlen: Optional[int] = (
|
|
None # seq len of the last cached mamba state
|
|
)
|
|
# 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
|
|
|
|
# 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
|
|
|
|
# Prefix info
|
|
# The indices to kv cache for the shared prefix.
|
|
self.prefix_indices: torch.Tensor = torch.empty((0,), dtype=torch.int64)
|
|
# Number of tokens to run prefill.
|
|
self.extend_input_len = 0
|
|
# The relative logprob_start_len in an extend batch
|
|
self.extend_logprob_start_len = 0
|
|
self.last_node: Any = None
|
|
self.last_host_node: Any = None
|
|
self.last_host_backup_node: Any = None
|
|
self.host_hit_length = 0
|
|
# Tokens loaded from storage backend (L3) during prefetch for this request
|
|
self.storage_hit_length = 0
|
|
# The node to lock until for swa radix tree lock ref
|
|
self.swa_uuid_for_lock: Optional[int] = None
|
|
# The prefix length that is inserted into the tree cache
|
|
self.cache_protected_len: int = 0
|
|
|
|
# Whether or not if it is chunked. It increments whenever
|
|
# it is chunked, and decrement whenever chunked request is
|
|
# processed.
|
|
self.is_chunked = 0
|
|
|
|
# For retraction
|
|
self.is_retracted = False
|
|
# Indicates if the req has ever been retracted.
|
|
self.retracted_stain = False
|
|
|
|
# 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.top_logprobs_num = top_logprobs_num
|
|
self.token_ids_logprob = token_ids_logprob
|
|
self.temp_scaled_logprobs = False
|
|
self.top_p_normalized_logprobs = False
|
|
|
|
# Logprobs (return values)
|
|
# True means the input logprob has been already sent to detokenizer.
|
|
self.input_logprob_sent: bool = False
|
|
self.input_token_logprobs_val: Optional[List[float]] = None
|
|
self.input_token_logprobs_idx: Optional[List[int]] = None
|
|
self.input_top_logprobs_val: Optional[List[float]] = None
|
|
self.input_top_logprobs_idx: Optional[List[int]] = None
|
|
self.input_token_ids_logprobs_val: Optional[List[float]] = None
|
|
self.input_token_ids_logprobs_idx: Optional[List[int]] = None
|
|
# 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.output_token_logprobs_val = []
|
|
self.output_token_logprobs_idx = []
|
|
# shape: (bs, k)
|
|
self.output_top_logprobs_val = []
|
|
self.output_top_logprobs_idx = []
|
|
# Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
|
|
self.output_token_ids_logprobs_val: List[
|
|
Union[List[float], torch.Tensor]
|
|
] = []
|
|
self.output_token_ids_logprobs_idx = []
|
|
else:
|
|
self.output_token_logprobs_val = self.output_token_logprobs_idx = (
|
|
self.output_top_logprobs_val
|
|
) = self.output_top_logprobs_idx = self.output_token_ids_logprobs_val = (
|
|
self.output_token_ids_logprobs_idx
|
|
) = None
|
|
self.hidden_states: 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
|
|
|
|
# capture routed experts
|
|
self.return_routed_experts = return_routed_experts
|
|
self.routed_experts: Optional[torch.Tensor] = (
|
|
None # cpu tensor: shape (seqlen, 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
|
|
)
|
|
|
|
# The number of verification forward passes in the speculative decoding.
|
|
# This is used to compute the average acceptance length per request.
|
|
self.spec_verify_ct = 0
|
|
|
|
# The number of accepted tokens in speculative decoding for this request.
|
|
# This is used to compute the acceptance rate and average acceptance length per request.
|
|
self.spec_accepted_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_acceptance_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
|
|
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
|
|
|
|
# the start index of the sent kv cache
|
|
# We want to send it chunk by chunk for chunked prefill.
|
|
# After every chunk forward, we do the following:
|
|
# kv_send(req.input_ids[req.start_send_idx:len(req.fill_ids)])
|
|
# start_send_idx = len(req.fill_ids)
|
|
self.start_send_idx: int = 0
|
|
|
|
# For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap
|
|
# This is because kv is not ready in `process_prefill_chunk`.
|
|
# We use `tmp_end_idx` to store the end index of the kv cache to send.
|
|
self.tmp_end_idx: int = -1
|
|
self.metadata_buffer_index: int = -1
|
|
|
|
# For Matryoshka embeddings
|
|
self.dimensions = dimensions
|
|
|
|
# For diffusion LLM
|
|
self.init_diffusion_llm(dllm_config)
|
|
|
|
# For hisparse
|
|
self.staging = False
|
|
|
|
@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_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_global_server_args().speculative_algorithm
|
|
return self.sampling_params.max_new_tokens == 0 and spec_alg is None
|
|
|
|
@property
|
|
def output_ids_through_stop(self) -> List[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 pop_committed_kv_cache(self) -> int:
|
|
"""Return the length of committed KV cache and mark them as freed."""
|
|
assert (
|
|
not self.kv_committed_freed
|
|
), f"Committed KV cache already freed ({self.kv_committed_len=})"
|
|
self.kv_committed_freed = True
|
|
return self.kv_committed_len
|
|
|
|
def pop_overallocated_kv_cache(self) -> Tuple[int, int]:
|
|
"""Return the range of over-allocated KV cache and mark them as freed."""
|
|
|
|
# NOTE: This function is called when there is over-allocation of KV cache.
|
|
# Over-allocation: we allocate more KV cache than the committed length.
|
|
# e.g., speculative decoding may allocate more KV cache than actually used.
|
|
assert (
|
|
not self.kv_overallocated_freed
|
|
), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv_allocated_len=}"
|
|
self.kv_overallocated_freed = True
|
|
return self.kv_committed_len, self.kv_allocated_len
|
|
|
|
def update_spec_acceptance_histogram(self, accepted_draft_tokens: int):
|
|
"""Update the speculative decoding acceptance histogram.
|
|
|
|
Args:
|
|
accepted_draft_tokens: Number of draft tokens accepted in this step.
|
|
"""
|
|
if len(self.spec_acceptance_histogram) <= accepted_draft_tokens:
|
|
self.spec_acceptance_histogram.extend(
|
|
[0] * (accepted_draft_tokens - len(self.spec_acceptance_histogram) + 1)
|
|
)
|
|
self.spec_acceptance_histogram[accepted_draft_tokens] += 1
|
|
|
|
def extend_image_inputs(self, image_inputs):
|
|
if self.multimodal_inputs is None:
|
|
self.multimodal_inputs = image_inputs
|
|
else:
|
|
self.multimodal_inputs.merge(image_inputs)
|
|
|
|
def finished(self) -> bool:
|
|
# Whether request reached finished condition
|
|
return self.finished_reason is not None
|
|
|
|
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.fill_ids = self.origin_input_ids + self.output_ids
|
|
|
|
input_len = len(self.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
|
|
|
|
# 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.return_logprob and self.logprob_start_len >= 0:
|
|
max_prefix_len = min(max_prefix_len, self.logprob_start_len)
|
|
max_prefix_len = max(max_prefix_len, 0)
|
|
token_ids = self.fill_ids[:max_prefix_len]
|
|
|
|
if tree_cache is not None:
|
|
if cow_mamba is None:
|
|
cow_mamba = tree_cache.supports_mamba()
|
|
match_result = tree_cache.match_prefix(
|
|
MatchPrefixParams(
|
|
key=RadixKey(token_ids=token_ids, extra_key=self.extra_key),
|
|
req=self,
|
|
cow_mamba=cow_mamba,
|
|
)
|
|
)
|
|
(
|
|
self.prefix_indices,
|
|
self.last_node,
|
|
self.last_host_node,
|
|
self.last_host_backup_node,
|
|
self.host_hit_length,
|
|
self.mamba_branching_seqlen,
|
|
) = (
|
|
match_result.device_indices,
|
|
match_result.last_device_node,
|
|
match_result.last_host_node,
|
|
match_result.last_host_backup_node,
|
|
match_result.host_hit_length,
|
|
match_result.mamba_branching_seqlen,
|
|
)
|
|
if match_result.cache_protected_len is not None:
|
|
self.cache_protected_len = match_result.cache_protected_len
|
|
else:
|
|
self.cache_protected_len = len(self.prefix_indices)
|
|
|
|
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)
|
|
)
|
|
)
|
|
|
|
self.set_extend_input_len(len(self.fill_ids) - len(self.prefix_indices))
|
|
|
|
# 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 tail_str(self) -> 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 ""
|
|
|
|
max_len_tail_str = max(
|
|
self.sampling_params.stop_str_max_len + 1,
|
|
self.sampling_params.stop_regex_max_len + 1,
|
|
)
|
|
|
|
tail_len = min(max_len_tail_str, len(self.output_ids))
|
|
return self.tokenizer.decode(self.output_ids[-tail_len:])
|
|
|
|
def check_match_stop_str_prefix(self) -> bool:
|
|
"""
|
|
Check if the suffix of tail_str overlaps with any stop_str prefix
|
|
"""
|
|
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 _check_str_based_finish(self):
|
|
if (
|
|
len(self.sampling_params.stop_strs) > 0
|
|
or len(self.sampling_params.stop_regex_strs) > 0
|
|
):
|
|
tail_str = self.tail_str()
|
|
|
|
# Check stop strings
|
|
if len(self.sampling_params.stop_strs) > 0:
|
|
for stop_str in self.sampling_params.stop_strs:
|
|
if stop_str in tail_str or stop_str in self.decoded_text:
|
|
self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
|
|
return True
|
|
|
|
# Check stop regex
|
|
if len(self.sampling_params.stop_regex_strs) > 0:
|
|
for stop_regex_str in self.sampling_params.stop_regex_strs:
|
|
if re.search(stop_regex_str, tail_str):
|
|
self.finished_reason = FINISHED_MATCHED_REGEX(
|
|
matched=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 > self.vocab_size or token_id < 0:
|
|
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 check_finished(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
|
|
|
|
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:
|
|
if self.grammar.is_terminated():
|
|
self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1])
|
|
return
|
|
|
|
new_accepted_tokens = self.output_ids[-new_accepted_len:]
|
|
|
|
if self._check_token_based_finish(new_accepted_tokens):
|
|
return
|
|
|
|
if self._check_vocab_boundary_finish(new_accepted_tokens):
|
|
return
|
|
|
|
if self._check_str_based_finish():
|
|
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_indices = torch.empty((0,), dtype=torch.int64)
|
|
self.routed_experts = None
|
|
self.last_node = None
|
|
self.swa_uuid_for_lock = None
|
|
self.extend_input_len = 0
|
|
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.extend_logprob_start_len = 0
|
|
self.is_chunked = 0
|
|
self.mamba_pool_idx = None
|
|
self.mamba_ping_pong_track_buffer = None
|
|
self.mamba_next_track_idx = None
|
|
self.mamba_last_track_seqlen = None
|
|
self.mamba_branching_seqlen = None
|
|
self.already_computed = 0
|
|
self.kv_allocated_len = 0
|
|
self.kv_committed_len = 0
|
|
self.kv_committed_freed = False
|
|
self.kv_overallocated_freed = False
|
|
self.swa_evicted_seqlen = 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 = []
|
|
|
|
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
|
|
token_indices = req_to_token_pool.req_to_token[
|
|
self.req_pool_idx, : self.seqlen - 1
|
|
]
|
|
self.kv_cache_cpu = token_to_kv_pool_allocator.get_cpu_copy(token_indices)
|
|
|
|
def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
|
|
token_indices = req_to_token_pool.req_to_token[
|
|
self.req_pool_idx, : self.seqlen - 1
|
|
]
|
|
token_to_kv_pool_allocator.load_cpu_copy(self.kv_cache_cpu, token_indices)
|
|
del self.kv_cache_cpu
|
|
|
|
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"Req Time Stats(rid={self.rid}{bootstrap_info}, input len={len(self.origin_input_ids)}, output len={len(self.output_ids)}, 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_extend_input_len(self, extend_input_len: int):
|
|
# Setting extend_input_len and computing the relative logprob_start_len in an extend batch
|
|
#
|
|
# 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
|
|
self.extend_input_len = extend_input_len
|
|
if self.logprob_start_len == -1:
|
|
logprob_start_len = len(self.fill_ids)
|
|
else:
|
|
# logprob_start_len should be at least the length of the prefix indices
|
|
logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices))
|
|
self.extend_logprob_start_len = min(
|
|
logprob_start_len - len(self.prefix_indices),
|
|
self.extend_input_len,
|
|
)
|
|
|
|
def set_finish_with_abort(self, error_msg: str):
|
|
if get_tensor_model_parallel_rank() == 0:
|
|
logger.error(f"{error_msg}, {self.rid=}")
|
|
self.multimodal_inputs = None
|
|
self.grammar = None
|
|
self.origin_input_ids = [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, HTTPStatus.BAD_REQUEST, "BadRequestError"
|
|
)
|
|
|
|
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=})"
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|
"""Store all information of a batch on the scheduler."""
|
|
|
|
# Request, memory pool, and cache
|
|
reqs: List[Req]
|
|
req_to_token_pool: ReqToTokenPool = None
|
|
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator = None
|
|
tree_cache: BasePrefixCache = None
|
|
is_hybrid_swa: bool = False
|
|
|
|
# Batch configs
|
|
model_config: ModelConfig = None
|
|
forward_mode: ForwardMode = None
|
|
enable_overlap: bool = False
|
|
# 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
|
|
|
|
# Sampling info
|
|
sampling_info: SamplingBatchInfo = None
|
|
|
|
# Batched arguments to model runner
|
|
input_ids: torch.Tensor = None # shape: [b], int64
|
|
input_embeds: torch.Tensor = None # shape: [b, hidden_size], float32
|
|
ne_token_table: torch.Tensor = None
|
|
token_type_ids: torch.Tensor = None # shape: [b], int64
|
|
req_pool_indices: torch.Tensor = None # shape: [b], int64
|
|
seq_lens: torch.Tensor = None # shape: [b], int64
|
|
seq_lens_cpu: torch.Tensor = None # shape: [b], int64
|
|
# The output locations of the KV cache
|
|
out_cache_loc: torch.Tensor = None # shape: [b], int64
|
|
output_ids: torch.Tensor = None # shape: [b], int64
|
|
|
|
# For hybrid GDN prefix cache
|
|
mamba_track_indices: torch.Tensor = None # shape: [b], int64
|
|
mamba_track_mask: torch.Tensor = None # shape: [b], bool
|
|
mamba_track_seqlens: torch.Tensor = None # shape: [b], int64
|
|
|
|
# For multimodal inputs
|
|
multimodal_inputs: Optional[List] = None
|
|
|
|
# The sum of all sequence lengths
|
|
seq_lens_sum: int = None
|
|
# The original sequence lengths, Qwen-1M related
|
|
orig_seq_lens: torch.Tensor = None # shape: [b], int32
|
|
|
|
# For DP attention
|
|
inner_idle_batch: Optional[ScheduleBatch] = None
|
|
global_num_tokens: Optional[List[int]] = None
|
|
global_num_tokens_for_logprob: Optional[List[int]] = None
|
|
is_extend_in_batch: bool = False
|
|
all_extend_in_batch: bool = False
|
|
can_run_dp_cuda_graph: bool = False
|
|
tbo_split_seq_index: Optional[int] = None
|
|
global_forward_mode: Optional[ForwardMode] = None
|
|
|
|
# For processing logprobs
|
|
return_logprob: bool = False
|
|
top_logprobs_nums: Optional[List[int]] = None
|
|
token_ids_logprobs: Optional[List[List[int]]] = None
|
|
|
|
# For logits and logprob post processing
|
|
temp_scaled_logprobs: bool = False
|
|
top_p_normalized_logprobs: bool = False
|
|
|
|
# For extend and mixed chunekd prefill
|
|
prefix_lens: List[int] = None
|
|
extend_lens: List[int] = None
|
|
extend_num_tokens: Optional[int] = None
|
|
decoding_reqs: List[Req] = None
|
|
extend_logprob_start_lens: List[int] = None
|
|
# It comes empty list if logprob is not required.
|
|
extend_input_logprob_token_ids: Optional[torch.Tensor] = None
|
|
|
|
# For encoder-decoder architectures
|
|
encoder_cached: Optional[List[bool]] = None
|
|
encoder_lens: Optional[torch.Tensor] = None
|
|
encoder_lens_cpu: Optional[List[int]] = None
|
|
encoder_out_cache_loc: Optional[torch.Tensor] = None
|
|
|
|
# For matryoshka embeddings
|
|
dimensions: Optional[list[int]] = None
|
|
|
|
# For split prefill
|
|
split_index: int = 0
|
|
split_prefill_finished: bool = False
|
|
split_forward_count: int = 1
|
|
split_forward_batch: ForwardBatch = None
|
|
seq_lens_cpu_cache: torch.Tensor = None
|
|
|
|
# Stream
|
|
has_stream: bool = False
|
|
|
|
# Has grammar
|
|
has_grammar: bool = False
|
|
|
|
# Device
|
|
device: str = "cuda"
|
|
|
|
# Speculative decoding
|
|
spec_algorithm: SpeculativeAlgorithm = None
|
|
# spec_info: Optional[SpecInput] = None
|
|
spec_info: Optional[SpecInput] = None
|
|
|
|
# Whether to return hidden states
|
|
return_hidden_states: bool = False
|
|
|
|
# Whether to return captured experts
|
|
return_routed_experts: bool = False
|
|
|
|
# Whether this batch is prefill-only (no token generation needed)
|
|
is_prefill_only: bool = False
|
|
|
|
# hicache pointer for synchronizing data loading from CPU to GPU
|
|
hicache_consumer_index: int = -1
|
|
|
|
# Diffusion LLM
|
|
dllm_config: Optional[DllmConfig] = None
|
|
|
|
# Metrics
|
|
dp_cooperation_info: Optional[DPCooperationInfo] = None
|
|
prefill_stats: Optional[PrefillStats] = None
|
|
|
|
# HiSparse
|
|
hisparse_coordinator: Optional[HiSparseCoordinator] = 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)
|
|
|
|
is_hybrid_swa = False
|
|
if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator):
|
|
is_hybrid_swa = True
|
|
|
|
return 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,
|
|
is_hybrid_swa=is_hybrid_swa,
|
|
model_config=model_config,
|
|
enable_overlap=enable_overlap,
|
|
return_logprob=return_logprob,
|
|
has_stream=any(req.stream for req in reqs),
|
|
has_grammar=any(req.grammar for req in reqs),
|
|
device=req_to_token_pool.device,
|
|
spec_algorithm=spec_algorithm,
|
|
return_hidden_states=any(req.return_hidden_states for req in reqs),
|
|
return_routed_experts=any(req.return_routed_experts for req in reqs),
|
|
is_prefill_only=all(req.is_prefill_only for req in reqs),
|
|
chunked_req=chunked_req,
|
|
dllm_config=dllm_config,
|
|
)
|
|
|
|
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 prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]):
|
|
_pin = is_pin_memory_available(self.device)
|
|
self.encoder_lens_cpu = []
|
|
self.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
|
|
self.encoder_lens_cpu.append(0)
|
|
self.encoder_cached.append(True)
|
|
else:
|
|
self.encoder_lens_cpu.append(im.num_image_tokens)
|
|
self.encoder_cached.append(
|
|
self.forward_mode.is_decode()
|
|
or len(req.prefix_indices) >= im.num_image_tokens
|
|
)
|
|
|
|
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 = []
|
|
for i, req in enumerate(self.reqs):
|
|
encoder_len = self.encoder_lens_cpu[i]
|
|
seq_lens[i] -= encoder_len
|
|
|
|
if len(req.prefix_indices) < encoder_len:
|
|
# NOTE: the encoder part should be considered as a whole
|
|
assert len(req.prefix_indices) == 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_input_len]
|
|
)
|
|
self.extend_lens[i] -= encoder_len
|
|
self.extend_num_tokens -= encoder_len
|
|
else:
|
|
decoder_out_cache_loc.append(
|
|
self.out_cache_loc[pt : pt + req.extend_input_len]
|
|
)
|
|
self.prefix_lens[i] -= encoder_len
|
|
|
|
pt += req.extend_input_len
|
|
|
|
# Reassign
|
|
self.input_ids = torch.tensor(
|
|
sum(input_ids, []), dtype=torch.int64, pin_memory=_pin
|
|
).to(self.device, non_blocking=True)
|
|
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}"
|
|
|
|
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.fill_ids[len(r.prefix_indices) :] for r in reqs]
|
|
extend_num_tokens = sum(len(ids) for ids in input_ids)
|
|
seq_lens = [len(r.fill_ids) for r in reqs]
|
|
orig_seq_lens = [max(len(r.fill_ids), len(r.origin_input_ids)) for r in reqs]
|
|
prefix_lens = [len(r.prefix_indices) for r in reqs]
|
|
extend_lens = [r.extend_input_len for r in reqs]
|
|
|
|
# For matryoshka embeddings
|
|
if self.model_config.is_matryoshka and any(
|
|
r.dimensions is not None for r in reqs
|
|
):
|
|
self.dimensions = [
|
|
r.dimensions if r.dimensions else self.model_config.hidden_size
|
|
for r in reqs
|
|
]
|
|
|
|
token_type_ids = [
|
|
r.token_type_ids for r in reqs if r.token_type_ids is not None
|
|
]
|
|
|
|
_pin = is_pin_memory_available(self.device)
|
|
input_ids_tensor = torch.tensor(
|
|
list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin
|
|
).to(self.device, non_blocking=True)
|
|
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)
|
|
|
|
token_type_ids_tensor = None
|
|
if len(token_type_ids) > 0:
|
|
token_type_ids_tensor = torch.tensor(
|
|
sum(token_type_ids, []), dtype=torch.int64, 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
|
|
|
|
# Allocate memory
|
|
out_cache_loc, req_pool_indices_tensor, req_pool_indices = alloc_for_extend(
|
|
self
|
|
)
|
|
|
|
# Set fields
|
|
input_embeds = []
|
|
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.req_pool_idx = req_pool_indices[i]
|
|
assert seq_len - pre_len == req.extend_input_len
|
|
|
|
req.extend_batch_idx += 1
|
|
|
|
# update req-level memory management fields
|
|
req.kv_committed_len = seq_len
|
|
req.kv_allocated_len = seq_len
|
|
|
|
# If input_embeds are available, store them
|
|
if req.input_embeds is not None:
|
|
# Slice to match extend_input_len — PrefillAdder truncates
|
|
# fill_ids/extend_input_len on chunk overflow but not input_embeds.
|
|
input_embeds.extend(
|
|
req.input_embeds[pre_len : pre_len + req.extend_input_len]
|
|
)
|
|
|
|
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:
|
|
# At this point, prefix_indices has been extended with host data
|
|
# via init_load_back in schedule_policy, so:
|
|
# - len(prefix_indices) = device_original + host_loaded
|
|
# - host_hit_length = total tokens from host cache (including storage-prefetched)
|
|
# - storage_hit_length = tokens loaded from storage backend (L3 hits)
|
|
# - device_portion = len(prefix_indices) - host_hit_length
|
|
#
|
|
# Storage hits are now tracked via scheduler after prefetch completes.
|
|
# storage_hit_length is set by scheduler.pop_prefetch_loaded_tokens()
|
|
host_total = req.host_hit_length
|
|
# Clamp storage to host_total to handle edge cases
|
|
storage_portion = min(host_total, req.storage_hit_length)
|
|
host_portion = host_total - storage_portion
|
|
device_portion = max(0, len(req.prefix_indices) - host_total)
|
|
|
|
req.cached_tokens_device = device_portion
|
|
req.cached_tokens_host = host_portion
|
|
req.cached_tokens_storage = storage_portion
|
|
req._cache_breakdown_computed = True
|
|
|
|
req.already_computed = seq_len
|
|
req.is_retracted = False
|
|
|
|
if get_global_server_args().enable_mamba_extra_buffer():
|
|
self._mamba_radix_cache_v2_req_prepare_for_extend(
|
|
req,
|
|
mamba_track_mask_cpu,
|
|
mamba_track_indices_cpu,
|
|
mamba_track_seqlens_cpu,
|
|
)
|
|
|
|
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]
|
|
# 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]
|
|
# fill_ids = [3, 4]
|
|
# extend_input_logprob_token_id = [4, 0]
|
|
global_start_idx, global_end_idx = (
|
|
len(req.prefix_indices),
|
|
len(req.fill_ids),
|
|
)
|
|
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_input_len - req.extend_logprob_start_len 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_input_len
|
|
- req.extend_logprob_start_len
|
|
- 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
|
|
|
|
self.input_ids = input_ids_tensor
|
|
self.req_pool_indices = req_pool_indices_tensor
|
|
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
|
|
)
|
|
for mm_input in multimodal_inputs:
|
|
if mm_input is None:
|
|
continue
|
|
for mm_item in mm_input.mm_items:
|
|
pixel_values = getattr(mm_item, "feature", None)
|
|
if isinstance(pixel_values, torch.Tensor):
|
|
mm_item.feature = pixel_values.to(self.device, non_blocking=True)
|
|
elif isinstance(pixel_values, CudaIpcTensorTransportProxy):
|
|
mm_item.feature = pixel_values.reconstruct_on_target_device(
|
|
torch.cuda.current_device()
|
|
)
|
|
# The reference by CudaIpcTensorTransportProxy was cut off,
|
|
# proactively delete to avoid slow gc.
|
|
del pixel_values
|
|
if get_global_server_args().language_only:
|
|
precomputed_embeddings = getattr(
|
|
mm_item, "precomputed_embeddings", None
|
|
)
|
|
if isinstance(precomputed_embeddings, torch.Tensor):
|
|
mm_item.precomputed_embeddings = precomputed_embeddings.to(
|
|
self.device, non_blocking=True
|
|
)
|
|
self.multimodal_inputs = multimodal_inputs
|
|
self.token_type_ids = token_type_ids_tensor
|
|
self.seq_lens_sum = sum(seq_lens)
|
|
|
|
if self.return_logprob:
|
|
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
|
|
self.token_ids_logprobs = [r.token_ids_logprob for r in reqs]
|
|
|
|
self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
|
|
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
|
|
|
|
if get_global_server_args().enable_mamba_extra_buffer():
|
|
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,
|
|
)
|
|
|
|
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,
|
|
mamba_track_mask_cpu: List[bool],
|
|
mamba_track_indices_cpu: List[int],
|
|
mamba_track_seqlens_cpu: List[int],
|
|
):
|
|
def _force_track_h(i: int) -> int:
|
|
assert i % FLA_CHUNK_SIZE == 0
|
|
# There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
|
|
# 1) aligned with FLA_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 FLA_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
|
|
|
|
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
|
|
mask = req.extend_input_len >= mamba_cache_chunk_size
|
|
mamba_track_mask_cpu.append(mask)
|
|
mamba_track_indices_cpu.append(
|
|
req.mamba_ping_pong_track_buffer[req.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 = len(req.prefix_indices) + req.extend_input_len
|
|
|
|
# mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to
|
|
# mamba radix cache to track which seqlen this mamba state should store at.
|
|
mamba_track_seqlen_aligned = (
|
|
len(req.prefix_indices)
|
|
+ (req.extend_input_len // mamba_cache_chunk_size)
|
|
* mamba_cache_chunk_size
|
|
)
|
|
|
|
# mamba_track_fla_chunk_aligned is the aligned seqlen based on FLA_CHUNK_SIZE
|
|
# If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which can be true when
|
|
# page_size > FLA_CHUNK_SIZE, we need to force the math calculation to retrieve the correct mamba state from h
|
|
# by _force_track_h()
|
|
mamba_track_fla_chunk_aligned = (
|
|
len(req.prefix_indices)
|
|
+ (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_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)
|
|
|
|
req.mamba_next_track_idx = (
|
|
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
|
|
req.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 - len(req.prefix_indices)
|
|
) % mamba_cache_chunk_size == 0
|
|
if (
|
|
req.mamba_branching_seqlen > len(req.prefix_indices)
|
|
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.mamba_last_track_seqlen = mamba_track_seqlen_aligned
|
|
mamba_track_seqlens_cpu.append(mamba_track_seqlen)
|
|
|
|
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()
|
|
|
|
for req in running_batch.reqs:
|
|
req.fill_ids = req.origin_input_ids + req.output_ids
|
|
req.set_extend_input_len(1)
|
|
|
|
input_ids = torch.cat([self.input_ids, running_batch.input_ids])
|
|
out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
|
|
|
|
self.merge_batch(running_batch)
|
|
self.input_ids = input_ids
|
|
self.out_cache_loc = out_cache_loc
|
|
|
|
# For overlap scheduler, the output_ids has one step delay
|
|
delta = 0 if self.enable_overlap else -1
|
|
|
|
# NOTE: prefix_indices is what has been cached, but we don't cache each decode step
|
|
self.prefix_lens.extend(
|
|
[
|
|
len(r.origin_input_ids) + len(r.output_ids) + delta
|
|
for r in running_batch.reqs
|
|
]
|
|
)
|
|
self.extend_lens.extend([1] * running_bs)
|
|
self.extend_num_tokens += running_bs
|
|
# TODO (lianmin): Revisit this. It should be seq_len - 1
|
|
self.extend_logprob_start_lens.extend([0] * running_bs)
|
|
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_committed_len % page_size == 0)
|
|
return new_pages * page_size
|
|
|
|
server_args = get_global_server_args()
|
|
len_per_topk = server_args.speculative_num_steps or 1
|
|
spec_topk = server_args.speculative_eagle_topk or 1
|
|
spec_tokens = server_args.speculative_num_draft_tokens
|
|
|
|
if page_size > 1 and spec_topk > 1:
|
|
# last partial page and ceil alignment
|
|
len_per_topk = ceil_align(len_per_topk + page_size, page_size)
|
|
spec_tokens = ceil_align(spec_tokens, page_size)
|
|
elif page_size > 1:
|
|
# only page alignment
|
|
len_per_topk = ceil_align(len_per_topk, page_size)
|
|
spec_tokens = ceil_align(spec_tokens, page_size)
|
|
|
|
num_tokens = max(len_per_topk * spec_topk, spec_tokens) * len(requests)
|
|
|
|
# v2 eagle has over-allocation
|
|
return num_tokens * (1 + self.is_spec_v2)
|
|
|
|
def check_decode_mem(self, selected_indices: Optional[List[int]] = None):
|
|
num_tokens = self.new_tokens_required_next_decode(selected_indices)
|
|
evict_from_tree_cache(self.tree_cache, num_tokens)
|
|
return self.token_to_kv_pool_allocator.available_size() >= num_tokens
|
|
|
|
def retract_all(self, server_args: ServerArgs):
|
|
retracted_reqs = self.reqs
|
|
for idx in range(len(self.reqs)):
|
|
self.release_req(idx, len(self.reqs) - idx, server_args)
|
|
|
|
self.filter_batch(retracted_reqs)
|
|
return retracted_reqs
|
|
|
|
def retract_decode(
|
|
self, server_args: ServerArgs
|
|
) -> Tuple[List[Req], float, List[Req]]:
|
|
"""Retract the decoding requests when there is not enough memory."""
|
|
sorted_indices = list(range(len(self.reqs)))
|
|
|
|
# TODO(lsyin): improve retraction policy for radix cache
|
|
# For spec decoding, filter_batch API can only filter
|
|
# requests from the back, so we can only retract from the back.
|
|
# TODO(sang): Clean up finish path and support better retract
|
|
# policy.
|
|
if not server_args.speculative_algorithm:
|
|
sorted_indices.sort(
|
|
key=lambda i: (
|
|
len(self.reqs[i].output_ids),
|
|
-len(self.reqs[i].origin_input_ids),
|
|
),
|
|
reverse=True,
|
|
)
|
|
|
|
retracted_reqs = []
|
|
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]
|
|
retracted_reqs.append(req)
|
|
# release memory and don't insert into the tree because we need the space instantly
|
|
self.release_req(idx, len(sorted_indices), server_args)
|
|
|
|
reqs_to_abort: List[Req] = []
|
|
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)
|
|
self.release_req(last_idx, 0, server_args)
|
|
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
|
|
total_decoded_tokens = sum(len(r.output_ids) for r in self.reqs)
|
|
total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in self.reqs)
|
|
|
|
new_estimate_ratio = (
|
|
total_decoded_tokens
|
|
+ envs.SGLANG_RETRACT_DECODE_STEPS.get() * len(self.reqs)
|
|
) / (
|
|
total_max_new_tokens + 1
|
|
) # avoid zero division
|
|
new_estimate_ratio = min(1.0, new_estimate_ratio)
|
|
|
|
return retracted_reqs, new_estimate_ratio, reqs_to_abort
|
|
|
|
def release_req(self, idx: int, remaing_req_count: int, server_args: ServerArgs):
|
|
req = self.reqs[idx]
|
|
|
|
if server_args.disaggregation_mode == "decode":
|
|
req.offload_kv_cache(
|
|
self.req_to_token_pool, self.token_to_kv_pool_allocator
|
|
)
|
|
# TODO (csy): for preempted requests, we may want to insert into the tree
|
|
release_kv_cache(req, self.tree_cache, is_insert=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(self.tree_cache, num_tokens)
|
|
|
|
req.reset_for_retract()
|
|
|
|
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.seq_lens_sum = 0
|
|
self.extend_num_tokens = 0
|
|
self.sampling_info = SamplingBatchInfo.from_schedule_batch(
|
|
self,
|
|
self.model_config.vocab_size,
|
|
)
|
|
|
|
@property
|
|
def is_spec_v2(self):
|
|
# FIXME: finally deprecate is_spec_v2
|
|
ret = self.enable_overlap and not self.spec_algorithm.is_none()
|
|
assert not ret or self.spec_algorithm.supports_spec_v2()
|
|
return ret
|
|
|
|
def prepare_for_decode(self):
|
|
self.forward_mode = ForwardMode.DECODE
|
|
bs = len(self.reqs)
|
|
# 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
|
|
|
|
# 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 hasattr(self, "nsa_cp_metadata") and self.nsa_cp_metadata is not None:
|
|
self.nsa_cp_metadata = None
|
|
|
|
if self.is_spec_v2:
|
|
# TODO(spec-v2): all spec v2 should go through this path
|
|
draft_input: EagleDraftInput = self.spec_info
|
|
draft_input.prepare_for_decode(self)
|
|
|
|
if not self.spec_algorithm.is_none():
|
|
# if spec decoding is used, the decode batch is prepared inside
|
|
# `forward_batch_speculative_generation` after running draft models.
|
|
return
|
|
|
|
if self.sampling_info.penalizer_orchestrator.is_required:
|
|
if self.enable_overlap:
|
|
# TODO: this can be slow, optimize this.
|
|
delayed_output_ids = torch.tensor(
|
|
[
|
|
(
|
|
req.output_ids[-1]
|
|
if len(req.output_ids)
|
|
else req.origin_input_ids[-1]
|
|
)
|
|
for req in self.reqs
|
|
],
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
)
|
|
self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
|
|
delayed_output_ids
|
|
)
|
|
else:
|
|
self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
|
|
self.output_ids.to(torch.int64)
|
|
)
|
|
|
|
# Update fields
|
|
self.input_ids = self.output_ids
|
|
self.output_ids = None
|
|
|
|
if self.model_config.is_encoder_decoder:
|
|
self.prepare_encoder_info_decode()
|
|
|
|
# Allocate memory
|
|
self.out_cache_loc = alloc_for_decode(self, token_per_req=1)
|
|
|
|
# Update req-level memory management fields
|
|
for req in self.reqs:
|
|
req.decode_batch_idx += 1
|
|
req.kv_committed_len += 1
|
|
req.kv_allocated_len += 1
|
|
|
|
# Update seq_lens after allocation
|
|
if self.enable_overlap:
|
|
# Do not use in-place operations in the overlap mode
|
|
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
|
|
else:
|
|
# A faster in-place version
|
|
self.seq_lens.add_(1)
|
|
self.seq_lens_cpu.add_(1)
|
|
self.orig_seq_lens.add_(1)
|
|
self.seq_lens_sum += bs
|
|
|
|
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,
|
|
)
|
|
|
|
if get_global_server_args().enable_mamba_extra_buffer():
|
|
if len(self.reqs) == 0:
|
|
self.mamba_track_indices = torch.empty(
|
|
(0,), dtype=torch.int64, device=self.device
|
|
)
|
|
else:
|
|
# already on device
|
|
all_buffers = torch.stack(
|
|
[req.mamba_ping_pong_track_buffer for req in self.reqs]
|
|
)
|
|
idx = (
|
|
torch.tensor(
|
|
[req.mamba_next_track_idx for req in self.reqs],
|
|
dtype=torch.int64,
|
|
pin_memory=True,
|
|
)
|
|
.unsqueeze(1)
|
|
.to(device=all_buffers.device, non_blocking=True)
|
|
)
|
|
self.mamba_track_indices = (
|
|
torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64)
|
|
)
|
|
|
|
# async H2D
|
|
self.mamba_track_mask = (
|
|
(self.seq_lens_cpu % get_global_server_args().mamba_track_interval == 0)
|
|
.pin_memory()
|
|
.to(device=self.device, non_blocking=True)
|
|
)
|
|
|
|
def maybe_wait_verify_done(self):
|
|
if self.is_spec_v2:
|
|
draft_input: EagleDraftInput = self.spec_info
|
|
if draft_input.verify_done is not None:
|
|
draft_input.verify_done.synchronize()
|
|
|
|
def filter_batch(
|
|
self,
|
|
chunked_req_to_exclude: Optional[Union[Req, List[Req]]] = None,
|
|
keep_indices: Optional[List[int]] = None,
|
|
# FIXME(lsyin): deprecate this API after spec v1 is deprecated
|
|
v1_spec_info_filtered: Optional[bool] = False,
|
|
):
|
|
# FIXME(lsyin): used here to get the correct seq_lens
|
|
# The batch has been launched but we need it verified to get correct next batch info
|
|
self.maybe_wait_verify_done()
|
|
|
|
if keep_indices is None:
|
|
if isinstance(chunked_req_to_exclude, Req):
|
|
chunked_req_to_exclude = [chunked_req_to_exclude]
|
|
elif chunked_req_to_exclude is None:
|
|
chunked_req_to_exclude = []
|
|
keep_indices = [
|
|
i
|
|
for i in range(len(self.reqs))
|
|
if not self.reqs[i].finished()
|
|
and self.reqs[i] not in chunked_req_to_exclude
|
|
]
|
|
|
|
if keep_indices is None or len(keep_indices) == 0:
|
|
# Filter out all requests
|
|
self.reqs = []
|
|
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.seq_lens = self.seq_lens[keep_indices_device]
|
|
self.seq_lens_cpu = self.seq_lens_cpu[keep_indices]
|
|
self.orig_seq_lens = self.orig_seq_lens[keep_indices_device]
|
|
self.out_cache_loc = None
|
|
self.seq_lens_sum = self.seq_lens.sum().item()
|
|
|
|
if self.output_ids is not None:
|
|
self.output_ids = self.output_ids[keep_indices_device]
|
|
|
|
self.mamba_track_indices = None
|
|
self.mamba_track_mask = None
|
|
self.mamba_track_seqlens = 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_stream = any(req.stream for req in self.reqs)
|
|
self.has_grammar = any(req.grammar for req in self.reqs)
|
|
|
|
self.sampling_info.filter_batch(keep_indices, keep_indices_device)
|
|
# NOTE: spec_info filtered before batch filtering only happens in:
|
|
# - Spec v1's verify phase
|
|
# - Only for decode batch (running_batch)
|
|
has_been_filtered = v1_spec_info_filtered and not self.is_spec_v2
|
|
|
|
if self.spec_info:
|
|
self.spec_info.filter_batch(
|
|
new_indices=keep_indices_device,
|
|
has_been_filtered=has_been_filtered,
|
|
)
|
|
|
|
def merge_batch(self, other: "ScheduleBatch"):
|
|
# In the regular scheduler path:
|
|
# 1) self is always prefill, whose seq_lens is not a future
|
|
# 2) other is always decode, which is finished in previous step
|
|
# so verify_done is already synced and this is a no-op.
|
|
# In disagg decode + overlap, merge_batch can be called before
|
|
# filter_batch, so running_batch.seq_lens may still be a forward_stream
|
|
# future. Synchronize here to avoid a cross-stream data race.
|
|
self.maybe_wait_verify_done()
|
|
|
|
# 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.extend(other.encoder_lens_cpu)
|
|
self.req_pool_indices = torch.cat(
|
|
[self.req_pool_indices, other.req_pool_indices]
|
|
)
|
|
self.seq_lens = torch.cat([self.seq_lens, other.seq_lens])
|
|
self.seq_lens_cpu = torch.cat([self.seq_lens_cpu, other.seq_lens_cpu])
|
|
self.orig_seq_lens = torch.cat([self.orig_seq_lens, other.orig_seq_lens])
|
|
self.out_cache_loc = None
|
|
self.seq_lens_sum += other.seq_lens_sum
|
|
if self.output_ids is not None:
|
|
self.output_ids = torch.cat([self.output_ids, other.output_ids])
|
|
self.mamba_track_indices = None
|
|
self.mamba_track_mask = None
|
|
self.mamba_track_seqlens = None
|
|
if self.return_logprob and other.return_logprob:
|
|
self.top_logprobs_nums.extend(other.top_logprobs_nums)
|
|
self.token_ids_logprobs.extend(other.token_ids_logprobs)
|
|
elif self.return_logprob:
|
|
self.top_logprobs_nums.extend([0] * len(other.reqs))
|
|
self.token_ids_logprobs.extend([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.extend(other.reqs)
|
|
if self.multimodal_inputs is not None:
|
|
self.multimodal_inputs.extend(other.multimodal_inputs)
|
|
|
|
self.return_logprob |= other.return_logprob
|
|
self.has_stream |= other.has_stream
|
|
self.has_grammar |= other.has_grammar
|
|
self.return_hidden_states |= other.return_hidden_states
|
|
|
|
if self.spec_info:
|
|
self.spec_info.merge_batch(other.spec_info)
|
|
|
|
def get_model_worker_batch(
|
|
self, seq_lens_cpu_cache: Optional[torch.Tensor] = None
|
|
) -> ModelWorkerBatch:
|
|
if self.forward_mode.is_decode_or_idle():
|
|
extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = None
|
|
else:
|
|
extend_seq_lens = self.extend_lens
|
|
extend_prefix_lens = self.prefix_lens
|
|
extend_logprob_start_lens = self.extend_logprob_start_lens
|
|
|
|
if self.sampling_info:
|
|
if self.has_grammar:
|
|
self.sampling_info.grammars = [req.grammar for req in self.reqs]
|
|
else:
|
|
self.sampling_info.grammars = None
|
|
|
|
seq_lens_cpu = (
|
|
seq_lens_cpu_cache if seq_lens_cpu_cache is not None else self.seq_lens_cpu
|
|
)
|
|
|
|
return ModelWorkerBatch(
|
|
forward_mode=self.forward_mode,
|
|
input_ids=self.input_ids,
|
|
req_pool_indices=self.req_pool_indices,
|
|
seq_lens=self.seq_lens,
|
|
orig_seq_lens=self.orig_seq_lens,
|
|
out_cache_loc=self.out_cache_loc,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
seq_lens_sum=self.seq_lens_sum,
|
|
return_logprob=self.return_logprob,
|
|
top_logprobs_nums=self.top_logprobs_nums,
|
|
token_ids_logprobs=self.token_ids_logprobs,
|
|
global_num_tokens=self.global_num_tokens,
|
|
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
|
|
is_extend_in_batch=self.is_extend_in_batch,
|
|
all_extend_in_batch=self.all_extend_in_batch,
|
|
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
|
|
tbo_split_seq_index=self.tbo_split_seq_index,
|
|
global_forward_mode=self.global_forward_mode,
|
|
extend_num_tokens=self.extend_num_tokens,
|
|
extend_seq_lens=extend_seq_lens,
|
|
extend_prefix_lens=extend_prefix_lens,
|
|
extend_logprob_start_lens=extend_logprob_start_lens,
|
|
multimodal_inputs=self.multimodal_inputs,
|
|
encoder_cached=self.encoder_cached,
|
|
encoder_lens=self.encoder_lens,
|
|
encoder_lens_cpu=self.encoder_lens_cpu,
|
|
encoder_out_cache_loc=self.encoder_out_cache_loc,
|
|
lora_ids=[req.lora_id for req in self.reqs],
|
|
sampling_info=self.sampling_info,
|
|
input_embeds=self.input_embeds,
|
|
ne_token_table=self.ne_token_table,
|
|
token_type_ids=self.token_type_ids,
|
|
spec_algorithm=self.spec_algorithm,
|
|
spec_info=self.spec_info,
|
|
hicache_consumer_index=self.hicache_consumer_index,
|
|
capture_hidden_mode=(
|
|
CaptureHiddenMode.FULL
|
|
if self.return_hidden_states
|
|
else (
|
|
getattr(
|
|
self.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL
|
|
)
|
|
if self.spec_info
|
|
else CaptureHiddenMode.NULL
|
|
)
|
|
),
|
|
extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
|
|
is_prefill_only=self.is_prefill_only,
|
|
dimensions=self.dimensions,
|
|
dllm_block_offsets=[req.dllm_block_offset for req in self.reqs],
|
|
dllm_config=self.dllm_config,
|
|
reqs=self.reqs,
|
|
has_grammar=self.has_grammar,
|
|
mamba_track_indices=self.mamba_track_indices,
|
|
mamba_track_mask=self.mamba_track_mask,
|
|
mamba_track_seqlens=self.mamba_track_seqlens,
|
|
)
|
|
|
|
def copy(self):
|
|
# Only contain fields that will be used by process_batch_result.
|
|
# Shallow-copy the reqs list so that in-place mutations (filter_batch,
|
|
# merge_batch) on the original don't corrupt this snapshot.
|
|
return ScheduleBatch(
|
|
reqs=self.reqs[:],
|
|
req_to_token_pool=self.req_to_token_pool,
|
|
req_pool_indices=self.req_pool_indices,
|
|
model_config=self.model_config,
|
|
forward_mode=self.forward_mode,
|
|
out_cache_loc=self.out_cache_loc,
|
|
return_logprob=self.return_logprob,
|
|
decoding_reqs=self.decoding_reqs,
|
|
spec_algorithm=self.spec_algorithm,
|
|
spec_info=self.spec_info,
|
|
global_num_tokens=self.global_num_tokens,
|
|
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
|
|
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
|
|
all_extend_in_batch=self.all_extend_in_batch,
|
|
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_mask=self.mamba_track_mask,
|
|
mamba_track_seqlens=self.mamba_track_seqlens,
|
|
dp_cooperation_info=self.dp_cooperation_info,
|
|
prefill_stats=self.prefill_stats,
|
|
)
|
|
|
|
def maybe_evict_swa(self):
|
|
if self.tree_cache.supports_swa():
|
|
sliding_window_size = self.tree_cache.sliding_window_size
|
|
server_args = get_global_server_args()
|
|
|
|
if (
|
|
self.forward_mode.is_decode()
|
|
and not server_args.disable_piecewise_cuda_graph
|
|
and not self.tree_cache.is_chunk_cache()
|
|
):
|
|
return
|
|
|
|
for idx, req in enumerate(self.reqs):
|
|
if self.forward_mode.is_decode():
|
|
# We set evict_swa condition here with two reasons:
|
|
# 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running.
|
|
# 2. Evict swa every window_size tokens to reduce the overhead.
|
|
if req.decode_batch_idx % sliding_window_size == 1:
|
|
self._evict_swa(req, req.seqlen - 1)
|
|
elif self.forward_mode.is_extend() and self.tree_cache.is_chunk_cache():
|
|
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 - server_args.chunked_prefill_size
|
|
if server_args.chunked_prefill_size > 0
|
|
else pre_len
|
|
)
|
|
self._evict_swa(req, pre_len)
|
|
else:
|
|
self._evict_swa(req, pre_len)
|
|
|
|
def _evict_swa(self, req: Req, pre_len: int):
|
|
assert self.tree_cache.supports_swa(), "prefix cache must support swa"
|
|
sliding_window_size = self.tree_cache.sliding_window_size
|
|
|
|
# For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
|
|
assert (
|
|
req.cache_protected_len % self.tree_cache.page_size == 0
|
|
), "cache_protected_len must be page aligned"
|
|
req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, req.cache_protected_len)
|
|
|
|
new_swa_evicted_seqlen = max(
|
|
req.swa_evicted_seqlen, pre_len - sliding_window_size
|
|
)
|
|
|
|
if self.tree_cache.page_size > 1:
|
|
new_swa_evicted_seqlen = (
|
|
new_swa_evicted_seqlen // self.tree_cache.page_size
|
|
) * self.tree_cache.page_size
|
|
|
|
if new_swa_evicted_seqlen > req.swa_evicted_seqlen:
|
|
free_slots = self.req_to_token_pool.req_to_token[
|
|
req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen
|
|
]
|
|
self.token_to_kv_pool_allocator.free_swa(free_slots)
|
|
req.swa_evicted_seqlen = new_swa_evicted_seqlen
|
|
|
|
def __str__(self):
|
|
return (
|
|
f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "
|
|
f"#req={(len(self.reqs))})"
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ModelWorkerBatch:
|
|
# The forward mode
|
|
forward_mode: ForwardMode
|
|
# The input ids
|
|
input_ids: torch.Tensor
|
|
# The indices of requests in the req_to_token_pool
|
|
req_pool_indices: torch.Tensor
|
|
# The sequence length
|
|
seq_lens: torch.Tensor
|
|
# The indices of output tokens in the token_to_kv_pool_allocator
|
|
out_cache_loc: torch.Tensor
|
|
# The sequence length tensor on CPU
|
|
seq_lens_cpu: Optional[torch.Tensor]
|
|
seq_lens_sum: int
|
|
|
|
# For logprob
|
|
return_logprob: bool
|
|
top_logprobs_nums: Optional[List[int]]
|
|
token_ids_logprobs: Optional[List[List[int]]]
|
|
|
|
# For DP attention
|
|
global_num_tokens: Optional[List[int]]
|
|
global_num_tokens_for_logprob: Optional[List[int]]
|
|
is_extend_in_batch: bool
|
|
all_extend_in_batch: bool
|
|
can_run_dp_cuda_graph: bool
|
|
tbo_split_seq_index: Optional[int]
|
|
global_forward_mode: Optional[ForwardMode]
|
|
|
|
# For extend
|
|
extend_num_tokens: Optional[int]
|
|
extend_seq_lens: Optional[List[int]]
|
|
extend_prefix_lens: Optional[List[int]]
|
|
extend_logprob_start_lens: Optional[List[int]]
|
|
extend_input_logprob_token_ids: Optional[torch.Tensor]
|
|
|
|
# For multimodal
|
|
multimodal_inputs: Optional[List[MultimodalInputs]]
|
|
|
|
# For encoder-decoder
|
|
encoder_cached: Optional[List[bool]]
|
|
encoder_lens: Optional[torch.Tensor]
|
|
encoder_lens_cpu: Optional[List[int]]
|
|
encoder_out_cache_loc: Optional[torch.Tensor]
|
|
|
|
# For LoRA
|
|
lora_ids: Optional[List[str]]
|
|
|
|
# Sampling info
|
|
sampling_info: SamplingBatchInfo
|
|
|
|
# The original sequence lengths, Qwen-1M related
|
|
orig_seq_lens: Optional[torch.Tensor] = None
|
|
|
|
# The input Embeds
|
|
input_embeds: Optional[torch.Tensor] = None
|
|
|
|
# token table for ngram embedding
|
|
ne_token_table: Optional[torch.Tensor] = None
|
|
|
|
# For corss-encoder model
|
|
token_type_ids: Optional[torch.Tensor] = None
|
|
|
|
# Speculative decoding
|
|
spec_algorithm: SpeculativeAlgorithm = None
|
|
|
|
spec_info: Optional[SpecInput] = None
|
|
|
|
# If set, the output of the batch contains the hidden states of the run.
|
|
capture_hidden_mode: CaptureHiddenMode = None
|
|
hicache_consumer_index: int = -1
|
|
|
|
# For matryoshka embeddings
|
|
dimensions: Optional[list[int]] = None
|
|
|
|
# Whether this batch is prefill-only (no token generation needed)
|
|
is_prefill_only: bool = False
|
|
|
|
# Diffusion LLM
|
|
dllm_block_offsets: Optional[List[int]] = None
|
|
dllm_config: Optional[DllmConfig] = None
|
|
|
|
# For constrained decoding
|
|
# FIXME(lsyin): remove this after fully overlap grammar
|
|
reqs: Optional[List[Req]] = None
|
|
has_grammar: bool = False
|
|
|
|
# For hidden states before normal
|
|
return_hidden_states_before_norm: bool = False
|
|
|
|
# For mamba state tracking
|
|
mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64
|
|
mamba_track_mask: Optional[torch.Tensor] = None # shape: [b], bool
|
|
mamba_track_seqlens: Optional[torch.Tensor] = None # shape: [b], int64
|