已合并
fix(mla) issue288: restrict MASK_TYPE_SWA_NORM for decode path #2078
ivanS创建于 4月30日
fix(mla) issue288: restrict MASK_TYPE_SWA_NORM for decode path #2078
已合并
共 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}; |
| 141 | return NO_ERROR; | 153 | return NO_ERROR; |
| 142 | } | 154 | } |
🚨 ModifyKernelGraph 中缺少 SWA_NORM maskType 的转换处理
参考依据
![]() ![]() | |||
| @@ -127,3 +127,20 @@ CaseNum |CaseName |OpName |OpParam | |||
| 127 | 126 |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 | 127 | 126 |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 |
| 128 | 127 |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 | 128 | 127 |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 |
| 129 | 128 |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 | 129 | 128 |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 | ||
⚠️ 多个测试用例因其他现有检查失败而非新增检查 多个测试用例(如 Case 130, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 145)期望返回 参考依据
![]() ![]() | |||
| 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_FLOAT | 8446 | return OpTypes.COMPUTE_FLOAT |
| 8447 | 8447 | ||
| 8448 | class MultiLatentAttentionOperation(DataGen): | 8448 | class MultiLatentAttentionOperation(DataGen): |
| 8449 | + | ||
| 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 | 8500 | ||
| 8450 | def get_op_type(op_params) -> OpTypes: | 8501 | def get_op_type(op_params) -> OpTypes: |
| 8451 | - return OpTypes.COMPUTE_FLOAT | 8502 | + return OpTypes.CV_FUSION |
| 8503 | + | ||
| 8504 | + | ||
| 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 | ||
| 8453 | class PagedCacheLoadOperation(DataGen): | 8521 | class PagedCacheLoadOperation(DataGen): |
| 8454 | 8522 | ||


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