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

[PD] Support PD with context parallel after refactor (#19504)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-28 13:11:15 +08:00
committed by GitHub
co-authored by Vladislav Nosivskoy
parent 0404ef9023
commit b01f3590be
5 changed files with 192 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+15 -8
View File
@@ -53,7 +53,7 @@ class TransferKVChunk:
room: int
prefill_kv_indices: npt.NDArray[np.int32]
index_slice: slice
is_last: bool
is_last_chunk: bool
prefill_aux_index: Optional[int]
state_indices: Optional[List[int]]
@@ -861,7 +861,7 @@ class MooncakeKVManager(CommonKVManager):
)
break
if kv_chunk.is_last:
if kv_chunk.is_last_chunk:
if kv_chunk.state_indices is not None:
self.maybe_send_extra(
req,
@@ -893,7 +893,7 @@ class MooncakeKVManager(CommonKVManager):
else:
# Dummy request means the decode instance is not used, so its status can be marked as success directly
# Dummy request does not need to sync status to decode endpoint
if kv_chunk.is_last and req.room in self.request_status:
if kv_chunk.is_last_chunk and req.room in self.request_status:
self.update_status(req.room, KVPoll.Success)
if (
@@ -1038,12 +1038,12 @@ class MooncakeKVManager(CommonKVManager):
bootstrap_room: int,
kv_indices: npt.NDArray[np.int32],
index_slice: slice,
is_last: bool,
is_last_chunk: bool,
aux_index: Optional[int] = None,
state_indices: Optional[List[int]] = None,
):
assert self.disaggregation_mode == DisaggregationMode.PREFILL
assert not is_last or (is_last and aux_index is not None)
assert not is_last_chunk or (is_last_chunk and aux_index is not None)
if (
bootstrap_room not in self.request_status
@@ -1072,7 +1072,7 @@ class MooncakeKVManager(CommonKVManager):
room=bootstrap_room,
prefill_kv_indices=kv_indices,
index_slice=index_slice,
is_last=is_last,
is_last_chunk=is_last_chunk,
prefill_aux_index=aux_index,
state_indices=state_indices,
)
@@ -1134,9 +1134,16 @@ class MooncakeKVSender(CommonKVSender):
):
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
self.curr_idx += len(kv_indices)
is_last = self.curr_idx == self.num_kv_indices
is_last_chunk = self.curr_idx == self.num_kv_indices
if not is_last:
if self.kv_mgr.is_dummy_cp_rank:
if not is_last_chunk:
return
else:
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
return
if not is_last_chunk:
self.kv_mgr.add_transfer_request(
self.bootstrap_room,
kv_indices,