diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 282a49ba8..3fe895ed1 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -288,12 +288,11 @@ class MooncakeKVManager(CommonKVManager): src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = ( self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) ) - kv_item_len = item_lens[0] layers_params = [ ( src_kv_ptrs[layer_id], dst_kv_ptrs[layer_id], - kv_item_len, + item_lens[layer_id], ) for layer_id in range(layers_current_pp_stage) ] @@ -301,19 +300,18 @@ class MooncakeKVManager(CommonKVManager): src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = ( self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) ) - kv_item_len = item_lens[0] layers_params = [ ( src_k_ptrs[layer_id], dst_k_ptrs[layer_id], - kv_item_len, + item_lens[layer_id], ) for layer_id in range(layers_current_pp_stage) ] + [ ( src_v_ptrs[layer_id], dst_v_ptrs[layer_id], - kv_item_len, + item_lens[layer_id], ) for layer_id in range(layers_current_pp_stage) ]