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_out = max(|x * smooth|) / dtype_max // per-token（逐行）
y = round(x * smooth / scale_out)
如果hasSmooth=1，会加上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 → KEY=8 (MOE_LARGE_SHAPE)
├─ groupNum > 0 → KEY=7 (MOE)
├─ innerLoopTimes > 0 → KEY=6 (LARGE_SHAPE)
├─ useDb & BF16 → KEY=2 (DB_BF16)
├─ useDb & FP16 → KEY=3 (DB_HALF)
├─ BF16 → KEY=0
└─ FP16 → KEY=1
测试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大小为例分析）
核间Tiling
size_t dimNum = xShape->GetStorageShape().GetDimNum() - 1; //获取最后一维数值 2048
uint64_t tempHeadCoreNum = 1;
for (size_t i = 0; i < dimNum; i++) {
tempHeadCoreNum *= xShape->GetStorageShape().GetDim(i); //获取前几维相乘数值（转2维） 4032
}
rowNum = tempHeadCoreNum; // 4032
// For 910B
rowNumPerMinTask = 1U;
scaleNumPerMinTask = 1U;
rowNumPerTask = std::max(rowNumPerMinTask, scaleNumPerMinTask); // 1
wholeTaskNum = rowNum / rowNumPerTask; // 每行看作一个task 4032
coreNum = std::max(std::min(vectorCoreNum, wholeTaskNum), ONE); // 这里vectorCoreNum=40
headCoreNum = rowNum % coreNum; // 4032 % 40 = 32
rowPerHeadCore = (rowNum + coreNum - 1U) / coreNum; // (4032 + 40 - 1) / 40 = 101
rowPerTailCore = rowNum / coreNum; // 4032 / 40 = 100
这里考虑我们vector核一共40个，4032行无法整除40，所以核间分配不完全均匀，这里头核32个，每个核要处理101行数据，尾核8个，每个核要处理100行数据。
验证：32×101 + 8×100 = 4032 ✓
核内Tiling
uint32_t alignedRowLen = AlignUp<NUM_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_cast<uint32_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 (IsSameType<xDtype, bfloat16_t>::value)
ProcessBf16(...);
else
ProcessFp16(...);
}
}
这里在循环内每次处理一行数据
// 每行循环内：
Div(temp, constScale, temp, MAX_VALUE_NUM); // ① V流水线：计算 constScale/max
event_t eventVS = GetTPipePtr()->FetchEventID(HardEvent::V_S);
SetFlag<HardEvent::V_S>(eventVS);
WaitFlag<HardEvent::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 = ...;
SetFlag<MTE3_S>(event_mte3_s);
if (i == 0) WaitFlag<MTE3_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});
PipeBarrier<PIPE_V>();
// 第二步：向量除法，一次计算 quantScale = constScale / max
Div(quantScaleTmp, constScale, brcbTmp, MAX_VALUE_NUM);
PipeBarrier<PIPE_V>();
// 第三步：向量乘法，stride参数实现"广播乘"
// src1步长=0意味着 quantScaleTmp[0] 被广播到 rowLen 个元素
Mul(tempCast, tempCast, quantScaleTmp, 64, (rowLen+63)>>6, {1, 1, 0, 8, 8, 0});
PipeBarrier<PIPE_V>();
// 第四步：两步Cast完成量化（float → int16 → half → output）
Cast(tempInt16, tempCast, RoundMode::CAST_RINT, rowLen);
PipeBarrier<PIPE_V>();
Cast(tempHalf, tempInt16, RoundMode::CAST_ROUND, rowLen);
PipeBarrier<PIPE_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 ──────────────►
WaitFlag<MTE2_V> ↓ 等DMA完成
Cast(smooth)
└─ DuplicateConst() ← 只有smooth结束后才能执行，只计算了constScale
Duplicate(constScale)
└─ ProcessQuant() ...
新执行顺序:
Process()
└─ DuplicateConst() Duplicate(constScale/constInvScale) ──►
└─ CopyInSmooth() MTE2搬运smooth ─────────────────────►
WaitFlag<MTE2_V> ↓
↑ 两段并行，Duplicate在MTE2搬运期间完成
└─ ProcessQuant() ...
Duplicate 是纯向量指令，不占用 MTE2 通道，提前到 smooth 搬运之前执行，可与 DMA 完全并行，隐藏了常量初始化的延迟。
新增 constInvScale
旧代码用 Div(temp, constScale, temp) 求 1/max，然后 scaleLocal.SetValue(i, 1/scale) 再做一次倒数。新代码在 DuplicateConst 中预填：
// INT8对称量化，constScale = 127.0
Duplicate<float>(constScale, DYNAMIC_QUANT_INT8_SYM_SCALE, MAX_VALUE_NUM);
Duplicate<float>(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 对比：
| 维度 | DynamicQuantDb |
DynamicQuantMultiRow |
|---|---|---|
| 流水模型 | Double-Buffer（TQue） | 单缓冲（TBuf） |
| 每次处理粒度 | 单行 × loop | 多行一批 × loop |
| ReduceMax | 每行标量路径 | WholeReduceMax 批量 |
| Scale广播 | 每行 Brcb+Div |
向量Div + BroadCast |
为何不用 Double-Buffer：小shape场景下，单次搬运量（几行 × rowLen字节）很小，MTE2搬运极快，DMA和计算之间几乎没有可以 overlap 的时间窗口，double-buffer的收益接近零，反而引入了额外的 TQue 管理开销。用 TBuf<VECCALC> 直接一次性分配所有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);
PipeBarrier<PIPE_V>();
}
// 如果有多段（rowLen > 64），把第1段起的所有段 max 逐步折入第0段
if (repeatTimes > 1) {
Max(calcTensor, calcTensor[BLOCK_SIZE_64], calcTensor,
BLOCK_SIZE_64, repeatTimes - 1, repeatParams);
PipeBarrier<PIPE_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);
PipeBarrier<PIPE_V>();
// 第二步：批量求每行最大值（内含尾部折叠预处理，见 ComputeReduceMax）
// 结果：scaleTmp[i] = Row_i 的绝对值最大值
ComputeReduceMax(ubRows);
PipeBarrier<PIPE_V>();
// 第三步：向量除法，一次算出所有行的量化scale = constScale / rowMax
// stride参数 {dst=1, src0=0, src1=1, dstRep=8, src0Rep=0, src1Rep=8}：
// src0(constScale) stride=0 → 广播同一个值给每行
// src1(scaleTmp) stride=8 → 每行的max间隔8个float（WholeReduceMax输出格式）
Div(quantScaleTmp, constScale, scaleTmp,
BLOCK_SIZE_64, CeilDiv(ubRows, BLOCK_SIZE_64),
{1, 0, 1, STRIDE_8, 0, STRIDE_8});
PipeBarrier<PIPE_V>();
// 第四步：将每行的单个 quantScale 广播到该行所有 rowLen 个位置
// [ubRows, 1] → [ubRows, rowLen]，沿 axis=1 广播
BroadCast<float, DIMENSION_2, 1>(absTmp, quantScaleTmp, dstShape1, srcShape1);
PipeBarrier<PIPE_V>();
// 第五步：向量乘法完成量化缩放，所有行 ubRows × rowLen 个元素同时处理
// absTmp[i][j] = castTmp[i][j] × quantScaleTmp[i]，结果范围 [-127, 127]
Mul(absTmp, castTmp, absTmp, ubNums);
PipeBarrier<PIPE_V>();
// 第六步：两步 Cast 完成量化（float → int16 → half）
// CAST_RINT：四舍五入到最近整数
Cast(absTmp.ReinterpretCast<int16_t>(), absTmp, RoundMode::CAST_RINT, ubNums);
PipeBarrier<PIPE_V>();
// int16 → half：为后续 CAST_TRUNC 截断到 int8/int4 做准备
Cast(absTmp.ReinterpretCast<half>(), absTmp.ReinterpretCast<int16_t>(),
RoundMode::CAST_ROUND, ubNums);
再看下Profiling：去掉 Double Buffer，节省一份 fp16 输入缓冲（2 bytes/elem），UB 常量下调，ubAvail 从 5 升至 6，同时一次vector计算内，对应MAX指令在一次计算中执行6次，其余指令均按计算逻辑单次执行。
优化效果
单算子性能提升约30%左右，符合预期
作为“人工智能6S店”的官方数字引擎，为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。
更多推荐
2
0- 0
已为社区贡献1条内容
所有评论(0)