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
parent 750084ae08
commit 6be65ae462

View File

@@ -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)
]