背景(需求)

quantize_add_layer_norm 是 (x1 + x2 + bias) → LayerNorm(mean/rstd, gamma/beta) → 静态量化 → y(INT8) + x(残差) 的融合算子,属于 LayerNorm-quant 家族:

x        = x1 + x2 + bias
mean     = sum_d(x) / N                       (N = 每行列数)
rstd     = 1 / sqrt(Var_d(x) + eps)
norm     = (x - mean) * rstd * gamma + beta
y(int8)  = round(norm / scales + zero_points)  per_channel(默认);mul_mode 为 ×scales;per_tensor 为 ×单标量
x_out    = x                                   残差输出(additional_output=false 时不写回)

此前该算子仅支持 910B / 910_93 / Kirin,950 缺位:def.cpp 无 AddConfig("ascend950")、算子目录内 grep ascend950|arch35|regbase 零命中、config/ascend950/ 目录不存在。同族 add_layer_norm_quant 已全量支持 950,补齐本算子即可让 LayerNorm-quant 家族(add_layer_norm_quant / quantize_add_layer_norm / dua_quantize_add_layer_norm)在 950 上完整可用,直接服务 W8A8 量化推理的归一化段。

实现进展

已由 MR #8234 实现(单提交,基于 upstream/master;初始 9b3f1af1d,2026-09-10 为过 CI pre-commit 门禁 amend 为 860058449 —— 当前等待人工评审,欢迎到 MR 检视。

以下展开 tiling 排产与 kernel 运算的具体做法。

tiling 策略

tiling 跑在 CPU 上、kernel 启动之前,回答:分几个核、每核几行、每轮往工作台(UB,约 256KB)搬多少,产出一张 12 字段的调度表 + 一把 tiling key。决策全部在 host 完成,kernel 只照表执行。

① shape 归一:x=[8192,4096]、gamma=[4096] → rows=8192、cols=4096(一行 = 一个归一化单元);cols 向上对齐到 32 的倍数(DMA 最小块 32B,int8 输出按 1B 定标)。

② 切行(双取整回填):

rowsPerCore     = ceil(rows / 核数)           // 粗分
usedCoreNum     = ceil(rows / rowsPerCore)    // 活少时算实际需要几核
rowsPerCore     = ceil(rows / usedCoreNum)    // 回填抹平
rowsPerTailCore = rows - rowsPerCore×(核数-1) // 末核兜余数

例:rows=8192、50 核 → 49 核各 164 行 + 末核 156 行;rows=10、50 核 → 只用 10 核(没有第②步会点亮 50 核、40 核空转)。

③ 排工作台,两种策略二选一(整行放得下优先 full_load):

工作台(UB,约 256KB)有限,一行数据可能就超了。分界线一条:一行能不能整个放上工作台?

full_load 优先的原因:两遍法数值路径与 910B 完全一致,golden 比对友好;welford 是"装不下"的诚实兜底,用"读两遍"换"装得下"。两个都排不下 → 明确报错,不静默算错。

策略① full_load:一张预算表定每轮行数。 把 UB 摆成固定区(不随行数变,先扣)+ 行区(每装一行都要花这么多):

┌─────────────────────── UB ≈ 256KB ────────────────────────┐
│ 固定区(先扣):                                            │
│   binaryAddUbSize  归并树缓冲      ← 行宽二次幂化后摊 32B 块 │
│   quantBufSize     量化参数整行    ← scales[+zp],常驻       │
│   weightBufSize    权重整行        ← gamma+beta[+1D bias]   │
│   UB_RESERVED      硬件保留 256B                            │
├───────────────────────────────────────────────────────────┤
│ 行区:rowStep 行 × (每行开销)                               │
│   x1+x2 各一份 + y(int8) + x 输出,全部 ×双缓冲               │
│   + [逐元素 bias 再加一份] + fp32 临时行 + 2 个标量位         │
└───────────────────────────────────────────────────────────┘
rowStep = 剩余UB ÷ 每行开销    → ≥1 就用本策略

跟着算一遍(x=[8192,4096]、bf16 主路、scales=fp32、广播 bias、无 zp):

固定区 = 归并树 128 + 量化参数 4096×1份(scales)×4B = 16384
       + 权重 4096×3份(γ+β+广播bias)×2B = 24576 + 硬件保留 256 → 合计 41344 B
