[P/D disagg] - support decode side radix cache (#19746)
This commit is contained in:
@@ -95,6 +95,12 @@ class BaseKVSender(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return 0
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0
|
||||
|
||||
@abstractmethod
|
||||
def poll(self) -> KVPoll:
|
||||
"""
|
||||
@@ -136,6 +142,7 @@ class BaseKVReceiver(ABC):
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List[int]] = None,
|
||||
decode_prefix_len: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Notify the prefill server about the kv indices, aux index, and state_indices.
|
||||
|
||||
@@ -141,6 +141,7 @@ class CommonKVManager(BaseKVManager):
|
||||
)
|
||||
self.register_to_bootstrap()
|
||||
self.transfer_infos = {}
|
||||
self.req_to_decode_prefix_len: Dict[int, int] = {}
|
||||
self.decode_kv_args_table = {}
|
||||
self.pp_group = get_pp_group()
|
||||
# If a timeout happens on the prefill side, it means prefill instances
|
||||
@@ -179,6 +180,12 @@ class CommonKVManager(BaseKVManager):
|
||||
return self.request_status[bootstrap_room]
|
||||
|
||||
def update_status(self, bootstrap_room: int, status: KVPoll):
|
||||
if (
|
||||
status == KVPoll.Failed
|
||||
and self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
and hasattr(self, "req_to_decode_prefix_len")
|
||||
):
|
||||
self.req_to_decode_prefix_len.pop(bootstrap_room, None)
|
||||
if bootstrap_room not in self.request_status:
|
||||
self.request_status[bootstrap_room] = status
|
||||
else:
|
||||
@@ -489,6 +496,12 @@ class CommonKVSender(BaseKVSender):
|
||||
f"CommonKVSender init with num_kv_indices: {num_kv_indices} and aux_index: {aux_index}"
|
||||
)
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return 0
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0
|
||||
|
||||
def send(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
|
||||
@@ -44,7 +44,6 @@ from sglang.srt.disaggregation.utils import (
|
||||
TransferBackend,
|
||||
get_kv_class,
|
||||
is_mla_backend,
|
||||
kv_to_page_indices,
|
||||
poll_and_all_reduce,
|
||||
poll_and_all_reduce_with_staging,
|
||||
prepare_abort,
|
||||
@@ -52,10 +51,18 @@ from sglang.srt.disaggregation.utils import (
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
||||
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
|
||||
from sglang.srt.managers.schedule_policy import match_prefix_for_req
|
||||
from sglang.srt.managers.utils import GenerationBatchResult
|
||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
||||
from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
BasePrefixCache,
|
||||
EvictParams,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import (
|
||||
kv_to_page_indices,
|
||||
page_align_floor,
|
||||
release_kv_cache,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
HybridLinearKVPool,
|
||||
HybridReqToTokenPool,
|
||||
@@ -68,6 +75,7 @@ from sglang.srt.observability.req_time_stats import (
|
||||
set_schedule_time_batch,
|
||||
set_time_batch,
|
||||
)
|
||||
from sglang.srt.utils import get_num_new_pages
|
||||
from sglang.srt.utils.network import NetworkAddress
|
||||
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
||||
|
||||
@@ -281,7 +289,7 @@ class DecodePreallocQueue:
|
||||
self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
|
||||
self.scheduler = scheduler
|
||||
self.transfer_queue = transfer_queue
|
||||
self.tree_cache = tree_cache # this is always a chunk cache
|
||||
self.tree_cache = tree_cache
|
||||
self.gloo_group = gloo_group
|
||||
self.tp_rank = tp_rank
|
||||
self.tp_size = tp_size
|
||||
@@ -448,6 +456,25 @@ class DecodePreallocQueue:
|
||||
|
||||
self.pending_reqs.append(decode_req)
|
||||
|
||||
def _match_prefix_and_lock(self, req: Req) -> Tuple[torch.Tensor, int]:
|
||||
"""
|
||||
Match a request against the decode-side radix cache, lock the matched
|
||||
node to prevent eviction, and return the matched prefix information.
|
||||
"""
|
||||
result = match_prefix_for_req(
|
||||
self.tree_cache,
|
||||
req,
|
||||
req.origin_input_ids,
|
||||
cow_mamba=self.tree_cache.supports_mamba(),
|
||||
include_req=True,
|
||||
)
|
||||
prefix_indices = result.device_indices
|
||||
last_device_node = result.last_device_node
|
||||
# Always lock to match aggregated scheduling behavior
|
||||
self.tree_cache.inc_lock_ref(last_device_node)
|
||||
|
||||
return prefix_indices, len(prefix_indices)
|
||||
|
||||
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]:
|
||||
prefill_info = self.kv_manager.prefill_info_table.get(_bootstrap_addr(req))
|
||||
# If None, it will go to the slow path and resolve prefill_info by _ensure_prefill_info then cache it
|
||||
@@ -736,14 +763,42 @@ class DecodePreallocQueue:
|
||||
# Memory estimation: don't add if the projected memory cannot be met
|
||||
# TODO: add new_token ratio
|
||||
origin_input_len = len(decode_req.req.origin_input_ids)
|
||||
if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
|
||||
# Match prefix against decode's radix cache.
|
||||
prefix_indices, prefix_len = self._match_prefix_and_lock(decode_req.req)
|
||||
# Align prefix_len down to page boundary so both prefill and
|
||||
# decode agree on the page-aligned split point for KV transfer.
|
||||
page_size = self.token_to_kv_pool_allocator.page_size
|
||||
if page_size > 1 and prefix_len % page_size != 0:
|
||||
prefix_len = page_align_floor(prefix_len, page_size)
|
||||
prefix_indices = prefix_indices[:prefix_len]
|
||||
|
||||
fill_len = origin_input_len + max(len(decode_req.req.output_ids) - 1, 0)
|
||||
required_alloc_tokens = self._required_alloc_tokens(
|
||||
fill_len=fill_len, prefix_len=prefix_len
|
||||
)
|
||||
# Matching may lock previously-evictable radix pages, so refresh
|
||||
# the admission budget against the post-lock pool state before we
|
||||
# decide whether this request still fits.
|
||||
allocatable_tokens = self._allocatable_tokens(
|
||||
retractable_tokens=retractable_tokens,
|
||||
count_retracted=True,
|
||||
extra_reserved_reqs=len(preallocated_reqs),
|
||||
)
|
||||
else:
|
||||
prefix_indices = None
|
||||
prefix_len = 0
|
||||
required_alloc_tokens = origin_input_len
|
||||
|
||||
required_tokens_for_request = (
|
||||
origin_input_len + self.num_reserved_decode_tokens
|
||||
required_alloc_tokens + self.num_reserved_decode_tokens
|
||||
)
|
||||
|
||||
if (
|
||||
max(
|
||||
required_tokens_for_request,
|
||||
origin_input_len
|
||||
- prefix_len
|
||||
+ min(
|
||||
decode_req.req.sampling_params.max_new_tokens,
|
||||
CLIP_MAX_NEW_TOKEN,
|
||||
@@ -752,26 +807,43 @@ class DecodePreallocQueue:
|
||||
)
|
||||
> allocatable_tokens
|
||||
):
|
||||
if prefix_len > 0:
|
||||
self.tree_cache.dec_lock_ref(decode_req.req.last_node)
|
||||
break
|
||||
if required_tokens_for_request > allocatable_tokens:
|
||||
if prefix_len > 0:
|
||||
self.tree_cache.dec_lock_ref(decode_req.req.last_node)
|
||||
break
|
||||
|
||||
allocatable_tokens -= required_tokens_for_request
|
||||
dst_kv_indices = self._pre_alloc(decode_req.req, prefix_indices, prefix_len)
|
||||
hisparse_req_budget -= 1
|
||||
dst_kv_indices = self._pre_alloc(decode_req.req)
|
||||
# Recompute from actual pool state for the next queue entry.
|
||||
# This accounts for page rounding and newly locked evictable cache.
|
||||
allocatable_tokens = self._allocatable_tokens(
|
||||
retractable_tokens=retractable_tokens,
|
||||
count_retracted=True,
|
||||
extra_reserved_reqs=len(preallocated_reqs) + 1,
|
||||
)
|
||||
decode_req.req.cache_protected_len = prefix_len
|
||||
|
||||
origin_input_len = len(decode_req.req.origin_input_ids)
|
||||
if self.scheduler.enable_hisparse:
|
||||
# Must cast to int32 for ZMQ serialization — from_zmq reads np.int32.
|
||||
# Must cast to int32 for ZMQ serialization -- from_zmq reads np.int32.
|
||||
kv_indices = (
|
||||
dst_kv_indices[:origin_input_len].cpu().numpy().astype(np.int32)
|
||||
dst_kv_indices[: origin_input_len - prefix_len]
|
||||
.cpu()
|
||||
.numpy()
|
||||
.astype(np.int32)
|
||||
)
|
||||
page_size = 1 # host pool page_size
|
||||
else:
|
||||
kv_indices_full = self.req_to_token_pool.req_to_token[
|
||||
decode_req.req.req_pool_idx
|
||||
][:origin_input_len]
|
||||
kv_indices = kv_indices_full.cpu().numpy()
|
||||
# Only send delta indices (beyond prefix) to prefill.
|
||||
kv_indices = (
|
||||
self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][
|
||||
prefix_len:origin_input_len
|
||||
]
|
||||
.cpu()
|
||||
.numpy()
|
||||
)
|
||||
page_size = self.token_to_kv_pool_allocator.page_size
|
||||
|
||||
# Prepare extra pool indices for hybrid models
|
||||
@@ -790,7 +862,7 @@ class DecodePreallocQueue:
|
||||
window_size = self.scheduler.sliding_window_size
|
||||
|
||||
window_start = max(0, seq_len - window_size)
|
||||
window_start = (window_start // page_size) * page_size
|
||||
window_start = page_align_floor(window_start, page_size)
|
||||
window_kv_indices_full = self.req_to_token_pool.req_to_token[
|
||||
decode_req.req.req_pool_idx, window_start:seq_len
|
||||
]
|
||||
@@ -821,7 +893,10 @@ class DecodePreallocQueue:
|
||||
assert decode_req.metadata_buffer_index is not None
|
||||
page_indices = kv_to_page_indices(kv_indices, page_size)
|
||||
decode_req.kv_receiver.send_metadata(
|
||||
page_indices, decode_req.metadata_buffer_index, state_indices
|
||||
page_indices,
|
||||
decode_req.metadata_buffer_index,
|
||||
state_indices,
|
||||
decode_prefix_len=prefix_len,
|
||||
)
|
||||
if (
|
||||
self.transfer_queue.enable_staging
|
||||
@@ -848,7 +923,10 @@ class DecodePreallocQueue:
|
||||
)
|
||||
|
||||
def _allocatable_tokens(
|
||||
self, retractable_tokens: Optional[int] = None, count_retracted: bool = True
|
||||
self,
|
||||
retractable_tokens: Optional[int] = None,
|
||||
count_retracted: bool = True,
|
||||
extra_reserved_reqs: int = 0,
|
||||
) -> int:
|
||||
need_space_for_single_req = (
|
||||
max(
|
||||
@@ -871,6 +949,10 @@ class DecodePreallocQueue:
|
||||
)
|
||||
else:
|
||||
available_size = self.token_to_kv_pool_allocator.available_size()
|
||||
# Include evictable decode-radix cache entries in the budget -- they
|
||||
# can be freed on demand before allocation.
|
||||
if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
|
||||
available_size += self.tree_cache.evictable_size()
|
||||
allocatable_tokens = available_size - max(
|
||||
# preserve some space for future decode
|
||||
self.num_reserved_decode_tokens
|
||||
@@ -878,6 +960,7 @@ class DecodePreallocQueue:
|
||||
len(self.scheduler.running_batch.reqs)
|
||||
+ len(self.transfer_queue.queue)
|
||||
+ len(self.scheduler.waiting_queue)
|
||||
+ extra_reserved_reqs
|
||||
),
|
||||
# make sure each request can finish if reach max_tokens with all other requests retracted
|
||||
need_space_for_single_req,
|
||||
@@ -904,20 +987,77 @@ class DecodePreallocQueue:
|
||||
)
|
||||
return allocatable_tokens
|
||||
|
||||
def _pre_alloc(self, req: Req) -> torch.Tensor:
|
||||
def _required_alloc_tokens(self, *, fill_len: int, prefix_len: int) -> int:
|
||||
page_size = self.token_to_kv_pool_allocator.page_size
|
||||
if page_size == 1:
|
||||
return fill_len - prefix_len
|
||||
|
||||
num_new_pages = get_num_new_pages(
|
||||
seq_lens=torch.tensor([fill_len], dtype=torch.int64),
|
||||
prefix_lens=torch.tensor([prefix_len], dtype=torch.int64),
|
||||
page_size=page_size,
|
||||
)
|
||||
return num_new_pages * page_size
|
||||
|
||||
def _pre_alloc(
|
||||
self,
|
||||
req: Req,
|
||||
prefix_indices: Optional[torch.Tensor] = None,
|
||||
prefix_len: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Pre-allocate the memory for req_to_token and token_kv_pool"""
|
||||
if prefix_len is None:
|
||||
prefix_len = 0
|
||||
|
||||
req_pool_indices = self.req_to_token_pool.alloc([req])
|
||||
|
||||
assert (
|
||||
req_pool_indices is not None
|
||||
), "req_pool_indices is full! There is a bug in memory estimation."
|
||||
|
||||
# Alloc all tokens for the prebuilt req (except for the reserved input token for decoding)
|
||||
fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
|
||||
req.kv_allocated_len = fill_len
|
||||
req.kv_committed_len = fill_len
|
||||
|
||||
if prefix_len > 0:
|
||||
self.req_to_token_pool.write(
|
||||
(req.req_pool_idx, slice(0, prefix_len)), prefix_indices
|
||||
)
|
||||
|
||||
# TODO(retraction): when retraction is implemented with radix cache
|
||||
# awareness, a retracted request should re-match the tree here
|
||||
# instead of re-allocating from scratch. See resume_retracted_reqs.
|
||||
delta_len = fill_len - prefix_len
|
||||
required_alloc_tokens = self._required_alloc_tokens(
|
||||
fill_len=fill_len, prefix_len=prefix_len
|
||||
)
|
||||
|
||||
# Evict cached entries if the pool doesn't have enough free pages.
|
||||
if (
|
||||
self.scheduler.server_args.disaggregation_decode_enable_radix_cache
|
||||
and self.token_to_kv_pool_allocator.available_size() < required_alloc_tokens
|
||||
):
|
||||
num_to_evict = (
|
||||
required_alloc_tokens - self.token_to_kv_pool_allocator.available_size()
|
||||
)
|
||||
result = self.tree_cache.evict(EvictParams(num_tokens=num_to_evict))
|
||||
if self.token_to_kv_pool_allocator.available_size() < required_alloc_tokens:
|
||||
logger.warning(
|
||||
f"Eviction insufficient: needed {required_alloc_tokens} tokens, "
|
||||
f"available {self.token_to_kv_pool_allocator.available_size()} "
|
||||
f"after evicting {result.num_tokens_evicted}/{num_to_evict} tokens. "
|
||||
f"evictable_size={self.tree_cache.evictable_size()}, "
|
||||
f"protected_size={self.tree_cache.protected_size()}, "
|
||||
f"fill_len={fill_len}, prefix_len={prefix_len}, delta_len={delta_len}, "
|
||||
f"page_size={self.token_to_kv_pool_allocator.page_size}, "
|
||||
f"req={req.rid}"
|
||||
)
|
||||
|
||||
if self.scheduler.enable_hisparse:
|
||||
# HiSparse is incompatible with decode-side L1 radix cache. Keep
|
||||
# this path on the upstream full-allocation semantics.
|
||||
assert prefix_len == 0
|
||||
|
||||
# Direct-to-host path: only allocate logical indices (no hisparse
|
||||
# device indices) and allocate host indices for RDMA destination.
|
||||
coordinator = self.scheduler.hisparse_coordinator
|
||||
@@ -930,7 +1070,7 @@ class DecodePreallocQueue:
|
||||
last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
|
||||
extend_num_tokens=fill_len,
|
||||
)
|
||||
# Allocate host indices for the RDMA transfer target
|
||||
# Allocate host indices for the RDMA transfer target.
|
||||
host_indices = coordinator.mem_pool_host.alloc(fill_len)
|
||||
if host_indices is None:
|
||||
raise RuntimeError(
|
||||
@@ -940,27 +1080,52 @@ class DecodePreallocQueue:
|
||||
host_indices = host_indices.to(device=coordinator.device)
|
||||
coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices
|
||||
elif self.token_to_kv_pool_allocator.page_size == 1:
|
||||
kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
|
||||
kv_loc = self.token_to_kv_pool_allocator.alloc(delta_len)
|
||||
else:
|
||||
device = self.token_to_kv_pool_allocator.device
|
||||
last_loc = (
|
||||
prefix_indices[-1:].to(dtype=torch.int64, device=device)
|
||||
if prefix_len > 0
|
||||
else torch.tensor([-1], dtype=torch.int64, device=device)
|
||||
)
|
||||
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
|
||||
prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
|
||||
prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
|
||||
prefix_lens=torch.tensor(
|
||||
[prefix_len], dtype=torch.int64, device=device
|
||||
),
|
||||
prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
|
||||
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
|
||||
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
|
||||
last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
|
||||
extend_num_tokens=fill_len,
|
||||
last_loc=last_loc,
|
||||
extend_num_tokens=delta_len,
|
||||
)
|
||||
|
||||
assert (
|
||||
kv_loc is not None
|
||||
), "KV cache is full! There is a bug in memory estimation."
|
||||
assert kv_loc is not None, (
|
||||
f"KV cache is full! Bug in memory estimation. "
|
||||
f"available={self.token_to_kv_pool_allocator.available_size()}, "
|
||||
f"evictable={self.tree_cache.evictable_size()}, "
|
||||
f"protected={self.tree_cache.protected_size()}, "
|
||||
f"required_alloc={required_alloc_tokens}, delta={delta_len}, "
|
||||
f"fill={fill_len}, prefix={prefix_len}, "
|
||||
f"page_size={self.token_to_kv_pool_allocator.page_size}, "
|
||||
f"req={req.rid}"
|
||||
)
|
||||
|
||||
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
|
||||
self.req_to_token_pool.write(
|
||||
(req.req_pool_idx, slice(prefix_len, prefix_len + len(kv_loc))), kv_loc
|
||||
)
|
||||
|
||||
# populate metadata
|
||||
req.fill_ids = req.origin_input_ids + req.output_ids
|
||||
req.set_extend_input_len(len(req.fill_ids))
|
||||
# Truncate fill_ids to kv_committed_len so cache_unfinished_req only
|
||||
# inserts committed KV into the radix tree. The last output token
|
||||
# hasn't had KV committed yet (fill_ids is 1 ahead).
|
||||
req.fill_ids = (req.origin_input_ids + req.output_ids)[: req.kv_committed_len]
|
||||
# Set prefix_indices so downstream consumers (init_next_round_input,
|
||||
# prepare_for_extend) see the correct prefix length. In the agg path
|
||||
# this is done inside init_next_round_input, but decode-disagg needs
|
||||
# allocation info before batch assembly so we set it here.
|
||||
req.prefix_indices = (
|
||||
prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64)
|
||||
)
|
||||
req.set_extend_input_len(len(req.fill_ids) - prefix_len)
|
||||
|
||||
# Return the transfer destination indices:
|
||||
if self.scheduler.enable_hisparse:
|
||||
@@ -1328,7 +1493,27 @@ class SchedulerDisaggregationDecodeMixin:
|
||||
# we can only add at least `num_not_used_batch` new batch to the running queue
|
||||
if i < num_not_used_batch:
|
||||
can_run_list.append(req)
|
||||
req.init_next_round_input(self.tree_cache)
|
||||
# Decode-radix path: do NOT re-match prefix here.
|
||||
# `pop_preallocated` already took a tree snapshot and used it
|
||||
# to (1) pre-allocate KV, (2) choose delta pages for transfer,
|
||||
# and (3) set cache_protected_len/last_node for correct frees.
|
||||
# Re-matching now can observe a newer tree (other reqs may have
|
||||
# inserted the same prefix) and overwrite cache_protected_len,
|
||||
# making `cache_unfinished_req` free the wrong range (leak).
|
||||
# Non-radix decode keeps the original behavior.
|
||||
tree_cache = (
|
||||
None
|
||||
if self.server_args.disaggregation_decode_enable_radix_cache
|
||||
else self.tree_cache
|
||||
)
|
||||
req.init_next_round_input(tree_cache)
|
||||
# Truncate fill_ids to kv_committed_len so cache_unfinished_req
|
||||
# only sees committed KV (fill_ids includes one uncommitted token).
|
||||
if req.kv_committed_len is not None:
|
||||
req.fill_ids = req.fill_ids[: req.kv_committed_len]
|
||||
req.set_extend_input_len(
|
||||
len(req.fill_ids) - len(req.prefix_indices)
|
||||
)
|
||||
else:
|
||||
waiting_queue.append(req)
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.common import maybe_cache_unfinished_req
|
||||
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
|
||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
|
||||
@@ -109,7 +110,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
||||
self.output_ids = []
|
||||
for req in self.reqs:
|
||||
self.output_ids.append(req.output_ids[-1])
|
||||
self.tree_cache.cache_unfinished_req(req)
|
||||
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||
if req.grammar is not None:
|
||||
# FIXME: this try-except block is for handling unexpected xgrammar issue.
|
||||
try:
|
||||
|
||||
@@ -27,6 +27,7 @@ class FakeKVManager(BaseKVManager):
|
||||
is_mla_backend: Optional[bool] = False,
|
||||
):
|
||||
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
||||
self.req_to_decode_prefix_len = {}
|
||||
|
||||
def register_to_bootstrap(self):
|
||||
pass
|
||||
@@ -41,6 +42,7 @@ class FakeKVSender(BaseKVSender):
|
||||
dest_tp_ranks: List[int],
|
||||
pp_rank: int,
|
||||
):
|
||||
self.kv_mgr = mgr
|
||||
self.has_sent = False
|
||||
|
||||
def poll(self) -> KVPoll:
|
||||
@@ -106,6 +108,7 @@ class FakeKVReceiver(BaseKVReceiver):
|
||||
kv_indices: list[int],
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List[int]] = None,
|
||||
decode_prefix_len: Optional[int] = None,
|
||||
):
|
||||
self.has_sent_metadata = True
|
||||
logger.debug(
|
||||
|
||||
@@ -1827,6 +1827,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List[int]] = None,
|
||||
decode_prefix_len: Optional[int] = None,
|
||||
):
|
||||
if self.bootstrap_infos is None:
|
||||
self.kv_mgr.record_failure(
|
||||
|
||||
@@ -1033,6 +1033,7 @@ class MoriKVReceiver(CommonKVReceiver):
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List[int]] = None,
|
||||
decode_prefix_len: Optional[int] = None,
|
||||
):
|
||||
if self.bootstrap_infos is None or self.bootstrap_room is None:
|
||||
return
|
||||
|
||||
@@ -44,8 +44,15 @@ class TransferInfo:
|
||||
dst_aux_index: int
|
||||
required_dst_info_num: int
|
||||
dst_state_indices: List[int]
|
||||
decode_prefix_len: Optional[int] = None # for decode radix cache
|
||||
|
||||
def is_dummy(self):
|
||||
# A transfer is "dummy" only for CP non-authoritative ranks.
|
||||
# When dst_kv_indices is empty due to a decode-side radix cache
|
||||
# full hit (decode_prefix_len > 0), the transfer is NOT dummy --
|
||||
# aux/state data still needs to be sent.
|
||||
if self.dst_kv_indices.size == 0 and self.decode_prefix_len:
|
||||
return False
|
||||
return self.dst_kv_indices.size == 0
|
||||
|
||||
@classmethod
|
||||
@@ -65,6 +72,9 @@ class TransferInfo:
|
||||
dst_aux_index=int(msg[5].decode("ascii")),
|
||||
required_dst_info_num=int(msg[6].decode("ascii")),
|
||||
dst_state_indices=dst_state_indices,
|
||||
decode_prefix_len=(
|
||||
int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
|
||||
), # hacky just add it into the message that will be sent
|
||||
)
|
||||
|
||||
|
||||
@@ -883,39 +893,44 @@ class NixlKVManager(CommonKVManager):
|
||||
assert len(chunked_dst_kv_indice) == len(kv_indices)
|
||||
assert req.agent_name in self.decode_kv_args_table
|
||||
|
||||
notif = (
|
||||
f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.engine_rank}"
|
||||
)
|
||||
decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size
|
||||
|
||||
if self.is_mla_backend or (decode_tp_size == self.attn_tp_size):
|
||||
kv_xfer_handle = self.send_kvcache(
|
||||
req.agent_name,
|
||||
kv_indices,
|
||||
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
||||
chunked_dst_kv_indice,
|
||||
self.decode_kv_args_table[req.agent_name].gpu_id,
|
||||
notif,
|
||||
)
|
||||
else:
|
||||
kv_xfer_handle = self.send_kvcache_slice(
|
||||
req.agent_name,
|
||||
kv_indices,
|
||||
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
||||
chunked_dst_kv_indice,
|
||||
self.decode_kv_args_table[req.agent_name].gpu_id,
|
||||
notif,
|
||||
prefill_tp_size=self.attn_tp_size,
|
||||
decode_tp_size=decode_tp_size,
|
||||
decode_tp_rank=self.decode_kv_args_table[
|
||||
req.agent_name
|
||||
].decode_tp_rank,
|
||||
dst_kv_item_len=self.decode_kv_args_table[
|
||||
req.agent_name
|
||||
].dst_kv_item_len,
|
||||
# Skip KV RDMA transfer when there are no pages to send
|
||||
# (e.g., decode-side radix cache matched the entire prefix).
|
||||
# Aux data is still sent below when is_last=True.
|
||||
if len(kv_indices) > 0:
|
||||
notif = (
|
||||
f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.pp_rank}"
|
||||
)
|
||||
|
||||
handles.append(kv_xfer_handle)
|
||||
if self.is_mla_backend or (decode_tp_size == self.attn_tp_size):
|
||||
kv_xfer_handle = self.send_kvcache(
|
||||
req.agent_name,
|
||||
kv_indices,
|
||||
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
||||
chunked_dst_kv_indice,
|
||||
self.decode_kv_args_table[req.agent_name].gpu_id,
|
||||
notif,
|
||||
)
|
||||
else:
|
||||
kv_xfer_handle = self.send_kvcache_slice(
|
||||
req.agent_name,
|
||||
kv_indices,
|
||||
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
||||
chunked_dst_kv_indice,
|
||||
self.decode_kv_args_table[req.agent_name].gpu_id,
|
||||
notif,
|
||||
prefill_tp_size=self.attn_tp_size,
|
||||
decode_tp_size=decode_tp_size,
|
||||
decode_tp_rank=self.decode_kv_args_table[
|
||||
req.agent_name
|
||||
].decode_tp_rank,
|
||||
dst_kv_item_len=self.decode_kv_args_table[
|
||||
req.agent_name
|
||||
].dst_kv_item_len,
|
||||
)
|
||||
|
||||
handles.append(kv_xfer_handle)
|
||||
# Only the last chunk we need to send the aux data.
|
||||
if is_last:
|
||||
if state_indices is not None:
|
||||
@@ -936,16 +951,24 @@ class NixlKVManager(CommonKVManager):
|
||||
handles.append(state_xfer_handle)
|
||||
|
||||
assert aux_index is not None
|
||||
# When no KV pages were sent (decode-side cache hit),
|
||||
# encode pp_rank in aux notif so receiver can mark
|
||||
# expected_kvs_per_pp[pp_rank] = 0.
|
||||
if len(kv_indices) == 0:
|
||||
aux_notif = f"{req.room}_aux_nokv_{self.kv_args.pp_rank}"
|
||||
else:
|
||||
aux_notif = f"{req.room}_aux"
|
||||
aux_xfer_handle = self.send_aux(
|
||||
req.agent_name,
|
||||
aux_index,
|
||||
self.decode_kv_args_table[req.agent_name].dst_aux_ptrs,
|
||||
req.dst_aux_index,
|
||||
f"{req.room}_aux",
|
||||
aux_notif,
|
||||
)
|
||||
handles.append(aux_xfer_handle)
|
||||
if is_last:
|
||||
del self.transfer_infos[bootstrap_room]
|
||||
self.req_to_decode_prefix_len.pop(bootstrap_room, None)
|
||||
return handles
|
||||
|
||||
def update_transfer_status(self):
|
||||
@@ -978,6 +1001,15 @@ class NixlKVManager(CommonKVManager):
|
||||
)
|
||||
elif components[1] == "aux":
|
||||
self.transfer_statuses[room].received_aux = True
|
||||
# Handle "nokv" marker: no KV pages were sent for
|
||||
# this pp_rank (decode-side radix cache hit).
|
||||
if len(components) > 3 and components[2] == "nokv":
|
||||
pp_rank = int(components[3])
|
||||
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = 0
|
||||
if self.transfer_statuses[room].num_pp_ranks_expected is None:
|
||||
self.transfer_statuses[room].num_pp_ranks_expected = (
|
||||
self.required_prefill_response_num_table.get(room, 1)
|
||||
)
|
||||
elif components[1] == "state":
|
||||
pp_rank = int(components[2]) if len(components) > 2 else 0
|
||||
self.transfer_statuses[room].received_state_per_pp.add(pp_rank)
|
||||
@@ -1019,6 +1051,14 @@ class NixlKVManager(CommonKVManager):
|
||||
].required_dst_info_num
|
||||
logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}")
|
||||
if len(self.transfer_infos[room]) == required_dst_info_num:
|
||||
self.req_to_decode_prefix_len[room] = next(
|
||||
(
|
||||
info.decode_prefix_len
|
||||
for info in self.transfer_infos[room].values()
|
||||
if info.decode_prefix_len is not None
|
||||
),
|
||||
0,
|
||||
)
|
||||
logger.debug(f"{room=} is bootstrapped")
|
||||
self.update_status(room, KVPoll.WaitingForInput)
|
||||
|
||||
@@ -1039,6 +1079,12 @@ class NixlKVSender(CommonKVSender):
|
||||
self.has_sent = False
|
||||
self.chunk_id = 0
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0 or last_chunk
|
||||
|
||||
def send(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
@@ -1113,6 +1159,7 @@ class NixlKVReceiver(CommonKVReceiver):
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List[int]] = None,
|
||||
decode_prefix_len: Optional[int] = None,
|
||||
):
|
||||
if self.bootstrap_infos is None:
|
||||
logger.error(
|
||||
@@ -1146,6 +1193,7 @@ class NixlKVReceiver(CommonKVReceiver):
|
||||
if not is_dummy and state_indices is not None
|
||||
else b""
|
||||
),
|
||||
str(decode_prefix_len or 0).encode("ascii"),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -37,8 +37,6 @@ from sglang.srt.disaggregation.utils import (
|
||||
TransferBackend,
|
||||
get_kv_class,
|
||||
is_mla_backend,
|
||||
kv_to_page_indices,
|
||||
kv_to_page_num,
|
||||
poll_and_all_reduce_attn_cp_tp_group,
|
||||
prepare_abort,
|
||||
)
|
||||
@@ -49,7 +47,12 @@ from sglang.srt.managers.schedule_batch import (
|
||||
Req,
|
||||
ScheduleBatch,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.mem_cache.common import (
|
||||
kv_to_page_indices,
|
||||
kv_to_page_num,
|
||||
maybe_cache_unfinished_req,
|
||||
release_kv_cache,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool
|
||||
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
|
||||
from sglang.srt.observability.req_time_stats import set_schedule_time_batch
|
||||
@@ -336,7 +339,7 @@ class PrefillBootstrapQueue:
|
||||
self.scheduler.tree_cache.release_aborted_request(req.rid)
|
||||
continue
|
||||
|
||||
# KV.WaitingForInput - init here
|
||||
# KV.WaitingForInput - decode is ready to receive. initialize the kv sender
|
||||
req.time_stats.set_bootstrap_done_time()
|
||||
num_kv_indices = len(req.origin_input_ids)
|
||||
if self.req_to_metadata_buffer_idx_allocator.available_size() == 0:
|
||||
@@ -347,7 +350,15 @@ class PrefillBootstrapQueue:
|
||||
)
|
||||
assert req.metadata_buffer_index is not None
|
||||
|
||||
num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size)
|
||||
# Cal number of pages to send
|
||||
# if decode has a cached prefix, we need to send the delta indices
|
||||
# otherwise, send the entire request
|
||||
decode_prefix_len = req.disagg_kv_sender.pop_decode_prefix_len()
|
||||
req.start_send_idx = decode_prefix_len
|
||||
num_kv_indices_to_send = num_kv_indices - decode_prefix_len
|
||||
num_pages = kv_to_page_num(
|
||||
num_kv_indices_to_send, self.token_to_kv_pool.page_size
|
||||
)
|
||||
req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
|
||||
|
||||
bootstrapped_reqs.append(req)
|
||||
@@ -525,7 +536,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
|
||||
# There is no output_ids for prefill
|
||||
req.output_ids.append(next_token_id)
|
||||
self.tree_cache.cache_unfinished_req(req) # update the tree and lock
|
||||
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||
self.disagg_prefill_inflight_queue.append(req)
|
||||
if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
|
||||
req.output_topk_p = batch.spec_info.topk_p[i]
|
||||
@@ -736,7 +747,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
chunked_req_to_exclude = set()
|
||||
if self.chunked_req:
|
||||
chunked_req_to_exclude.add(self.chunked_req)
|
||||
self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
|
||||
maybe_cache_unfinished_req(self.chunked_req, self.tree_cache, chunked=True)
|
||||
if self.enable_overlap:
|
||||
# Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
|
||||
self.chunked_req.tmp_end_idx = min(
|
||||
@@ -781,12 +792,20 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
# if not the last chunk and the last page is partial, delay the last partial page to the next send
|
||||
end_idx = end_idx - end_idx % page_size
|
||||
|
||||
if end_idx < start_idx:
|
||||
logger.debug(
|
||||
"send_kv_chunk skip: rid=%s start_send_idx=%s end_idx=%s",
|
||||
req.rid,
|
||||
start_idx,
|
||||
end_idx,
|
||||
)
|
||||
return
|
||||
|
||||
kv_indices = (
|
||||
self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
|
||||
.cpu()
|
||||
.numpy()
|
||||
)
|
||||
req.start_send_idx = end_idx
|
||||
state_indices = None
|
||||
if last_chunk:
|
||||
self.disagg_metadata_buffers.set_buf(req)
|
||||
@@ -833,9 +852,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
state_indices = kv_to_page_indices(state_indices, page_size)
|
||||
|
||||
page_indices = kv_to_page_indices(kv_indices, page_size)
|
||||
if len(page_indices) == 0:
|
||||
logger.info(
|
||||
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
|
||||
)
|
||||
if not req.disagg_kv_sender.should_send_kv_chunk(len(page_indices), last_chunk):
|
||||
return
|
||||
req.disagg_kv_sender.send(page_indices, state_indices)
|
||||
req.start_send_idx = end_idx
|
||||
|
||||
@@ -439,26 +439,6 @@ def get_kv_class(
|
||||
raise ValueError(f"Unsupported transfer backend: {transfer_backend}")
|
||||
|
||||
|
||||
#########################
|
||||
# KV Pages
|
||||
#########################
|
||||
|
||||
|
||||
def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
|
||||
# 1. The page is guaranteed to be full except the last page.
|
||||
# 2. page index = kv_index // page_size
|
||||
# The return vector is kv_indices[::page_size] // page_size
|
||||
if page_size == 1: # shortcut
|
||||
return kv_indices
|
||||
|
||||
return kv_indices[::page_size] // page_size
|
||||
|
||||
|
||||
def kv_to_page_num(num_kv_indices: int, page_size: int):
|
||||
# ceil(num_kv_indices / page_size)
|
||||
return (num_kv_indices + page_size - 1) // page_size
|
||||
|
||||
|
||||
def page_indices_to_cp_rank_page_indices(
|
||||
page_indices: np.ndarray,
|
||||
total_pages: int,
|
||||
|
||||
@@ -1990,8 +1990,9 @@ def _execute_server_warmup(server_args: ServerArgs):
|
||||
)
|
||||
if res.status_code == 200:
|
||||
logger.info(
|
||||
f"End of prefill disaggregation mode warmup with status {res.status_code}, resp: {res.json()}"
|
||||
f"Disaggregation warmup request completed with status {res.status_code}, resp: {res.json()}"
|
||||
)
|
||||
logger.info("End of disaggregation warmup")
|
||||
_global_state.tokenizer_manager.server_status = ServerStatus.Up
|
||||
else:
|
||||
logger.info(
|
||||
|
||||
@@ -54,7 +54,7 @@ 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.disaggregation.utils import FAKE_BOOTSTRAP_HOST, 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
|
||||
@@ -868,6 +868,7 @@ class Req(ReqDllmMixin):
|
||||
self.bootstrap_host: str = bootstrap_host
|
||||
self.bootstrap_port: Optional[int] = bootstrap_port
|
||||
self.bootstrap_room: Optional[int] = bootstrap_room
|
||||
self.skip_radix_cache_insert = bootstrap_host == FAKE_BOOTSTRAP_HOST
|
||||
self.disagg_kv_sender: Optional[BaseKVSender] = None
|
||||
|
||||
self.routed_dp_rank: Optional[int] = routed_dp_rank
|
||||
@@ -1229,6 +1230,7 @@ class Req(ReqDllmMixin):
|
||||
self.prefix_indices = torch.empty((0,), dtype=torch.int64)
|
||||
self.routed_experts = None
|
||||
self.last_node = None
|
||||
self.cache_protected_len = 0
|
||||
self.swa_uuid_for_lock = None
|
||||
self.extend_input_len = 0
|
||||
self.is_retracted = True
|
||||
|
||||
@@ -77,6 +77,42 @@ IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD = int(
|
||||
IGNORE_EOS_RESERVE_TOKENS = 1
|
||||
|
||||
|
||||
def match_prefix_for_req(
|
||||
tree_cache: BasePrefixCache,
|
||||
req: Req,
|
||||
token_ids: Optional[List[int]] = None,
|
||||
*,
|
||||
cow_mamba: bool = False,
|
||||
include_req: bool = False,
|
||||
):
|
||||
if token_ids is None:
|
||||
token_ids = req.origin_input_ids + req.output_ids
|
||||
|
||||
match_result = tree_cache.match_prefix(
|
||||
MatchPrefixParams(
|
||||
key=RadixKey(token_ids=token_ids, extra_key=req.extra_key),
|
||||
cow_mamba=cow_mamba,
|
||||
req=req if include_req else None,
|
||||
)
|
||||
)
|
||||
(
|
||||
req.prefix_indices,
|
||||
req.last_node,
|
||||
req.last_host_node,
|
||||
req.host_hit_length,
|
||||
) = (
|
||||
match_result.device_indices,
|
||||
match_result.last_device_node,
|
||||
match_result.last_host_node,
|
||||
match_result.host_hit_length,
|
||||
)
|
||||
if match_result.mamba_branching_seqlen is not None:
|
||||
req.mamba_branching_seqlen = match_result.mamba_branching_seqlen
|
||||
if match_result.cache_protected_len is not None:
|
||||
req.cache_protected_len = match_result.cache_protected_len
|
||||
return match_result
|
||||
|
||||
|
||||
class CacheAwarePolicy(Enum):
|
||||
"""Scheduling policies that are aware of the tree cache."""
|
||||
|
||||
@@ -195,23 +231,7 @@ class SchedulePolicy:
|
||||
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,
|
||||
)
|
||||
match_result = match_prefix_for_req(self.tree_cache, r, prefix_ids)
|
||||
|
||||
# NOTE(sang): This logic is for in-batch prefix caching;
|
||||
# If there are more than 1 request that have small matching prefix from
|
||||
|
||||
@@ -185,7 +185,7 @@ from sglang.srt.managers.scheduler_update_weights_mixin import (
|
||||
)
|
||||
from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
|
||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||
from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache
|
||||
from sglang.srt.mem_cache.radix_cache import RadixCache
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
|
||||
from sglang.srt.model_loader.utils import get_resolved_model_impl
|
||||
@@ -800,6 +800,24 @@ class Scheduler(
|
||||
"Transformers backend to avoid multimodal prefix-cache mismatches."
|
||||
)
|
||||
|
||||
# Decode radix cache is unsupported with hybrid SWA/SSM models —
|
||||
# these use specialized memory pools incompatible with the
|
||||
# prefix-match-and-lock allocation path.
|
||||
if (
|
||||
server_args.disaggregation_decode_enable_radix_cache
|
||||
and server_args.disaggregation_mode == "decode"
|
||||
):
|
||||
if self.is_hybrid_swa:
|
||||
raise ValueError(
|
||||
"--disaggregation-decode-enable-radix-cache is incompatible "
|
||||
"with sliding window attention (SWA) models"
|
||||
)
|
||||
if self.is_hybrid_ssm:
|
||||
raise ValueError(
|
||||
"--disaggregation-decode-enable-radix-cache is incompatible "
|
||||
"with Mamba/SSM models"
|
||||
)
|
||||
|
||||
effective_chunked_prefill_size = server_args.chunked_prefill_size
|
||||
if self.model_config.is_multimodal and uses_transformers_backend:
|
||||
effective_chunked_prefill_size = None
|
||||
@@ -2344,7 +2362,7 @@ class Scheduler(
|
||||
self.handle_embedding_request(tokenized_req)
|
||||
|
||||
def stash_chunked_request(self, req: Req):
|
||||
self.tree_cache.cache_unfinished_req(req, chunked=True)
|
||||
maybe_cache_unfinished_req(req, self.tree_cache, chunked=True)
|
||||
|
||||
def _build_hisparse_decode_batch(self, reqs):
|
||||
"""Build a ScheduleBatch for hisparse requests transitioning from staging to decode."""
|
||||
|
||||
@@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import (
|
||||
Req,
|
||||
ScheduleBatch,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache
|
||||
from sglang.srt.server_args import MIS_DELIMITER_TOKEN_ID, get_global_server_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -87,6 +87,9 @@ class SchedulerOutputProcessorMixin:
|
||||
|
||||
def process_batch_result_prebuilt(self: Scheduler, batch: ScheduleBatch):
|
||||
assert self.disaggregation_mode == DisaggregationMode.DECODE
|
||||
use_free_group = self.server_args.disaggregation_decode_enable_radix_cache
|
||||
if use_free_group:
|
||||
self.token_to_kv_pool_allocator.free_group_begin()
|
||||
for req in batch.reqs:
|
||||
req.time_stats.set_decode_prebuilt_finish_time()
|
||||
req.check_finished()
|
||||
@@ -98,6 +101,8 @@ class SchedulerOutputProcessorMixin:
|
||||
|
||||
# Note: Logprobs should be handled on the prefill engine.
|
||||
self.stream_output(batch.reqs, batch.return_logprob)
|
||||
if use_free_group:
|
||||
self.token_to_kv_pool_allocator.free_group_end()
|
||||
|
||||
def maybe_collect_routed_experts(self: Scheduler, req: Req):
|
||||
"""Collect routed experts for a finished request."""
|
||||
@@ -199,7 +204,7 @@ class SchedulerOutputProcessorMixin:
|
||||
release_kv_cache(req, self.tree_cache)
|
||||
req.time_stats.set_completion_time()
|
||||
elif not batch.decoding_reqs or req not in batch.decoding_reqs:
|
||||
self.tree_cache.cache_unfinished_req(req)
|
||||
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||
if self.enable_hisparse:
|
||||
self.hisparse_coordinator.admit_request_into_staging(req)
|
||||
|
||||
@@ -333,7 +338,7 @@ class SchedulerOutputProcessorMixin:
|
||||
release_kv_cache(req, self.tree_cache)
|
||||
req.time_stats.set_completion_time()
|
||||
else:
|
||||
self.tree_cache.cache_unfinished_req(req)
|
||||
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||
else:
|
||||
# being chunked reqs' prefill is not finished
|
||||
req.is_chunked -= 1
|
||||
|
||||
@@ -30,6 +30,13 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChunkCache(BasePrefixCache):
|
||||
"""
|
||||
ChunkCache is used when radix cache is disabled.
|
||||
|
||||
That includes standard chunked-prefill setups and the decode side of P/D
|
||||
disaggregation when decode radix cache is not enabled.
|
||||
"""
|
||||
|
||||
def __init__(self, params: CacheInitParams):
|
||||
self.req_to_token_pool = params.req_to_token_pool
|
||||
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
@@ -26,6 +27,29 @@ MAMBA_STATE_PER_REQ_NO_CACHE = 1
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
|
||||
# The page is guaranteed to be full except the last page.
|
||||
if page_size == 1:
|
||||
return kv_indices
|
||||
|
||||
return kv_indices[::page_size] // page_size
|
||||
|
||||
|
||||
def kv_to_page_num(num_kv_indices: int, page_size: int):
|
||||
return (num_kv_indices + page_size - 1) // page_size
|
||||
|
||||
|
||||
def page_align_floor(length: int, page_size: int) -> int:
|
||||
return (length // page_size) * page_size
|
||||
|
||||
|
||||
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
|
||||
if getattr(req, "skip_radix_cache_insert", False):
|
||||
return
|
||||
|
||||
tree_cache.cache_unfinished_req(req, **kwargs)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def write_req_to_token_pool_triton(
|
||||
req_to_token_ptr, # [max_batch, max_context_len]
|
||||
@@ -553,7 +577,10 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
|
||||
req.mamba_pool_idx = None
|
||||
return
|
||||
|
||||
tree_cache.cache_finished_req(req, is_insert=is_insert)
|
||||
tree_cache.cache_finished_req(
|
||||
req,
|
||||
is_insert=is_insert and not getattr(req, "skip_radix_cache_insert", False),
|
||||
)
|
||||
|
||||
# StreamingSession.cache_finished_req handles speculative tail trim
|
||||
# and bookkeeping flag sync internally, then sets req_pool_idx = None.
|
||||
|
||||
@@ -419,8 +419,8 @@ class RadixCache(BasePrefixCache):
|
||||
Returns:
|
||||
MatchResult: ``device_indices`` is a 1-D ``torch.int64`` tensor of
|
||||
the concatenated KV cache indices corresponding to the longest
|
||||
cached prefix (may be length 0). ``last_device_node`` and
|
||||
``last_host_node`` (currently the same) are the tree node objects
|
||||
cached prefix (may be length 0).
|
||||
``last_device_node`` and ``last_host_node`` (currently the same) are the tree node objects
|
||||
representing the terminal node of the matched prefix. This method
|
||||
may mutate internal structure by splitting an existing node if the
|
||||
match ends inside a stored segment.
|
||||
@@ -516,10 +516,9 @@ class RadixCache(BasePrefixCache):
|
||||
result = self.insert(
|
||||
InsertParams(key=radix_key, value=values, priority=priority)
|
||||
)
|
||||
new_prefix_len = result.prefix_len
|
||||
# Free the duplicates that were already in the tree
|
||||
self.token_to_kv_pool_allocator.free(
|
||||
kv_indices[req.cache_protected_len : new_prefix_len]
|
||||
kv_indices[req.cache_protected_len : result.prefix_len]
|
||||
)
|
||||
else:
|
||||
self.token_to_kv_pool_allocator.free(
|
||||
@@ -530,7 +529,8 @@ class RadixCache(BasePrefixCache):
|
||||
self.token_to_kv_pool_allocator.free(kv_indices[key_len:])
|
||||
|
||||
# Remove req slot release the cache lock
|
||||
self.dec_lock_ref(req.last_node)
|
||||
if req.last_node is not None:
|
||||
self.dec_lock_ref(req.last_node)
|
||||
|
||||
def cache_unfinished_req(self, req: Req, chunked=False):
|
||||
"""Cache request when it is unfinished."""
|
||||
|
||||
@@ -850,7 +850,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase):
|
||||
|
||||
Returns a dict with latency_ms, total_mb, speed_gb_s if computable, else None.
|
||||
"""
|
||||
from sglang.srt.disaggregation.utils import kv_to_page_num
|
||||
from sglang.srt.mem_cache.common import kv_to_page_num
|
||||
|
||||
result = {}
|
||||
|
||||
|
||||
@@ -721,6 +721,7 @@ class ServerArgs:
|
||||
disaggregation_transfer_backend: str = "mooncake"
|
||||
disaggregation_bootstrap_port: int = 8998
|
||||
disaggregation_ib_device: Optional[str] = None
|
||||
disaggregation_decode_enable_radix_cache: bool = False
|
||||
disaggregation_decode_enable_offload_kvcache: bool = False
|
||||
num_reserved_decode_tokens: int = 512 # used for decode kv cache offload in PD
|
||||
# FIXME: hack to reduce ITL when decode bs is small
|
||||
@@ -786,6 +787,9 @@ class ServerArgs:
|
||||
# Validate SSL arguments early (before dummy-model short-circuit).
|
||||
self._handle_ssl_validation()
|
||||
|
||||
# Validate PD disaggregation flags early (before dummy-model short-circuit).
|
||||
self._handle_pd_disaggregation()
|
||||
|
||||
if self.model_path.lower() in ["none", "dummy"]:
|
||||
# Skip for dummy models
|
||||
return
|
||||
@@ -873,9 +877,6 @@ class ServerArgs:
|
||||
# Handle model loading format.
|
||||
self._handle_load_format()
|
||||
|
||||
# Handle PD disaggregation.
|
||||
self._handle_pd_disaggregation()
|
||||
|
||||
# Handle Encoder disaggregation.
|
||||
self._handle_encoder_disaggregation()
|
||||
|
||||
@@ -3729,14 +3730,39 @@ class ServerArgs:
|
||||
|
||||
def _handle_pd_disaggregation(self):
|
||||
if self.disaggregation_mode == "decode":
|
||||
self.disable_radix_cache = True
|
||||
logger.warning("KV cache is forced as chunk cache for decode server")
|
||||
if self.enable_mamba_extra_buffer():
|
||||
logger.warning(
|
||||
"Mamba extra_buffer is disabled because decode disaggregation "
|
||||
"currently forces chunk cache. Falling back to no_buffer."
|
||||
)
|
||||
self.mamba_scheduler_strategy = "no_buffer"
|
||||
if self.disaggregation_decode_enable_radix_cache:
|
||||
if self.enable_hisparse:
|
||||
raise ValueError(
|
||||
"--disaggregation-decode-enable-radix-cache is incompatible "
|
||||
"with --enable-hisparse"
|
||||
)
|
||||
if self.disaggregation_transfer_backend != "nixl":
|
||||
raise ValueError(
|
||||
"--disaggregation-decode-enable-radix-cache currently "
|
||||
"requires --disaggregation-transfer-backend nixl"
|
||||
)
|
||||
if self.speculative_algorithm is not None:
|
||||
raise ValueError(
|
||||
"--disaggregation-decode-enable-radix-cache is incompatible "
|
||||
"with speculative decoding "
|
||||
f"(--speculative-algorithm {self.speculative_algorithm})"
|
||||
)
|
||||
if self.enable_dp_attention:
|
||||
logger.warning(
|
||||
"EXPERIMENTAL: Decode radix cache with DP attention. "
|
||||
"Requires prefix-aware DP rank routing for optimal cache hits."
|
||||
)
|
||||
self.disable_radix_cache = False
|
||||
logger.warning("EXPERIMENTAL: Radix cache is enabled for decode server")
|
||||
else:
|
||||
self.disable_radix_cache = True
|
||||
logger.warning("KV cache is forced as chunk cache for decode server")
|
||||
if self.enable_mamba_extra_buffer():
|
||||
logger.warning(
|
||||
"Mamba extra_buffer is disabled because decode disaggregation "
|
||||
"currently forces chunk cache. Falling back to no_buffer."
|
||||
)
|
||||
self.mamba_scheduler_strategy = "no_buffer"
|
||||
|
||||
elif self.disaggregation_mode == "prefill":
|
||||
assert (
|
||||
@@ -6410,6 +6436,11 @@ class ServerArgs:
|
||||
"or multiple comma-separated devices (e.g., --disaggregation-ib-device mlx5_0,mlx5_1). "
|
||||
"Default is None, which triggers automatic device detection when mooncake backend is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--disaggregation-decode-enable-radix-cache",
|
||||
action="store_true",
|
||||
help="Enable radix cache on decode server (PD mode). Caches KV prefixes to avoid redundant transfers. Requires --disaggregation-transfer-backend nixl and is incompatible with --enable-hisparse.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--disaggregation-decode-enable-offload-kvcache",
|
||||
action="store_true",
|
||||
|
||||
Reference in New Issue
Block a user