dynamic_quant算子优化实践

​作者​昇腾实战派​知识地图​https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述昇腾DynamicQuant算子可以实现动态逐 token 量化将浮点张量FP16/BF16量化为低精度整数INT8/INT4或浮点8位FP8同时输出每行的量化 scale和可选的 offset。本文通过深入分析dynamic_quant算子实现识别出流水线同步开销和单行处理粒度是性能瓶颈并针对性地提出了向量指令优化与多行批量处理方案。算子分析dynamic_quant功能概述DynamicQuant 实现动态逐 token 量化将浮点张量FP16/BF16量化为低精度整数INT8/INT4或浮点8位FP8同时输出每行的量化 scale和可选的 offset。量化公式对称scale_outmax(|x*smooth|)/dtype_max//per-token逐行 yround(x*smooth/scale_out)如果hasSmooth1会加上smooth操作拆解为伪代码tmp_x Cast(x, fp32) //先转为fp32进行后续计算 tmp_x Abs(tmp_x) // x.shape [N, D] max_x ReduceMax(tmp_x) //计算行最大值 [N] scaleInv Div(max_x dtype_max) // [N] 计算scale的倒数 tmp_x Mul(tmp_x, scaleInv) // [N, D] fp32 tmp_x Cast(tmp_x, half) // 先转为half才能转为int8 y Cast(tmp_x, int_8)量化公式非对称scale_out (max(x) - min(x)) / (2 * dtype_max) offset_out dtype_max - max(x) / scale_out y round(x / scale_out offset_out)代码逻辑分析通用路径 SetTilingData() ├─ groupNum 0 innerLoopTimes 0 → KEY8 (MOE_LARGE_SHAPE) ├─ groupNum 0 → KEY7 (MOE) ├─ innerLoopTimes 0 → KEY6 (LARGE_SHAPE) ├─ useDb BF16 → KEY2 (DB_BF16) ├─ useDb FP16 → KEY3 (DB_HALF) ├─ BF16 → KEY0 └─ FP16 → KEY1测试shape输入(x)输入type输出scale/ y输出type[1,2048]BF16[1];[1,2048]FLOAT;INT8[1008,64]BF16[1008];[1008,64]FLOAT;INT8[1008,2048]BF16[1008];[1008,2048]FLOAT;INT8[1008,4096]BF16[1008];[1008,4096]FLOAT;INT8[1008,8192]BF16[1008];[1008,8192]FLOAT;INT8[4032,2048]BF16[4032];[4032,2048]FLOAT;INT8按照上述shape根据代码中已有Tiling逻辑会走到dynamic_quant_db.h这里我们来分析一下这里的代码host侧如何Tiling以4032 * 2048的shape大小为例分析核间Tilingsize_t dimNumxShape-GetStorageShape().GetDimNum()-1;//获取最后一维数值 2048uint64_ttempHeadCoreNum1;for(size_t i0;idimNum;i){tempHeadCoreNum*xShape-GetStorageShape().GetDim(i);//获取前几维相乘数值转2维 4032}rowNumtempHeadCoreNum;// 4032// For 910BrowNumPerMinTask1U;scaleNumPerMinTask1U;rowNumPerTaskstd::max(rowNumPerMinTask,scaleNumPerMinTask);// 1wholeTaskNumrowNum/rowNumPerTask;// 每行看作一个task 4032coreNumstd::max(std::min(vectorCoreNum,wholeTaskNum),ONE);// 这里vectorCoreNum40headCoreNumrowNum%coreNum;// 4032 % 40 32rowPerHeadCore(rowNumcoreNum-1U)/coreNum;// (4032 40 - 1) / 40 101rowPerTailCorerowNum/coreNum;// 4032 / 40 100这里考虑我们vector核一共40个4032行无法整除40所以核间分配不完全均匀这里头核32个每个核要处理101行数据尾核8个每个核要处理100行数据。验证32×101 8×100 4032✓核内Tilinguint32_t alignedRowLen AlignUpNUM_SIXTEEN(rowLen); //向上对齐16位 2048 if (xDtype ge::DT_BF16) { if (hasSmooth) { ubPerRow BF16_DB_SMOOTH_UB_SIZE * alignedRowLen; // 19 * 2048 38912 ubPerRowNew BF16_SMOOTH_UB_SIZE * alignedRowLen; // 17 * 2048 34816 } else { ubPerRow BF16_DB_UB_SIZE * alignedRowLen; // 13 * 2048 26624 ubPerRowNew BF16_UB_SIZE * alignedRowLen; // 11 * 2048 22528 }uint64_t maxUseUbSize ubSize - RESERVED_LENGTH; // 192 *1024 - 1024 195328 uint32_t ubAvail static_castuint32_t(maxUseUbSize / ubPerRowNew); // 195328 / 34816 5 bool useDb true; /* 省略 */ SetTilingKey(context, xDtype, useDb); // DynamicQuantDb tilingData.set_coreNum(coreNum); // 40 tilingData.set_rowLen(rowLen); // 2048 tilingData.set_headCoreNum(headCoreNum); //32 tilingData.set_rowPerHeadCore(rowPerHeadCore); //101 tilingData.set_rowPerTailCore(rowPerTailCore); //100 tilingData.set_multiRowNumHeadCore(std::min({COMPARE_INT, ubAvail, rowPerHeadCore})); // min{255,5,101}5 tilingData.set_multiRowNumTailCore(std::min({COMPARE_INT, ubAvail, rowPerTailCore})); // min{255,5,100}5 /* 省略 */例如core0要处理101行数据每次循环处理5行数据那么一共会处理 101 / 5 21次但是最后一次会单独处理一行数据存在一定的尾部开销见下图第一个Profiling末尾处。kernel侧主体结构// 存在两个几乎完全相同的函数仅类型细节不同 void ProcessBf16(smoothLocal, multiRow, loopNum) { /* ~65行 */ } void ProcessFp16(smoothLocal, multiRow, loopNum) { /* ~55行 */ } void Process(...) { if (isAsymmetrical) { ProcessAsymmetric(...); } else { if constexpr (IsSameTypexDtype, bfloat16_t::value) ProcessBf16(...); else ProcessFp16(...); } }这里在循环内每次处理一行数据// 每行循环内 Div(temp, constScale, temp, MAX_VALUE_NUM); // ① V流水线计算 constScale/max event_t eventVS GetTPipePtr()-FetchEventID(HardEvent::V_S); SetFlagHardEvent::V_S(eventVS); WaitFlagHardEvent::V_S(eventVS); // ← ★停顿点1等V流水线结果刷到标量寄存器 scale temp.GetValue(0); // ② 标量读出S流水线 scaleLocal.SetValue(i, 1/scale); // ③ 标量写入scale数组S流水线 Muls(tempCast, tempCast, scale, tilingData_.rowLen); // ④ V流水线向量×标量V操作但依赖②的结果 // 第0行还需等MTE3_S完成才能写scale event_t event_mte3_s ...; SetFlagMTE3_S(event_mte3_s); if (i 0) WaitFlagMTE3_S(...); // ← ★停顿点2等上次DMA写出完成两个停顿点的代价V_S同步V流水线必须把结果刷到标量寄存器所有未完成的V指令必须drain代价约几十个cycle。Muls本身是向量指令但它需要的参数scale是通过GetValue从标量流水线读出的所以整条链路Div → V_S停顿 → GetValue → Muls被串行化了MTE3_S同步SetValue写 scaleLocal需要等上一次 DMA 写出完成引入 MTE3→Scalar 的等待通过Profiling来看由于搬运时间比Vector计算时间少很多而且实际由于搬运时间已经提前因此开启double buffer这类操作实际没有太大的性能收益无法做到计算和搬运的掩盖我们希望的理想流水线是vector计算占比尽可能提高中间不要有断流现象。通过指令流水来看V_S这个同步是可以消除的可以通过广播操作将GetValue及SetValue这类Scaler计算消除。优化方案–Vector指令优化消除V_S同步Brcb操作替换GetValue针对dynamic_quant_db.h减少Scaler操作消除V_S同步但是现在vector计算的时候仍然是每个核处理一行仅是核内的计算优化这里实测性能收益不高单算子仅3us的收益。// 第一步把 max 值广播到8个连续位置满足向量指令最小粒度要求 Brcb(brcbTmp, temp, 1, {1, 8}); PipeBarrierPIPE_V(); // 第二步向量除法一次计算 quantScale constScale / max Div(quantScaleTmp, constScale, brcbTmp, MAX_VALUE_NUM); PipeBarrierPIPE_V(); // 第三步向量乘法stride参数实现广播乘 // src1步长0意味着 quantScaleTmp[0] 被广播到 rowLen 个元素 Mul(tempCast, tempCast, quantScaleTmp, 64, (rowLen63)6, {1, 1, 0, 8, 8, 0}); PipeBarrierPIPE_V(); // 第四步两步Cast完成量化float → int16 → half → output Cast(tempInt16, tempCast, RoundMode::CAST_RINT, rowLen); PipeBarrierPIPE_V(); Cast(tempHalf, tempInt16, RoundMode::CAST_ROUND, rowLen); PipeBarrierPIPE_V(); // 第五步用 mask 向量乘法写 scale uint64_t mask[2] {1ULL (rowIndex 7), 0ULL}; Mul(scaleLocal[rowIndex ~7], brcbTmp, constInvScale, mask, 1, {1, 1, 0, 8, 8, 0}); // ↑ mask只激活第 rowIndex%8 位避免标量SetValue Cast(outLocal[rowIndex * outAlignLen], tempHalf, RoundMode::CAST_TRUNC, rowLen);DuplicateConst前移 新增constInvScale常量前移的意义旧执行顺序: Process() └─ CopyInSmooth() MTE2搬运smooth ──────────────► WaitFlagMTE2_V ↓ 等DMA完成 Cast(smooth) └─ DuplicateConst() ← 只有smooth结束后才能执行只计算了constScale Duplicate(constScale) └─ ProcessQuant() ... 新执行顺序: Process() └─ DuplicateConst() Duplicate(constScale/constInvScale) ──► └─ CopyInSmooth() MTE2搬运smooth ─────────────────────► WaitFlagMTE2_V ↓ ↑ 两段并行Duplicate在MTE2搬运期间完成 └─ ProcessQuant() ...Duplicate是纯向量指令不占用 MTE2 通道提前到 smooth 搬运之前执行可与 DMA 完全并行隐藏了常量初始化的延迟。新增constInvScale旧代码用Div(temp, constScale, temp)求1/max然后scaleLocal.SetValue(i, 1/scale)再做一次倒数。新代码在DuplicateConst中预填// INT8对称量化constScale 127.0Duplicatefloat(constScale,DYNAMIC_QUANT_INT8_SYM_SCALE,MAX_VALUE_NUM);Duplicatefloat(constInvScale,float(1)/DYNAMIC_QUANT_INT8_SYM_SCALE,MAX_VALUE_NUM);// constScale [127.0, 127.0, 127.0, 127.0, 127.0, 127.0, 127.0, 127.0]// constInvScale [0.00787, 0.00787, ..., 0.00787] (1/127)后续mask Mul中scaleLocal[i] brcbTmp[i] * constInvScale[i] max[i] * (1/127.0) max[i] / 127.0把浮点除法 标量倒数替换为向量乘法并消除了V_S同步。多行批量处理新增dynamic_quant_multi_row.h针对小shape场景UB可以放下多行数据的场景进行优化。这里的核心思想是与其让每个核独立处理一行算子头开销 × 行数不如让每个核一次性把多行装入UB消除流水切换开销用批量向量指令处理所有行与dynamic_quant_db对比维度DynamicQuantDbDynamicQuantMultiRow流水模型Double-BufferTQue单缓冲TBuf每次处理粒度单行 × loop多行一批 × loopReduceMax每行标量路径WholeReduceMax批量Scale广播每行BrcbDiv向量Div BroadCast为何不用 Double-Buffer小shape场景下单次搬运量几行 × rowLen字节很小MTE2搬运极快DMA和计算之间几乎没有可以 overlap 的时间窗口double-buffer的收益接近零反而引入了额外的 TQue 管理开销。用TBufVECCALC直接一次性分配所有UB让向量计算尽快开始更合适。代码流程这里只有Max指令会循环多次其他指令直接处理多行数据从而获得性能收益// ───────────────────────────────────────────────────────────── // ComputeReduceMax批量求每行绝对值最大值 // 输入absTmp[ubRows × rowLen]已取绝对值 // 输出scaleTmp[ubRows]每行一个max值 // ───────────────────────────────────────────────────────────── // 第一步对每行做尾部折叠预处理满足 WholeReduceMax 要求 rowLen 是64整数倍 for (uint32_t i 0; i ubRows; i) { calcTensor absTmp[i * rowLen]; uint32_t repeatTimes rowLen / BLOCK_SIZE_64; // 完整的64元素块数 uint32_t remainElements rowLen % BLOCK_SIZE_64; // 不足64的尾部元素数 // 如果有尾巴rowLen % 64 ! 0把尾部 max 折入行首第0段 if (remainElements 0) { Max(calcTensor, calcTensor, calcTensor[repeatTimes * BLOCK_SIZE_64], remainElements, 1, repeatParams); PipeBarrierPIPE_V(); } // 如果有多段rowLen 64把第1段起的所有段 max 逐步折入第0段 if (repeatTimes 1) { Max(calcTensor, calcTensor[BLOCK_SIZE_64], calcTensor, BLOCK_SIZE_64, repeatTimes - 1, repeatParams); PipeBarrierPIPE_V(); } } // 第二步一条指令批量求所有行的 max结果紧排存入 scaleTmp[0..ubRows-1] // rowLen/8 每行的 repeat 次数WholeReduceMax 每次处理8个float WholeReduceMax(scaleTmp, absTmp, BLOCK_SIZE_64, ubRows, 1, 1, rowLen / STRIDE_8, ReduceOrder::ORDER_ONLY_VALUE); // ───────────────────────────────────────────────────────────── // ComputeQuantization批量量化所有行 // 输入castTmp[ubRows × rowLen]fp16/bf16已Cast为float // 输出absTmp量化后的half值等待后续Cast到int8/int4 // ───────────────────────────────────────────────────────────── // 第一步对整批数据取绝对值 Abs(absTmp, castTmp, ubNums); PipeBarrierPIPE_V(); // 第二步批量求每行最大值内含尾部折叠预处理见 ComputeReduceMax // 结果scaleTmp[i] Row_i 的绝对值最大值 ComputeReduceMax(ubRows); PipeBarrierPIPE_V(); // 第三步向量除法一次算出所有行的量化scale constScale / rowMax // stride参数 {dst1, src00, src11, dstRep8, src0Rep0, src1Rep8} // src0(constScale) stride0 → 广播同一个值给每行 // src1(scaleTmp) stride8 → 每行的max间隔8个floatWholeReduceMax输出格式 Div(quantScaleTmp, constScale, scaleTmp, BLOCK_SIZE_64, CeilDiv(ubRows, BLOCK_SIZE_64), {1, 0, 1, STRIDE_8, 0, STRIDE_8}); PipeBarrierPIPE_V(); // 第四步将每行的单个 quantScale 广播到该行所有 rowLen 个位置 // [ubRows, 1] → [ubRows, rowLen]沿 axis1 广播 BroadCastfloat, DIMENSION_2, 1(absTmp, quantScaleTmp, dstShape1, srcShape1); PipeBarrierPIPE_V(); // 第五步向量乘法完成量化缩放所有行 ubRows × rowLen 个元素同时处理 // absTmp[i][j] castTmp[i][j] × quantScaleTmp[i]结果范围 [-127, 127] Mul(absTmp, castTmp, absTmp, ubNums); PipeBarrierPIPE_V(); // 第六步两步 Cast 完成量化float → int16 → half // CAST_RINT四舍五入到最近整数 Cast(absTmp.ReinterpretCastint16_t(), absTmp, RoundMode::CAST_RINT, ubNums); PipeBarrierPIPE_V(); // int16 → half为后续 CAST_TRUNC 截断到 int8/int4 做准备 Cast(absTmp.ReinterpretCasthalf(), absTmp.ReinterpretCastint16_t(), RoundMode::CAST_ROUND, ubNums);再看下Profiling去掉 Double Buffer节省一份 fp16 输入缓冲2 bytes/elemUB 常量下调ubAvail 从 5 升至 6同时一次vector计算内对应MAX指令在一次计算中执行6次其余指令均按计算逻辑单次执行。优化效果单算子性能提升约30%左右符合预期