已合并
fix(mla) issue288: restrict MASK_TYPE_SWA_NORM for decode path #2078
fix(mla) issue288: restrict MASK_TYPE_SWA_NORM for decode path #2078
已合并
ivanS创建于 4月30日
4 个文件变更+112-1
@@ -109,6 +109,14 @@ static bool ParamCheck(const infer::MultiLatentAttentionParam &opParam)
109 ATB_LOG(ERROR) << "only mtp(CALC_TYPE_SPEC) support mask";109 ATB_LOG(ERROR) << "only mtp(CALC_TYPE_SPEC) support mask";
110 return false;110 return false;
111 }111 }
112+ if ((opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC ||
113+ opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC_AND_RING) &&
114+ opParam.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_SWA_NORM) {
115+ ATB_LOG(ERROR) << "MLA decode (CALC_TYPE_SPEC) does not support SWA_NORM mask. "
116+ << "Kernel SoftmaxStage1 only handles mask_type=3 (LOOK_AHEAD) and 4 (MASK_FREE). "
117+ << "Use CALC_TYPE_PREFILL for SWA_NORM or set maskType to MASK_FREE.";
118+ return false;
119+ }
112 if ((opParam.cacheMode == infer::MultiLatentAttentionParam::CacheMode::INT8_NZCACHE) &&120 if ((opParam.cacheMode == infer::MultiLatentAttentionParam::CacheMode::INT8_NZCACHE) &&
113 (opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_RING ||121 (opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_RING ||
114 opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC_AND_RING) &&122 opParam.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC_AND_RING) &&
@@ -130,6 +138,12 @@ static bool ParamCheck(const infer::MultiLatentAttentionParam &opParam)
130 return false;138 return false;
131 }139 }
132 }140 }
141+ if (opParam.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_SWA_NORM) {
142+ if (opParam.windowSize <= 0) {
143+ ATB_LOG(ERROR) << "windowSize in swa mode should be greater than 0";
144+ return false;
145+ }
146+ }
133 return true;147 return true;
134}148}
135 149 
@@ -76,6 +76,14 @@ Status MultiLatentAttentionOpsRunner::SetupKernelGraph(const OpsTensorPack &opsT
76 asdParam.tor = param_.qkScale;76 asdParam.tor = param_.qkScale;
77 asdParam.kvHead = param_.kvHeadNum;77 asdParam.kvHead = param_.kvHeadNum;
78 asdParam.isRing = isRing ? 1 : 0;78 asdParam.isRing = isRing ? 1 : 0;
79+ asdParam.windowSize = static_cast<int32_t>(param_.windowSize);
80+ if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_SPEC) {
81+ asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_LOOK_AHEAD;
82+ } else if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_MASK_FREE) {
83+ asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_MASK_FREE;
84+ } else if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_CAUSAL_MASK) {
85+ asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_CAUSAL_MASK;
86+ }
79 mlaNode.opDesc = {0, "MLAOperation", asdParam};87 mlaNode.opDesc = {0, "MLAOperation", asdParam};
80 mlaNode.inTensors = {&query, &queryRope, &kvCache, &kvCacheRope, &blockTables, mask, qkDescale, pvDescale};88 mlaNode.inTensors = {&query, &queryRope, &kvCache, &kvCacheRope, &blockTables, mask, qkDescale, pvDescale};
81 if (!isRing) {89 if (!isRing) {
@@ -131,12 +139,16 @@ Status MultiLatentAttentionOpsRunner::ModifyKernelGraph(const OpsTensorPack &ops
131 asdParam.maskUseStatus = newParam_.maskUseStatus;139 asdParam.maskUseStatus = newParam_.maskUseStatus;
132 asdParam.isRing = param_.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_RING || param_.140 asdParam.isRing = param_.calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_RING || param_.
133 calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC_AND_RING;141 calcType == infer::MultiLatentAttentionParam::CalcType::CALC_TYPE_SPEC_AND_RING;
142+ asdParam.windowSize = static_cast<int32_t>(param_.windowSize);
134 if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_SPEC) {143 if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_SPEC) {
135 asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_LOOK_AHEAD;144 asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_LOOK_AHEAD;
136 }145 }
137 if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_MASK_FREE) {146 if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_MASK_FREE) {
138 asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_MASK_FREE;147 asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_MASK_FREE;
139 }148 }
149+ if (param_.maskType == infer::MultiLatentAttentionParam::MaskType::MASK_TYPE_CAUSAL_MASK) {
150+ asdParam.maskType = AtbOps::OpParam::MLA::MaskType::MASK_TYPE_CAUSAL_MASK;
151+ }
140 mlaNode.opDesc = {0, "MLAOperation", asdParam};152 mlaNode.opDesc = {0, "MLAOperation", asdParam};
atomgit-bot
atomgit-botatomgit-bot4月30日

