[P/D disagg] - support decode side radix cache (#19746)

This commit is contained in:
ishandhanani
2026-05-01 21:55:34 +08:00
committed by GitHub
parent 4197c55968
commit 5b7ce417d0
24 changed files with 993 additions and 138 deletions
@@ -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],
+218 -33
View File
@@ -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
+78 -30
View File
@@ -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"),
]
)
+29 -12
View File
@@ -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
-20
View File
@@ -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,
+2 -1
View File
@@ -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(
+3 -1
View File
@@ -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
+37 -17
View File
@@ -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
+20 -2
View File
@@ -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
+28 -1
View File
@@ -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.
+5 -5
View File
@@ -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 = {}
+42 -11
View File
@@ -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",