剩余   = 262144 − 41344 = 220800 B
每行开销 = (x1 的 2B×2份 + y 的 1B + x 输出的 2B)×双缓冲×4096 = 57344
         + fp32 临时行 4B×4096 + 2 个标量位 = 16392  → 共 73736 B/行
rowStep = 220800 ÷ 73736 = 2.99 → 每轮 2 行(每核 164 行 ÷ 2 = 82 轮,双缓冲让搬运与计算重叠)

两本账分开记:广播 bias 整行常驻、不随行数涨,并入权重区(weightTensorNums_ +1);逐元素 bias 随行增长,才进每行开销(full_load 记在 inOutCols、welford 记在 elewiseSliceNums)——谁也不漏、谁也不重。

策略② welford:按"每列开销"记账。 full_load 排不下时,固定装 1 行 × colsPerLoop 列,预算按每列花多少字节算(所有切片 ×双缓冲):

量化参数 = 2缓冲 × 4B × 1份(scales)                    = 8  B/列
权重     = 2缓冲 × 2B × 3份(γ+β+bias)                 = 12 B/列
主体     = 2缓冲 × 2B × 3份(x1/x2/x输出) + 2缓冲×1B(y) = 14 B/列
标量+树  = 4B × 3(mean/var 2个标量 + 归约树≤1切片)     = 12 B/列
                                                        ────────
                                                        46 B/列

跟着算一遍(超宽行:cols=131072、fp16、广播 bias、无 zp):

colsPerLoop   = (262144−256) ÷ 46 = 5693 → 向下对齐到 64 的倍数 = 5632 列
                (寄存器按整条装载,切片是 64 的倍数才不浪费)
colsLoopCount = ceil(131072 / 5632) = 24 轮
colsTail      = 131072 % 5632 = 1536(最后一轮的尾巴)

两个精修动作:均分重算——cols 恰能被轮数整除时,把 colsPerLoop 改成均分值,每轮等长、没有残缺尾巴,kernel 收尾走无分支的"对齐版"快路;对齐 64——让寄存器按整条装载。

④ 归并树参数:binaryAddNum / K / LastNum:求均值要把一行(或一片)加起来,调度表带三个"归并树形状"字段,把元素个数换算成寄存器级二叉归并树的参数:

full_load 例:binaryAddNum = 4096 的二次幂化 = 2048
             2048 ÷ 64(一条寄存器装 64 个数)= 32 条 ≤ 64 → K=0、LastNum=32
             → 末层一把硬件 ReduceSum 把 32 条并成 1 条
welford 例: binaryAddNum = 5632 的二次幂化 = 4096 → 64 条 → 恰好一层,K=0、LastNum=64

K = ⌈log₂(部分和条数 ÷ 64)⌉:条数 >64 时每归并一级条数减半,K 级后恰好压进一条寄存器(64 个标量),末层 ReduceSum 收头。常规隐藏宽度(≤8192)全是 K=0,行宽每大 64 倍多一级(16384→K=1、65536→K=3)。设计动机:跨寄存器通道的 ReduceSum 又贵又慢,元素级 Add 便宜——先用便宜的加法把数据折半,贵的归约只在末层用一次。

⑤ 组钥匙:key = 8000 + 策略(0/100) + bias(1 逐元素/2 广播) + 量化模式(0 mul / 10 div / 20 per_tensor);上例 bf16+广播+per_channel → 8012。

kernel 运算(NPU 侧)

全链 fp32:x1/x2/bias(bf16/fp16)装载时各自 Cast 升 fp32 再相加,mean/var/仿射/量化全程 fp32,仅两头转换(残差 x 出口 Cast 降回原 dtype;y 的 int8 本是终点)。

full_load 两遍法求 mean/var:先扫一遍求均值、再扫一遍求方差,加法顺序与 910B 完全一致。行内求和使用"对折 + 归并树"。背景:整行数据摊在若干条向量寄存器里(每条 64 个数),硬件加法分两种身价:

第①步 对折:整行从折痕 binaryAddNum(严格小于 cols 的最大二次幂)切开,
     前段第 i 条与后段第 i 条 → Add(Q_i, R_i)      ← 便宜的元素级加
                     ↓ 折后条数减半
第②步 条内归约:每条部分和条 ReduceSum(64 路 → 1 个标量)→ 部分和进归并缓冲
第③步 末层收尾:部分和 ≤64 个时装进 1 条寄存器,mask=LastNum 一把 ReduceSum → mean
     (超宽行先做 K 级二分归并,条数逐级减半压到 ≤64 再收头)

