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:
@@ -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)
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user