minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: support Kimi Linear PD disaggregation with DCP (#32837)

Browse Source 
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
Khoa Pham
2026-07-31 02:14:09 -07:00
committed by GitHub
co-authored by Yangmin Li
parent 33c27d8e7f
commit 2573190b93
13 changed files with 1084 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/pd_disaggregation_hook.py
+25
View File
@@ -26,6 +26,31 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None:
"with MC_FORCE_TCP=1 (TCP transport, no RDMA)"
)
if server_args.disaggregation_mode == "prefill" and server_args.dcp_size > 1:
logger.warning(
"DCP on a PD prefill server is supported when prefill and decode "
"use the same DCP layout, but it usually adds communication "
"overhead without improving prefill performance."
)
if server_args.disaggregation_mode == "decode" and server_args.dcp_size > 1:
if server_args.disaggregation_transfer_backend not in ("mooncake", "nixl"):
raise ValueError(
"PD decode DCP requires --disaggregation-transfer-backend "
"mooncake or nixl, got "
f"{server_args.disaggregation_transfer_backend!r}."
)
if server_args.disaggregation_decode_enable_radix_cache:
raise ValueError(
"PD decode DCP currently requires chunk cache; "
"--disaggregation-decode-enable-radix-cache is not supported."
)
if server_args.enable_hierarchical_cache:
raise ValueError(
"PD decode DCP currently requires chunk cache; "
"--enable-hierarchical-cache is not supported."
)
if server_args.disaggregation_mode == "decode":
if server_args.disaggregation_decode_enable_radix_cache:
if server_args.enable_hisparse:
python/sglang/srt/disaggregation/base/conn.py
+1
View File
@@ -135,6 +135,7 @@ class BaseKVSender(ABC):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
"""
Send the kv cache at the given kv indices and the extra cache/state at the given indices to the decoder server.
python/sglang/srt/disaggregation/common/conn.py
+57 -5
View File
@@ -163,10 +163,13 @@ class CommonKVManager(BaseKVManager):
self.bootstrap_host = server_args.host
self.bootstrap_port = server_args.disaggregation_bootstrap_port
self.dist_init_addr = server_args.dist_init_addr
self.attn_tp_size = get_parallel().attn_tp_size
self.attn_tp_rank = get_parallel().attn_tp_rank
self.attn_cp_size = get_parallel().attn_cp_size
self.attn_cp_rank = get_parallel().attn_cp_rank
parallel = get_parallel()
self.attn_tp_size = parallel.attn_tp_size
self.attn_tp_rank = parallel.attn_tp_rank
self.attn_cp_size = parallel.attn_cp_size
self.attn_cp_rank = parallel.attn_cp_rank
self.dcp_size = server_args.dcp_size
self.dcp_rank = parallel.dcp_rank if self.dcp_size > 1 else 0
self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank()
self.system_dp_size = (
@@ -268,6 +271,39 @@ class CommonKVManager(BaseKVManager):
f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
)
def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool:
if self.dcp_size == dst_dcp_size:
if self.dcp_rank != dst_dcp_rank:
raise RuntimeError(
"PD peers must connect matching DCP ranks, got "
f"prefill={self.dcp_rank}, decode={dst_dcp_rank}"
)
return False
if (
self.dcp_size == 1
and dst_dcp_size > 1
and (self.is_mla_backend or self.is_hybrid_mla_backend)
):
return True
raise RuntimeError(
f"Unsupported PD DCP topology: {self.dcp_size} -> {dst_dcp_size}"
)
def prepare_dcp_token_item_lens(self, dst_page_item_lens: List[int]) -> List[int]:
page_size = self.kv_args.page_size
src_token_lens = [
item_len // page_size for item_len in self.kv_args.kv_item_lens
]
dst_token_lens = [item_len // page_size for item_len in dst_page_item_lens]
if src_token_lens != dst_token_lens:
raise RuntimeError(
"PD DCP source/destination KV geometry differs: "
f"src={src_token_lens}, dst={dst_token_lens}"
)
return src_token_lens
def check_status(self, bootstrap_room: int) -> KVPoll:
return self.request_status[bootstrap_room]
@@ -469,7 +505,11 @@ class CommonKVManager(BaseKVManager):
info: PrefillServerInfo = None
try:
url = f"http://{bootstrap_addr}/route?prefill_dp_rank={-1}&prefill_cp_rank={-1}&target_tp_rank={-1}&target_pp_rank={-1}"
url = (
f"http://{bootstrap_addr}/route?"
f"prefill_dp_rank={-1}&prefill_cp_rank={-1}&"
f"target_tp_rank={-1}&target_pp_rank={-1}"
)
response = requests.get(url, timeout=5)
if response.status_code == 200:
data = response.json()
@@ -501,6 +541,17 @@ class CommonKVManager(BaseKVManager):
f"Both servers must use the same --kv-cache-dtype value."
)
if self.dcp_size > 1:
if not (self.is_mla_backend or self.is_hybrid_mla_backend):
raise RuntimeError(
"PD decode DCP requires an MLA or hybrid-MLA KV pool."
)
if info.attn_cp_size != 1:
raise RuntimeError(
"PD decode DCP currently requires prefill attention CP=1, "
f"got {info.attn_cp_size}."
)
self._resolve_rank_mapping(info)
self.prefill_info_table[bootstrap_addr] = info
logger.debug(f"Prefill parallel info for [{bootstrap_addr}]: {info}")
@@ -1137,6 +1188,7 @@ class CommonKVSender(BaseKVSender):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
pass
python/sglang/srt/disaggregation/common/utils.py
+69
View File
@@ -25,6 +25,7 @@ class TransferKVChunk:
prefill_aux_index: Optional[int]
state_indices: Optional[List]
chunk_id: Optional[int] = None
num_kv_tokens: Optional[int] = None
trace_ctx: Union[TraceReqContext, TraceNullContext] = dataclasses.field(
default_factory=TraceNullContext
)
@@ -127,3 +128,71 @@ def group_concurrent_contiguous(
dst_groups = [g.tolist() for g in dst_groups]
return src_groups, dst_groups
@dataclasses.dataclass(frozen=True)
class DCPTokenTransferPlan:
src_token_indices: npt.NDArray[np.int64]
dst_token_indices: npt.NDArray[np.int64]
def build_dcp_token_transfer_plan(
src_page_indices: npt.NDArray[np.int32],
dst_page_indices: npt.NDArray[np.int32],
*,
physical_page_size: int,
dcp_size: int,
dcp_rank: int,
src_page_offset: int = 0,
decode_prefix_len: int = 0,
num_kv_tokens: Optional[int] = None,
) -> DCPTokenTransferPlan:
virtual_page_size = physical_page_size * dcp_size
if decode_prefix_len % virtual_page_size != 0:
raise ValueError(
"PD DCP transfer requires decode_prefix_len to align to the virtual "
f"DCP page size ({virtual_page_size}), got {decode_prefix_len}"
)
src_pages = np.asarray(src_page_indices, dtype=np.int64)
dst_pages = np.asarray(dst_page_indices, dtype=np.int64)
source_capacity = src_pages.size * physical_page_size
if num_kv_tokens is None:
num_kv_tokens = source_capacity
if not 0 <= num_kv_tokens <= source_capacity:
raise ValueError(
"num_kv_tokens must fit in the provided source pages, "
f"got tokens={num_kv_tokens}, capacity={source_capacity}"
)
if src_pages.size == 0:
empty = np.empty((0,), dtype=np.int64)
return DCPTokenTransferPlan(empty, empty.copy())
chunk_start = decode_prefix_len + src_page_offset * physical_page_size
first_owned_offset = (dcp_rank - chunk_start) % dcp_size
owned_offsets = np.arange(
first_owned_offset, num_kv_tokens, dcp_size, dtype=np.int64
)
src_token_indices = (
src_pages[owned_offsets // physical_page_size] * physical_page_size
+ owned_offsets % physical_page_size
)
relative_positions = src_page_offset * physical_page_size + owned_offsets
dst_local_offsets = relative_positions // dcp_size
dst_page_ordinals = dst_local_offsets // physical_page_size
if dst_page_ordinals.size and (
dst_pages.size == 0 or int(dst_page_ordinals.max()) >= dst_pages.size
):
required_pages = int(dst_page_ordinals.max()) + 1
raise ValueError(
"Insufficient destination DCP pages: "
f"required={required_pages}, provided={dst_pages.size}, "
f"src_page_offset={src_page_offset}, dcp_rank={dcp_rank}"
)
dst_token_indices = (
dst_pages[dst_page_ordinals] * physical_page_size
+ dst_local_offsets % physical_page_size
)
return DCPTokenTransferPlan(src_token_indices, dst_token_indices)
python/sglang/srt/disaggregation/fake/conn.py
+1
View File
@@ -78,6 +78,7 @@ class FakeKVSender(BaseKVSender):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
self.has_sent = True
logger.debug(
python/sglang/srt/disaggregation/mooncake/conn.py
+179 -16
View File
@@ -34,6 +34,7 @@ from sglang.srt.disaggregation.common.utils import (
AuxDataCodec,
FastQueue,
TransferKVChunk,
build_dcp_token_transfer_plan,
group_concurrent_contiguous,
pack_int_lists,
unpack_int_lists,
@@ -45,6 +46,7 @@ from sglang.srt.disaggregation.utils import (
DisaggregationMode,
build_transfer_entry_pairs,
compute_mamba_state_slice_byte_blocks,
resolve_dcp_dst_entry_indices,
)
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
from sglang.srt.environ import envs
@@ -131,6 +133,10 @@ class KVArgsRegisterInfo:
dst_state_dim_per_tensor: List[List[int]]
dst_kv_layer_ids: List[int]
dst_state_layer_ids: List[List[int]]
dst_dcp_size: int = 1
dst_dcp_rank: int = 0
requires_dcp_relayout: bool = False
dcp_token_item_lens: Optional[List[int]] = None
# Note: always put the staging field at the final (since the staging field is optional and contains multiple inputs)
staging: Optional[StagingRegisterInfo] = None
@@ -163,6 +169,13 @@ class KVArgsRegisterInfo:
if len(msg) > 13 and msg[13] != b""
else []
),
# msg[14:16] belong to the staging field below; DCP trails it.
dst_dcp_size=(
int(msg[16].decode("ascii")) if len(msg) > 16 and msg[16] != b"" else 1
),
dst_dcp_rank=(
int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 14),
)
@@ -749,6 +762,106 @@ class MooncakeKVManager(CommonKVManager):
dst_layer_ids=dst_layer_ids,
)
def send_kvcache_dcp(
self,
mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int32],
*,
dcp_token_item_lens: List[int],
dst_dcp_size: int,
dst_dcp_rank: int,
src_page_offset: int,
decode_prefix_len: int,
num_kv_tokens: int,
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: List[int],
) -> int:
if num_kv_tokens is None:
raise ValueError("PD DCP transfer requires num_kv_tokens")
physical_page_size = self.kv_args.page_size
plan = build_dcp_token_transfer_plan(
prefill_kv_indices,
dst_kv_indices,
physical_page_size=physical_page_size,
dcp_size=dst_dcp_size,
dcp_rank=dst_dcp_rank,
src_page_offset=src_page_offset,
decode_prefix_len=decode_prefix_len,
num_kv_tokens=num_kv_tokens,
)
if plan.src_token_indices.size == 0:
return 0
src_layer_ids = self.kv_args.kv_layer_ids
if src_layer_ids or dst_layer_ids:
dst_indices = resolve_dcp_dst_entry_indices(
src_layer_ids,
dst_layer_ids,
len(self.kv_args.kv_data_ptrs),
len(dst_kv_ptrs),
)
src_kv_ptrs = self.kv_args.kv_data_ptrs
dst_kv_ptrs = [dst_kv_ptrs[j] for j in dst_indices]
else:
src_kv_ptrs, dst_kv_ptrs, _ = self.get_mla_kv_ptrs_with_pp(
self.kv_args.kv_data_ptrs,
dst_kv_ptrs,
)
layers_current_pp_stage = len(src_kv_ptrs)
src_groups, dst_groups = group_concurrent_contiguous(
plan.src_token_indices,
plan.dst_token_indices,
)
layers_params = [
(
src_kv_ptrs[layer_id],
dst_kv_ptrs[layer_id],
dcp_token_item_lens[layer_id],
)
for layer_id in range(layers_current_pp_stage)
]
def set_transfer_blocks(
src_ptr: int, dst_ptr: int, token_item_len: int
) -> List[Tuple[int, int, int]]:
return [
(
src_ptr + int(src_group[0]) * token_item_len,
dst_ptr + int(dst_group[0]) * token_item_len,
len(src_group) * token_item_len,
)
for src_group, dst_group in zip(src_groups, dst_groups)
]
def process_layer(src_ptr: int, dst_ptr: int, token_item_len: int) -> int:
return self._transfer_data(
mooncake_session_id,
set_transfer_blocks(src_ptr, dst_ptr, token_item_len),
)
if self.enable_custom_mem_pool:
futures = [
executor.submit(process_layer, src_ptr, dst_ptr, token_item_len)
for src_ptr, dst_ptr, token_item_len in layers_params
]
for future in concurrent.futures.as_completed(futures):
status = future.result()
if status != 0:
for pending in futures:
pending.cancel()
return status
return 0
transfer_blocks = []
for src_ptr, dst_ptr, token_item_len in layers_params:
transfer_blocks.extend(
set_transfer_blocks(src_ptr, dst_ptr, token_item_len)
)
return self._transfer_data(mooncake_session_id, transfer_blocks)
def send_kvcache_slice(
self,
mooncake_session_id: str,
@@ -1437,23 +1550,33 @@ class MooncakeKVManager(CommonKVManager):
)
break
chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
if len(chunked_dst_kv_indice) < len(
kv_chunk.prefill_kv_indices
):
logger.warning(
f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
)
kv_chunk.prefill_kv_indices = kv_chunk.prefill_kv_indices[
: len(chunked_dst_kv_indice)
]
target_rank_registration_info: KVArgsRegisterInfo = (
self.decode_kv_args_table[req.mooncake_session_id]
)
is_dcp_transfer = (
target_rank_registration_info.requires_dcp_relayout
)
if is_dcp_transfer:
chunked_dst_kv_indice = req.dst_kv_indices
else:
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
if len(chunked_dst_kv_indice) < len(
kv_chunk.prefill_kv_indices
):
logger.warning(
f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
)
kv_chunk.prefill_kv_indices = (
kv_chunk.prefill_kv_indices[
: len(chunked_dst_kv_indice)
]
)
skip_kv, skip_state = self._get_dsa_cache_transfer_skip_flags(
target_rank_registration_info
)
@@ -1463,6 +1586,27 @@ class MooncakeKVManager(CommonKVManager):
or skip_kv
):
ret = 0
elif is_dcp_transfer:
dcp_token_item_lens = (
target_rank_registration_info.dcp_token_item_lens
)
assert dcp_token_item_lens is not None
ret = self.send_kvcache_dcp(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
dcp_token_item_lens=dcp_token_item_lens,
dst_dcp_size=target_rank_registration_info.dst_dcp_size,
dst_dcp_rank=target_rank_registration_info.dst_dcp_rank,
src_page_offset=kv_chunk.index_slice.start or 0,
decode_prefix_len=req.decode_prefix_len or 0,
num_kv_tokens=kv_chunk.num_kv_tokens,
executor=executor,
dst_layer_ids=(
target_rank_registration_info.dst_kv_layer_ids
),
)
elif (
self.is_mla_backend
or self.is_hybrid_mla_backend
@@ -1670,9 +1814,19 @@ class MooncakeKVManager(CommonKVManager):
continue
mooncake_session_id = waiting_req_bytes[3].decode("ascii")
if room == "None":
self.decode_kv_args_table[mooncake_session_id] = (
KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
decode_kv_args = KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout(
decode_kv_args.dst_dcp_size,
decode_kv_args.dst_dcp_rank,
)
if decode_kv_args.requires_dcp_relayout:
decode_kv_args.dcp_token_item_lens = (
self.prepare_dcp_token_item_lens(
[decode_kv_args.dst_kv_item_len]
* len(self.kv_args.kv_item_lens)
)
)
self.decode_kv_args_table[mooncake_session_id] = decode_kv_args
with self.session_lock:
if mooncake_session_id in self.failed_sessions:
self.failed_sessions.remove(mooncake_session_id)
@@ -1786,6 +1940,7 @@ class MooncakeKVManager(CommonKVManager):
is_last_chunk: bool,
aux_index: Optional[int] = None,
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
trace_ctx: Optional[Union[TraceReqContext, TraceNullContext]] = None,
):
assert self.disaggregation_mode == DisaggregationMode.PREFILL
@@ -1824,6 +1979,7 @@ class MooncakeKVManager(CommonKVManager):
is_last_chunk=is_last_chunk,
prefill_aux_index=aux_index,
state_indices=state_indices,
num_kv_tokens=num_kv_tokens,
trace_ctx=trace_ctx,
)
)
@@ -1904,6 +2060,7 @@ class MooncakeKVSender(CommonKVSender):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
kv_indices, index_slice, is_last_chunk, should_skip = (
self._prepare_send_indices(kv_indices, state_indices)
@@ -1917,6 +2074,7 @@ class MooncakeKVSender(CommonKVSender):
kv_indices,
index_slice,
False,
num_kv_tokens=num_kv_tokens,
trace_ctx=self.trace_ctx.copy_for_thread(),
)
else:
@@ -1927,6 +2085,7 @@ class MooncakeKVSender(CommonKVSender):
True,
aux_index=self.aux_index,
state_indices=state_indices,
num_kv_tokens=num_kv_tokens,
trace_ctx=self.trace_ctx.copy_for_thread(),
)
self._record_transfer_indices(kv_indices, state_indices)
@@ -2030,6 +2189,8 @@ class MooncakeKVReceiver(CommonKVReceiver):
dst_tp_rank = str(tp_rank).encode("ascii")
dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
dst_kv_item_len = str(kv_item_len).encode("ascii")
dst_dcp_size = str(self.kv_mgr.dcp_size).encode("ascii")
dst_dcp_rank = str(self.kv_mgr.dcp_rank).encode("ascii")
if (
self.kv_mgr.enable_staging
and self.kv_mgr._staging_ctx.allocator is not None
@@ -2062,6 +2223,8 @@ class MooncakeKVReceiver(CommonKVReceiver):
packed_state_layer_ids,
packed_staging_base_ptr,
staging_total_size_str,
dst_dcp_size,
dst_dcp_rank,
]
)
except zmq.ZMQError:
python/sglang/srt/disaggregation/mori/conn.py
+1
View File
@@ -1417,6 +1417,7 @@ class MoriKVSender(CommonKVSender):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
kv_indices, index_slice, is_last_chunk, should_skip = (
self._prepare_send_indices(kv_indices, state_indices)
python/sglang/srt/disaggregation/nixl/conn.py
+158 -24
View File
@@ -32,6 +32,7 @@ from sglang.srt.disaggregation.common.staging_handler import (
from sglang.srt.disaggregation.common.utils import (
FastQueue,
TransferKVChunk,
build_dcp_token_transfer_plan,
group_concurrent_contiguous,
pack_int_lists,
unpack_int_lists,
@@ -40,6 +41,7 @@ from sglang.srt.disaggregation.utils import (
DisaggregationMode,
build_transfer_entry_pairs,
compute_mamba_state_slice_byte_blocks,
resolve_dcp_dst_entry_indices,
)
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
@@ -158,6 +160,7 @@ class TransferInfo:
required_dst_info_num: int
dst_state_indices: List[List[int]]
decode_prefix_len: Optional[int] = None # for decode radix cache
is_dummy_rank: Optional[bool] = None
# NOTE: optional staging field; populated via STAGING_RSP. Keep at the
# end so positional construction in from_zmq() continues to work.
staging: Optional[StagingTransferInfo] = None
@@ -167,6 +170,8 @@ class TransferInfo:
# 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.is_dummy_rank is not None:
return self.is_dummy_rank
if self.dst_kv_indices.size == 0 and self.decode_prefix_len:
return False
return self.dst_kv_indices.size == 0
@@ -189,6 +194,11 @@ class TransferInfo:
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
is_dummy_rank=(
bool(int(msg[9].decode("ascii")))
if len(msg) > 9 and msg[9] != b""
else None
),
)
@@ -211,6 +221,11 @@ class KVArgsRegisterInfo:
dst_kv_item_len: int
dst_kv_item_lens: list[int]
dst_kv_layer_ids: list[int] = dataclasses.field(default_factory=list)
dst_dcp_size: int = 1
dst_dcp_rank: int = 0
requires_dcp_relayout: bool = False
dcp_token_item_lens: Optional[List[int]] = None
dcp_dst_region_indices: Optional[List[int]] = None
dst_num_slots: Optional[int] = None
dst_state_item_lens: List[List[int]] = dataclasses.field(default_factory=list)
dst_state_dim_per_tensor: List[List[int]] = dataclasses.field(default_factory=list)
@@ -277,6 +292,13 @@ class KVArgsRegisterInfo:
dst_kv_item_len=dst_kv_item_len,
dst_kv_item_lens=dst_kv_item_lens,
dst_kv_layer_ids=dst_kv_layer_ids,
# msg[19:21] are the layer-id frames above; DCP trails them.
dst_dcp_size=(
int(msg[21].decode("ascii")) if len(msg) > 21 and msg[21] != b"" else 1
),
dst_dcp_rank=(
int(msg[22].decode("ascii")) if len(msg) > 22 and msg[22] != b"" else 0
),
dst_num_slots=dst_num_slots,
dst_state_item_lens=dst_state_item_lens,
dst_state_dim_per_tensor=dst_state_dim_per_tensor,
@@ -973,6 +995,27 @@ class NixlKVManager(CommonKVManager):
assert self.src_mem_kind is not None
src_mem_kind = self.src_mem_kind
decode_only_spec_dec = n_dst > n_src
if peer_info.requires_dcp_relayout:
dst_indices = resolve_dcp_dst_entry_indices(
self.kv_args.kv_layer_ids,
peer_info.dst_kv_layer_ids,
n_src,
n_dst,
)
peer_info.dcp_dst_region_indices = dst_indices
dst_kv_mem_kinds = [peer_info.dst_kv_mem_kinds[j] for j in dst_indices]
dst_kv_item_lens = [peer_info.dst_kv_item_lens[j] for j in dst_indices]
dst_mem_kind = _homogeneous_kv_mem_kind(
dst_kv_mem_kinds,
"PD DCP destination",
)
peer_info.dst_homogeneous_mem_kind = dst_mem_kind
peer_info.dcp_token_item_lens = self.prepare_dcp_token_item_lens(
dst_kv_item_lens
)
return
if (
self.is_mla_backend
or self.is_hybrid_mla_backend
@@ -1110,20 +1153,28 @@ class NixlKVManager(CommonKVManager):
len(kv_chunk.prefill_kv_indices) > 0
and self.kv_args.kv_data_ptrs
):
chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
if len(chunked_dst_kv_indice) < len(
kv_chunk.prefill_kv_indices
):
logger.warning(
f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
)
kv_chunk.prefill_kv_indices = kv_chunk.prefill_kv_indices[
: len(chunked_dst_kv_indice)
is_dcp_transfer = dst_info.requires_dcp_relayout
if is_dcp_transfer:
chunked_dst_kv_indice = req.dst_kv_indices
else:
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
if len(chunked_dst_kv_indice) < len(
kv_chunk.prefill_kv_indices
):
logger.warning(
f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
)
kv_chunk.prefill_kv_indices = (
kv_chunk.prefill_kv_indices[
: len(chunked_dst_kv_indice)
]
)
src_prefill_kv_indices = kv_chunk.prefill_kv_indices
notif = (
@@ -1164,7 +1215,18 @@ class NixlKVManager(CommonKVManager):
break
if kv_xfer_handle is None:
if (
if is_dcp_transfer:
kv_xfer_handle = self.send_kvcache_dcp(
req.agent_name,
src_prefill_kv_indices,
dst_info,
chunked_dst_kv_indice,
src_page_offset=kv_chunk.index_slice.start or 0,
decode_prefix_len=req.decode_prefix_len or 0,
num_kv_tokens=kv_chunk.num_kv_tokens,
notif=notif,
)
elif (
self.is_mla_backend
or self.is_hybrid_mla_backend
or decode_tp_size == self.attn_tp_size
@@ -1228,15 +1290,17 @@ class NixlKVManager(CommonKVManager):
if kv_chunk.prefill_aux_index is None:
raise RuntimeError("Missing aux index for last chunk")
# A no-KV notification still identifies its PP source.
# Empty non-final chunks do not consume chunk IDs, so a
# final no-KV chunk_id equals the prior KV chunk count.
aux_notif = f"{req.room}_aux"
if (
len(kv_chunk.prefill_kv_indices) == 0
or not self.kv_args.kv_data_ptrs
):
aux_notif = (
f"{req.room}_aux_nokv_{self.transfer_source_rank}"
aux_notif += (
f"_nokv_{self.transfer_source_rank}"
f"_{kv_chunk.chunk_id}"
)
else:
aux_notif = f"{req.room}_aux"
aux_xfer_handle = self.send_aux(
req.agent_name,
kv_chunk.prefill_aux_index,
@@ -1354,6 +1418,9 @@ class NixlKVManager(CommonKVManager):
if agent_name in self.decode_kv_args_table:
logger.info(f"Peer {agent_name} was already registered, ignoring.")
return
decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout(
decode_kv_args.dst_dcp_size, decode_kv_args.dst_dcp_rank
)
self.decode_kv_args_table[agent_name] = decode_kv_args
self.agent.add_remote_agent(decode_kv_args.agent_metadata)
if self.disaggregation_mode == DisaggregationMode.PREFILL:
@@ -1373,6 +1440,7 @@ class NixlKVManager(CommonKVManager):
src_mem_kind: str = "VRAM",
dst_mem_kind: str = "VRAM",
force_flat: bool = False,
bypass_prepped: bool = False,
):
"""Generic KV cache transfer supporting both MHA and MLA architectures.
Used by both send_kvcache and maybe_send_extra.
@@ -1382,7 +1450,8 @@ class NixlKVManager(CommonKVManager):
index) whose per-layer list must not be half-split into K/V."""
# Prepped path (KV only; state transfers use the non-prepped path below).
if (
src_data_ptrs is self.kv_args.kv_data_ptrs
not bypass_prepped
and src_data_ptrs is self.kv_args.kv_data_ptrs
and "" in self.prep_handles
and peer_name in self.prep_handles
):
@@ -1550,6 +1619,63 @@ class NixlKVManager(CommonKVManager):
dst_mem_kind=dst_mem_kind,
)
def send_kvcache_dcp(
self,
peer_name: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_info: KVArgsRegisterInfo,
dst_kv_indices: npt.NDArray[np.int32],
*,
src_page_offset: int,
decode_prefix_len: int,
num_kv_tokens: int,
notif: str,
):
if self.src_mem_kind is None:
raise RuntimeError("Missing NIXL source KV memory kind")
if dst_info.dst_homogeneous_mem_kind is None:
raise RuntimeError("Missing NIXL destination KV memory kind")
if num_kv_tokens is None:
raise ValueError("PD DCP transfer requires num_kv_tokens")
physical_page_size = self.kv_args.page_size
plan = build_dcp_token_transfer_plan(
prefill_kv_indices,
dst_kv_indices,
physical_page_size=physical_page_size,
dcp_size=dst_info.dst_dcp_size,
dcp_rank=dst_info.dst_dcp_rank,
src_page_offset=src_page_offset,
decode_prefix_len=decode_prefix_len,
num_kv_tokens=num_kv_tokens,
)
if plan.src_token_indices.size == 0:
self.agent.send_notif(peer_name, notif.encode("ascii"))
return None
token_item_lens = dst_info.dcp_token_item_lens
assert token_item_lens is not None
dst_kv_ptrs = [
dst_info.dst_kv_ptrs[dst_idx] for dst_idx in dst_info.dcp_dst_region_indices
]
# Prepared handles encode page-level offsets, while DCP relayout needs
# flat descriptors for the selected token rows.
return self._send_kvcache_generic(
peer_name=peer_name,
src_data_ptrs=self.kv_args.kv_data_ptrs,
dst_data_ptrs=dst_kv_ptrs,
item_lens=token_item_lens,
prefill_data_indices=plan.src_token_indices,
dst_data_indices=plan.dst_token_indices,
dst_gpu_id=dst_info.gpu_id,
notif=notif,
src_mem_kind=self.src_mem_kind,
dst_mem_kind=dst_info.dst_homogeneous_mem_kind,
force_flat=True,
bypass_prepped=True,
)
def send_kvcache_mixed(
self,
peer_name: str,
@@ -2243,6 +2369,7 @@ class NixlKVManager(CommonKVManager):
chunk_id: int,
aux_index: Optional[int] = None,
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
assert self.disaggregation_mode == DisaggregationMode.PREFILL
assert not is_last_chunk or (is_last_chunk and aux_index is not None)
@@ -2273,6 +2400,7 @@ class NixlKVManager(CommonKVManager):
chunk_id=chunk_id,
prefill_aux_index=aux_index,
state_indices=state_indices,
num_kv_tokens=num_kv_tokens,
)
)
return None
@@ -2313,8 +2441,8 @@ class NixlKVManager(CommonKVManager):
elif tag == "stg":
self._handle_stg_notification(components, room)
elif tag == "aux":
# main's "nokv" marker (decode-side radix cache hit):
# mark expected_kvs_per_pp[pp_rank] = 0 for this rank.
# Main's "nokv" marker carries the number of earlier KV
# chunks expected from this PP rank.
self._handle_aux_notification(room, components)
elif tag == "state":
pp_rank = int(components[2]) if len(components) > 2 else 0
@@ -2342,15 +2470,16 @@ class NixlKVManager(CommonKVManager):
Notification tag layouts:
aux: {room}_aux -> 2 fields
aux (nokv): {room}_aux_nokv_{pp_rank} -> 4 fields
(decode-side radix cache hit; this pp_rank sent
no KV pages, so expected_kvs_per_pp[pp_rank] = 0)
aux (nokv): {room}_aux_nokv_{pp_rank}_{expected} -> 5 fields
(the last chunk had no KV pages for this rank;
`expected` is the number of prior KV chunks)
"""
self.transfer_statuses[room].received_aux = True
# main's "nokv" marker (decode-side radix cache hit, see #19746).
if len(components) > 3 and components[2] == "nokv":
pp_rank = int(components[3])
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = 0
expected = int(components[4]) if len(components) > 4 else 0
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = expected
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)
@@ -2597,6 +2726,7 @@ class NixlKVSender(CommonKVSender):
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
if self._send_failed:
return
@@ -2620,6 +2750,7 @@ class NixlKVSender(CommonKVSender):
self.chunk_id,
self.aux_index,
state_indices,
num_kv_tokens,
)
self._record_transfer_indices(kv_indices, state_indices)
self.chunk_id += 1
@@ -2737,6 +2868,7 @@ class NixlKVReceiver(CommonKVReceiver):
str(self.required_dst_info_num).encode("ascii"),
packed_state_indices,
str(decode_prefix_len or 0).encode("ascii"),
str(int(is_dummy)).encode("ascii"),
]
)
except zmq.ZMQError:
@@ -2865,6 +2997,8 @@ class NixlKVReceiver(CommonKVReceiver):
packed_kv_item_lens,
packed_state_layer_ids,
packed_kv_layer_ids,
str(self.kv_mgr.dcp_size).encode("ascii"),
str(self.kv_mgr.dcp_rank).encode("ascii"),
]
)
except zmq.ZMQError:
python/sglang/srt/disaggregation/prefill.py
+12 -3
View File
@@ -332,7 +332,8 @@ class PrefillBootstrapQueue:
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
num_kv_indices_to_send,
self.scheduler.token_to_kv_pool_allocator.page_size,
)
req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
req.pending_bootstrap = False
@@ -1082,7 +1083,11 @@ class SchedulerDisaggregationPrefillMixin:
cached_end = len(req.prefix_indices) - req.host_hit_length
if cached_end <= req.start_send_idx:
return
assert cached_end % self.token_to_kv_pool_allocator.page_size == 0
if cached_end % self.token_to_kv_pool_allocator.page_size != 0:
# DCP radix hits can end on a logical cache-page boundary that is
# not a complete physical DCP page. The regular final send covers
# the full range; only skip this optional early-send optimization.
return
# Early-send issues the KV read before this step's forward is enqueued,
# but under overlap scheduling the PRIOR step's prefill forward may still
# be writing these prefix pages on forward_stream. Record a completion
@@ -1235,7 +1240,11 @@ class SchedulerDisaggregationPrefillMixin:
page_indices = kv_to_page_indices(kv_indices, page_size)
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.disagg_kv_sender.send(
page_indices,
state_indices,
num_kv_tokens=end_idx - start_idx,
)
req.start_send_idx = end_idx
def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None:
python/sglang/srt/disaggregation/utils.py
+25
View File
@@ -934,6 +934,31 @@ def build_transfer_entry_pairs(
return [(i, i) for i in range(n_src)]
def resolve_dcp_dst_entry_indices(
src_layer_ids: List[int],
dst_layer_ids: List[int],
n_src: int,
n_dst: int,
) -> List[int]:
"""Destination entry index for each local KV entry, for a DCP relayout.
DCP re-splits the KV by context while PP re-splits it by layer, so the two
index spaces only line up when neither peer is pipelined. Both backends
need the same resolution, hence the shared helper.
"""
if not src_layer_ids and not dst_layer_ids:
# Legacy/non-PP layout. n_dst may exceed n_src when the decode side
# runs speculative decoding and the prefill side does not.
return list(range(n_src))
# A one-sided mapping is rejected by build_transfer_entry_pairs itself.
return [
j
for _, j in build_transfer_entry_pairs(
src_layer_ids, dst_layer_ids, n_src, n_dst
)
]
def append_state_component(
kv_args: KVArgs,
state_type: StateType,