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