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

[PD] Minor code cleanup for mooncake backend (#18279)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-05 17:38:09 +08:00
committed by GitHub
parent 079fc8f3c5
commit afae4c7178
1 changed files with 27 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+27 -26
View File
@@ -315,8 +315,10 @@ class MooncakeKVManager(CommonKVManager):
# Use correct item lengths for K and V separately # Use correct item lengths for K and V separately
if layers_current_pp_stage > len(dst_k_ptrs): if layers_current_pp_stage > len(dst_k_ptrs):
logger.error( logger.error(
f"layers_current_pp_stage is out of range: {layers_current_pp_stage=}, {len(dst_k_ptrs)}" "Prefill transfer kvcache error, layers_current_pp_stage is out of range: "
f"layers_current_pp_stage={layers_current_pp_stage}, len(dst_k_ptrs)={len(dst_k_ptrs)}"
) )
return -1
layers_params = [ layers_params = [
( (
src_k_ptrs[layer_id], src_k_ptrs[layer_id],
@@ -373,13 +375,12 @@ class MooncakeKVManager(CommonKVManager):
for f in futures: for f in futures:
f.cancel() f.cancel()
return status return status
return 0
else: else:
# Combining all layers' params in one batch transfer is more efficient # Combining all layers' params in one batch transfer is more efficient
# compared to using multiple threads # compared to using multiple threads
return process_layers(layers_params) return process_layers(layers_params)
return 0
def send_kvcache( def send_kvcache(
self, self,
mooncake_session_id: str, mooncake_session_id: str,
@@ -401,9 +402,9 @@ class MooncakeKVManager(CommonKVManager):
def send_kvcache_slice( def send_kvcache_slice(
self, self,
mooncake_session_id: str, mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int64], prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int], dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int64], dst_kv_indices: npt.NDArray[np.int32],
dst_tp_rank: int, dst_tp_rank: int,
dst_attn_tp_size: int, dst_attn_tp_size: int,
dst_kv_item_len: int, dst_kv_item_len: int,
@@ -422,7 +423,6 @@ class MooncakeKVManager(CommonKVManager):
src_kv_item_len = self.kv_args.kv_item_lens[0] src_kv_item_len = self.kv_args.kv_item_lens[0]
dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size
num_kv_heads = self.kv_args.kv_head_num num_kv_heads = self.kv_args.kv_head_num
num_layers = len(self.kv_args.kv_data_ptrs)
page_size = self.kv_args.page_size page_size = self.kv_args.page_size
# Calculate head distribution # Calculate head distribution
@@ -464,30 +464,31 @@ class MooncakeKVManager(CommonKVManager):
) )
return -1 return -1
prefill_kv_indices_reshaped = prefill_kv_indices.astype(np.int64).reshape(-1, 1) prefill_page_indices = prefill_kv_indices.reshape(-1, 1)
dst_kv_indices_reshaped = dst_kv_indices.astype(np.int64).reshape(-1, 1) decode_page_indices = dst_kv_indices.reshape(-1, 1)
token_offsets = np.arange(page_size, dtype=np.int64).reshape(1, -1) tokens_per_page = np.arange(page_size, dtype=np.int32).reshape(1, -1)
bytes_per_token_on_prefill = src_kv_item_len // page_size bytes_per_token_on_prefill = src_kv_item_len // page_size
bytes_per_token_on_decode = dst_kv_item_len // page_size bytes_per_token_on_decode = dst_kv_item_len // page_size
src_token_offsets_base = ( src_token_slot_offsets = (
token_offsets * bytes_per_token_on_prefill + src_head_slice_offset tokens_per_page * bytes_per_token_on_prefill + src_head_slice_offset
) )
dst_token_offsets_base = ( dst_token_slot_offsets = (
token_offsets * bytes_per_token_on_decode + dst_head_slice_offset tokens_per_page * bytes_per_token_on_decode + dst_head_slice_offset
) )
def process_layer_tp_aware(ptrs): def process_layer_tp_aware(src_layer_ptr, dst_layer_ptr):
src_ptr, dst_ptr = ptrs src_page_base_addrs = src_layer_ptr + prefill_page_indices * src_kv_item_len
src_page_starts = src_ptr + prefill_kv_indices_reshaped * src_kv_item_len dst_page_base_addrs = dst_layer_ptr + decode_page_indices * dst_kv_item_len
dst_page_starts = dst_ptr + dst_kv_indices_reshaped * dst_kv_item_len src_slice_addrs = src_page_base_addrs + src_token_slot_offsets
src_addrs = src_page_starts + src_token_offsets_base dst_slice_addrs = dst_page_base_addrs + dst_token_slot_offsets
dst_addrs = dst_page_starts + dst_token_offsets_base
src_addr_list = src_addrs.reshape(-1).tolist() src_addr_list = src_slice_addrs.reshape(-1).tolist()
if not src_addr_list: if not src_addr_list:
# Nothing to transfer for this layer.
return 0 return 0
dst_addr_list = dst_addrs.reshape(-1).tolist() dst_addr_list = dst_slice_addrs.reshape(-1).tolist()
total_chunks = len(src_addr_list) total_slices = len(src_addr_list)
length_list = [heads_bytes_per_token_to_send] * total_chunks length_list = [heads_bytes_per_token_to_send] * total_slices
return self.engine.batch_transfer_sync( return self.engine.batch_transfer_sync(
mooncake_session_id, src_addr_list, dst_addr_list, length_list mooncake_session_id, src_addr_list, dst_addr_list, length_list
) )
@@ -495,11 +496,11 @@ class MooncakeKVManager(CommonKVManager):
futures = [] futures = []
for i in range(layers_current_pp_stage): for i in range(layers_current_pp_stage):
futures.append( futures.append(
executor.submit(process_layer_tp_aware, (src_k_ptrs[i], dst_k_ptrs[i])) executor.submit(process_layer_tp_aware, src_k_ptrs[i], dst_k_ptrs[i])
) )
for i in range(layers_current_pp_stage): for i in range(layers_current_pp_stage):
futures.append( futures.append(
executor.submit(process_layer_tp_aware, (src_v_ptrs[i], dst_v_ptrs[i])) executor.submit(process_layer_tp_aware, src_v_ptrs[i], dst_v_ptrs[i])
) )
for future in concurrent.futures.as_completed(futures): for future in concurrent.futures.as_completed(futures):
@@ -985,7 +986,7 @@ class MooncakeKVManager(CommonKVManager):
elif status == KVPoll.Failed: elif status == KVPoll.Failed:
self.record_failure( self.record_failure(
bootstrap_room, bootstrap_room,
f"Failed to get kvcache from prefill instance, it might be dead", "Failed to get kvcache from prefill instance, it might be dead",
) )
self.update_status(bootstrap_room, status) self.update_status(bootstrap_room, status)