⚠️ ModifyKernelGraph 缺少 windowSize 校验

SetupKernelGraph 调用前会在 ParamCheck 中校验 SWA 模式的 windowSize,但 ModifyKernelGraph 中设置 SWA maskType 时没有校验 windowSize。如果 param_ 在 Setup 之后被修改(设置 maskType=SWA 但 windowSize<=0),会绕过校验。存在数据一致性风险。

参考依据

  • 类别: bug
  • 置信度: 75%
likedislike
ivanS
5月6日 评论:
141 return NO_ERROR;153 return NO_ERROR;
142}154}
atomgit-bot
atomgit-botatomgit-bot5月6日

🚨 ModifyKernelGraph 中缺少 SWA_NORM maskType 的转换处理

SetupKernelGraph 函数中新增了 MASK_TYPE_SPECMASK_TYPE_MASK_FREEMASK_TYPE_CAUSAL_MASK 三种 maskType 的转换,但 ModifyKernelGraph 函数中新增的转换逻辑与之不一致。当 maskType == MASK_TYPE_SWA_NORM 时,SetupKernelGraph 会保持原值(因为不在新增的转换分支中),但 ModifyKernelGraph 同样不会进入任何转换分支,导致 asdParam.maskType 可能使用未正确设置的值。对于 decode path (CALC_TYPE_SPEC),在 ParamCheck 中已限制 SWA_NORM,但对于其他情况可能导致运行时行为异常。

参考依据

  • 类别: bug
  • 置信度: 95%
likedislike
ivanS
5月6日 评论:
@@ -127,3 +127,20 @@ CaseNum |CaseName |OpName |OpParam
127126 |MultiLatentAttentionDecoderWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":2,"cacheMode":1,"headNum":32,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;nd;nd;nd;nd|32,32,512;32,32,64;64,64,1,512;64,64,1,64;32,2;32|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM127126 |MultiLatentAttentionDecoderWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":2,"cacheMode":1,"headNum":32,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;nd;nd;nd;nd|32,32,512;32,32,64;64,64,1,512;64,64,1,64;32,2;32|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM
128127 |MultiLatentAttentionInt8NzWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":0,"cacheMode":2,"headNum":32,"kvHeadNum":1} | 8 |int8;float16;int8;float16;int32;int32;float;float |nd;nd;fractal_nz;fractal_nz;nd;nd;nd;nd|10,32,512;10,32,64;64,16,1,32;64,4,1,16;10,2;10;32;32|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM128127 |MultiLatentAttentionInt8NzWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":0,"cacheMode":2,"headNum":32,"kvHeadNum":1} | 8 |int8;float16;int8;float16;int32;int32;float;float |nd;nd;fractal_nz;fractal_nz;nd;nd;nd;nd|10,32,512;10,32,64;64,16,1,32;64,4,1,16;10,2;10;32;32|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM
129128 |MultiLatentAttentionNzWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":0,"cacheMode":3,"headNum":32,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;fractal_nz;fractal_nz;nd;nd|10,32,512;10,32,64;64,32,1234,16;64,4,1234,16;10,2;10|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM129128 |MultiLatentAttentionNzWrongBlockSize |MultiLatentAttentionOperation |{"maskType":0,"calcType":0,"cacheMode":3,"headNum":32,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;fractal_nz;fractal_nz;nd;nd|10,32,512;10,32,64;64,32,1234,16;64,4,1234,16;10,2;10|2 |float16;float16|nd;nd|32,32,512;32,32,1|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:ERROR_INVALID_TENSOR_DIM
130+129 |MultiLatentAttentionDecodeSWAWindowSizeZero |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":16,"kvHeadNum":1,"windowSize":0} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|32,32,512;32,32,64;64,128,1,512;64,128,1,64;32,2;32;512,512;32|1 |float16|nd|32,32,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
atomgit-bot
atomgit-botatomgit-bot5月6日

