Star137
243
代码
代码
Issues74
Pull Requests182
项目讨论
Wiki
分析
Star137
243
已合并
[atb-models] [bug-fix] Fix Quant+LoRA+FlashComm combination bugs #1080
gengli8创建于 6月27日
[atb-models] [bug-fix] Fix Quant+LoRA+FlashComm combination bugs #1080
已合并
gengli8创建于 6月27日
从forcekeng/MindIE-LLMfix_lora_flashcomm合入到Ascend/MindIE-LLMdev
共 1 个文件变更+119-144
Mexamples/atb_models/atb_framework/operations/fusion/linear/linear.cpp+119-144
@@ -10,62 +10,56 @@
10 * See the Mulan PSL v2 for more details.10 * See the Mulan PSL v2 for more details.
11 */11 */
12 12 
13+#include "operations/fusion/linear/linear.h"
14+ 
13#include <atb/atb_infer.h>15#include <atb/atb_infer.h>
16+ 
14#include <cmath>17#include <cmath>
18+ 
15#include "atb_speed/log.h"19#include "atb_speed/log.h"
16#include "atb_speed/utils/check_util.h"20#include "atb_speed/utils/check_util.h"
17-#include "operations/aclnn/ops/w8a16_operation.h"21+#include "operations/aclnn/ops/dynamic_quant_operation.h"
22+#include "operations/aclnn/ops/grouped_matmul_operation.h"
23+#include "operations/aclnn/ops/w16a16_operation.h"
24+#include "operations/aclnn/ops/w16a16sc_operation.h"
18#include "operations/aclnn/ops/w4a16_operation.h"25#include "operations/aclnn/ops/w4a16_operation.h"
19#include "operations/aclnn/ops/w4a8_operation.h"26#include "operations/aclnn/ops/w4a8_operation.h"
27+#include "operations/aclnn/ops/w8a16_operation.h"
20#include "operations/aclnn/ops/w8a8_operation.h"28#include "operations/aclnn/ops/w8a8_operation.h"
21-#include "operations/aclnn/ops/w16a16_operation.h"
22-#include "operations/aclnn/ops/grouped_matmul_operation.h"
23-#include "operations/aclnn/ops/dynamic_quant_operation.h"
24-#include "operations/aclnn/ops/w16a16sc_operation.h"
25#include "operations/aclnn/utils/utils.h"29#include "operations/aclnn/utils/utils.h"
26#include "operations/fusion/utils.h"30#include "operations/fusion/utils.h"
27-#include "operations/fusion/linear/linear.h"
28 31 
29namespace atb_speed {32namespace atb_speed {
30namespace common {33namespace common {
31 34 
32// 是否为matmulBackend开启下LINEAR_W8A8_QUANT、LINEAR_W8A8_DEQUANT场景35// 是否为matmulBackend开启下LINEAR_W8A8_QUANT、LINEAR_W8A8_DEQUANT场景
33-bool IsAclnnPerTensor(const FusionLinearParam &param)36+bool IsAclnnPerTensor(const FusionLinearParam &param) {
34-{
35 return param.matmulBackend == atb_speed::common::OpBackend::ACLNN &&37 return param.matmulBackend == atb_speed::common::OpBackend::ACLNN &&
36- (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DEQUANT);38+ (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DEQUANT);
37}39}
38 40 
39// 是否使用aclnn的QuantBatchMatmul接口41// 是否使用aclnn的QuantBatchMatmul接口
40-bool UseQuantBatchMatmul(const FusionLinearParam &param)42+bool UseQuantBatchMatmul(const FusionLinearParam &param) {
41-{
42 // All机型: dynamic、pdmix43 // All机型: dynamic、pdmix
43- return IsAclnnPerTensor(param) || \44+ return IsAclnnPerTensor(param) || param.quantType == LINEAR_W8A8_DYNAMIC_QUANT ||
44- param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || \45+ param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT ||
45- param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT || \
46- param.quantType == LINEAR_W4A8_DYNAMIC_QUANT || \
47 param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT;46 param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT;
48}47}
49 48 
50-std::map<std::string, std::vector<std::string>> GetLinearInTensorCandidates()49+std::map<std::string, std::vector<std::string>> GetLinearInTensorCandidates() {
51-{
52 std::map<std::string, std::vector<std::string>> linearInTensorCandidates = {50 std::map<std::string, std::vector<std::string>> linearInTensorCandidates = {
53- {"default", {51+ {"default", {"in_input", "in_weight", "in_scale", "in_offset", "in_descale", "in_bias", "in_compress_idx"}},
54- "in_input", "in_weight", "in_scale", "in_offset", "in_descale", "in_bias", "in_compress_idx"}
55- },
56 {"lora", {"in_group_list", "in_lora_a", "in_lora_b"}},52 {"lora", {"in_group_list", "in_lora_a", "in_lora_b"}},
57 {"lora_with_mask", {"in_im_mask"}},53 {"lora_with_mask", {"in_im_mask"}},
58 {"dynamic_quant", {"dynamic_input_scale"}},54 {"dynamic_quant", {"dynamic_input_scale"}},
59 {"add_swiglu_quant_sacle_in", {"swiglu_quant_input_scale"}},55 {"add_swiglu_quant_sacle_in", {"swiglu_quant_input_scale"}},
60- {"flash_comm", {56+ {"flash_comm",
61- "send_counts", "sdispls", "send_count", "recv_counts", "rdispls", "recv_count", "fake_ag_shape"}57+ {"send_counts", "sdispls", "send_count", "recv_counts", "rdispls", "recv_count", "fake_ag_shape"}},
62- },
63 };58 };
64 return linearInTensorCandidates;59 return linearInTensorCandidates;
65}60}
66 61 
67-std::map<std::string, std::vector<std::string>> GetLinearIntermediateTensorCandidates()62+std::map<std::string, std::vector<std::string>> GetLinearIntermediateTensorCandidates() {
68-{
69 std::map<std::string, std::vector<std::string>> linearIntermediateTensorCandidates = {63 std::map<std::string, std::vector<std::string>> linearIntermediateTensorCandidates = {
70 {"quant_input", {"intermediate_quant_input"}},64 {"quant_input", {"intermediate_quant_input"}},
71 {"lora", {"intermediate_base_linear_out", "intermediate_lora_a_out", "intermediate_lora_b_out"}},65 {"lora", {"intermediate_base_linear_out", "intermediate_lora_a_out", "intermediate_lora_b_out"}},
@@ -77,10 +71,8 @@ std::map<std::string, std::vector<std::string>> GetLinearIntermediateTensorCandi
77 return linearIntermediateTensorCandidates;71 return linearIntermediateTensorCandidates;
78}72}
79 73 
80-std::map<std::string, uint32_t> ConstructLinearTensorMap(74+std::map<std::string, uint32_t> ConstructLinearTensorMap(const FusionLinearParam &param, uint32_t &inTensorNum,
81- const FusionLinearParam &param,75+ uint32_t &outTensorNum, uint32_t &internalTensorNum) {
82- uint32_t &inTensorNum, uint32_t &outTensorNum, uint32_t &internalTensorNum)
83-{
84 auto linearInTensorCandidates = GetLinearInTensorCandidates();76 auto linearInTensorCandidates = GetLinearInTensorCandidates();
85 auto linearIntermediateTensorCandidates = GetLinearIntermediateTensorCandidates();77 auto linearIntermediateTensorCandidates = GetLinearIntermediateTensorCandidates();
86 78 
@@ -91,17 +83,17 @@ std::map<std::string, uint32_t> ConstructLinearTensorMap(
91 // 添加默认的Tensor83 // 添加默认的Tensor
92 AddTensorToList(linearInTensorCandidates, "default", inTensorList);84 AddTensorToList(linearInTensorCandidates, "default", inTensorList);
93 85 
94- if (!param.enableSwigluQuant || (param.quantType != LINEAR_W8A8_DYNAMIC_DEQUANT86+ if (!param.enableSwigluQuant ||
95- && param.quantType != LINEAR_W4A8_DYNAMIC_DEQUANT)) {87+ (param.quantType != LINEAR_W8A8_DYNAMIC_DEQUANT && param.quantType != LINEAR_W4A8_DYNAMIC_DEQUANT)) {
96 // 添加额外的中间Tensor88 // 添加额外的中间Tensor
97- if (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_SC_QUANT89+ if (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_SC_QUANT ||
98- || ((param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT)90+ ((param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) &&
99- && !param.enableSwiGLUQuantForSharedExperts)) {91+ !param.enableSwiGLUQuantForSharedExperts)) {
100 AddTensorToList(linearIntermediateTensorCandidates, "quant_input", intermediateTensorList);92 AddTensorToList(linearIntermediateTensorCandidates, "quant_input", intermediateTensorList);
101 }93 }
102 // 添加动态量化中间Tensor94 // 添加动态量化中间Tensor
103- if ((param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT)95+ if ((param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) &&
104- && !param.enableSwiGLUQuantForSharedExperts) {96+ !param.enableSwiGLUQuantForSharedExperts) {
105 AddTensorToList(linearIntermediateTensorCandidates, "dynamic_quant", intermediateTensorList);97 AddTensorToList(linearIntermediateTensorCandidates, "dynamic_quant", intermediateTensorList);
106 }98 }
107 }99 }
@@ -119,8 +111,7 @@ std::map<std::string, uint32_t> ConstructLinearTensorMap(
119 AddTensorToList(linearInTensorCandidates, "flash_comm", inTensorList);111 AddTensorToList(linearInTensorCandidates, "flash_comm", inTensorList);
120 AddTensorToList(linearIntermediateTensorCandidates, "flashComm", intermediateTensorList);112 AddTensorToList(linearIntermediateTensorCandidates, "flashComm", intermediateTensorList);
121 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) {113 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) {
122- AddTensorToList(linearIntermediateTensorCandidates, "flashComm_dynamic_quant",114+ AddTensorToList(linearIntermediateTensorCandidates, "flashComm_dynamic_quant", intermediateTensorList);
123- intermediateTensorList);
124 }115 }
125 }116 }
126 inTensorNum = inTensorList.size();117 inTensorNum = inTensorList.size();
@@ -131,8 +122,7 @@ std::map<std::string, uint32_t> ConstructLinearTensorMap(
131}122}
132 123 
133int64_t AddElewiseQuant(atb::GraphParam &opGraph, const FusionLinearParam &param,124int64_t AddElewiseQuant(atb::GraphParam &opGraph, const FusionLinearParam &param,
134- std::map<std::string, uint32_t> &tensorMap)125+ std::map<std::string, uint32_t> &tensorMap) {
135-{
136 if (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_SC_QUANT) {126 if (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_SC_QUANT) {
137 // quant127 // quant
138 atb::Node inputQuantNode;128 atb::Node inputQuantNode;
@@ -146,8 +136,8 @@ int64_t AddElewiseQuant(atb::GraphParam &opGraph, const FusionLinearParam &param
146 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) {136 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) {
147 atb::Node inputDynamicQuantNode;137 atb::Node inputDynamicQuantNode;
148 inputDynamicQuantNode.inTensorIds = GetTensorIdxList(tensorMap, {"in_input"});138 inputDynamicQuantNode.inTensorIds = GetTensorIdxList(tensorMap, {"in_input"});
149- inputDynamicQuantNode.outTensorIds = GetTensorIdxList(tensorMap, {"intermediate_quant_input",139+ inputDynamicQuantNode.outTensorIds =
150- "intermediate_input_scale"});140+ GetTensorIdxList(tensorMap, {"intermediate_quant_input", "intermediate_input_scale"});
151 inputDynamicQuantNode.operation = new atb_speed::common::DynamicQuantOperation("DynamicQuantNode");141 inputDynamicQuantNode.operation = new atb_speed::common::DynamicQuantOperation("DynamicQuantNode");
152 opGraph.nodes.push_back(inputDynamicQuantNode);142 opGraph.nodes.push_back(inputDynamicQuantNode);
153 }143 }
@@ -155,17 +145,19 @@ int64_t AddElewiseQuant(atb::GraphParam &opGraph, const FusionLinearParam &param
155}145}
156 146 
157int64_t AddAllGather(atb::GraphParam &opGraph, const FusionLinearParam &param,147int64_t AddAllGather(atb::GraphParam &opGraph, const FusionLinearParam &param,
158- std::map<std::string, uint32_t> &tensorMap)148+ std::map<std::string, uint32_t> &tensorMap) {
159-{
160 atb::Node allGatherVNode;149 atb::Node allGatherVNode;
161 atb::infer::AllGatherVParam allGatherVParam;150 atb::infer::AllGatherVParam allGatherVParam;
162 allGatherVParam.rank = param.flashCommParallelInfo.rank;151 allGatherVParam.rank = param.flashCommParallelInfo.rank;
163 allGatherVParam.rankSize = param.flashCommParallelInfo.worldSize;152 allGatherVParam.rankSize = param.flashCommParallelInfo.worldSize;
164 allGatherVParam.backend = param.flashCommParallelInfo.backend;153 allGatherVParam.backend = param.flashCommParallelInfo.backend;
165 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(allGatherVParam, &allGatherVNode.operation));154 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(allGatherVParam, &allGatherVNode.operation));
166- allGatherVNode.inTensorIds = {GetTensorIdx(155+ allGatherVNode.inTensorIds = {
167- tensorMap, (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_QUANT156+ GetTensorIdx(tensorMap, (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_QUANT ||
168- || param.quantType == LINEAR_W8A8_SC_QUANT) ? "intermediate_quant_input" : "in_input")};157+ param.quantType == LINEAR_W8A8_SC_QUANT) &&
158+ !param.supportLora
159+ ? "intermediate_quant_input"
160+ : "in_input")};
169 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "recv_count"));161 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "recv_count"));
170 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "send_counts"));162 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "send_counts"));
171 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "sdispls"));163 allGatherVNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "sdispls"));
@@ -176,12 +168,13 @@ int64_t AddAllGather(atb::GraphParam &opGraph, const FusionLinearParam &param,
176 opGraph.nodes.push_back(allGatherVNode);168 opGraph.nodes.push_back(allGatherVNode);
177 CHECK_OPERATION_STATUS_RETURN(common::AddDapEventsAfterComm(opGraph));169 CHECK_OPERATION_STATUS_RETURN(common::AddDapEventsAfterComm(opGraph));
178 170 
179- if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) {171+ if ((param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) &&
172+ !param.supportLora) {
180 atb::Node allGatherInputScaleNode;173 atb::Node allGatherInputScaleNode;
181 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(allGatherVParam, &allGatherInputScaleNode.operation));174 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(allGatherVParam, &allGatherInputScaleNode.operation));
182- allGatherInputScaleNode.inTensorIds = {GetTensorIdx(175+ allGatherInputScaleNode.inTensorIds = {GetTensorIdx(tensorMap, param.quantType == LINEAR_W8A8_DYNAMIC_QUANT
183- tensorMap, param.quantType == LINEAR_W8A8_DYNAMIC_QUANT176+ ? "intermediate_input_scale"
184- ? "intermediate_input_scale" : "dynamic_input_scale")};177+ : "dynamic_input_scale")};
185 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "send_count"));178 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "send_count"));
186 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "recv_counts"));179 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "recv_counts"));
187 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "rdispls"));180 allGatherInputScaleNode.inTensorIds.push_back(GetTensorIdx(tensorMap, "rdispls"));
@@ -195,12 +188,10 @@ int64_t AddAllGather(atb::GraphParam &opGraph, const FusionLinearParam &param,
195}188}
196 189 
197int64_t AddAclNNWeightQuantBatchMatmul(atb::Node &linearNode, const FusionLinearParam &param,190int64_t AddAclNNWeightQuantBatchMatmul(atb::Node &linearNode, const FusionLinearParam &param,
198- std::map<std::string, uint32_t> &tensorMap)191+ std::map<std::string, uint32_t> &tensorMap) {
199-{192+ linearNode.inTensorIds = GetTensorIdxList(
200- linearNode.inTensorIds = GetTensorIdxList(tensorMap, {193+ tensorMap,
201- param.enableFlashComm ? "intermediate_allgather_out" : "in_input",194+ {param.enableFlashComm ? "intermediate_allgather_out" : "in_input", "in_weight", "in_scale", "in_offset"});
202- "in_weight", "in_scale", "in_offset"
203- });
204 AclNNWeightQuantBatchMatmulParam aclnnParam;195 AclNNWeightQuantBatchMatmulParam aclnnParam;
205 aclnnParam.transposeB = param.transposeType == TRANSPOSE;196 aclnnParam.transposeB = param.transposeType == TRANSPOSE;
206 if (param.hasBias) {197 if (param.hasBias) {
@@ -221,23 +212,22 @@ int64_t AddAclNNWeightQuantBatchMatmul(atb::Node &linearNode, const FusionLinear
221}212}
222 213 
223int64_t AddW4A8Matmul(atb::Node &linearNode, const FusionLinearParam &param,214int64_t AddW4A8Matmul(atb::Node &linearNode, const FusionLinearParam &param,
224- std::map<std::string, uint32_t> &tensorMap)215+ std::map<std::string, uint32_t> &tensorMap) {
225-{
226 const bool containingQuant = param.quantType == LINEAR_W4A8_DYNAMIC_QUANT;216 const bool containingQuant = param.quantType == LINEAR_W4A8_DYNAMIC_QUANT;
227 AclNNW4A8Param aclnnParam;217 AclNNW4A8Param aclnnParam;
228 std::string key;218 std::string key;
229 if (param.enableSwigluQuant) {219 if (param.enableSwigluQuant) {
230 key = "in_input";220 key = "in_input";
231 } else {221 } else {
232- key = (containingQuant && !param.enableSwiGLUQuantForSharedExperts) ?222+ key = (containingQuant && !param.enableSwiGLUQuantForSharedExperts) ? "intermediate_quant_input" : "in_input";
233- "intermediate_quant_input" : "in_input";
234 }223 }
235 std::string inputScaleKey;224 std::string inputScaleKey;
236 if (param.enableSwigluQuant && containingQuant) {225 if (param.enableSwigluQuant && containingQuant) {
237 inputScaleKey = "intermediate_quant_input_scale";226 inputScaleKey = "intermediate_quant_input_scale";
238 } else {227 } else {
239- inputScaleKey = !containingQuant ? "dynamic_input_scale" : param.enableSwiGLUQuantForSharedExperts ?228+ inputScaleKey = !containingQuant ? "dynamic_input_scale"
240- "swiglu_quant_input_scale" : "intermediate_input_scale";229+ : param.enableSwiGLUQuantForSharedExperts ? "swiglu_quant_input_scale"
230+ : "intermediate_input_scale";
241 }231 }
242 std::vector<std::string> tensorNames = {key, "in_weight", inputScaleKey, "in_scale", "in_bias"};232 std::vector<std::string> tensorNames = {key, "in_weight", inputScaleKey, "in_scale", "in_bias"};
243 linearNode.inTensorIds = GetTensorIdxList(tensorMap, tensorNames);233 linearNode.inTensorIds = GetTensorIdxList(tensorMap, tensorNames);
@@ -249,8 +239,7 @@ int64_t AddW4A8Matmul(atb::Node &linearNode, const FusionLinearParam &param,
249}239}
250 240 
251int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &param,241int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &param,
252- std::map<std::string, uint32_t> &tensorMap)242+ std::map<std::string, uint32_t> &tensorMap) {
253-{
254 AclNNQuantMatmulParam aclnnQuantMatmulParam;243 AclNNQuantMatmulParam aclnnQuantMatmulParam;
255 aclnnQuantMatmulParam.transposeB = param.transposeType == TRANSPOSE;244 aclnnQuantMatmulParam.transposeB = param.transposeType == TRANSPOSE;
256 std::string key;245 std::string key;
@@ -258,19 +247,19 @@ int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &para
258 key = "intermediate_allgather_out";247 key = "intermediate_allgather_out";
259 } else {248 } else {
260 key = (param.quantType == LINEAR_W8A8_QUANT ||249 key = (param.quantType == LINEAR_W8A8_QUANT ||
261- (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT250+ (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT && !param.enableSwiGLUQuantForSharedExperts))
262- && !param.enableSwiGLUQuantForSharedExperts)) ?251+ ? "intermediate_quant_input"
263- "intermediate_quant_input" : "in_input";252+ : "in_input";
264 }253 }
265- std::string inScaleKey = (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DEQUANT) ?254+ std::string inScaleKey =
266- "in_descale" : "in_scale";255+ (param.quantType == LINEAR_W8A8_QUANT || param.quantType == LINEAR_W8A8_DEQUANT) ? "in_descale" : "in_scale";
267 std::string inputScaleKey;256 std::string inputScaleKey;
268 if (param.enableFlashComm) {257 if (param.enableFlashComm) {
269 inputScaleKey = "intermediate_allgather_input_scale_out";258 inputScaleKey = "intermediate_allgather_input_scale_out";
270 } else {259 } else {
271- inputScaleKey = param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT ?260+ inputScaleKey = param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT ? "dynamic_input_scale"
272- "dynamic_input_scale" : param.enableSwiGLUQuantForSharedExperts ? "swiglu_quant_input_scale" :261+ : param.enableSwiGLUQuantForSharedExperts ? "swiglu_quant_input_scale"
273- "intermediate_input_scale";262+ : "intermediate_input_scale";
274 }263 }
275 std::vector<std::string> tensorNames = {key, "in_weight", inScaleKey};264 std::vector<std::string> tensorNames = {key, "in_weight", inScaleKey};
276 // per token265 // per token
@@ -287,8 +276,8 @@ int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &para
287 linearNode.inTensorIds = GetTensorIdxList(tensorMap, tensorNames);276 linearNode.inTensorIds = GetTensorIdxList(tensorMap, tensorNames);
288 ATB_SPEED_LOG_DEBUG("tensorNames: " << tensorNames << "; inTensorIds: " << linearNode.inTensorIds);277 ATB_SPEED_LOG_DEBUG("tensorNames: " << tensorNames << "; inTensorIds: " << linearNode.inTensorIds);
289 linearNode.inTensorReshapeFuncs.resize(linearNode.inTensorIds.size());278 linearNode.inTensorReshapeFuncs.resize(linearNode.inTensorIds.size());
290- linearNode.inTensorReshapeFuncs[0] = [=](const atb::Dims &oldShape, atb::Dims &newShape) { // 1: input279+ linearNode.inTensorReshapeFuncs[0] = [=](const atb::Dims &oldShape, atb::Dims &newShape) { // 1: input
291- newShape.dimNum = 2; // dimNum: 2280+ newShape.dimNum = 2; // dimNum: 2
292 // 开启TURBO_ATTN时, w8a8_pdmix或pertoken场景, canndev算子, input会多出一维(2维)281 // 开启TURBO_ATTN时, w8a8_pdmix或pertoken场景, canndev算子, input会多出一维(2维)
293 if (oldShape.dimNum == NUM3) {282 if (oldShape.dimNum == NUM3) {
294 newShape.dims[DIM0] = oldShape.dims[DIM0] * oldShape.dims[DIM1];283 newShape.dims[DIM0] = oldShape.dims[DIM0] * oldShape.dims[DIM1];
@@ -297,8 +286,8 @@ int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &para
297 };286 };
298 // dynamic的inputScaleKey转换287 // dynamic的inputScaleKey转换
299 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) {288 if (param.quantType == LINEAR_W8A8_DYNAMIC_QUANT || param.quantType == LINEAR_W8A8_DYNAMIC_DEQUANT) {
300- linearNode.inTensorReshapeFuncs[3] = [=](const atb::Dims &oldShape, atb::Dims &newShape) { // 3: 3号scale289+ linearNode.inTensorReshapeFuncs[3] = [=](const atb::Dims &oldShape, atb::Dims &newShape) { // 3: 3号scale
301- newShape.dimNum = 1; // dimNum: 1290+ newShape.dimNum = 1; // dimNum: 1
302 // 开启TURBO_ATTN时, canndev算子, scale会多出一维(2维)291 // 开启TURBO_ATTN时, canndev算子, scale会多出一维(2维)
303 newShape.dims[0] = oldShape.dimNum == NUM2 ? oldShape.dims[0] * oldShape.dims[1] : oldShape.dims[0];292 newShape.dims[0] = oldShape.dimNum == NUM2 ? oldShape.dims[0] * oldShape.dims[1] : oldShape.dims[0];
304 };293 };
@@ -310,11 +299,9 @@ int64_t AddAclNNQuantMatmul(atb::Node &linearNode, const FusionLinearParam &para
310}299}
311 300 
312int64_t AddAclNNMatmul(atb::Node &linearNode, const FusionLinearParam &param,301int64_t AddAclNNMatmul(atb::Node &linearNode, const FusionLinearParam &param,
313- std::map<std::string, uint32_t> &tensorMap)302+ std::map<std::string, uint32_t> &tensorMap) {
314-{303+ linearNode.inTensorIds =
315- linearNode.inTensorIds = GetTensorIdxList(tensorMap, {304+ GetTensorIdxList(tensorMap, {(param.enableFlashComm) ? "intermediate_allgather_out" : "in_input", "in_weight"});
316- (param.enableFlashComm) ?
317- "intermediate_allgather_out" : "in_input", "in_weight"});
318 AclNNMatmulParam aclnnMatmulParam;305 AclNNMatmulParam aclnnMatmulParam;
319 aclnnMatmulParam.transposeB = param.transposeType == TRANSPOSE;306 aclnnMatmulParam.transposeB = param.transposeType == TRANSPOSE;
320 if (param.hasBias) {307 if (param.hasBias) {
@@ -326,19 +313,17 @@ int64_t AddAclNNMatmul(atb::Node &linearNode, const FusionLinearParam &param,
326}313}
327 314 
328int64_t AddW16A16SCMatmul(atb::Node &linearNode, const FusionLinearParam &param,315int64_t AddW16A16SCMatmul(atb::Node &linearNode, const FusionLinearParam &param,
329- std::map<std::string, uint32_t> &tensorMap)316+ std::map<std::string, uint32_t> &tensorMap) {
330-{317+ linearNode.inTensorIds =
331- linearNode.inTensorIds = GetTensorIdxList(tensorMap, {318+ GetTensorIdxList(tensorMap, {(param.enableFlashComm) ? "intermediate_allgather_out" : "in_input", "in_weight",
332- (param.enableFlashComm) ?319+ "in_bias", "in_compress_idx"});
333- "intermediate_allgather_out" : "in_input", "in_weight", "in_bias", "in_compress_idx"});
334 AclNNW16A16SCParam aclnnw16a16scParam;320 AclNNW16A16SCParam aclnnw16a16scParam;
335 linearNode.operation = new atb_speed::common::W16A16SCOperation("W16A16SCLinearNode", aclnnw16a16scParam);321 linearNode.operation = new atb_speed::common::W16A16SCOperation("W16A16SCLinearNode", aclnnw16a16scParam);
336 return atb::NO_ERROR;322 return atb::NO_ERROR;
337}323}
338 324 
339int64_t AddAclNNLinear(atb::Node &linearNode, const FusionLinearParam &param,325int64_t AddAclNNLinear(atb::Node &linearNode, const FusionLinearParam &param,
340- std::map<std::string, uint32_t> &tensorMap)326+ std::map<std::string, uint32_t> &tensorMap) {
341-{
342 if (param.quantType == LINEAR_W4A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT) {327 if (param.quantType == LINEAR_W4A8_DYNAMIC_QUANT || param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT) {
343 CHECK_OPERATION_STATUS_RETURN(AddW4A8Matmul(linearNode, param, tensorMap));328 CHECK_OPERATION_STATUS_RETURN(AddW4A8Matmul(linearNode, param, tensorMap));
344 return atb::NO_ERROR;329 return atb::NO_ERROR;
@@ -361,13 +346,12 @@ int64_t AddAclNNLinear(atb::Node &linearNode, const FusionLinearParam &param,
361 CHECK_OPERATION_STATUS_RETURN(AddW16A16SCMatmul(linearNode, param, tensorMap));346 CHECK_OPERATION_STATUS_RETURN(AddW16A16SCMatmul(linearNode, param, tensorMap));
362 return atb::NO_ERROR;347 return atb::NO_ERROR;
363 }348 }
364- 349+ 
365 return atb::NO_ERROR;350 return atb::NO_ERROR;
366}351}
367 352 
368int64_t AddLinear(atb::GraphParam &opGraph, const FusionLinearParam &param,353int64_t AddLinear(atb::GraphParam &opGraph, const FusionLinearParam &param,
369- std::map<std::string, uint32_t> &tensorMap)354+ std::map<std::string, uint32_t> &tensorMap) {
370-{
371 atb::Node linearNode;355 atb::Node linearNode;
372 atb::infer::LinearParam linearParam;356 atb::infer::LinearParam linearParam;
373 int matmulBackend = param.matmulBackend;357 int matmulBackend = param.matmulBackend;
@@ -397,9 +381,8 @@ int64_t AddLinear(atb::GraphParam &opGraph, const FusionLinearParam &param,
397 } else {381 } else {
398 key = param.quantType == LINEAR_W8A8_SC_DEQUANT ? "in_input" : "intermediate_quant_input";382 key = param.quantType == LINEAR_W8A8_SC_DEQUANT ? "in_input" : "intermediate_quant_input";
399 }383 }
400- linearNode.inTensorIds = GetTensorIdxList(tensorMap, {384+ linearNode.inTensorIds =
401- key, "in_weight", "in_bias", "in_descale", "in_compress_idx"385+ GetTensorIdxList(tensorMap, {key, "in_weight", "in_bias", "in_descale", "in_compress_idx"});
402- });
403 opGraph.nodes.push_back(linearNode);386 opGraph.nodes.push_back(linearNode);
404 return atb::NO_ERROR;387 return atb::NO_ERROR;
405 }388 }
@@ -431,31 +414,30 @@ int64_t AddLinear(atb::GraphParam &opGraph, const FusionLinearParam &param,
431 linearNode.inTensorIds = GetTensorIdxList(tensorMap, {key, "in_weight"});414 linearNode.inTensorIds = GetTensorIdxList(tensorMap, {key, "in_weight"});
432 } else {415 } else {
433 linearParam.hasBias = true;416 linearParam.hasBias = true;
434- linearNode.inTensorIds = GetTensorIdxList(tensorMap, {417+ linearNode.inTensorIds = GetTensorIdxList(tensorMap, {key, "in_weight", "in_bias", "in_descale"});
435- key, "in_weight", "in_bias", "in_descale"
436- });
437 }418 }
438 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(linearParam, &linearNode.operation));419 CHECK_OPERATION_STATUS_RETURN(atb::CreateOperation(linearParam, &linearNode.operation));
439 } else {420 } else {
440 // AclNN Linear (NO_QUANT)421 // AclNN Linear (NO_QUANT)
441 CHECK_OPERATION_STATUS_RETURN(AddAclNNLinear(linearNode, param, tensorMap));422 CHECK_OPERATION_STATUS_RETURN(AddAclNNLinear(linearNode, param, tensorMap));
442 }423 }
443- 424+ 
444 opGraph.nodes.push_back(linearNode);425 opGraph.nodes.push_back(linearNode);
445 426 
446 return atb::NO_ERROR;427 return atb::NO_ERROR;
447}428}
448 429 
449-atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **operation)430+atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **operation) {
450-{
451 atb::GraphParam opGraph;431 atb::GraphParam opGraph;
452- opGraph.name = param.quantType == NO_QUANT ? "LinearNoQuant" : \432+ opGraph.name = param.quantType == NO_QUANT ? "LinearNoQuant"
453- param.quantType == LINEAR_W8A8_DEQUANT || param.quantType == LINEAR_W8A8_SC_DEQUANT ? "LinearDequantOnly" : \433+ : param.quantType == LINEAR_W8A8_DEQUANT || param.quantType == LINEAR_W8A8_SC_DEQUANT
454- param.quantType == LINEAR_W16A16_SC ? "LinearW16A16SC" : \434+ ? "LinearDequantOnly"
455- param.quantType == W8A16 ? "LinearW8A16" : \435+ : param.quantType == LINEAR_W16A16_SC ? "LinearW16A16SC"
456- param.quantType == W4A16 ? "LinearW4A16" : "LinearQuant";436+ : param.quantType == W8A16 ? "LinearW8A16"
457- std::map<std::string, uint32_t> tensorMap = ConstructLinearTensorMap(437+ : param.quantType == W4A16 ? "LinearW4A16"
458- param, opGraph.inTensorNum, opGraph.outTensorNum, opGraph.internalTensorNum);438+ : "LinearQuant";
439+ std::map<std::string, uint32_t> tensorMap =
440+ ConstructLinearTensorMap(param, opGraph.inTensorNum, opGraph.outTensorNum, opGraph.internalTensorNum);
459 441 
460 if (param.transposeType == TRANSPOSE_INVALID) {442 if (param.transposeType == TRANSPOSE_INVALID) {
461 ATB_SPEED_LOG_ERROR("param.transposeType is invalid");443 ATB_SPEED_LOG_ERROR("param.transposeType is invalid");
@@ -463,8 +445,9 @@ atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **
463 }445 }
464 // dense层: enableSwiGLUQuantForSharedExperts 不开446 // dense层: enableSwiGLUQuantForSharedExperts 不开
465 // down层: 1) 不开 2) 开、非down 3)开、down、非DYNAMIC_DEQUANT447 // down层: 1) 不开 2) 开、非down 3)开、down、非DYNAMIC_DEQUANT
466- if (!param.enableSwiGLUQuantForSharedExperts && (!param.enableSwigluQuant \448+ if (!param.enableSwiGLUQuantForSharedExperts &&
467- || (param.quantType != LINEAR_W8A8_DYNAMIC_DEQUANT && param.quantType != LINEAR_W4A8_DYNAMIC_DEQUANT))) {449+ (!param.enableSwigluQuant ||
450+ (param.quantType != LINEAR_W8A8_DYNAMIC_DEQUANT && param.quantType != LINEAR_W4A8_DYNAMIC_DEQUANT))) {
468 CHECK_OPERATION_STATUS_RETURN(AddElewiseQuant(opGraph, param, tensorMap));451 CHECK_OPERATION_STATUS_RETURN(AddElewiseQuant(opGraph, param, tensorMap));
469 if (param.enableFlashComm) {452 if (param.enableFlashComm) {
470 CHECK_OPERATION_STATUS_RETURN(AddAllGather(opGraph, param, tensorMap));453 CHECK_OPERATION_STATUS_RETURN(AddAllGather(opGraph, param, tensorMap));
@@ -473,8 +456,8 @@ atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **
473 if (param.enableCVOverlap) {456 if (param.enableCVOverlap) {
474 CHECK_OPERATION_STATUS_RETURN(atb_speed::common::CreateRecordWithoutNodeId(457 CHECK_OPERATION_STATUS_RETURN(atb_speed::common::CreateRecordWithoutNodeId(
475 opGraph, atb_speed::EventAction::PUSH, atb_speed::common::VECTOR_CONTROL));458 opGraph, atb_speed::EventAction::PUSH, atb_speed::common::VECTOR_CONTROL));
476- CHECK_OPERATION_STATUS_RETURN(atb_speed::common::CreateWaitWithoutNodeId(459+ CHECK_OPERATION_STATUS_RETURN(atb_speed::common::CreateWaitWithoutNodeId(opGraph, atb_speed::EventAction::PUSH,
477- opGraph, atb_speed::EventAction::PUSH, atb_speed::common::CUBE_CONTROL));460+ atb_speed::common::CUBE_CONTROL));
478 }461 }
479 CHECK_OPERATION_STATUS_RETURN(AddLinear(opGraph, param, tensorMap));462 CHECK_OPERATION_STATUS_RETURN(AddLinear(opGraph, param, tensorMap));
480 463 
@@ -499,21 +482,22 @@ atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **
499 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(biasIdx).shape.dims[0];482 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(biasIdx).shape.dims[0];
500 } else if (param.quantType == W4A16) {483 } else if (param.quantType == W4A16) {
501 if (param.transposeType == TransposeType::TRANSPOSE) {484 if (param.transposeType == TransposeType::TRANSPOSE) {
502- outTensorDescs.at(0).shape.dims[outDimSize - 1] = \485+ outTensorDescs.at(0).shape.dims[outDimSize - 1] =
503 inTensorDescs.at(weightIdx).shape.dims[0]; // 0: n维shape486 inTensorDescs.at(weightIdx).shape.dims[0]; // 0: n维shape
504 } else {487 } else {
505- outTensorDescs.at(0).shape.dims[outDimSize - 1] = \488+ outTensorDescs.at(0).shape.dims[outDimSize - 1] =
506 CheckIntMulOverFlow(inTensorDescs.at(weightIdx).shape.dims[1], 2); // 1, 2: 最后一维shape * 2489 CheckIntMulOverFlow(inTensorDescs.at(weightIdx).shape.dims[1], 2); // 1, 2: 最后一维shape * 2
507 }490 }
508 } else if (param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) {491 } else if (param.quantType == LINEAR_W4A8_DYNAMIC_DEQUANT || param.quantType == LINEAR_W4A8_DYNAMIC_QUANT) {
509- outTensorDescs.at(0).shape.dims[outDimSize - 1] = \492+ outTensorDescs.at(0).shape.dims[outDimSize - 1] =
510 CheckIntMulOverFlow(inTensorDescs.at(weightIdx).shape.dims[1], 8); // 8: [m, k] @ [k, n//8] -> [m, n]493 CheckIntMulOverFlow(inTensorDescs.at(weightIdx).shape.dims[1], 8); // 8: [m, k] @ [k, n//8] -> [m, n]
511- } else if (inTensorDescs.at(weightIdx).shape.dimNum == 3) { // 3: dimNum494+ } else if (inTensorDescs.at(weightIdx).shape.dimNum == 3) { // 3: dimNum
512 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(weightIdx).shape.dims[nDim + 1];495 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(weightIdx).shape.dims[nDim + 1];
513- } else if (param.enEin && inTensorDescs.at(weightIdx).shape.dimNum == 4) { // 4: dimNum496+ } else if (param.enEin && inTensorDescs.at(weightIdx).shape.dimNum == 4) { // 4: dimNum
514- outTensorDescs.at(0).shape.dims[outDimSize - 1] = param.transposeType == TransposeType::TRANSPOSE ? \497+ outTensorDescs.at(0).shape.dims[outDimSize - 1] =
515- inTensorDescs.at(weightIdx).shape.dims[2] : // 2: dimNum498+ param.transposeType == TransposeType::TRANSPOSE ? inTensorDescs.at(weightIdx).shape.dims[2]
516- inTensorDescs.at(weightIdx).shape.dims[1] * inTensorDescs.at(weightIdx).shape.dims[3]; // 3: dimNum499+ : // 2: dimNum
500+ inTensorDescs.at(weightIdx).shape.dims[1] * inTensorDescs.at(weightIdx).shape.dims[3]; // 3: dimNum
517 } else {501 } else {
518 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(weightIdx).shape.dims[nDim];502 outTensorDescs.at(0).shape.dims[outDimSize - 1] = inTensorDescs.at(weightIdx).shape.dims[nDim];
519 }503 }
@@ -524,10 +508,8 @@ atb::Status CreateFusionLinear(const FusionLinearParam &param, atb::Operation **
524 return atb::NO_ERROR;508 return atb::NO_ERROR;
525}509}
526 510 
527-std::map<std::string, uint32_t> ConstructLinearWithLoraTensorMap(511+std::map<std::string, uint32_t> ConstructLinearWithLoraTensorMap(const FusionLinearParam &param, uint32_t &inTensorNum,
528- const FusionLinearParam &param,512+ uint32_t &outTensorNum, uint32_t &internalTensorNum) {
529- uint32_t &inTensorNum, uint32_t &outTensorNum, uint32_t &internalTensorNum)
530-{
531 auto linearInTensorCandidates = GetLinearInTensorCandidates();513 auto linearInTensorCandidates = GetLinearInTensorCandidates();
532 auto linearIntermediateTensorCandidates = GetLinearIntermediateTensorCandidates();514 auto linearIntermediateTensorCandidates = GetLinearIntermediateTensorCandidates();
533 515 
@@ -561,8 +543,7 @@ std::map<std::string, uint32_t> ConstructLinearWithLoraTensorMap(
561 return GetTensorMap(inTensorList, outTensorList, intermediateTensorList);543 return GetTensorMap(inTensorList, outTensorList, intermediateTensorList);
562}544}
563 545 
564-int64_t AddImMask(atb::GraphParam &opGraph, std::map<std::string, uint32_t> &tensorMap)546+int64_t AddImMask(atb::GraphParam &opGraph, std::map<std::string, uint32_t> &tensorMap) {
565-{
566 atb::Node mulNode;547 atb::Node mulNode;
567 atb::infer::ElewiseParam mulParam;548 atb::infer::ElewiseParam mulParam;
568 mulParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_MUL;549 mulParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_MUL;
@@ -573,9 +554,8 @@ int64_t AddImMask(atb::GraphParam &opGraph, std::map<std::string, uint32_t> &ten
573 return atb::NO_ERROR;554 return atb::NO_ERROR;
574}555}
575 556 
576-int64_t AddLoraA(atb::GraphParam &opGraph, const FusionLinearParam &param,557+int64_t AddLoraA(atb::GraphParam &opGraph, const FusionLinearParam &param, std::map<std::string, uint32_t> &tensorMap,
577- std::map<std::string, uint32_t> &tensorMap, bool enableFlashComm)558+ bool enableFlashComm) {
578-{
579 // 添加Lora A559 // 添加Lora A
580 atb::Node loraALinearNode;560 atb::Node loraALinearNode;
581 if (param.loraEnableGMM) {561 if (param.loraEnableGMM) {
@@ -607,9 +587,7 @@ int64_t AddLoraA(atb::GraphParam &opGraph, const FusionLinearParam &param,
607 return atb::NO_ERROR;587 return atb::NO_ERROR;
608}588}
609 589 
610-int64_t AddLoraB(atb::GraphParam &opGraph, const FusionLinearParam &param,590+int64_t AddLoraB(atb::GraphParam &opGraph, const FusionLinearParam &param, std::map<std::string, uint32_t> &tensorMap) {
611- std::map<std::string, uint32_t> &tensorMap)
612-{
613 // 添加Lora B591 // 添加Lora B
614 atb::Node loraBLinearNode;592 atb::Node loraBLinearNode;
615 if (param.loraEnableGMM) {593 if (param.loraEnableGMM) {
@@ -635,11 +613,10 @@ int64_t AddLoraB(atb::GraphParam &opGraph, const FusionLinearParam &param,
635 return atb::NO_ERROR;613 return atb::NO_ERROR;
636}614}
637 615 
638-atb::Status CreateFusionLinearWithLora(const FusionLinearParam &param, atb::Operation **operation)616+atb::Status CreateFusionLinearWithLora(const FusionLinearParam &param, atb::Operation **operation) {
639-{
640 atb::GraphParam opGraph;617 atb::GraphParam opGraph;
641- std::map<std::string, uint32_t> tensorMap = ConstructLinearWithLoraTensorMap(618+ std::map<std::string, uint32_t> tensorMap =
642- param, opGraph.inTensorNum, opGraph.outTensorNum, opGraph.internalTensorNum);619+ ConstructLinearWithLoraTensorMap(param, opGraph.inTensorNum, opGraph.outTensorNum, opGraph.internalTensorNum);
643 opGraph.name = "LinearWithLora";620 opGraph.name = "LinearWithLora";
644 621 
645 if (param.enableFlashComm) {622 if (param.enableFlashComm) {
@@ -653,10 +630,9 @@ atb::Status CreateFusionLinearWithLora(const FusionLinearParam &param, atb::Oper
653 baseLinearParam.loraEnableGMM = false;630 baseLinearParam.loraEnableGMM = false;
654 baseLinearParam.enableFlashComm = false;631 baseLinearParam.enableFlashComm = false;
655 CHECK_OPERATION_STATUS_RETURN(CreateFusionLinear(baseLinearParam, &baseLinearNode.operation));632 CHECK_OPERATION_STATUS_RETURN(CreateFusionLinear(baseLinearParam, &baseLinearNode.operation));
656- baseLinearNode.inTensorIds = GetTensorIdxList(tensorMap, {633+ baseLinearNode.inTensorIds =
657- param.enableFlashComm ? "intermediate_allgather_out" : "in_input",634+ GetTensorIdxList(tensorMap, {param.enableFlashComm ? "intermediate_allgather_out" : "in_input", "in_weight",
658- "in_weight", "in_scale", "in_offset", "in_descale", "in_bias", "in_compress_idx"635+ "in_scale", "in_offset", "in_descale", "in_bias", "in_compress_idx"});
659- });
660 baseLinearNode.outTensorIds = {GetTensorIdx(tensorMap, "intermediate_base_linear_out")};636 baseLinearNode.outTensorIds = {GetTensorIdx(tensorMap, "intermediate_base_linear_out")};
661 opGraph.nodes.push_back(baseLinearNode);637 opGraph.nodes.push_back(baseLinearNode);
662 638 
@@ -686,13 +662,12 @@ atb::Status CreateFusionLinearWithLora(const FusionLinearParam &param, atb::Oper
686 return atb::NO_ERROR;662 return atb::NO_ERROR;
687}663}
688 664 
689-atb::Status FusionLinear(const FusionLinearParam &param, atb::Operation **operation)665+atb::Status FusionLinear(const FusionLinearParam &param, atb::Operation **operation) {
690-{
691 if (param.supportLora) {666 if (param.supportLora) {
692 return CreateFusionLinearWithLora(param, operation);667 return CreateFusionLinearWithLora(param, operation);
693 } else {668 } else {
694 return CreateFusionLinear(param, operation);669 return CreateFusionLinear(param, operation);
695 }670 }
696}671}
697-} // namespace common672+} // namespace common
698-} // namespace atb_speed673+} // namespace atb_speed