leavelet
6a1e862f48
专题 S4 (design: docs_internal/perf/prefill-compute-intensity-plan.md S4, amended). Under FCFS a cold request joining a warm-led batch turns a 1-2s cache-hit forward into a 5-10s one, splitting the warm work into the 新-cache-新 pattern. The policy prevents exactly that one thing: - WARM candidates always admit (into a cold-led batch they are free density — the cold extend dominates the forward anyway). - COLD admits into an empty or cold-led batch (small colds co-batch today; the FCFS head always starts a batch so the queue keeps moving). - COLD into a WARM-led batch is skipped, bounded by a per-pass window (W=16 skips), a head defer count (K=3 passes) and an age bound (T=5s). On any bound the scan STOPS instead of force-admitting: the cold waits for the same forward either way, but leads its own clean batch next pass instead of polluting this one. The skip is strictly post-match / pre-admit (after init_next_round_input, before add_one_req): no lock, no allocation, no budget mutation to unwind, and re-matching a skipped candidate next pass is exactly what the scan already does after a cap rejection. Classification is the in-scan match result (device prefix + host hit vs a 64-token floor) — under FCFS+L2 no pre-scan signal exists, so this adds zero matching work for inspected candidates. Disabled wholesale under priority scheduling (the skip must not reorder across priority classes). Three amendments vs the design draft, reasoned in the decision-table docstring: cold+cold-led admits (STOP would regress today's small-cold co-batching); starved heads STOP rather than force-admit (clean batch boundaries at identical latency); priority interaction handled by disabling rather than per-request comparison. Decision logic is a pure function with table + bounds unit tests (28/28 adder suite green). Default OFF. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2538 lines
100 KiB
Python
2538 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
|
|
# Cache-affinity scheduling (plan doc S4): consecutive passes this
|
|
# request was deferred as the FCFS head while a warm batch formed.
|
|
self.affinity_defer_count = 0
|
|
self.prefix_match_deferred_by_pending_backup = False
|
|
self.cp_hicache_prepared_backup = None
|
|
# 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)
|
|
self.prefix_match_deferred_by_pending_backup = (
|
|
match_result.pending_backup_deferred_node is not None
|
|
)
|
|
|
|
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,
|
|
status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
|
|
err_type: str = "BadRequestError",
|
|
):
|
|
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, status_code, err_type)
|
|
|
|
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
|
|
)
|
|
),
|
|
capture_draft_hidden_states=False,
|
|
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
|
|
capture_draft_hidden_states: bool = False
|
|
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
|