diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index f7d4092d8..2b9ddde75 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -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. diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 79dfc79c4..3c916229a 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -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], diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 0c46f8104..e54e711ce 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 342296585..a74f01420 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -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: diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 60bf5465b..073faecfb 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -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( diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 0ea63af73..1272f3a3b 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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( diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 70154f9e9..6226c19df 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -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 diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 005d5b05c..41ab0fa97 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -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"), ] ) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index f946a5c28..140ed72a9 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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 diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 4b2e8c209..43c323058 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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, diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index e2f005337..bcf8a28a4 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4c472c31a..ee6c1c77b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index aa72acfe5..fea14cc98 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 86b4b51bb..31d3aa022 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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.""" diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 504d562d1..044deafdd 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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 diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 5d58bcde5..6641bcfa0 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index c3c01e8ba..56fccddc9 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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. diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 06fec7547..a707a4b24 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -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.""" diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index f8b3f3233..0710fce85 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -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 = {} diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 35be1c510..8e30444ea 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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", diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 347c87b10..de882d944 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -395,6 +395,9 @@ install_extra_deps() { fi $PIP_CMD install ${MOONCAKE_PKG} ${EXTRA_NVIDIA_SPECS} py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX + # Best-effort NIXL install for decode-radix disaggregation coverage. + $PIP_CMD install nixl $PIP_INSTALL_SUFFIX || echo "Warning: nixl install failed; continuing without nixl" + if [ "$IS_BLACKWELL" != "1" ]; then git clone --branch v0.5 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git $PIP_CMD install -e lmms-eval/ $PIP_INSTALL_SUFFIX diff --git a/test/registered/distributed/test_disaggregation_decode_radix_cache.py b/test/registered/distributed/test_disaggregation_decode_radix_cache.py new file mode 100644 index 000000000..c92c48fcd --- /dev/null +++ b/test/registered/distributed/test_disaggregation_decode_radix_cache.py @@ -0,0 +1,83 @@ +import time +import unittest + +import requests + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.cache_hit_kit import run_multiturn_cache_hit_test +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + is_in_ci, + try_cached_model, +) + +register_cuda_ci(est_time=120, suite="stage-c-test-8-gpu-h20") + + +def _has_nixl(): + try: + import nixl._api # noqa: F401 + except ImportError: + return False + return True + + +@unittest.skipUnless( + is_in_ci() or _has_nixl(), + "NIXL is required for decode radix cache disaggregation coverage.", +) +class TestDisaggregationDecodeRadixCache(PDDisaggregationServerBase): + extra_decode_args = ["--disaggregation-decode-enable-radix-cache"] + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST) + cls.transfer_backend = ["--disaggregation-transfer-backend", "nixl"] + cls.launch_all() + + def _assert_process_healthy(self, name, process, url): + self.assertIsNotNone(process, f"{name} process was not started") + self.assertIsNone( + process.poll(), + f"{name} exited unexpectedly with code {process.returncode}", + ) + response = requests.get(f"{url}/health", timeout=10) + response.raise_for_status() + + def test_decode_radix_cache_hits_and_workers_stay_alive(self): + decode_info = requests.get(f"{self.decode_url}/server_info", timeout=10).json() + self.assertFalse( + decode_info.get("disable_radix_cache", True), + "decode server did not enable radix cache", + ) + + result = run_multiturn_cache_hit_test( + base_url=self.base_url, + model_path=self.model, + num_clients=4, + num_rounds=3, + request_length=384, + output_length=64, + max_parallel=4, + ) + self.assertGreater( + result["overall"]["total_cached_tokens"], + 0, + "expected decode radix cache to reuse at least some tokens", + ) + + # Give the schedulers a short idle window so any post-request leak/crash + # paths have a chance to surface before the liveness checks below. + time.sleep(5) + + self._assert_process_healthy("load balancer", self.process_lb, self.lb_url) + self._assert_process_healthy("prefill", self.process_prefill, self.prefill_url) + self._assert_process_healthy("decode", self.process_decode, self.decode_url) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py new file mode 100644 index 000000000..9415e7258 --- /dev/null +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -0,0 +1,386 @@ +""" +Unit tests for lock_ref correctness in decode disagg radix cache scenarios. + +Verifies that inc_lock_ref / dec_lock_ref are balanced across the four +transfer scenarios identified in PR #19746: + +1. Incremental transfer & success (prefix match > 0) + inc_lock_ref(pop_preallocated) -> dec+inc(cache_unfinished_req) -> dec(cache_finished_req) + +2. Full transfer & success (prefix match == 0, full KV transferred) + inc_lock_ref(get_new_prebuilt_batch) -> dec+inc(cache_unfinished_req) -> dec(cache_finished_req) + +3. Incremental transfer & failure (prefix match > 0, transfer fails) + inc_lock_ref(pop_preallocated) -> dec(cache_finished_req via release_kv_cache is_insert=False) + +4. Full transfer & failure (prefix match == 0, transfer fails) + no inc_lock_ref -> dec(root_node) is no-op since root lock_ref starts at 1 + +Usage: + python -m pytest test/registered/unit/mem_cache/test_decode_radix_lock_ref.py -v +""" + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-small") + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.mem_cache.base_prefix_cache import ( + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey + + +def _make_cache_with_pools(page_size=1): + """Create a RadixCache with mock pools sufficient for cache_unfinished/finished_req.""" + mock_allocator = MagicMock() + mock_allocator.device = torch.device("cpu") + + # req_to_token pool: stores kv indices per request slot + max_seq_len = 64 + max_batch = 4 + req_to_token = torch.zeros(max_batch, max_seq_len, dtype=torch.int64) + + mock_pool = MagicMock() + mock_pool.req_to_token = req_to_token + mock_pool.write = lambda idx_tuple, values: req_to_token.__setitem__( + idx_tuple, values + ) + + cache = RadixCache.create_simulated( + mock_allocator=mock_allocator, page_size=page_size + ) + cache.req_to_token_pool = mock_pool + return cache, req_to_token + + +class MockReq: + """Minimal mock Req with fields needed by cache_unfinished/finished_req.""" + + def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): + self.fill_ids = list(fill_ids) + self.origin_input_ids = ( + list(fill_ids[:-1]) if len(fill_ids) > 1 else list(fill_ids) + ) + self.output_ids = [fill_ids[-1]] if len(fill_ids) > 1 else [] + self.req_pool_idx = req_pool_idx + self.cache_protected_len = cache_protected_len + self.last_node = last_node + self.extra_key = None + self.prefix_indices = torch.empty(0, dtype=torch.int64) + self.priority = 0 + self.kv_committed_len = len(fill_ids) + self.kv_allocated_len = len(fill_ids) + self.kv_committed_freed = False + + def pop_committed_kv_cache(self): + self.kv_committed_freed = True + return self.kv_committed_len + + def pop_overallocated_kv_cache(self): + return (self.kv_committed_len, self.kv_allocated_len) + + +def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): + return MockReq(fill_ids, req_pool_idx, cache_protected_len, last_node) + + +class TestDecodeLockRefScenarios(unittest.TestCase): + """Test lock_ref balance across decode transfer scenarios.""" + + def _populate_prefix(self, cache, prefix_ids, prefix_values): + """Insert a prefix into the tree so future requests can match it.""" + cache.insert( + InsertParams( + key=RadixKey(prefix_ids), + value=torch.tensor(prefix_values, dtype=torch.int64), + ) + ) + + def test_incremental_transfer_success(self): + """Scenario 1: prefix match > 0, transfer succeeds. + + Flow: inc_lock_ref(pop_preallocated) + -> dec_lock_ref + inc_lock_ref(cache_unfinished_req) + -> dec_lock_ref(cache_finished_req) + """ + cache, req_to_token = _make_cache_with_pools() + + # Pre-populate a prefix [1,2,3] in the tree + prefix = [1, 2, 3] + prefix_vals = [10, 20, 30] + self._populate_prefix(cache, prefix, prefix_vals) + + # Match prefix (simulates _match_prefix_and_lock in pop_preallocated) + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) + matched_node = result.last_device_node + prefix_len = len(result.device_indices) + self.assertEqual(prefix_len, 3) + + # Step 1: inc_lock_ref (pop_preallocated locks the matched node) + cache.inc_lock_ref(matched_node) + self.assertGreater(matched_node.lock_ref, 0) + + # Simulate _pre_alloc: write prefix + new tokens to req_to_token + full_ids = [1, 2, 3, 4, 5] # prefix + 2 new tokens + full_vals = [10, 20, 30, 40, 50] + req_to_token[0, : len(full_vals)] = torch.tensor(full_vals, dtype=torch.int64) + + req = _make_req( + fill_ids=full_ids, + req_pool_idx=0, + cache_protected_len=prefix_len, + last_node=matched_node, + ) + + # Step 2: cache_unfinished_req (dec old lock, inc new lock) + cache.cache_unfinished_req(req) + + # Step 3: cache_finished_req with is_insert=True (dec lock) + cache.cache_finished_req(req) + + # Verify: all non-root nodes should have lock_ref == 0 + # (root always has lock_ref == 1) + self.assertEqual(cache.root_node.lock_ref, 1) + self.assertEqual(cache.protected_size(), 0) + # The evictable size should equal total inserted tokens + self.assertEqual(cache.evictable_size(), len(full_ids)) + + def test_full_transfer_success(self): + """Scenario 2: no prefix match, full KV transferred, succeeds. + + Flow: inc_lock_ref(root, via init_next_round_input/get_new_prebuilt_batch) + -> dec_lock_ref + inc_lock_ref(cache_unfinished_req) + -> dec_lock_ref(cache_finished_req) + """ + cache, req_to_token = _make_cache_with_pools() + + # No prefix in tree -- match returns root + full_ids = [10, 20, 30] + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids))) + matched_node = result.last_device_node + self.assertEqual(len(result.device_indices), 0) # no match + # matched_node is root + + root_lock_before = cache.root_node.lock_ref + # Step 1: inc_lock_ref on root (simulates get_new_prebuilt_batch) + # Note: inc/dec_lock_ref skip the root node (while node != root_node), + # so this is a no-op. Root always keeps lock_ref=1. + cache.inc_lock_ref(matched_node) + self.assertEqual(cache.root_node.lock_ref, root_lock_before) # no-op on root + + # Write full KV to pool + full_vals = [100, 200, 300] + req_to_token[0, : len(full_vals)] = torch.tensor(full_vals, dtype=torch.int64) + + req = _make_req( + fill_ids=full_ids, + req_pool_idx=0, + cache_protected_len=0, + last_node=matched_node, + ) + + # Step 2: cache_unfinished_req (dec root=no-op, inc new leaf) + cache.cache_unfinished_req(req) + + # Step 3: cache_finished_req (dec leaf) + cache.cache_finished_req(req) + + # Root lock unchanged, all nodes unlocked + self.assertEqual(cache.root_node.lock_ref, root_lock_before) + self.assertEqual(cache.protected_size(), 0) + self.assertEqual(cache.evictable_size(), len(full_ids)) + + def test_incremental_transfer_failure(self): + """Scenario 3: prefix match > 0, transfer fails. + + Flow: inc_lock_ref(pop_preallocated) + -> dec_lock_ref(cache_finished_req via release_kv_cache is_insert=False) + """ + cache, req_to_token = _make_cache_with_pools() + + # Pre-populate prefix + prefix = [1, 2, 3] + prefix_vals = [10, 20, 30] + self._populate_prefix(cache, prefix, prefix_vals) + + # Match and lock + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) + matched_node = result.last_device_node + prefix_len = len(result.device_indices) + + cache.inc_lock_ref(matched_node) + # Prefix tokens should now be protected (locked) + self.assertGreater(cache.protected_size(), 0) + + # Simulate _pre_alloc with additional tokens + full_ids = [1, 2, 3, 4, 5] + full_vals = [10, 20, 30, 40, 50] + req_to_token[0, : len(full_vals)] = torch.tensor(full_vals, dtype=torch.int64) + + req = _make_req( + fill_ids=full_ids, + req_pool_idx=0, + cache_protected_len=prefix_len, + last_node=matched_node, + ) + + # Transfer fails -> cache_finished_req with is_insert=False + # This frees delta tokens and dec_lock_ref on last_node + cache.cache_finished_req(req, is_insert=False) + + # The prefix node should be unlocked (back to evictable) + self.assertEqual(cache.root_node.lock_ref, 1) + self.assertEqual(cache.protected_size(), 0) + # Prefix tokens should still be in tree and evictable + self.assertEqual(cache.evictable_size(), len(prefix)) + + def test_full_transfer_failure(self): + """Scenario 4: no prefix match, transfer fails. + + Flow: _match_prefix_and_lock sets last_node=root and calls + inc_lock_ref(root) which is a no-op. On failure, + cache_finished_req calls dec_lock_ref(root) which is also + a no-op. Net: balanced. + """ + cache, req_to_token = _make_cache_with_pools() + + root_lock_before = cache.root_node.lock_ref + + # No prefix in tree -- match returns root (simulates _match_prefix_and_lock) + full_ids = [10, 20, 30] + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids))) + matched_node = result.last_device_node + self.assertIs(matched_node, cache.root_node) + + # inc_lock_ref(root) is a no-op + cache.inc_lock_ref(matched_node) + self.assertEqual(cache.root_node.lock_ref, root_lock_before) + + full_vals = [100, 200, 300] + req_to_token[0, : len(full_vals)] = torch.tensor(full_vals, dtype=torch.int64) + + # last_node = root (as set by _match_prefix_and_lock) + req = _make_req( + fill_ids=full_ids, + req_pool_idx=0, + cache_protected_len=0, + last_node=matched_node, + ) + + # Transfer fails -> cache_finished_req with is_insert=False + # dec_lock_ref(root) is a no-op + cache.cache_finished_req(req, is_insert=False) + + # Root lock unchanged, nothing protected or evictable + self.assertEqual(cache.root_node.lock_ref, root_lock_before) + self.assertEqual(cache.protected_size(), 0) + self.assertEqual(cache.evictable_size(), 0) + + def test_pop_preallocated_rechecks_budget_after_lock(self): + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + + req = MagicMock() + req.rid = "req-1" + req.origin_input_ids = list(range(8)) + req.output_ids = [99] + req.last_node = object() + req.finished_reason = None + req.cache_protected_len = 0 + req.sampling_params.max_new_tokens = 16 + + decode_req = MagicMock() + decode_req.req = req + decode_req.waiting_for_input = True + + queue.queue = [decode_req] + queue.pending_reqs = [] + queue.retracted_queue = [] + queue.num_reserved_decode_tokens = 0 + queue._resolve_pending_reqs = MagicMock() + queue._update_handshake_waiters = MagicMock() + queue._match_prefix_and_lock = MagicMock( + return_value=(torch.arange(4, dtype=torch.int64), 4) + ) + queue._pre_alloc = MagicMock( + side_effect=AssertionError("_pre_alloc should not run") + ) + queue.transfer_queue = MagicMock(queue=[], enable_staging=False) + queue.tree_cache = MagicMock() + queue.tree_cache.dec_lock_ref = MagicMock() + queue.req_to_token_pool = MagicMock() + queue.req_to_token_pool.available_size.return_value = 1 + queue.req_to_metadata_buffer_idx_allocator = MagicMock() + queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1 + queue.token_to_kv_pool_allocator = MagicMock() + queue.token_to_kv_pool_allocator.page_size = 4 + + running_batch = MagicMock() + running_batch.reqs = [] + server_args = MagicMock() + server_args.disaggregation_decode_enable_radix_cache = True + scheduler = MagicMock() + scheduler.running_batch = running_batch + scheduler.server_args = server_args + scheduler.enable_hisparse = False + scheduler.waiting_queue = [] + scheduler.last_batch = None + scheduler.stream_output = MagicMock() + queue.scheduler = scheduler + + # Initial budget says the request fits; post-lock budget says it does not. + queue._allocatable_tokens = MagicMock(side_effect=[8, 3]) + + preallocated, failed = queue.pop_preallocated() + + self.assertEqual(preallocated, []) + self.assertEqual(failed, []) + queue._pre_alloc.assert_not_called() + queue.tree_cache.dec_lock_ref.assert_called_once_with(req.last_node) + self.assertEqual(queue._allocatable_tokens.call_count, 2) + + def test_repeated_incremental_no_leak(self): + """Multiple incremental transfers shouldn't leak lock_refs.""" + cache, req_to_token = _make_cache_with_pools() + + prefix = [1, 2, 3] + prefix_vals = [10, 20, 30] + self._populate_prefix(cache, prefix, prefix_vals) + + for iteration in range(5): + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) + matched_node = result.last_device_node + prefix_len = len(result.device_indices) + + cache.inc_lock_ref(matched_node) + + suffix_token = 40 + iteration + full_ids = prefix + [suffix_token] + full_vals = prefix_vals + [100 + iteration] + req_to_token[0, : len(full_vals)] = torch.tensor( + full_vals, dtype=torch.int64 + ) + + req = _make_req( + fill_ids=full_ids, + req_pool_idx=0, + cache_protected_len=prefix_len, + last_node=matched_node, + ) + + cache.cache_unfinished_req(req) + cache.cache_finished_req(req) + + # After all iterations, root lock should be 1, no protected nodes + self.assertEqual(cache.root_node.lock_ref, 1) + self.assertEqual(cache.protected_size(), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index a83a2c693..a62591158 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -47,6 +47,22 @@ class TestLoadBalanceMethod(unittest.TestCase): server_args = ServerArgs(model_path="dummy", disaggregation_mode="decode") self.assertEqual(server_args.load_balance_method, "round_robin") + def test_pd_decode_radix_cache_rejects_hisparse(self): + with self.assertRaises(ValueError) as context: + ServerArgs( + model_path="dummy", + disaggregation_mode="decode", + disaggregation_decode_enable_radix_cache=True, + disaggregation_transfer_backend="nixl", + enable_hisparse=True, + ) + + self.assertIn( + "--disaggregation-decode-enable-radix-cache is incompatible with " + "--enable-hisparse", + str(context.exception), + ) + class TestPortArgs(unittest.TestCase): @patch("sglang.srt.server_args.get_free_port")