Fix target MLA with eagle3 support for PD disaggregation (#13555)

Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com>
Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
This commit is contained in:
Mike Qiu
2025-11-21 12:16:05 +08:00
committed by GitHub
co-authored by Mike_Qiu
parent 750084ae08
commit 6be65ae462
@@ -288,12 +288,11 @@ class MooncakeKVManager(CommonKVManager):
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = ( src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
) )
kv_item_len = item_lens[0]
layers_params = [ layers_params = [
( (
src_kv_ptrs[layer_id], src_kv_ptrs[layer_id],
dst_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) 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 = ( 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) self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
) )
kv_item_len = item_lens[0]
layers_params = [ layers_params = [
( (
src_k_ptrs[layer_id], src_k_ptrs[layer_id],
dst_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) for layer_id in range(layers_current_pp_stage)
] + [ ] + [
( (
src_v_ptrs[layer_id], src_v_ptrs[layer_id],
dst_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) for layer_id in range(layers_current_pp_stage)
] ]