907 lines
35 KiB
Python
907 lines
35 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
|
|
from sglang.srt.utils import get_bool_env_var
|
|
|
|
_ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 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.
|
|
# ==============================================================================
|
|
"""Request scheduler policy"""
|
|
|
|
import os
|
|
import random
|
|
from collections import Counter, defaultdict
|
|
from contextlib import contextmanager
|
|
from enum import Enum, auto
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
|
|
|
|
import torch
|
|
|
|
from sglang.srt.dllm.config import DllmConfig
|
|
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
|
|
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
|
from sglang.srt.mem_cache.base_prefix_cache import (
|
|
BasePrefixCache,
|
|
InsertParams,
|
|
MatchPrefixParams,
|
|
)
|
|
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
|
|
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
|
|
# Clip the estimation of max_new_tokens for the request whose max_new_tokens is very large.
|
|
# This can prevent the server from being too conservative.
|
|
# Note that this only clips the estimation in the scheduler but does not change the stop
|
|
# condition. The request can still generate tokens until it hits the unclipped max_new_tokens.
|
|
CLIP_MAX_NEW_TOKENS = int(
|
|
os.environ.get("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", "4096")
|
|
)
|
|
|
|
# Threshold for in-batch prefix cache.
|
|
# If a request has a matched prefix length (against existing cache) less than this value,
|
|
# the scheduler runs the in-batch prefix caching check for this request.
|
|
# If we set it to -1, it means we disable in-batch prefix caching.
|
|
IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD = int(
|
|
os.environ.get("IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD", "32")
|
|
)
|
|
|
|
# Threshold for in-batch prefix cache.
|
|
# If a request has a matched prefix length (within the waiting queue) larger than this value,
|
|
# the scheduler deprioritizes this request
|
|
IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD = int(
|
|
os.environ.get("IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD", "32")
|
|
)
|
|
|
|
|
|
IGNORE_EOS_RESERVE_TOKENS = 1
|
|
|
|
|
|
class CacheAwarePolicy(Enum):
|
|
"""Scheduling policies that are aware of the tree cache."""
|
|
|
|
LPM = "lpm" # longest prefix match
|
|
DFS_WEIGHT = "dfs-weight" # depth-first search weighting
|
|
|
|
|
|
class CacheAgnosticPolicy(Enum):
|
|
"""Scheduling policies that are not aware of the tree cache."""
|
|
|
|
FCFS = "fcfs" # first come first serve
|
|
LOF = "lof" # longest output first
|
|
RANDOM = "random"
|
|
ROUTING_KEY = "routing-key" # prioritize by routing key frequency in running batch
|
|
|
|
|
|
class SchedulePolicy:
|
|
Policy = Union[CacheAwarePolicy, CacheAgnosticPolicy]
|
|
|
|
def __init__(
|
|
self,
|
|
policy: str,
|
|
tree_cache: BasePrefixCache,
|
|
enable_hierarchical_cache: bool,
|
|
enable_priority_scheduling: bool,
|
|
schedule_low_priority_values_first: bool,
|
|
):
|
|
self.policy = self._validate_and_adjust_policy(policy, tree_cache)
|
|
self.tree_cache = tree_cache
|
|
self.enable_hierarchical_cache = enable_hierarchical_cache
|
|
self.enable_priority_scheduling = enable_priority_scheduling
|
|
self.schedule_low_priority_values_first = schedule_low_priority_values_first
|
|
self.priority_sign = 1 if schedule_low_priority_values_first else -1
|
|
|
|
# It is used to find the matching prefix for in-batch prefix caching.
|
|
self.waiting_queue_radix_tree = RadixCache.create_simulated()
|
|
|
|
def calc_priority(
|
|
self, waiting_queue: List[Req], running_batch: Optional[ScheduleBatch] = None
|
|
) -> bool:
|
|
if self.policy == CacheAgnosticPolicy.FCFS:
|
|
if self.enable_priority_scheduling:
|
|
SchedulePolicy._sort_by_priority_and_fcfs(
|
|
waiting_queue, self.priority_sign
|
|
)
|
|
return False
|
|
|
|
policy = self._determine_active_policy(waiting_queue)
|
|
|
|
prefix_computed = False
|
|
if isinstance(policy, CacheAwarePolicy):
|
|
prefix_computed = True
|
|
temporary_deprioritized = self._compute_prefix_matches(
|
|
waiting_queue, policy
|
|
)
|
|
if policy == CacheAwarePolicy.LPM:
|
|
SchedulePolicy._sort_by_longest_prefix(
|
|
waiting_queue, temporary_deprioritized
|
|
)
|
|
elif policy == CacheAwarePolicy.DFS_WEIGHT:
|
|
SchedulePolicy._sort_by_dfs_weight(waiting_queue, self.tree_cache)
|
|
else:
|
|
raise ValueError(f"Unknown CacheAware Policy: {policy=}")
|
|
else:
|
|
if policy == CacheAgnosticPolicy.FCFS:
|
|
pass
|
|
elif policy == CacheAgnosticPolicy.LOF:
|
|
SchedulePolicy._sort_by_longest_output(
|
|
waiting_queue,
|
|
self.enable_priority_scheduling,
|
|
self.priority_sign,
|
|
)
|
|
elif policy == CacheAgnosticPolicy.RANDOM:
|
|
SchedulePolicy._sort_randomly(waiting_queue)
|
|
elif policy == CacheAgnosticPolicy.ROUTING_KEY:
|
|
if running_batch is not None:
|
|
SchedulePolicy._sort_by_routing_key(waiting_queue, running_batch)
|
|
else:
|
|
raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}")
|
|
return prefix_computed
|
|
|
|
def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy:
|
|
if self.policy == CacheAwarePolicy.LPM and len(waiting_queue) > 128:
|
|
# Turn off the expensive prefix matching and sorting when the #queue is large.
|
|
return CacheAgnosticPolicy.FCFS
|
|
return self.policy
|
|
|
|
def _validate_and_adjust_policy(
|
|
self, policy: str, tree_cache: BasePrefixCache
|
|
) -> Policy:
|
|
"""
|
|
Validates the policy and adjusts it if necessary based on tree cache settings.
|
|
"""
|
|
try:
|
|
policy_enum = CacheAwarePolicy(policy)
|
|
if getattr(tree_cache, "disable", True):
|
|
# If tree_cache is disabled, using CacheAgnosticPolicy policy
|
|
return CacheAgnosticPolicy.FCFS
|
|
return policy_enum
|
|
except ValueError:
|
|
try:
|
|
return CacheAgnosticPolicy(policy)
|
|
except ValueError:
|
|
raise ValueError(f"Unknown schedule_policy: {policy=}")
|
|
|
|
def _compute_prefix_matches(
|
|
self, waiting_queue: List[Req], policy: CacheAwarePolicy
|
|
) -> Set[int]:
|
|
"""
|
|
Computes and caches the matching prefixes for requests in the waiting queue,
|
|
and handles in-batch prefix caching logic.
|
|
"""
|
|
temporary_deprioritized: Set[int] = set()
|
|
self.waiting_queue_radix_tree.reset()
|
|
|
|
for r in waiting_queue:
|
|
prefix_ids = r.origin_input_ids + r.output_ids
|
|
extra_key = r.extra_key
|
|
# NOTE: the prefix_indices must always be aligned with last_node
|
|
match_result = self.tree_cache.match_prefix(
|
|
MatchPrefixParams(
|
|
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key)
|
|
)
|
|
)
|
|
(
|
|
r.prefix_indices,
|
|
r.last_node,
|
|
r.last_host_node,
|
|
r.host_hit_length,
|
|
) = (
|
|
match_result.device_indices,
|
|
match_result.last_device_node,
|
|
match_result.last_host_node,
|
|
match_result.host_hit_length,
|
|
)
|
|
|
|
# NOTE(sang): This logic is for in-batch prefix caching;
|
|
# If there are more than 1 request that have small matching prefix from
|
|
# existing cache, but all those requests share the same prefix, we prefer
|
|
# to schedule only one of them so that we can increase the cache hit rate.
|
|
# We prefer to set IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD > 0 because too small
|
|
# threshold means we cannot use in-batch prefix caching for short prefixes.
|
|
# It is kind of common when the engine is long running (e.g., imagine the prefix "the").
|
|
if len(r.prefix_indices) <= IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD:
|
|
match_result = self.waiting_queue_radix_tree.match_prefix(
|
|
MatchPrefixParams(
|
|
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key)
|
|
)
|
|
)
|
|
in_batch_matching_prefixes = match_result.device_indices
|
|
if (
|
|
len(in_batch_matching_prefixes)
|
|
>= IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD
|
|
):
|
|
temporary_deprioritized.add(r.rid)
|
|
else:
|
|
# Insert with a dummy key
|
|
self.waiting_queue_radix_tree.insert(
|
|
InsertParams(
|
|
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key),
|
|
value=torch.empty(len(prefix_ids), dtype=torch.bool),
|
|
)
|
|
)
|
|
return temporary_deprioritized
|
|
|
|
@staticmethod
|
|
def _sort_by_longest_prefix(
|
|
waiting_queue: List[Req], temporary_deprioritized: Set[int]
|
|
) -> None:
|
|
"""Sorts the waiting queue based on the longest prefix match."""
|
|
waiting_queue.sort(
|
|
key=lambda r: (
|
|
-len(r.prefix_indices)
|
|
if r.rid not in temporary_deprioritized
|
|
else float("inf")
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _sort_by_dfs_weight(
|
|
waiting_queue: List[Req], tree_cache: BasePrefixCache
|
|
) -> None:
|
|
"""Sorts the waiting queue based on a depth-first search weighting."""
|
|
last_node_to_reqs = defaultdict(list)
|
|
for req in waiting_queue:
|
|
last_node_to_reqs[req.last_node].append(req)
|
|
|
|
node_to_weight = defaultdict(int)
|
|
for node in last_node_to_reqs:
|
|
node_to_weight[node] = len(last_node_to_reqs[node])
|
|
SchedulePolicy._calc_weight(tree_cache.root_node, node_to_weight)
|
|
|
|
waiting_queue.clear()
|
|
SchedulePolicy._get_dfs_priority(
|
|
tree_cache.root_node,
|
|
node_to_weight,
|
|
last_node_to_reqs,
|
|
waiting_queue,
|
|
)
|
|
|
|
@staticmethod
|
|
def _sort_by_longest_output(
|
|
waiting_queue: List[Req],
|
|
enable_priority_scheduling: bool,
|
|
priority_sign: int,
|
|
) -> None:
|
|
"""Sorts the waiting queue based on the longest output (max_new_tokens). If using priority scheduling, sort by priority first."""
|
|
if enable_priority_scheduling:
|
|
waiting_queue.sort(
|
|
key=lambda x: (
|
|
x.priority * priority_sign,
|
|
-x.sampling_params.max_new_tokens,
|
|
)
|
|
)
|
|
else:
|
|
waiting_queue.sort(key=lambda x: -x.sampling_params.max_new_tokens)
|
|
|
|
@staticmethod
|
|
def _sort_randomly(waiting_queue: List[Req]) -> None:
|
|
"""Shuffles the waiting queue randomly."""
|
|
random.shuffle(waiting_queue)
|
|
|
|
@staticmethod
|
|
def _sort_by_priority_and_fcfs(
|
|
waiting_queue: List[Req], priority_sign: int
|
|
) -> None:
|
|
"""Sorts the waiting queue based on the request priority then received titmestamp."""
|
|
waiting_queue.sort(
|
|
key=lambda x: (
|
|
x.priority * priority_sign,
|
|
x.time_stats.wait_queue_entry_time,
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _sort_by_routing_key(
|
|
waiting_queue: List[Req], running_batch: ScheduleBatch
|
|
) -> None:
|
|
"""Sorts waiting queue by routing key frequency in running batch."""
|
|
routing_key_counts = Counter(
|
|
r.routing_key for r in running_batch.reqs if r.routing_key
|
|
)
|
|
|
|
if _ROUTING_KEY_POLICY_DEBUG_LOG:
|
|
waiting_keys_before = [r.routing_key for r in waiting_queue]
|
|
logger.info(
|
|
f"routing_key_counts={dict(routing_key_counts)}, "
|
|
f"waiting_keys_before={waiting_keys_before}"
|
|
)
|
|
|
|
if not routing_key_counts:
|
|
return
|
|
|
|
def sort_key(req: Req):
|
|
key = req.routing_key
|
|
if key and key in routing_key_counts:
|
|
count = routing_key_counts[key]
|
|
return (0, -count, key)
|
|
else:
|
|
return (1, 0, key or "")
|
|
|
|
waiting_queue.sort(key=sort_key)
|
|
|
|
if _ROUTING_KEY_POLICY_DEBUG_LOG:
|
|
waiting_keys_after = [r.routing_key for r in waiting_queue]
|
|
logger.info(f"waiting_keys_after={waiting_keys_after}")
|
|
|
|
@staticmethod
|
|
def _calc_weight(cur_node: TreeNode, node_to_weight: Dict[TreeNode, int]) -> None:
|
|
for child in cur_node.children.values():
|
|
SchedulePolicy._calc_weight(child, node_to_weight)
|
|
node_to_weight[cur_node] += node_to_weight[child]
|
|
|
|
@staticmethod
|
|
def _get_dfs_priority(
|
|
cur_node: TreeNode,
|
|
node_to_priority: Dict[TreeNode, int],
|
|
last_node_to_reqs: Dict[TreeNode, List[Req]],
|
|
q: List,
|
|
) -> None:
|
|
children = [child for child in cur_node.children.values()]
|
|
children.sort(key=lambda x: -node_to_priority[x])
|
|
for child in children:
|
|
SchedulePolicy._get_dfs_priority(
|
|
child, node_to_priority, last_node_to_reqs, q
|
|
)
|
|
q.extend(last_node_to_reqs[cur_node])
|
|
|
|
|
|
class AddReqResult(Enum):
|
|
CONTINUE = auto() # Continue to add requests
|
|
NO_TOKEN = auto() # No token left
|
|
OTHER = auto() # Other reasons to stop adding requests
|
|
|
|
|
|
class PrefillAdder:
|
|
def __init__(
|
|
self,
|
|
page_size: int,
|
|
tree_cache: BasePrefixCache,
|
|
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
|
|
running_batch: ScheduleBatch,
|
|
new_token_ratio: float,
|
|
rem_input_tokens: int,
|
|
rem_chunk_tokens: Optional[int],
|
|
mixed_with_decode_tokens: int = 0,
|
|
priority_scheduling_preemption_threshold: int = 0,
|
|
max_prefill_bs: int = 0,
|
|
max_running_requests: Optional[int] = None,
|
|
prefill_max_requests: Optional[int] = None,
|
|
prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
|
|
dllm_config: Optional[DllmConfig] = None,
|
|
):
|
|
self.page_size = page_size
|
|
self.tree_cache = tree_cache
|
|
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
|
|
self.running_batch = running_batch
|
|
self.new_token_ratio = new_token_ratio
|
|
self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
|
|
self.rem_chunk_tokens = rem_chunk_tokens
|
|
self.dllm_config = dllm_config
|
|
|
|
if self.dllm_config is not None:
|
|
self._init_dllm_meta(dllm_config)
|
|
|
|
if self.rem_chunk_tokens is not None:
|
|
self.rem_chunk_tokens -= mixed_with_decode_tokens
|
|
self.rem_total_token_offset = mixed_with_decode_tokens
|
|
self.cur_rem_token_offset = mixed_with_decode_tokens
|
|
|
|
self.req_states = None
|
|
self.can_run_list = []
|
|
self.preempt_list = []
|
|
self.new_chunked_req = None
|
|
self.log_hit_tokens = 0
|
|
# TODO(lsyin): report the real input tokens excluding page alignment
|
|
self.log_input_tokens = 0
|
|
|
|
if running_batch is not None:
|
|
self.rem_total_token_offset += sum(
|
|
[
|
|
self._get_running_request_total_token_offset(r)
|
|
for r in running_batch.reqs
|
|
]
|
|
)
|
|
|
|
self.is_hybrid_swa = isinstance(
|
|
self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
|
|
)
|
|
self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba()
|
|
|
|
self.priority_scheduling_preemption_threshold = (
|
|
priority_scheduling_preemption_threshold
|
|
)
|
|
self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split()
|
|
self.max_running_requests = max_running_requests
|
|
self.prefill_max_requests = prefill_max_requests
|
|
self.prefill_delayer_single_pass = prefill_delayer_single_pass
|
|
self.max_prefill_bs = max_prefill_bs
|
|
|
|
def _init_dllm_meta(self, dllm_config: DllmConfig):
|
|
self.dllm_block_size = dllm_config.block_size
|
|
max_running_reqs = dllm_config.max_running_requests
|
|
|
|
self.rem_dllm_tokens = max_running_reqs * self.dllm_block_size
|
|
|
|
def _get_running_request_total_token_offset(self, req: Req) -> int:
|
|
return (
|
|
min(
|
|
(req.sampling_params.max_new_tokens - len(req.output_ids)),
|
|
CLIP_MAX_NEW_TOKENS,
|
|
)
|
|
* self.new_token_ratio
|
|
)
|
|
|
|
@property
|
|
def rem_total_tokens(self):
|
|
if self.is_hybrid_swa:
|
|
available_and_evictable = min(
|
|
self.token_to_kv_pool_allocator.full_available_size()
|
|
+ self.tree_cache.full_evictable_size(),
|
|
self.token_to_kv_pool_allocator.swa_available_size()
|
|
+ self.tree_cache.swa_evictable_size(),
|
|
)
|
|
elif self.is_hybrid_ssm_cache:
|
|
available_and_evictable = (
|
|
self.token_to_kv_pool_allocator.available_size()
|
|
+ self.tree_cache.full_evictable_size()
|
|
)
|
|
else:
|
|
available_and_evictable = (
|
|
self.token_to_kv_pool_allocator.available_size()
|
|
+ self.tree_cache.evictable_size()
|
|
)
|
|
return available_and_evictable - self.rem_total_token_offset
|
|
|
|
@property
|
|
def cur_rem_tokens(self):
|
|
if self.is_hybrid_swa:
|
|
available_and_evictable = min(
|
|
self.token_to_kv_pool_allocator.full_available_size()
|
|
+ self.tree_cache.full_evictable_size(),
|
|
self.token_to_kv_pool_allocator.swa_available_size()
|
|
+ self.tree_cache.swa_evictable_size(),
|
|
)
|
|
elif self.is_hybrid_ssm_cache:
|
|
available_and_evictable = (
|
|
self.token_to_kv_pool_allocator.available_size()
|
|
+ self.tree_cache.full_evictable_size()
|
|
)
|
|
else:
|
|
available_and_evictable = (
|
|
self.token_to_kv_pool_allocator.available_size()
|
|
+ self.tree_cache.evictable_size()
|
|
)
|
|
|
|
return available_and_evictable - self.cur_rem_token_offset
|
|
|
|
def ceil_paged_tokens(self, tokens: int) -> int:
|
|
return -(-tokens // self.page_size) * self.page_size
|
|
|
|
def budget_state(self):
|
|
if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0:
|
|
return AddReqResult.NO_TOKEN
|
|
|
|
if self.rem_input_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
if self.dllm_config is not None:
|
|
if self.rem_dllm_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
else:
|
|
if self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
return AddReqResult.CONTINUE
|
|
|
|
def _update_prefill_budget(
|
|
self, prefix_len: int, extend_input_len: int, max_new_tokens: int
|
|
):
|
|
# TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative
|
|
extend_input_len = self.ceil_paged_tokens(extend_input_len)
|
|
|
|
# alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit
|
|
page_overhead = self.page_size
|
|
self.rem_total_token_offset += extend_input_len + max_new_tokens + page_overhead
|
|
self.cur_rem_token_offset += extend_input_len + page_overhead
|
|
self.rem_input_tokens -= extend_input_len
|
|
|
|
if self.dllm_config is not None:
|
|
self.rem_dllm_tokens -= extend_input_len
|
|
elif self.rem_chunk_tokens is not None:
|
|
self.rem_chunk_tokens -= extend_input_len
|
|
|
|
self.log_hit_tokens += prefix_len
|
|
self.log_input_tokens += extend_input_len
|
|
|
|
def _get_dllm_remain_tokens(self) -> int:
|
|
_rem_tokens = min(
|
|
self.rem_dllm_tokens,
|
|
self.dllm_block_size,
|
|
int(self.rem_total_tokens),
|
|
)
|
|
if _rem_tokens <= 0:
|
|
_rem_tokens = self.rem_dllm_tokens
|
|
|
|
return _rem_tokens
|
|
|
|
def _add_dllm_req(self, req: Req, prefix_len: int):
|
|
# FIXME: consider the case when rem_dllm_tokens < dllm_block_size,
|
|
# the diffusion unmask process may have some problems
|
|
# Make sure at least one page is available
|
|
trunc_len = (
|
|
min(self.rem_dllm_tokens, self.dllm_block_size)
|
|
// self.page_size
|
|
* self.page_size
|
|
)
|
|
|
|
req.extend_input_len = trunc_len
|
|
req.fill_ids = req.fill_ids[: prefix_len + trunc_len]
|
|
|
|
self.can_run_list.append(req)
|
|
|
|
self._update_prefill_budget(prefix_len, trunc_len, 0)
|
|
|
|
def _req_inc_lock_ref(self, req: Req):
|
|
if self.is_hybrid_swa:
|
|
swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node)
|
|
req.swa_uuid_for_lock = swa_uuid_for_lock
|
|
else:
|
|
self.tree_cache.inc_lock_ref(req.last_node)
|
|
|
|
def add_dllm_staging_req(self, req: Req):
|
|
assert self.dllm_config is not None
|
|
_rem_tokens = self._get_dllm_remain_tokens()
|
|
|
|
if _rem_tokens <= 0:
|
|
return AddReqResult.NO_TOKEN
|
|
|
|
# Truncate input length to available tokens and update request metadata
|
|
truncated = req.extend_input_len > _rem_tokens
|
|
req.extend_input_len = min(req.extend_input_len, _rem_tokens)
|
|
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
|
|
self.can_run_list.append(req)
|
|
|
|
# Update budget: reserve max_new_tokens only if not truncated
|
|
max_new_tokens = (
|
|
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
|
|
if not truncated
|
|
else 0
|
|
)
|
|
self._update_prefill_budget(0, req.extend_input_len, max_new_tokens)
|
|
|
|
# Return based on remaining token availability
|
|
return (
|
|
AddReqResult.NO_TOKEN
|
|
if self._get_dllm_remain_tokens() <= 0
|
|
else AddReqResult.CONTINUE
|
|
)
|
|
|
|
def add_chunked_req(self, req: Req):
|
|
if self.dllm_config is not None:
|
|
_rem_tokens = self._get_dllm_remain_tokens()
|
|
else:
|
|
_rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
|
|
# The chunked_req must be added to the list; otherwise, it will cause a memory leak.
|
|
# Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
|
|
if _rem_tokens <= 0:
|
|
_rem_tokens = self.rem_chunk_tokens
|
|
|
|
truncated = req.extend_input_len > _rem_tokens
|
|
req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
|
|
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
|
|
self.can_run_list.append(req)
|
|
self._update_prefill_budget(
|
|
0,
|
|
req.extend_input_len,
|
|
(
|
|
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
|
|
if not truncated
|
|
else 0
|
|
),
|
|
)
|
|
|
|
# Return if chunked prefill not finished
|
|
return req if truncated else None
|
|
|
|
@contextmanager
|
|
def _lock_node(self, last_node: TreeNode):
|
|
try:
|
|
if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
|
|
swa_uuid_for_lock = self.tree_cache.inc_lock_ref(last_node)
|
|
else:
|
|
self.tree_cache.inc_lock_ref(last_node)
|
|
yield None
|
|
finally:
|
|
if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
|
|
self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock)
|
|
else:
|
|
self.tree_cache.dec_lock_ref(last_node)
|
|
|
|
def add_one_req_ignore_eos(self, req: Req):
|
|
# Early exit if no enough tokens for the input tokens
|
|
if self.ceil_paged_tokens(req.extend_input_len) > min(
|
|
self.cur_rem_tokens, self.rem_total_tokens
|
|
):
|
|
return AddReqResult.NO_TOKEN
|
|
|
|
def add_req_state(r, insert_sort=False):
|
|
new_token_ratio = (
|
|
1.0 if r.sampling_params.ignore_eos else self.new_token_ratio
|
|
)
|
|
tokens_left = r.sampling_params.max_new_tokens * new_token_ratio - len(
|
|
r.output_ids
|
|
)
|
|
tokens_occupied = len(r.origin_input_ids) + len(r.output_ids)
|
|
|
|
if tokens_left <= 0:
|
|
return
|
|
|
|
if not insert_sort:
|
|
self.req_states.append((tokens_left, tokens_occupied))
|
|
else:
|
|
i = 0
|
|
for i in range(len(self.req_states)):
|
|
if tokens_left <= self.req_states[i][0]:
|
|
break
|
|
self.req_states.insert(i, (tokens_left, tokens_occupied))
|
|
|
|
if self.req_states is None:
|
|
self.req_states = []
|
|
add_req_state(req)
|
|
if self.running_batch is not None:
|
|
for r in self.running_batch.reqs:
|
|
add_req_state(r)
|
|
for r in self.can_run_list:
|
|
add_req_state(r)
|
|
self.req_states.sort(key=lambda x: x[0])
|
|
else:
|
|
add_req_state(req, insert_sort=True)
|
|
|
|
if not self.is_hybrid_swa:
|
|
# Skip this logic for swa. The SWA has different memory management, and
|
|
# this mechanism is underestimating the memory usage.
|
|
cur_rem_tokens = self.cur_rem_tokens - self.ceil_paged_tokens(
|
|
req.extend_input_len
|
|
)
|
|
tokens_freed = 0
|
|
for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
|
|
# tokens_left gives a reservative calculation as the last token is not stored
|
|
bs = len(self.req_states) - i
|
|
min_free_tokens = cur_rem_tokens + tokens_freed - tokens_left * bs
|
|
# reserve tokens for corner cases
|
|
if min_free_tokens <= IGNORE_EOS_RESERVE_TOKENS * bs:
|
|
return AddReqResult.NO_TOKEN
|
|
tokens_freed += tokens_occupied
|
|
|
|
if self.dllm_config is not None:
|
|
if self.rem_dllm_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
self._add_dllm_req(req, 0)
|
|
elif (
|
|
self.rem_chunk_tokens is None # chunked prefill is disabled
|
|
or req.extend_input_len <= self.rem_chunk_tokens # it is the last chunk
|
|
):
|
|
# Non-chunked prefill
|
|
self.can_run_list.append(req)
|
|
self._update_prefill_budget(
|
|
0,
|
|
req.extend_input_len,
|
|
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS),
|
|
)
|
|
else:
|
|
if self.rem_chunk_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
# Chunked prefill
|
|
trunc_len = self.rem_chunk_tokens
|
|
|
|
req.set_extend_input_len(trunc_len)
|
|
req.fill_ids = req.fill_ids[:trunc_len]
|
|
self.can_run_list.append(req)
|
|
self.new_chunked_req = req
|
|
self._update_prefill_budget(0, trunc_len, 0)
|
|
|
|
return self.budget_state()
|
|
|
|
def add_one_req(
|
|
self, req: Req, has_chunked_req: bool, truncation_align_size: Optional[int]
|
|
):
|
|
if (self.prefill_delayer_single_pass is not None) and (
|
|
not self.prefill_delayer_single_pass.negotiate_should_allow_prefill(
|
|
local_prefillable=True,
|
|
running_batch=self.running_batch.batch_size(),
|
|
max_prefill_bs=self.max_prefill_bs,
|
|
max_running_requests=self.max_running_requests,
|
|
)
|
|
):
|
|
return AddReqResult.OTHER
|
|
# TODO support cp with multiple requests
|
|
# Enabling context parallelism currently presents precision issues;
|
|
# therefore, the prefill-batch setting is temporarily set to 1.
|
|
if self.nsa_prefill_cp_in_seq_split and len(self.can_run_list) >= 1:
|
|
return AddReqResult.OTHER
|
|
|
|
if (x := self.prefill_max_requests) is not None and len(self.can_run_list) >= x:
|
|
return AddReqResult.OTHER
|
|
|
|
if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True):
|
|
return self.add_one_req_ignore_eos(req)
|
|
|
|
total_tokens = req.extend_input_len + min(
|
|
max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
|
|
CLIP_MAX_NEW_TOKENS,
|
|
)
|
|
|
|
# adjusting the input_tokens based on host_hit_length and page_size
|
|
real_input_tokens = req.extend_input_len - req.host_hit_length
|
|
real_input_tokens = self.ceil_paged_tokens(real_input_tokens)
|
|
prefix_len = len(req.prefix_indices)
|
|
|
|
if total_tokens >= self.rem_total_tokens:
|
|
return AddReqResult.NO_TOKEN
|
|
|
|
if real_input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0:
|
|
return AddReqResult.OTHER
|
|
|
|
with self._lock_node(req.last_node):
|
|
# self.rem_total_tokens may decrease after the lock acquisition
|
|
if total_tokens >= self.rem_total_tokens:
|
|
return AddReqResult.NO_TOKEN
|
|
|
|
if req.host_hit_length > 0:
|
|
new_indices, req.last_node = self.tree_cache.init_load_back(
|
|
req.last_host_node, req.host_hit_length
|
|
)
|
|
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
|
|
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
|
prefix_len = len(req.prefix_indices)
|
|
req.cache_protected_len = prefix_len
|
|
|
|
input_tokens = self.ceil_paged_tokens(req.extend_input_len)
|
|
|
|
if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0:
|
|
return AddReqResult.OTHER
|
|
|
|
if self.dllm_config is not None:
|
|
if self.rem_dllm_tokens <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
assert (
|
|
truncation_align_size is None
|
|
), "truncation_align_size is not supported for dllm prefill"
|
|
|
|
self._add_dllm_req(req, prefix_len)
|
|
self._req_inc_lock_ref(req)
|
|
elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
|
|
# Non-chunked prefill
|
|
self.can_run_list.append(req)
|
|
|
|
self._req_inc_lock_ref(req)
|
|
self._update_prefill_budget(
|
|
prefix_len,
|
|
input_tokens,
|
|
min(
|
|
req.sampling_params.max_new_tokens,
|
|
CLIP_MAX_NEW_TOKENS,
|
|
),
|
|
)
|
|
else:
|
|
# Make sure at least one page is available
|
|
trunc_len = self.rem_chunk_tokens // self.page_size * self.page_size
|
|
|
|
if trunc_len <= 0:
|
|
return AddReqResult.OTHER
|
|
|
|
# When truncation align size is set, we want to assert that the prefill prefix length is multiple of truncation align size
|
|
# A typical use case is when deterministic inference is enabled with flashinfer attention backend,
|
|
# we need the prefill prefix length to be multiple of attention split size
|
|
if truncation_align_size is not None:
|
|
if trunc_len < truncation_align_size:
|
|
return AddReqResult.OTHER
|
|
else:
|
|
trunc_len = truncation_align_size * (
|
|
trunc_len // truncation_align_size
|
|
)
|
|
|
|
# Chunked prefill
|
|
req.set_extend_input_len(trunc_len)
|
|
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len]
|
|
|
|
self.can_run_list.append(req)
|
|
self.new_chunked_req = req
|
|
|
|
self._req_inc_lock_ref(req)
|
|
self._update_prefill_budget(prefix_len, trunc_len, 0)
|
|
|
|
return self.budget_state()
|
|
|
|
def preempt_to_schedule(self, req: Req, server_args: ServerArgs) -> bool:
|
|
"""
|
|
Preempt running requests to serve the new request if the priority threshold is met and token count sum is verified.
|
|
Returns True if preemption was committed, and the new request can be scheduled.
|
|
"""
|
|
# Iterate running requests to find preemptible requests
|
|
priority_sign = 1 if server_args.schedule_low_priority_values_first else -1
|
|
|
|
# NOTE: A request finishes in two phases:
|
|
# 1) check_finished + release_kv_cache (in process_batch_result)
|
|
# 2) filter out of batch (in get_next_batch_to_run / update_running_batch)
|
|
# Preemption runs between these two phases (inside get_new_batch_prefill),
|
|
# so running_batch may still contain requests whose KV cache is already freed.
|
|
# We must skip them here to avoid a double-free on release_req.
|
|
valid_running_reqs = (
|
|
r
|
|
for r in self.running_batch.reqs
|
|
if r not in self.preempt_list and not r.finished()
|
|
)
|
|
|
|
sorted_valid_running_reqs = sorted(
|
|
valid_running_reqs,
|
|
key=lambda x: (
|
|
x.priority * (-priority_sign),
|
|
-x.time_stats.wait_queue_entry_time,
|
|
),
|
|
)
|
|
|
|
preemptible_reqs = []
|
|
min_tokens_to_remove = (
|
|
req.extend_input_len
|
|
+ min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
|
|
- self.rem_total_tokens
|
|
)
|
|
for running_req in sorted_valid_running_reqs:
|
|
# Priority difference needs to meet the threshold to be preemptible.
|
|
priority_diff = (req.priority - running_req.priority) * (-priority_sign)
|
|
|
|
if priority_diff > self.priority_scheduling_preemption_threshold:
|
|
preemptible_reqs.append(running_req)
|
|
min_tokens_to_remove -= self._get_running_request_total_token_offset(
|
|
running_req
|
|
)
|
|
if min_tokens_to_remove <= 0:
|
|
break
|
|
else:
|
|
break
|
|
|
|
# Check max token count limit can be met
|
|
if len(preemptible_reqs) == 0 or min_tokens_to_remove > 0:
|
|
return False
|
|
|
|
# Preempt running requests. Release allocated resources for immediate usage.
|
|
preemptible_reqs = set(preemptible_reqs)
|
|
keep_indices = []
|
|
release_counter = 0
|
|
for i, running_req in enumerate(self.running_batch.reqs):
|
|
if running_req in preemptible_reqs:
|
|
self.rem_total_token_offset -= (
|
|
self._get_running_request_total_token_offset(running_req)
|
|
)
|
|
release_counter += 1
|
|
self.running_batch.release_req(
|
|
i, len(self.running_batch.reqs) - release_counter, server_args
|
|
)
|
|
else:
|
|
keep_indices.append(i)
|
|
self.running_batch.filter_batch(keep_indices=keep_indices)
|
|
self.preempt_list.extend(preemptible_reqs)
|
|
return True
|