Star137
242
代码
代码
Issues83
Pull Requests180
项目讨论
Wiki
分析
Star137
242
已合并
[Feature] MindIE支持Prefix Cache叠加SP #21
chuyuelin创建于 2025年12月23日
[Feature] MindIE支持Prefix Cache叠加SP #21
已合并
chuyuelin创建于 2025年12月23日
从chuyuelin/MindIE-LLMprefixcache_sp_bugfix合入到Ascend/MindIE-LLMdev
共 7 个文件变更+101-25
Mexamples/atb_models/atb_framework/models/deepseekv2/layer/decoder_layer.cpp+9-1
@@ -128,6 +128,9 @@ std::map<std::string, std::vector<std::string>> GetDeepseekV2LayerInTensorCandid
128 {"prefixcache_cp", {128 {"prefixcache_cp", {
129 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"129 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"
130 }},130 }},
131+ {"prefixcache_sp", {
132+ "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx"
133+ }},
131 {"prefixcache_c8", {"in_history_compressed_kv_int"}},134 {"prefixcache_c8", {"in_history_compressed_kv_int"}},
132 {"dense_tp", {135 {"dense_tp", {
133 "in_dense_tp_padding_idx", "in_dense_tp_mlp_out_idx",136 "in_dense_tp_padding_idx", "in_dense_tp_mlp_out_idx",
@@ -230,6 +233,8 @@ std::vector<std::string> ConstructIntensorList(const DecoderLayerParam &param)
230 }233 }
231 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {234 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {
232 atb_speed::common::AddTensorToList(deepseekV2InTensorCandidates, "prefixcache_cp", inTensorList);235 atb_speed::common::AddTensorToList(deepseekV2InTensorCandidates, "prefixcache_cp", inTensorList);
236+ } else if (param.mapping.Get(base::ATTN_INNER_SP).IsEnabled()) {
237+ atb_speed::common::AddTensorToList(deepseekV2InTensorCandidates, "prefixcache_sp", inTensorList);
233 }238 }
234 }239 }
235 if (param.hasDenseTp) {240 if (param.hasDenseTp) {
@@ -407,7 +412,7 @@ void SetAttnCpParam(
407 atb_speed::common::LatentAttentionParam<atb::infer::RmsNormParam> &latentAttentionParam,412 atb_speed::common::LatentAttentionParam<atb::infer::RmsNormParam> &latentAttentionParam,
408 const DecoderLayerParam &param)413 const DecoderLayerParam &param)
409{414{
410- if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {415+ if (param.mapping.Get(base::ATTN_CP).IsEnabled() || param.mapping.Get(base::ATTN_INNER_SP).IsEnabled()) {
411 latentAttentionParam.contextParallelInfo = param.mapping.Get(base::ATTN_CP);416 latentAttentionParam.contextParallelInfo = param.mapping.Get(base::ATTN_CP);
412 latentAttentionParam.ringMLAParam.headNum = param.numAttentionHeadsPerRank;417 latentAttentionParam.ringMLAParam.headNum = param.numAttentionHeadsPerRank;
413 latentAttentionParam.ringMLAParam.kvHeadNum = param.numAttentionHeadsPerRank;418 latentAttentionParam.ringMLAParam.kvHeadNum = param.numAttentionHeadsPerRank;
@@ -567,6 +572,9 @@ int64_t SetAttention(atb::GraphParam &opGraph, const DecoderLayerParam &param,
567 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {572 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {
568 atb_speed::common::AddTensorToList(573 atb_speed::common::AddTensorToList(
569 GetDeepseekV2LayerInTensorCandidates(), "prefixcache_cp", attnInTensorNames);574 GetDeepseekV2LayerInTensorCandidates(), "prefixcache_cp", attnInTensorNames);
575+ } else if (param.mapping.Get(base::ATTN_INNER_SP).IsEnabled()) {
576+ atb_speed::common::AddTensorToList(
577+ GetDeepseekV2LayerInTensorCandidates(), "prefixcache_sp", attnInTensorNames);
570 }578 }
571 }579 }
572 if (param.enableFA3 && param.enableKvQuantLayer) {580 if (param.enableFA3 && param.enableKvQuantLayer) {
Mexamples/atb_models/atb_framework/models/deepseekv2/model/decoder_model.cpp+18-0
@@ -429,6 +429,9 @@ std::map<std::string, std::vector<std::string>> GetDeepseekV2ModelInTensorCandid
429 {"prefixcache_cp", {429 {"prefixcache_cp", {
430 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"430 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"
431 }},431 }},
432+ {"prefixcache_sp", {
433+ "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx"
434+ }},
432 {"prefixcache_c8", {"in_history_compressed_kv_int"}},435 {"prefixcache_c8", {"in_history_compressed_kv_int"}},
433 {"dense_tp", {436 {"dense_tp", {
434 "in_dense_tp_padding_idx_model", "in_dense_tp_mlp_out_idx_model",437 "in_dense_tp_padding_idx_model", "in_dense_tp_mlp_out_idx_model",
@@ -475,6 +478,8 @@ void DecoderModel::ConstructInTensorMap()
475 }478 }
476 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {479 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {
477 atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_cp", this->inTensorMap);480 atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_cp", this->inTensorMap);
481+ } else if (param.mapping.Get(base::ATTN_INNER_SP).IsEnabled()) {
482+ atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_sp", this->inTensorMap);
478 }483 }
479 }484 }
480 if (param.enableDenseTp) {485 if (param.enableDenseTp) {
@@ -1265,6 +1270,18 @@ atb::Status DecoderModel::AddPrefixCacheCpHostWeight(atb_speed::Model::Node &lay
1265 return atb::NO_ERROR;1270 return atb::NO_ERROR;
1266}1271}
1267 1272 
1273+atb::Status DecoderModel::AddPrefixCacheSpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId)
1274+{
1275+ if (param.enablePrefixCache && param.mapping.Get(base::ATTN_INNER_SP).IsEnabled() \
1276+ && (!param.mapping.Get(base::ATTN_CP).IsEnabled())) {
1277+ layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
1278+ atb_speed::common::GetTensorIdx(this->inTensorMap, "in_kv_cache_padding_idx"));
1279+ layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
1280+ atb_speed::common::GetTensorIdx(this->inTensorMap, "in_kv_cache_unpadding_idx"));
1281+ }
1282+ return atb::NO_ERROR;
1283+}
1284+ 
1268atb::Status DecoderModel::AddParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId)1285atb::Status DecoderModel::AddParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId)
1269{1286{
1270 layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(1287 layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
@@ -1369,6 +1386,7 @@ atb::Status DecoderModel::AddLayerHostWeight(atb_speed::Model::Node &layerNode,
1369 AddParallelHostWeight(layerNode, inTensorId);1386 AddParallelHostWeight(layerNode, inTensorId);
1370 AddPrefixCacheHostWeight(layerNode, inTensorId);1387 AddPrefixCacheHostWeight(layerNode, inTensorId);
1371 AddPrefixCacheCpHostWeight(layerNode, inTensorId);1388 AddPrefixCacheCpHostWeight(layerNode, inTensorId);
1389+ AddPrefixCacheSpHostWeight(layerNode, inTensorId);
1372 AddDenseTpHostWeight(layerNode, inTensorId, layerId);1390 AddDenseTpHostWeight(layerNode, inTensorId, layerId);
1373 // new inTensors please add here before1391 // new inTensors please add here before
1374 AddExpertHostWeight(layerNode, inTensorId);1392 AddExpertHostWeight(layerNode, inTensorId);
Mexamples/atb_models/atb_framework/models/deepseekv2/model/decoder_model.h+1-0
@@ -110,6 +110,7 @@ private:
110 atb::Status AddSequenceParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);110 atb::Status AddSequenceParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
111 atb::Status AddPrefixCacheHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);111 atb::Status AddPrefixCacheHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
112 atb::Status AddPrefixCacheCpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);112 atb::Status AddPrefixCacheCpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
113+ atb::Status AddPrefixCacheSpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
113 atb::Status AddExpertHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);114 atb::Status AddExpertHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
114 atb::Status AddDenseTpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);115 atb::Status AddDenseTpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);
115 atb::Status AddNodesBeforeLayer() override;116 atb::Status AddNodesBeforeLayer() override;
Mexamples/atb_models/atb_framework/models/deepseekv2/model/mtp_decoder_model.cpp+18-0
@@ -87,6 +87,9 @@ static std::map<std::string, std::vector<std::string>> GetDeepseekV2ModelInTenso
87 {"prefixcache_cp", {87 {"prefixcache_cp", {
88 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"88 "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"
89 }},89 }},
90+ {"prefixcache_sp", {
91+ "in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx"
92+ }},
90 {"prefixcache_c8", {"in_history_compressed_kv_int"}},93 {"prefixcache_c8", {"in_history_compressed_kv_int"}},
91 {"dense_tp", {94 {"dense_tp", {
92 "in_dense_tp_padding_idx_model", "in_dense_tp_mlp_out_idx_model",95 "in_dense_tp_padding_idx_model", "in_dense_tp_mlp_out_idx_model",
@@ -126,6 +129,8 @@ void MtpDecoderModel::ConstructInTensorMap()
126 }129 }
127 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {130 if (param.mapping.Get(base::ATTN_CP).IsEnabled()) {
128 atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_cp", this->inTensorMap);131 atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_cp", this->inTensorMap);
132+ } else if (param.mapping.Get(base::ATTN_INNER_SP).IsEnabled()) {
133+ atb_speed::common::AssignTensorIdx(deepseekV2ModelInTensorCandidates, "prefixcache_sp", this->inTensorMap);
129 }134 }
130 }135 }
131 if (param.enableDenseTp) {136 if (param.enableDenseTp) {
@@ -682,6 +687,18 @@ atb::Status MtpDecoderModel::AddPrefixCacheCpHostWeight(atb_speed::Model::Node &
682 return atb::NO_ERROR;687 return atb::NO_ERROR;
683}688}
684 689 
690+atb::Status MtpDecoderModel::AddPrefixCacheSpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId)
691+{
692+ if (param.enablePrefixCache && param.mapping.Get(base::ATTN_INNER_SP).IsEnabled() \
693+ && (!param.mapping.Get(base::ATTN_CP).IsEnabled())) {
694+ layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
695+ atb_speed::common::GetTensorIdx(this->inTensorMap, "in_kv_cache_padding_idx"));
696+ layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
697+ atb_speed::common::GetTensorIdx(this->inTensorMap, "in_kv_cache_unpadding_idx"));
698+ }
699+ return atb::NO_ERROR;
700+}
701+ 
685atb::Status MtpDecoderModel::AddLayerHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId)702atb::Status MtpDecoderModel::AddLayerHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId)
686{703{
687 layerNode.inTensors.at(inTensorId++) = &graph_.internalTensors.at(704 layerNode.inTensors.at(inTensorId++) = &graph_.internalTensors.at(
@@ -721,6 +738,7 @@ atb::Status MtpDecoderModel::AddLayerHostWeight(atb_speed::Model::Node &layerNod
721 AddParallelHostWeight(layerNode, inTensorId);738 AddParallelHostWeight(layerNode, inTensorId);
722 AddPrefixCacheHostWeight(layerNode, inTensorId);739 AddPrefixCacheHostWeight(layerNode, inTensorId);
723 AddPrefixCacheCpHostWeight(layerNode, inTensorId);740 AddPrefixCacheCpHostWeight(layerNode, inTensorId);
741+ AddPrefixCacheSpHostWeight(layerNode, inTensorId);
724 AddDenseTpHostWeight(layerNode, inTensorId, layerId);742 AddDenseTpHostWeight(layerNode, inTensorId, layerId);
725 if (param.enableEPWB) {743 if (param.enableEPWB) {
726 layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(744 layerNode.inTensors.at(inTensorId++) = &graph_.inTensors.at(
Mexamples/atb_models/atb_framework/models/deepseekv2/model/mtp_decoder_model.h+1-0
@@ -39,6 +39,7 @@ private:
39 atb::Status AddSingleLayer(uint32_t layerId) override;39 atb::Status AddSingleLayer(uint32_t layerId) override;
40 atb::Status AddParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);40 atb::Status AddParallelHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
41 atb::Status AddPrefixCacheCpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);41 atb::Status AddPrefixCacheCpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
42+ atb::Status AddPrefixCacheSpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
42 atb::Status AddPrefixCacheHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);43 atb::Status AddPrefixCacheHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId);
43 atb::Status AddLayerHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);44 atb::Status AddLayerHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);
44 atb::Status AddDenseTpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);45 atb::Status AddDenseTpHostWeight(atb_speed::Model::Node &layerNode, size_t &inTensorId, int layerId);
Mexamples/atb_models/atb_framework/models/deepseekv2/operation/latent_attention.cpp+19-5
@@ -74,6 +74,8 @@ std::map<std::string, std::vector<std::string>> GetLatentAttnInTensorCandidates(
74 },74 },
75 {"prefixcache_cp", {"in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"}75 {"prefixcache_cp", {"in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx", "in_kv_cache_len"}
76 },76 },
77+ {"prefixcache_sp", {"in_kv_cache_padding_idx", "in_kv_cache_unpadding_idx"}
78+ },
77 {"prefixcache_c8", {"in_history_compressed_kv_int"}79 {"prefixcache_c8", {"in_history_compressed_kv_int"}
78 },80 },
79 };81 };
@@ -140,6 +142,11 @@ std::map<std::string, std::vector<std::string>> GetLatentAttnIntermediateTensorC
140 "temp_v_proj_b", "intermediate_k_nope_history", "rope_k_o_repeat_history",142 "temp_v_proj_b", "intermediate_k_nope_history", "rope_k_o_repeat_history",
141 "intermediate_v_mha_history", "intermediate_history_kv_rope_concat_padding",143 "intermediate_v_mha_history", "intermediate_history_kv_rope_concat_padding",
142 "intermediate_history_kv_rope_concat_allgather_unpadding"}},144 "intermediate_history_kv_rope_concat_allgather_unpadding"}},
145+ {"prefixcache_sp", {"intermediate_history_kv_rope_concat", "intermediate_history_kv_rope_concat_allgather",
146+ "intermediate_kv_cp_history", "rope_k_o_cp_history", "intermediate_history_kv_rope_concat_padding",
147+ "intermediate_history_kv_rope_concat_allgather_unpadding", "temp_v_proj_b", "intermediate_k_nope",
148+ "intermediate_v_mha", "intermediate_k_nope_history", "rope_k_o_repeat_history",
149+ "intermediate_v_mha_history", "rope_k_o_repeat", "cur_lse", "cache_lse"}}
143 };150 };
144 return latentAttnIntermediateTensorCandidates;151 return latentAttnIntermediateTensorCandidates;
145}152}
@@ -162,6 +169,8 @@ std::map<std::string, uint32_t> ConstructTensorMap(const LatentAttentionParam<No
162 }169 }
163 if (param.contextParallelInfo.IsEnabled()) {170 if (param.contextParallelInfo.IsEnabled()) {
164 AddTensorToList(latentAttnInTensorCandidates, "prefixcache_cp", inTensorList);171 AddTensorToList(latentAttnInTensorCandidates, "prefixcache_cp", inTensorList);
172+ } else if (param.hasAttnInnerSp) {
173+ AddTensorToList(latentAttnInTensorCandidates, "prefixcache_sp", inTensorList);
165 }174 }
166 }175 }
167 if (EnableFA3Quant(param)) { // 添加FA3特性的Tensor176 if (EnableFA3Quant(param)) { // 添加FA3特性的Tensor
@@ -198,6 +207,8 @@ std::map<std::string, uint32_t> ConstructTensorMap(const LatentAttentionParam<No
198 if (param.enablePrefixCache) {207 if (param.enablePrefixCache) {
199 if (param.contextParallelInfo.IsEnabled()) {208 if (param.contextParallelInfo.IsEnabled()) {
200 AddTensorToList(latentAttnIntermediateTensorCandidates, "prefixcache_cp", intermediateTensorList);209 AddTensorToList(latentAttnIntermediateTensorCandidates, "prefixcache_cp", intermediateTensorList);
210+ } else if (param.hasAttnInnerSp) {
211+ AddTensorToList(latentAttnIntermediateTensorCandidates, "prefixcache_sp", intermediateTensorList);
201 } else {212 } else {
202 AddTensorToList(latentAttnIntermediateTensorCandidates, "prefixcache", intermediateTensorList);213 AddTensorToList(latentAttnIntermediateTensorCandidates, "prefixcache", intermediateTensorList);
203 }214 }
@@ -943,7 +954,8 @@ atb::Status AddLAttnKRopeRepeatHistoryNode(const LatentAttentionParam<NormParamT
943 newShape.dims[1] = 1;954 newShape.dims[1] = 1;
944 newShape.dims[2] = oldShape.dims[1]; // 2:dim id955 newShape.dims[2] = oldShape.dims[1]; // 2:dim id
945 };956 };
946- if (param.enablePrefixCache && param.contextParallelInfo.IsEnabled() && param.isPrefill) {957+ if ((param.enablePrefixCache && param.contextParallelInfo.IsEnabled() && param.isPrefill) || \
958+ (param.hasAttnInnerSp && param.enablePrefixCache)) {
947 keyRepeatNode.inTensorIds[0] = GetTensorIdx(tensorMap, "rope_k_o_cp_history");959 keyRepeatNode.inTensorIds[0] = GetTensorIdx(tensorMap, "rope_k_o_cp_history");
948 }960 }
949 keyRepeatNode.operation = new atb_speed::common::RepeatOperation("RepeatNode", kvRepeatParam);961 keyRepeatNode.operation = new atb_speed::common::RepeatOperation("RepeatNode", kvRepeatParam);
@@ -1117,7 +1129,8 @@ atb::Status AddLAttnKProjBHistoryNode(const LatentAttentionParam<NormParamType>
1117 newShape.dims[0] = oldShape.dims[0] * oldShape.dims[1];1129 newShape.dims[0] = oldShape.dims[0] * oldShape.dims[1];
1118 newShape.dims[1] = oldShape.dims[2]; // 2: dim id1130 newShape.dims[1] = oldShape.dims[2]; // 2: dim id
1119 };1131 };
1120- if (param.contextParallelInfo.IsEnabled() && param.isPrefill) {1132+ if ((param.contextParallelInfo.IsEnabled() && param.isPrefill) || \
1133+ (param.hasAttnInnerSp && param.enablePrefixCache)) {
1121 kProjBNode.inTensorIds[0] = GetTensorIdx(tensorMap, "intermediate_kv_cp_history");1134 kProjBNode.inTensorIds[0] = GetTensorIdx(tensorMap, "intermediate_kv_cp_history");
1122 }1135 }
1123 1136 
@@ -1224,7 +1237,8 @@ atb::Status AddLAttnVProjBHistoryNode(const LatentAttentionParam<NormParamType>
1224 newShape.dims[0] = oldShape.dims[0] * oldShape.dims[1];1237 newShape.dims[0] = oldShape.dims[0] * oldShape.dims[1];
1225 newShape.dims[1] = oldShape.dims[2]; // 2: dim id1238 newShape.dims[1] = oldShape.dims[2]; // 2: dim id
1226 };1239 };
1227- if (param.contextParallelInfo.IsEnabled() && param.isPrefill) {1240+ if ((param.contextParallelInfo.IsEnabled() && param.isPrefill) || \
1241+ (param.hasAttnInnerSp && param.enablePrefixCache)) {
1228 vProjBNode.inTensorIds[0] = GetTensorIdx(tensorMap, "intermediate_kv_cp_history");1242 vProjBNode.inTensorIds[0] = GetTensorIdx(tensorMap, "intermediate_kv_cp_history");
1229 }1243 }
1230 1244 
@@ -1617,7 +1631,7 @@ atb::Status PreprocessKV(const LatentAttentionParam<NormParamType> &param,
1617 if (EnableFA3Quant(param)) {1631 if (EnableFA3Quant(param)) {
1618 CHECK_OPERATION_STATUS_RETURN(AddKVHistoryDequantNode(param, opGraph, tensorMap));1632 CHECK_OPERATION_STATUS_RETURN(AddKVHistoryDequantNode(param, opGraph, tensorMap));
1619 }1633 }
1620- if (param.contextParallelInfo.IsEnabled()) {1634+ if (param.contextParallelInfo.IsEnabled() || param.hasAttnInnerSp) {
1621 CHECK_OPERATION_STATUS_RETURN(AddKAllGatherCpNode(param, opGraph, tensorMap, true));1635 CHECK_OPERATION_STATUS_RETURN(AddKAllGatherCpNode(param, opGraph, tensorMap, true));
1622 }1636 }
1623 CHECK_OPERATION_STATUS_RETURN(AddLAttnKProjBHistoryNode(param, opGraph, tensorMap));1637 CHECK_OPERATION_STATUS_RETURN(AddLAttnKProjBHistoryNode(param, opGraph, tensorMap));
@@ -1761,7 +1775,7 @@ atb::Status AddRingMLAEncoderNode(
1761 ringMLaParam.kvHeadNum = param.selfAttentionParam.kvHeadNum;1775 ringMLaParam.kvHeadNum = param.selfAttentionParam.kvHeadNum;
1762 ringMLaParam.qkScale = param.selfAttentionParam.qkScale;1776 ringMLaParam.qkScale = param.selfAttentionParam.qkScale;
1763 ringMLaParam.maskType = atb::infer::RingMLAParam::MaskType::MASK_TYPE_TRIU;1777 ringMLaParam.maskType = atb::infer::RingMLAParam::MaskType::MASK_TYPE_TRIU;
1764- 1778+ 
1765 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(ringMLaParam, &ringAttentionNode.operation));1779 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(ringMLaParam, &ringAttentionNode.operation));
1766 ringAttentionNode.inTensorIds = {1780 ringAttentionNode.inTensorIds = {
1767 GetTensorIdx(tensorMap, "nope_q"),1781 GetTensorIdx(tensorMap, "nope_q"),
Mexamples/atb_models/atb_llm/models/deepseekv2/flash_causal_deepseekv2.py+35-19
@@ -1413,30 +1413,43 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1413 1413 
1414 self.ring_cur_seqlen = torch.stack([q_seqlen, q_seqlen]).pin_memory()1414 self.ring_cur_seqlen = torch.stack([q_seqlen, q_seqlen]).pin_memory()
1415 self.ring_cur_seqlen = self.ring_cur_seqlen.to(self.device, non_blocking=True)1415 self.ring_cur_seqlen = self.ring_cur_seqlen.to(self.device, non_blocking=True)
1416- self.ring_cache_seqlen = torch.stack([q_seqlen, prefix_seqlen]).pin_memory()1416+ if self.mapping.has_attn_inner_sp() and not self.mapping.has_attn_cp():
1417- self.ring_cache_seqlen = self.ring_cache_seqlen.to(self.device, non_blocking=True)1417+ prefix_sp_lens = [i - j for i, j in zip(input_lens, q_lens)]
1418- self.acl_param = json.dumps({1418+ prefix_sp_seqlen = torch.from_numpy(np.array(prefix_sp_lens)).to(torch.int32)
1419- SEQUENCE_LENGTH: input_lens,1419+ self.ring_cache_seqlen = torch.stack([q_seqlen, prefix_sp_seqlen]).pin_memory()
1420- Q_LEN: q_lens if q_lens is not None else [],1420+ self.ring_cache_seqlen = self.ring_cache_seqlen.to(self.device, non_blocking=True)
1421- "ringCurSeqlen": q_lens + q_lens,1421+ self.acl_param = json.dumps({
1422- "ringCacheSeqlen": q_lens + self.prefix_lens,1422+ SEQUENCE_LENGTH: input_lens,
1423- })1423+ Q_LEN: q_lens if q_lens is not None else [],
1424+ "ringCurSeqlen": q_lens + q_lens,
1425+ "ringCacheSeqlen": q_lens + prefix_sp_lens,
1426+ })
1427+ else:
1428+ self.ring_cache_seqlen = torch.stack([q_seqlen, prefix_seqlen]).pin_memory()
1429+ self.ring_cache_seqlen = self.ring_cache_seqlen.to(self.device, non_blocking=True)
1430+ self.acl_param = json.dumps({
1431+ SEQUENCE_LENGTH: input_lens,
1432+ Q_LEN: q_lens if q_lens is not None else [],
1433+ "ringCurSeqlen": q_lens + q_lens,
1434+ "ringCacheSeqlen": q_lens + self.prefix_lens,
1435+ })
1424 1436 
1425 def prepare_paddingidx_for_prefixcache_contextparallel(self, **kwargs):1437 def prepare_paddingidx_for_prefixcache_contextparallel(self, **kwargs):
1426 sp_computed_slots_padding_idx = kwargs.get("sp_computed_slots_padding_idx", None)1438 sp_computed_slots_padding_idx = kwargs.get("sp_computed_slots_padding_idx", None)
1427 computed_slots_order = kwargs.get("sp_computed_slots_order", None)1439 computed_slots_order = kwargs.get("sp_computed_slots_order", None)
1428 self.kv_cache_padding_idx = sp_computed_slots_padding_idx1440 self.kv_cache_padding_idx = sp_computed_slots_padding_idx
1429 self.kv_cache_unpadding_idx = computed_slots_order1441 self.kv_cache_unpadding_idx = computed_slots_order
1430- acl_param = json.loads(self.acl_param)1442+ if self.mapping.has_attn_cp():
1431- q_lens = kwargs.get("q_lens", None)1443+ acl_param = json.loads(self.acl_param)
1432- chunk_lengths = [x // 2 for x in q_lens]1444+ q_lens = kwargs.get("q_lens", None)
1433- q_seqlen = torch.from_numpy(np.array(chunk_lengths)).to(torch.int32) # new tokens1445+ chunk_lengths = [x // 2 for x in q_lens]
1434- all_rank_prefix_lens = kwargs.get("all_rank_prefix_lens", None)1446+ q_seqlen = torch.from_numpy(np.array(chunk_lengths)).to(torch.int32) # new tokens
1435- all_rank_prefix_seqlen = torch.from_numpy(np.array(all_rank_prefix_lens)).to(torch.int32) # cache tokens1447+ all_rank_prefix_lens = kwargs.get("all_rank_prefix_lens", None)
1436- self.kv_cache_len = torch.stack([q_seqlen, all_rank_prefix_seqlen]).pin_memory()1448+ all_rank_prefix_seqlen = torch.from_numpy(np.array(all_rank_prefix_lens)).to(torch.int32) # cache tokens
1437- self.kv_cache_len = self.kv_cache_len.to(self.device, non_blocking=True)1449+ self.kv_cache_len = torch.stack([q_seqlen, all_rank_prefix_seqlen]).pin_memory()
1438- acl_param["kvCachelen"] = chunk_lengths + all_rank_prefix_lens1450+ self.kv_cache_len = self.kv_cache_len.to(self.device, non_blocking=True)
1439- self.acl_param = json.dumps(acl_param)1451+ acl_param["kvCachelen"] = chunk_lengths + all_rank_prefix_lens
1452+ self.acl_param = json.dumps(acl_param)
1440 1453 
1441 def prepare_paddingidx_for_contextparallel(self, input_ids):1454 def prepare_paddingidx_for_contextparallel(self, input_ids):
1442 input_length = len(input_ids) # The length of each sp_rank input sequence (batch*seq) is the same.1455 input_length = len(input_ids) # The length of each sp_rank input sequence (batch*seq) is the same.
@@ -1822,7 +1835,7 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1822 if self.mapping.has_attn_cp():1835 if self.mapping.has_attn_cp():
1823 self.prepare_cp_prefill_inputs(input_ids, input_lengths, q_lens)1836 self.prepare_cp_prefill_inputs(input_ids, input_lengths, q_lens)
1824 1837 
1825- if self.has_prefixcache and self.mapping.has_attn_cp():1838+ if self.has_prefixcache and (self.mapping.has_attn_cp() or self.mapping.has_attn_inner_sp()):
1826 self.prepare_paddingidx_for_prefixcache_contextparallel(**kwargs)1839 self.prepare_paddingidx_for_prefixcache_contextparallel(**kwargs)
1827 for b_i in range(block_tables.shape[0]): # 存在部分命中， 没有block的 rank 要添加一个用于只读的无用块，否则pagegloadcache算子会报错1840 for b_i in range(block_tables.shape[0]): # 存在部分命中， 没有block的 rank 要添加一个用于只读的无用块，否则pagegloadcache算子会报错
1828 if block_tables[b_i][0] == -1:1841 if block_tables[b_i][0] == -1:
@@ -1836,6 +1849,9 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1836 if self.has_prefixcache and self.mapping.has_attn_cp():1849 if self.has_prefixcache and self.mapping.has_attn_cp():
1837 self.acl_encoder_operation_inputs.extend([self.kv_cache_padding_idx, \1850 self.acl_encoder_operation_inputs.extend([self.kv_cache_padding_idx, \
1838 self.kv_cache_unpadding_idx, self.kv_cache_len])1851 self.kv_cache_unpadding_idx, self.kv_cache_len])
1852+ elif self.has_prefixcache and self.mapping.has_attn_inner_sp():
1853+ self.acl_encoder_operation_inputs.extend([self.kv_cache_padding_idx, \
1854+ self.kv_cache_unpadding_idx])
1839 if self.mapping.enable_dense_tp: # new padding idx please add here before1855 if self.mapping.enable_dense_tp: # new padding idx please add here before
1840 self.acl_encoder_operation_inputs.append(self.dense_tp_padding_idx)1856 self.acl_encoder_operation_inputs.append(self.dense_tp_padding_idx)
1841 self.acl_encoder_operation_inputs.append(self.dense_gather_mlpout_idx)1857 self.acl_encoder_operation_inputs.append(self.dense_gather_mlpout_idx)