⚠️ 多个测试用例因其他现有检查失败而非新增检查

多个测试用例(如 Case 130, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 145)期望返回 ERROR_INVALID_PARAM,但实际失败原因可能与预期的新增检查无关。例如:Case 130 使用 calcType=0 (CALC_TYPE_PREFILL) 与 maskType=4 (SWA_NORM) 组合,现有代码第101-107行检查 "only mtp(CALC_TYPE_SPEC) support mask",因此该用例实际上在第101行就返回了,无法测试新增的第141-147行 windowSize 检查。类似问题存在于多个测试用例中,测试设计不够精确。

参考依据

  • 类别: general
  • 置信度: 85%
likedislike
ivanS
5月6日 评论:
131+130 |MultiLatentAttentionDecodeUndefinedSWA |MultiLatentAttentionOperation |{"maskType":4,"calcType":0,"cacheMode":1,"headNum":16,"kvHeadNum":1,"windowSize":4096} | 7 |float16;float16;float16;float16;int32;int32;float16 |nd;nd;nd;nd;nd;nd;nd|32,32,512;32,32,64;64,128,1,512;64,128,1,64;32,2;32;512,512|1 |float16|nd|32,32,512|random;random;random;random;random;random;random|-5,5;-5,5;-5,5;-5,5;-5,5;-5,5;-5,5| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
132+131 |MultiLatentAttentionDecodeSWAPositive |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":8,"kvHeadNum":1,"qkScale":0.0416667,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,8,512;1,8,64;4,128,1,512;4,128,1,64;1,2;1;1,128;1|1 |float16|nd|1,8,512|customize;customize;customize;customize;customize;customize;customize;customize|-0.1,0.1;-0.1,0.1;-0.1,0.1;-0.1,0.1;-0.1,0.1;-0.1,0.1;-0.1,0.1;-0.1,0.1| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
133+132 |MultiLatentAttentionDecodeSpecAndRingSWAWindowZero |MultiLatentAttentionOperation |{"maskType":4,"calcType":3,"cacheMode":1,"headNum":16,"kvHeadNum":1,"windowSize":0} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32;512,512;32|2 |float16;float16|nd;nd|32,16,512;32,16,1|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
134+133 |MultiLatentAttentionDecodeRingSWA |MultiLatentAttentionOperation |{"maskType":4,"calcType":2,"cacheMode":1,"headNum":16,"kvHeadNum":1,"windowSize":4096} | 7 |float16;float16;float16;float16;int32;int32;float16 |nd;nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32;512,512|2 |float16;float16|nd;nd|32,16,512;32,16,1|random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
135+134 |MultiLatentAttentionDecodeUndefinedMaskFree |MultiLatentAttentionOperation |{"maskType":2,"calcType":0,"cacheMode":1,"headNum":16,"kvHeadNum":1} | 7 |float16;float16;float16;float16;int32;int32;float16 |nd;nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32;512,512|1 |float16|nd|32,16,512|random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
136+135 |MultiLatentAttentionDecodeRingMaskFree |MultiLatentAttentionOperation |{"maskType":2,"calcType":2,"cacheMode":1,"headNum":16,"kvHeadNum":1} | 7 |float16;float16;float16;float16;int32;int32;float16 |nd;nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32;512,512|2 |float16;float16|nd;nd|32,16,512;32,16,1|random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
137+136 |MultiLatentAttentionDecodeUndefinedCausal |MultiLatentAttentionOperation |{"maskType":3,"calcType":0,"cacheMode":1,"headNum":16,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32|1 |float16|nd|32,16,512|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
138+137 |MultiLatentAttentionDecodeRingCausal |MultiLatentAttentionOperation |{"maskType":3,"calcType":2,"cacheMode":1,"headNum":16,"kvHeadNum":1} | 6 |float16;float16;float16;float16;int32;int32 |nd;nd;nd;nd;nd;nd|32,16,512;32,16,64;64,128,1,512;64,128,1,64;32,2;32|2 |float16;float16|nd;nd|32,16,512;32,16,1|random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
139+138 |MultiLatentAttentionDecodeSWAHead8 |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":8,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,8,512;1,8,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,8,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
140+139 |MultiLatentAttentionDecodeSWAHead16 |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":16,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,16,512;1,16,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,16,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
141+140 |MultiLatentAttentionDecodeSWAHead32 |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":32,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,32,512;1,32,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,32,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
142+141 |MultiLatentAttentionDecodeSWAHead64 |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":64,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,64,512;1,64,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,64,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
143+142 |MultiLatentAttentionDecodeSWAHead128 |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":128,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,128,512;1,128,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,128,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
144+143 |MultiLatentAttentionDecodeSpecAndRingSWA |MultiLatentAttentionOperation |{"maskType":4,"calcType":3,"cacheMode":1,"headNum":8,"kvHeadNum":1,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,8,512;1,8,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|2 |float16;float16|nd;nd|1,8,512;1,8,1|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
145+144 |MultiLatentAttentionDecodeSpecMaskFree |MultiLatentAttentionOperation |{"maskType":2,"calcType":1,"cacheMode":1,"headNum":16,"kvHeadNum":1} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,16,512;1,16,64;1,128,1,512;1,128,1,64;1,1;1;1,128;1|1 |float16|nd|1,16,512|random;random;random;random;random;random;random;random|-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100;-100,100| | | | | |Ascend910B |I:NO_ERROR
146+145 |MultiLatentAttentionDecodeSWAPositiveNOERROR |MultiLatentAttentionOperation |{"maskType":4,"calcType":1,"cacheMode":1,"headNum":8,"kvHeadNum":1,"qkScale":0.0416667,"windowSize":64} | 8 |float16;float16;float16;float16;int32;int32;float16;int32 |nd;nd;nd;nd;nd;nd;nd;nd|1,8,512;1,8,64;4,128,1,512;4,128,1,64;1,2;1;1,128;1|1 |float16|nd|1,8,512|customize;customize;customize;customize;customize;customize;customize;customize|-5,5;-5,5;-5,5;-5,5;-5,5;-5,5;-5,5;-5,5| | | | | |Ascend910B |C:ERROR_INVALID_PARAM
@@ -8446,9 +8446,77 @@ class MlaPreprocessOperation(DataGen):
8446 return OpTypes.COMPUTE_FLOAT8446 return OpTypes.COMPUTE_FLOAT
8447 8447 
8448class MultiLatentAttentionOperation(DataGen):8448class MultiLatentAttentionOperation(DataGen):
8449+ @staticmethod
8450+ def customize(shapes, i, datatype, format, data_gen_ranges, op_params):
8451+ if i != 0:
8452+ if i < len(MultiLatentAttentionOperation.in_tensors):
8453+ t = MultiLatentAttentionOperation.in_tensors[i]
8454+ if t.numel() > 0:
8455+ return torch_npu.npu_format_cast(t, format_dict[format])
8456+ json_data = json.loads(op_params)
8457+ mask_type = json_data.get("maskType", 0)
8458+ num_tokens, num_heads, head_dim_q = shapes[0]
8459+ rope_dim = shapes[1][2]
8460+ num_blocks, block_size, kv_heads, kv_head_dim = shapes[2]
8461+ batch, max_blocks = shapes[4]
8462+ rope_head_dim = shapes[3][3]
8463+ 
8464+ low = float(data_gen_ranges.split(',')[0])
8465+ high = float(data_gen_ranges.split(',')[1])
8466+ query_np = np.random.uniform(low, high, size=(num_tokens, num_heads, head_dim_q)).astype(np.float32)
8467+ query_rope_np = np.random.uniform(low, high, size=(num_tokens, num_heads, rope_dim)).astype(np.float32)
8468+ kv_np = np.random.uniform(low, high, size=(num_blocks, block_size, kv_heads, kv_head_dim)).astype(np.float32)
8469+ kv_rope_np = np.random.uniform(low, high, size=(num_blocks, block_size, kv_heads, rope_head_dim)).astype(np.float32)
8470+ bt_np = np.zeros((batch, max_blocks), dtype=np.int32)
8471+ for b in range(batch):
8472+ for j in range(max_blocks):
8473+ idx = b * max_blocks + j
8474+ bt_np[b, j] = idx if idx < num_blocks else 0
8475+ ctx_np = np.full((batch,), block_size, dtype=np.int32)
8476+ qsl_np = np.ones((batch,), dtype=np.int32)
8477+ 
8478+ base_shape_idx = 6
8479+ has_mask = mask_type != 0
8480+ if has_mask:
8481+ m_shapes = shapes[base_shape_idx]
8482+ m_np = np.random.uniform(low, high, size=(m_shapes[0], m_shapes[1])).astype(np.float32)
8483+ base_shape_idx = 7
8484+ 
8485+ tensor_list = [
8486+ torch.from_numpy(query_np).to(dtype_dict[datatype]).npu(),
8487+ torch.from_numpy(query_rope_np).to(dtype_dict[datatype]).npu(),
8488+ torch.from_numpy(kv_np).to(dtype_dict[datatype]).npu(),
8489+ torch.from_numpy(kv_rope_np).to(dtype_dict[datatype]).npu(),
8490+ torch.from_numpy(bt_np).npu(),
8491+ torch.from_numpy(ctx_np).npu(),
8492+ ]
8493+ if has_mask:
8494+ tensor_list.append(torch.from_numpy(m_np).to(dtype_dict[datatype]).npu())
8495+ tensor_list.append(torch.from_numpy(qsl_np).npu())
8496+ 
8497+ MultiLatentAttentionOperation.in_tensors = tensor_list
8498+ return tensor_list[0]
8499+ 
8449 @staticmethod8500 @staticmethod
8450 def get_op_type(op_params) -> OpTypes:8501 def get_op_type(op_params) -> OpTypes:
8451- return OpTypes.COMPUTE_FLOAT8502+ return OpTypes.CV_FUSION
8503+ 
8504+ @staticmethod
8505+ def case_preprocess(op_params, operation, input_tensor_list):
8506+ json_data = json.loads(op_params)
8507+ host_dict = {}
8508+ host_dict["contextLens"] = input_tensor_list[5].tolist()
8509+ calc_type = json_data.get("calcType", 0)
8510+ if calc_type in [1, 3]:
8511+ mask_type = json_data.get("maskType", 0)
8512+ qseq_idx = 6 if mask_type == 0 else 7
8513+ host_dict["qSeqlen"] = input_tensor_list[qseq_idx].tolist()
8514+ if "maskType" in json_data:
8515+ host_dict["maskType"] = json_data["maskType"]
8516+ if "cacheMode" in json_data and json_data["cacheMode"] == 2:
8517+ host_dict["cacheType"] = 1
8518+ run_param = json.dumps(host_dict)
8519+ operation.set_varaintpack_param(run_param)
8452 8520 
8453class PagedCacheLoadOperation(DataGen):8521class PagedCacheLoadOperation(DataGen):
8454 @staticmethod8522 @staticmethod