GroupedMatmul算子性能调优案例

案例介绍

本案例对分组Matmul即GroupedMatmul算子的per-token量化场景进行性能分析和优化，GroupedMatmul算子计算过程（通过python代码表达）为：

offset = 0
for i in range(g):
mmOut = x[offset:offset + group[i]] * weight[i] + bias[i]
y[offset:offset + group[i]] = Gelu(mmOut * scale[i] * pertokenScale[offset:offset + group[i]])
offset += group[i]

验证平台为 Atlas A2 训练系列产品/Atlas 800I A2 推理产品 。

优化分析以如下算子规格为例：

主要介绍以下优化方法：

• 对于Vector占比较高（Vector Bound）的场景，将AI Core中的Cube和Vector启动比例设置为1:2；
• 优化CV并行流水，减少Cube和Vector间的空闲等待时间；
• 优化Vector计算流水，提高Vector并行计算速度。

获取性能数据

固定8核测试，即当前性能和后续优化tiling中blockDim固定设置为8。

通过msProf算子调优工具获取算子性能数据：

• 获取真实环境执行的性能数据（指令的cycle占比数据ArithmeticUtilization.csv），包含各个流水的占比情况；
• 获取仿真性能数据（指令流水图），包含各个流水的占用区间，可观察流水间依赖情况，从而优化并行效率。

分析主要瓶颈点

固定8核进行测试的情况下，通过msprof op命令获取指令的cycle占比数据如下：

通过msprof op simulator获取到的指令流水图如下图所示：

结合上述两种数据（真实数据和仿真数据）进行性能分析：

• Vector计算bound，当前为减少核启动开销设置为1:1；
• 实际优化过程中，对上述问题进行优化、Vector占比下降后，Cube和Vector各自都有间隙，相互之间都有等待耗时；

  

• Vector没有开启double buffer，计算和数据搬运部分没有并行。

  

设计优化方案

• 将AI Core中的Cube和Vector启动比例设置为1:2。每次Cube输出的数据，由两个Vector并行计算对应的反量化和激活函数；在Vector的循环里，Vector0和Vector1交替进行计算（前提条件，循环次数不为1）。代码示例如下：

  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

  uint32_t vecCount = 0; uint32_t taskRation = GetTaskRation(); for (uint32_t offsetN = 0; offsetN < curCubeSingleN; offsetN += mnConfig.baseN) { if (unlikely(offsetN + mnConfig.baseN >= curCubeSingleN)) { curVecBaseN = curCubeSingleN - offsetN; } uint32_t alignBaseN = Ceil(curVecBaseN, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block DataCopyScale(curVecBaseN, alignBaseN, scaleOffset + offsetN); uint32_t curVecBaseM = vecBaseM; uint64_t mmOutOffset = mnConfig.workSpaceOffset + offsetN * mnConfig.baseM; CrossCoreWaitFlag(SYNC_AIC_TO_AIV); for (uint32_t offsetM = 0; offsetM < curCubeSingleM; offsetM += vecBaseM) { vecCount++; if (vecCount % taskRation != subBlockIdx) { continue; // Vector0和Vector1交替进行计算 } if (unlikely(offsetM + vecBaseM >= curCubeSingleM)) { curVecBaseM = curCubeSingleM - offsetM; } // 使用AscendDequant接口做perchannel反量化 LocalTensor<cT::T> mmOutLocal = vecInQueue.AllocTensor<cT::T>(); DataCopyPad2D(mmOutLocal, mmOutGm[mmOutOffset + offsetM * curVecBaseN], curVecBaseM, curVecBaseN, curVecBaseN); vecInQueue.EnQue(mmOutLocal); ComputeDequantAndActivate(mnConfig, curVecBaseM, alignBaseN, curVecBaseN, offsetM); LocalTensor<DTYPE_Y> yLocal = vecOutQueue.DeQue<DTYPE_Y>(); DataCopyPad2D(yGm[outOffset + offsetM * tiling->n + offsetN], yLocal, curVecBaseM, curVecBaseN, alignBaseN, tiling->n); vecOutQueue.FreeTensor(yLocal); } ...

• Cube和Vector启动比例设置为1:2后，出现Cube和Vector各自都有间隙、相互之间都有等待耗时的情况。分析原因是因为Vector和Cube存在使用一份workspace进行数据传递的场景，通过4份workspace的方案进行优化：host按4倍baseM * baseN申请workspace，Cube在计算前可以跳过前4轮的等待。

  if ASCEND_IS_AIC {
      if (cubeCount >= tiling->parallNum) {  // tiling->parallNum设置为4
          CrossCoreWaitFlag(SYNC_AIV_TO_AIC);
      }
      mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k);
      mm.SetSingleShape(curSingleM, curSingleN, tiling->k);
      mm.SetTensorA(xGm[xOffset]);
      auto weightSlice = weightGm[weightOffset];
      if (mnConfig.blockDimM == 1) {
          weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
      }
      mm.SetTensorB(weightSlice);
      uint64_t worskspaceOffset = mnConfig.workSpaceOffset;
      while (mm.Iterate()) {
          mm.GetTensorC(mmOutGm[worskspaceOffset], 0, true);
          CrossCoreSetFlag<2, PIPE_FIX>(SYNC_AIC_TO_AIV);
          worskspaceOffset += (mnConfig.baseM * mnConfig.baseN);
      }
  }
  cubeCount++;

• Vector开启double buffer，InitBuffer指定分配内存块个数为2。

  pipe->InitBuffer(scaleInQueue, 2, tiling->mmTilingData.baseN * sizeof(DTYPE_SCALE));
  pipe->InitBuffer(perTokenScaleInQueue, 2, tiling->mmTilingData.baseM * sizeof(float));
  pipe->InitBuffer(vecInQueue, 2, tiling->ubCalSize * sizeof(cT::T));
  pipe->InitBuffer(vecOutQueue, 2, tiling->ubCalSize * sizeof(DTYPE_Y));

验证优化方案性能收益

• 将AI Core中的Cube和Vector启动比例设置为1:2后，执行总耗时从218.1us下降为154.2us。指令流水图显示Cube间等待变小。

  

• 如上图所示，Vector已经不处于bound状态，但Cube和Vector都有间隙，没有用满（上述两个箭头的位置）。分析原因如下：

  Vector在等Cube输出的数据，Cube需要等Vector计算完释放workspace以存放下一轮的计算结果，当前为了让Cube、Vector流水并行，workspace用了两份空间：

  

  因为Vector和Cube存在使用一份workspace进行数据传递的场景，存在数据依赖，所以会有等待的间隔。

  可以采用4份workspace进行优化：

  

  优化后，总耗时由154.2us下降为131.8us。指令流水图显示Vector、Cube各自间隙明显减小。

  

• Vector开启double buffer，优化后执行总耗时从131.8us下降为128.1us。

  

总结

• 在Vector为主要瓶颈点时，将AI Core中的Cube和Vector启动比例设置为1:2；
• Cube、Vector时间接近，且两者都有因相互等待导致的间隙时，采用4份workspace优化；
• 观察数据搬运是否与计算相互掩盖，多轮计算没有数据依赖，且buffer够大时，开启double buffer，增加并行效率。