经济账(4096 例):朴素 ReduceSum×64 + 末层×1 = 65 条贵指令;对折 Add×32 + ReduceSum×32 + 1 = 32 条便宜 + 33 条贵。4096 个数加成 1 个共 4095 次加法、两种做法次数相同——省的不是次数,是把恰好一半的加法从贵引擎搬进便宜引擎。第二遍算 var 重读 x 的 fp32 副本,同一套树形再用一遍(三个归并树参数即为此下发)。

welford 单遍流式统计(宽表路径):一行放不下就切段流式——读一段、更新一次统计量、扔掉。数据流:

原行 x[0..cols-1]
   ↓ 切段(J 段 × 每段 P 列;P = colsPerLoop,由 UB 预算定,通常几千)
段0: x[0..P-1]    段1: x[P..2P-1]    …    段(J-1)
   ↓ 按段内列位置抽给 P 个统计器(一条向量指令同时滚 64 个,P/64 条指令滚一段)
第 0 列统计器: x[0],   x[P],     x[2P],    …   ← 每个统计器跨段收 J 个数
第 1 列统计器: x[1],   x[P+1],   x[2P+1],  …
   ↓ 每个统计器独立跑三行更新(n = 已见样本数 = 段号+1)
(mean_0, M2_0), (mean_1, M2_1), …, (mean_{P-1}, M2_{P-1})
   ↓ Chan 合并(收尾走"对折+归并树"同款归并)
整行 mean、var → rstd = 1/√(var+eps)

三行更新:

delta1 = x − mean(旧);  mean += (1/n)·delta1;  M2 += delta1 × (x − mean(新))

收尾用 Chan 合并公式把 P 份部分统计并成整行:

整行 mean = Σ mean_p / P
整行 var  = Σ [ M2_p + J·(mean_p − 整行mean)² ] / cols

数字验证(行 [1,3,5,7],P=2、J=2):切段 [1,3] / [5,7],滚出 mean=[3,5]、M2=[8,8];合并 mean=(3+5)/2=4 ✓,var=(8+2×1 + 8+2×1)/4 = 5 ✓——不加 J·delta² 修正会算成 16/4=4 ≠ 5,漏的正是列间差异。J·delta_p² 是"列间差异"修正(总方差 = 组内 + 组间;M2_p 只记了组内,M2 为离差平方和、从没除过任何东西)。数值上全程只累积小偏差的乘积、不碰大数平方和,这是不用 E[x²]−E[x²] 朴素一遍法的原因。

量化收尾:per_channel 走高精度向量除(norm ÷ scales),mul/per_tensor 走乘;标量 scale 用"GM 拷 1 元素 + DIST_BRC 广播装载"整条进向量乘。round 由三步链完成(RINT 四舍五入 → fp16 中转 → int8 截断,舍入模式定死在 CastTrait)。welford 路径统计完重读一遍输入做量化(切片未驻留,带宽换内存);x 残差在统计遍顺手写回。

源码级复用 add_layer_norm 的 arch35 共享头与 norm_common 归约原语(编译期 include,运行时零依赖),op_host/CMakeLists 声明 DEPENDENCIES(norm 族 39 算子惯例)。

验证

真实 Ascend 950 环境(蓝区,CANN 9.0.0-beta.2);910B/910_93 全矩阵回归(910_93 为编译 + 注册验证);2026-09-07 环境重置后对最终提交重新 clone 复验,以下均为终态:

• UT:host tiling 新增 12 例 + 既有 2 例回归 = 14/14 PASSED;kernel 3 例(full_load / fast-path / welford × 3 量化模式)跑通冒烟 —— c310 模拟器不模拟量化系 regbase 指令流,数值比对交由真卡 ST(同仓 rms_norm_quant_v2 等同此限制)。
• 精度:golden 生成/比对脚本随算子交付(gen_data.py / compare_data.py);判据:浮点 x 严格容差,INT8 y 允许 ±1 舍入抖动且占比 ≤2%(welford 与两遍法加法顺序不同的固有差异),差 2 判错。
• binary:恰好 6 个 .o 与 binary.json 6 条一一对应,binary_info_config 注册正确。
• 回归:910b host 14 + kernel 6 无回归;910_93 编译通过。

关联

• 实现 MR:#8234(feat(norm): quantize_add_layer_norm 增加 Ascend 950 (arch35/regbase) 支持)