chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[PD] Enable all CP ranks for KVCache transfer (#19765)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-03-03 19:35:21 +08:00
committed by GitHub
parent 365ca1edb5
commit facde4c6d3
6 changed files with 163 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -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:

31
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -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:

18
python/sglang/srt/disaggregation/mori/conn.py
View File

@@ -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,

19
python/sglang/srt/disaggregation/nixl/conn.py
View File

@@ -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,

83
python/sglang/srt/disaggregation/utils.py
View File

@@ -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
#########################

1
python/sglang/srt/environ.py
View File

@@ -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.