From facde4c6d3f3b820c00f053fe9a45c9a086deaf2 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Tue, 3 Mar 2026 19:35:21 +0800 Subject: [PATCH] [PD] Enable all CP ranks for KVCache transfer (#19765) Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 32 ++++--- .../srt/disaggregation/mooncake/conn.py | 31 +++++-- python/sglang/srt/disaggregation/mori/conn.py | 18 +++- python/sglang/srt/disaggregation/nixl/conn.py | 19 ++++- python/sglang/srt/disaggregation/utils.py | 83 ++++++++++++++++++- python/sglang/srt/environ.py | 1 + 6 files changed, 163 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 16cf8b601..bc21f8882 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -110,6 +110,9 @@ class CommonKVManager(BaseKVManager): self.pp_size = server_args.pp_size self.pp_rank = self.kv_args.pp_rank self.local_ip = get_local_ip_auto() + self.enable_all_cp_ranks_for_transfer = ( + envs.SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER.get() + ) # bind zmq socket context = zmq.Context() @@ -124,9 +127,12 @@ class CommonKVManager(BaseKVManager): self.failure_lock = threading.Lock() if self.disaggregation_mode == DisaggregationMode.PREFILL: - # TODO(shangming): Fix me when we support MHA/GQA + CP, or when we utilize all cp ranks for KV transfer in CP mode. + # When SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER is True, all CP ranks + # participate in KV transfer; Otherwise only CP rank 0 sends. self.is_dummy_cp_rank = ( - is_mla_backend and self.attn_cp_size > 1 and self.attn_cp_rank != 0 + not self.enable_all_cp_ranks_for_transfer + and self.attn_cp_size > 1 + and self.attn_cp_rank != 0 ) self.register_to_bootstrap() self.transfer_infos = {} @@ -395,6 +401,9 @@ class CommonKVSender(BaseKVSender): def init(self, num_kv_indices: int, aux_index: Optional[int] = None): self.num_kv_indices = num_kv_indices self.aux_index = aux_index + logger.debug( + f"CommonKVSender init with num_kv_indices: {num_kv_indices} and aux_index: {aux_index}" + ) def send( self, @@ -488,23 +497,22 @@ class CommonKVReceiver(BaseKVReceiver): self.prefill_info.attn_tp_size // self.kv_mgr.attn_tp_size ) - # Decode cp size should be equal to prefill cp size or 1 - assert ( - self.kv_mgr.attn_cp_size == self.prefill_info.attn_cp_size - or self.kv_mgr.attn_cp_size == 1 - ), ( - f"Decode cp size ({self.kv_mgr.attn_cp_size}) should be equal to prefill cp size ({self.prefill_info.attn_cp_size}) or 1", + # Decode cp size should be equal to 1 + assert self.kv_mgr.attn_cp_size == 1, ( + f"Decode cp size ({self.kv_mgr.attn_cp_size}) should be equal to 1", ) if self.kv_mgr.attn_cp_size == self.prefill_info.attn_cp_size: + # This means that the prefill cp size is 1 + assert self.prefill_info.attn_cp_size == 1, ( + f"When prefill cp size is 1, attn cp size should be 1, but got {self.kv_mgr.attn_cp_size}", + ) self.target_cp_ranks = [self.kv_mgr.attn_cp_rank] else: self.target_cp_ranks = [ rank for rank in range(self.prefill_info.attn_cp_size) ] - # TODO(shangming): Support KVCache transfer for multiple prefill cp ranks -> 1 decode cp rank - # For now, we handle the control plane in advance, but we need to support the data plane in the future. - if self.kv_mgr.is_mla_backend: - # For MLA: we only need to retrieve KVCache from the first CP rank now + if not self.kv_mgr.enable_all_cp_ranks_for_transfer: + # Only retrieve from prefill CP rank 0 when not using all ranks self.target_cp_ranks = self.target_cp_ranks[:1] self.required_prefill_response_num *= 1 else: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index a223a2578..f42ebfffc 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -28,7 +28,10 @@ from sglang.srt.disaggregation.common.utils import ( from sglang.srt.disaggregation.mooncake.utils import ( check_mooncake_custom_mem_pool_enabled, ) -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + filter_kv_indices_for_cp_rank, +) from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs @@ -779,7 +782,12 @@ class MooncakeKVManager(CommonKVManager): ) polls = [] dst_ranks_infos = [] - local_rank = self.attn_tp_rank * self.pp_size + self.pp_rank + # Unique id per prefill sender so decode's response set size matches expected_response_num. + prefill_unique_rank = ( + self.attn_tp_rank * (self.pp_size * self.attn_cp_size) + + self.pp_rank * self.attn_cp_size + + self.attn_cp_rank + ) for req in reqs_to_be_processed: if not req.is_dummy: # Early exit if the request has failed @@ -795,7 +803,7 @@ class MooncakeKVManager(CommonKVManager): req.dst_port, req.room, KVPoll.Failed, - local_rank, + prefill_unique_rank, ) break @@ -857,7 +865,7 @@ class MooncakeKVManager(CommonKVManager): req.dst_port, req.room, KVPoll.Failed, - local_rank, + prefill_unique_rank, ) break @@ -888,7 +896,11 @@ class MooncakeKVManager(CommonKVManager): self.update_status(req.room, status) for endpoint, dst_port, room in dst_ranks_infos: self.sync_status_to_decode_endpoint( - endpoint, dst_port, room, status, local_rank + endpoint, + dst_port, + room, + status, + prefill_unique_rank, ) else: # Dummy request means the decode instance is not used, so its status can be marked as success directly @@ -1136,7 +1148,14 @@ class MooncakeKVSender(CommonKVSender): self.curr_idx += len(kv_indices) is_last_chunk = self.curr_idx == self.num_kv_indices - if self.kv_mgr.is_dummy_cp_rank: + # Special handling for cp + if self.kv_mgr.enable_all_cp_ranks_for_transfer: + kv_indices, index_slice = filter_kv_indices_for_cp_rank( + self.kv_mgr, + kv_indices, + index_slice, + ) + elif self.kv_mgr.is_dummy_cp_rank: if not is_last_chunk: return else: diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 054068530..ff7d3a056 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -32,7 +32,10 @@ from sglang.srt.disaggregation.common.conn import ( CommonKVSender, ) from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + filter_kv_indices_for_cp_rank, +) from sglang.srt.server_args import ServerArgs from sglang.srt.utils.common import ( format_tcp_address, @@ -865,6 +868,19 @@ class MoriKVSender(CommonKVSender): self.curr_idx += len(kv_indices) is_last = self.curr_idx == self.num_kv_indices + # Special handling for cp + if self.kv_mgr.enable_all_cp_ranks_for_transfer: + kv_indices, index_slice = filter_kv_indices_for_cp_rank( + self.kv_mgr, + kv_indices, + index_slice, + ) + elif self.kv_mgr.is_dummy_cp_rank: + if not is_last: + return + else: + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success) + return statuses, infos = self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 0adc412de..764fd9e42 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -20,7 +20,10 @@ from sglang.srt.disaggregation.common.conn import ( CommonKVSender, ) from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + filter_kv_indices_for_cp_rank, +) from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs @@ -905,6 +908,20 @@ class NixlKVSender(CommonKVSender): self.curr_idx += len(kv_indices) is_last = self.curr_idx == self.num_kv_indices + # Special handling for cp + if self.kv_mgr.enable_all_cp_ranks_for_transfer: + kv_indices, index_slice = filter_kv_indices_for_cp_rank( + self.kv_mgr, + kv_indices, + index_slice, + ) + elif self.kv_mgr.is_dummy_cp_rank: + if not is_last: + return + else: + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success) + return + new_xfer_handles = self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index fe1c71294..b7b3b0238 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -5,7 +5,7 @@ import random from collections import deque from contextlib import nullcontext from enum import Enum -from typing import TYPE_CHECKING, Literal, Optional, Type, overload +from typing import TYPE_CHECKING, Literal, Optional, Tuple, Type, overload import numpy as np import torch @@ -424,6 +424,87 @@ def kv_to_page_num(num_kv_indices: int, page_size: int): return (num_kv_indices + page_size - 1) // page_size +def page_indices_to_cp_rank_page_indices( + page_indices: np.ndarray, + total_pages: int, + cp_rank: int, + cp_size: int, +) -> np.ndarray: + """ + Filter page_indices (which are *global* page ids in the KV pool) to those + belonging to the given CP rank for this request. + + For a single request, its pages occupy a contiguous global range + [first_page, first_page + total_pages). We first compute the local + split [0, total_pages) across cp_size ranks, then shift that local + range by first_page back into the global page id space and take + the intersection with page_indices. + + Returns: + Subset of page_indices that fall in this rank's global + [start_page, end_page) slice for the given CP rank. + """ + if cp_size <= 1: + return page_indices + + if page_indices.size == 0: + return np.asarray(page_indices) + + first_page = int(page_indices.min()) + base = total_pages // cp_size + rem = total_pages % cp_size + + if rem == 0: + local_start = cp_rank * base + local_end = local_start + base + else: + local_start = cp_rank * base + min(cp_rank, rem) + n_pages = base + (1 if cp_rank < rem else 0) + local_end = local_start + n_pages + + # Map back to global page ids. + start_page = first_page + local_start + end_page = first_page + local_end + + mask = (page_indices >= start_page) & (page_indices < end_page) + return np.asarray(page_indices)[mask] + + +def filter_kv_indices_for_cp_rank( + kv_mgr: CommonKVManager, kv_indices: np.ndarray, index_slice: slice +) -> Tuple[np.ndarray, slice]: + """Filters kv_indices and index_slice for the current CP rank.""" + total_pages = len(kv_indices) + cp_rank = kv_mgr.attn_cp_rank + cp_size = kv_mgr.attn_cp_size + + rank_page_indices = page_indices_to_cp_rank_page_indices( + page_indices=kv_indices, + total_pages=total_pages, + cp_rank=cp_rank, + cp_size=cp_size, + ) + + if rank_page_indices.size == 0: + new_kv_indices = kv_indices[:0] + new_index_slice = slice(index_slice.start, index_slice.start) + else: + mask = np.isin(kv_indices, rank_page_indices) + if not mask.any(): + new_kv_indices = kv_indices[:0] + new_index_slice = slice(index_slice.start, index_slice.start) + else: + first_pos = int(mask.argmax()) + last_pos = len(mask) - int(mask[::-1].argmax()) + + new_kv_indices = kv_indices[first_pos:last_pos] + new_index_slice = slice( + index_slice.start + first_pos, + index_slice.start + last_pos, + ) + return new_kv_indices, new_index_slice + + ######################### # Misc ######################### diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e0cb3f46f..aa5db62d5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -242,6 +242,7 @@ class Envs: SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") + SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period.