在 Hexagon DSP 上,激活函数通常不是模型的主要耗时,但 exp、tanh、倒数等非线性运算的实现成本较高:它们往往需要近似计算、迭代求逆以及 FP16/FP32 数据类型转换,不容易映射成简洁的 HVX 向量指令。
这次我们针对 MNN Hexagon 后端的 FP16 激活路径做了一轮优化,核心思路是使用 PWL(Piecewise Linear,分段线性近似),以若干段直线替代复杂的非线性计算,并按 HVX 的 FP16 算术、vlut16 查表和 128 字节向量布局共同设计分段。该方案覆盖 Sigmoid、Tanh、GELU、SiLU 和 MulSiLU,最终使 MulSiLU 单算子的 DSP 耗时降低 52.47%、加速 2.10 倍;Qwen3-0.6B 的 BINARY_ELEMENTWISE 分类耗时在 prefill 和 decode 阶段分别降低 42.79% 和 27.48%。
本文记录从原始实现、PWL 映射、分段搜索、HVX 内核到精度与性能验证的完整过程,相关实现见 MNN PR #4700。
一、为什么激活函数在 DSP 上仍然昂贵
几个常见激活函数的定义如下:
\[\operatorname{Sigmoid}(x)=\frac{1}{1+e^{-x}}\] \[\operatorname{SiLU}(x)=x\cdot\operatorname{Sigmoid}(x)\] \[\operatorname{GELU}(x)\approx\frac{x}{2}\left(1+\tanh\left[\sqrt{\frac{2}{\pi}}\left(x+0.044715x^3\right)\right]\right)\]这些公式在 CPU 标量代码里只是一次函数调用,但在 DSP 向量内核中会展开为一串操作。优化前的 MulSiLU 路径大致需要:
- 把 FP16 输入拆分并转换为 FP32 向量;
- 用
exp2近似计算指数项; - 计算分母并通过倒数初值和 Newton 迭代求逆;
- 完成 Sigmoid、SiLU 和上投影分支的乘法;
- 再把 FP32 结果转换回 FP16。
Sigmoid、Tanh、GELU 和 SiLU 的 HVX 路径虽然已经是向量化实现,但它们仍然包含指数、倒数迭代、多次乘加或复合函数展开。对于本身只有逐元素计算的算子,这些数学步骤和 FP16/FP32 往返就是主要成本。
图中的核心变化可以概括为:不再实时计算完整的非线性函数,而是先判断输入属于哪一段,再查出该段的斜率和偏置,最后执行一次向量乘加。
二、PWL 是什么
PWL 的基本思想是把非线性曲线切成若干区间,并在每个区间内用一条直线近似:
\[y=f(x)\approx a_i x+b_i,\qquad x\in[e_i,e_{i+1})\]如果使用区间端点连线,初始系数可以写成:
\[a_i=\frac{f(e_{i+1})-f(e_i)}{e_{i+1}-e_i},\qquad b_i=f(e_i)-a_i e_i\]运行时只需要完成三件事:
- 根据
x确定分段编号i; - 从查找表取出
a[i]和b[i]; - 用 QF16 乘加计算
y=a[i]×x+b[i]。
这里有一个容易混淆的细节:“一次查表选段、两次查表取系数”特指默认的 learned8 SiLU 路径。 Sigmoid、Tanh 和 GELU 默认采用 companded16,它们直接利用 FP16 的指数和尾数位生成分段编号,然后只需两次 vlut16 读取 a、b。因此更通用、也更准确的描述是:
选段(位运算或 LUT) → 两次 LUT 取 a/b → 一次 QF16 乘加
为什么这种形式适合 HVX
Hexagon HVX 的一个向量是 128 字节,可以同时容纳 64 个 FP16 元素。PWL 路径中的符号提取、比较、查表和乘加都可以按 64 个 lane 并行执行,而且主体计算始终留在 FP16/QF16 向量域。
实现使用 vlut16 读取 FP16 系数。由于该指令会分别消费 halfword 的高、低字节,内核把 4 bit 分段编号复制到两个字节,再取 lookup 结果的低向量,从而保留全部 64 个 FP16 lane。系数表按 128 字节对齐和填充,既保证向量加载安全,也让编译器有机会把表加载移出元素循环。
三、分段不是越多越好
PWL 最直接的实现是在固定区间内等宽切分。我们最初使用 uniform32:在 SiLU 的 [0,8] 正半轴上按 0.25 的宽度切成 32 段。它精度直观,但需要两组 vlut16 bank;确定高低 bank、读取更多表项以及额外选择都会消耗指令和寄存器。
随后实现了 companded16,按函数曲率非均匀分配 16 段:
| 输入范围 | 分段宽度 | 原因 |
|---|---|---|
[0,2) | 0.25 | 曲率大,保留密集分段 |
[2,4) | 0.5 | 曲线开始变平 |
[4,8) | 1.0 | 接近饱和,允许更宽分段 |
这样一个 16-entry vlut16 bank 就能覆盖 [0,8]。Tanh 和 GELU 只需覆盖到 4,因此使用前 12 段即可。
但对于 LLM 中更重要的 SiLU/MulSiLU,我们继续做了面向硬件约束的分段搜索,得到默认的 learned8:
[0, 0.25), [0.25, 0.5), [0.5, 1), [1, 1.5),
[1.5, 3.5), [3.5, 5), [5, 6), [6, 8)
这里的 “learned” 不是在线训练一个模型,而是在 CPU 上搜索满足精度门限、同时适合 HVX 索引编码的分段边界和 FP16 系数。搜索器会模拟:
- FP32 输入进入 DSP 前的 FP16 量化;
- 斜率和偏置存为 FP16 后的误差;
- QF16 乘加到 FP16 输出的最终舍入;
- 与 DSP 完全一致的 FP16 位状态编码;
- 分段边界附近最容易出现的峰值误差。
learned8 不再做通用区间搜索,而是把 FP16 的指数和尾数高位压缩为 16 种状态,再用一张 16-entry LUT 映射到 8 个分段:
状态到分段:0, 0, 1, 1, 2, 2, 3, 4, 4, 4, 4, 5, 5, 6, 7, 7
于是每个向量的主要步骤变为:
FP16 位状态 → 1 次 vlut16 得到 segment
→ 2 次 vlut16 得到 slope / bias
→ QF16: y = slope × |x| + bias
需要注意,8 段并不意味着 DSP skeleton 中只保存 8 个紧凑的 half。HVX 查找表仍需按完整向量对齐和填充,所以“数学分段更少”不必然等于“二进制一定更小”;这里真正获得的是更短的数据通路。
四、利用对称性只拟合一半曲线
对称关系可以让 PWL 只处理正半轴,减少表项,并把负半轴恢复为简单的向量加减:
\[\operatorname{Sigmoid}(-x)=1-\operatorname{Sigmoid}(x)\] \[\tanh(-x)=-\tanh(x)\] \[\operatorname{SiLU}(-x)=\operatorname{SiLU}(x)-x\] \[\operatorname{GELU}(-x)=\operatorname{GELU}(x)-x\]内核先提取符号位并对 |x| 查表,算出正半轴结果,再通过 predicate 选择正负结果。曲线已经进入饱和区时则不再查表外推:
| 算子 | PWL 范围 | 正向饱和值 | 负向饱和值 |
|---|---|---|---|
| Sigmoid | |x| < 8 | 1 | 0 |
| Tanh | |x| < 4 | 1 | -1 |
| GELU | |x| < 4 | x | 0 |
| SiLU | |x| < 8 | x | 0 |
这种“正半轴近似 + 恒等式恢复 + 区间外饱和”的组合,比为整个实数域分别保存正负系数更适合向量实现。
MulSiLU 的融合路径
LLM 的门控 MLP 常见形式是:
\[\operatorname{MulSiLU}(u,g)=u\cdot\operatorname{SiLU}(g)\]优化后的 Binary 和 Loop 内核共享 learned8 SiLU 向量函数,在同一个 HVX 循环中先得到 SiLU(g),紧接着与 u 做 FP16 乘法并写回。这样既避免重复实现,也不会把中间 SiLU 结果落到额外缓冲区。
为什么 Log 没有强行使用 PWL
这轮优化同时把 Log 改为 HVX log2(x) × ln(2) 路径,但它不属于 PWL。Log 的输入动态范围和误差分布与饱和型激活不同,已有的向量 log2 helper 更合适。这个取舍也说明:目标不是把所有非线性函数统一成一种算法,而是为每种函数选择更适合硬件的数据通路。
五、正确性:CPU 模拟不能替代真机测试
近似函数优化最容易出现两类问题:一类是数学误差超限,另一类是数学模型正确但 HVX lane 映射、查表字节语义或 FP16 舍入不一致。
因此验证分为两层。
1. CPU 全量模拟
pwl_search.py 会遍历全部有限 FP16 输入,并模拟系数量化、QF16 乘加、输出舍入、对称恢复与饱和逻辑。对于 learned8 SiLU,还额外使用与真机单测相同的 FP32 输入网格,先量化成 FP16 再执行近似,但保留 FP32 数学结果作为参考。
learned8 SiLU 遍历全部有限 FP16 输入时的最大绝对误差为 0.00632850;按 Host 到 DSP 的 FP32→FP16 输入路径计算时,最大绝对误差为 0.00722693。
2. Hexagon 真机单测
真机测试使用 8193 个元素,刻意让长度不能被 64 个 FP16 lane 整除,以同时覆盖完整 HVX 向量和标量尾部。输入包含:
-
[-12,12]的密集网格; - 每个 0.25 边界的
edge-0.01、edge、edge+0.01; -
±0、±4、±8、±12、±100等特殊点。
在一台 v79 设备上的最大绝对误差如下:
| 算子 | 最大绝对误差 | 测试阈值 |
|---|---|---|
| Sigmoid | 0.00236678 | 0.005 |
| Tanh | 0.00627482 | 0.009 |
| SiLU | 0.00754023 | 0.008 |
| GELU | 0.00613671 | 0.009 |
| Log(HVX log2 路径) | 0.00332212 | 0.02 |
| MulSiLU | 0.07119751 | 0.08 |
这些专项测试只在选择 Hexagon 后端时执行;其他后端运行通用 run_test 时会自动跳过,避免把 DSP 专项测试误当成跨后端算子测试。
六、性能结果
性能对比在同一台 v79 设备上进行,保持 Host library、runtime、模型和测试参数完全一致,只替换 DSP skeleton。PWL 前原始实现、companded16 和 learned8 使用交替顺序运行并取 DSP 耗时中位数。
| 测试项 | PWL 前原始实现 | companded16 | learned8 | learned8 耗时降低 | 加速比 |
|---|---|---|---|---|---|
| MulSiLU 单算子(262144 个元素) | 10.2770 ms | 5.2995 ms | 4.8850 ms | 52.47% | 2.10× |
Qwen3-0.6B BINARY_ELEMENTWISE prefill | 56.0425 ms | 34.7880 ms | 32.0605 ms | 42.79% | 1.75× |
Qwen3-0.6B BINARY_ELEMENTWISE decode | 49.3120 ms | 37.1125 ms | 35.7615 ms | 27.48% | 1.38× |
MulSiLU 单算子从 10.2770 ms 降到 4.8850 ms,说明去掉 FP32 转换、指数和倒数迭代后,PWL 确实显著缩短了热点数据通路。模型 profile 中,prefill 的收益高于 decode;但两个阶段的 BINARY_ELEMENTWISE 分类耗时都明显下降。
这里报告的是 DSP 单算子与算子分类时间,用于直接观察激活内核的变化;它不等同于端到端 token throughput,后者还会受到 MatMul、Attention、调度和数据搬运等其他部分影响。
七、这次优化的几个经验
1. PWL 是软硬件协同问题,不只是拟合问题
只看数学误差,增加分段通常更容易逼近原函数;但在 HVX 上,更多分段可能意味着更多 LUT bank、选择指令、寄存器压力和表加载。最终需要共同优化的是:
分段边界 + FP16 系数量化 + 索引编码 + LUT 布局 + 向量指令数
2. 非均匀分段比盲目增加表项更有效
激活函数的误差敏感区域集中在零点附近和曲率较大的区间。把表项均匀铺满整个输入范围会浪费饱和区的容量。companded16 和 learned8 都在做同一件事:把有限的分段预算留给真正需要它的区域。
3. 对称性和饱和区也是算法的一部分
如果只把 PWL 理解成 y=ax+b,就会忽略一半优化空间。只拟合正半轴、用恒等式恢复负半轴,并在大输入区间直接返回极限值,既缩小查表范围,也减少误差累积。
4. 模拟器必须复现真实数据类型
用 FP64 拟合得到的漂亮误差,不代表 FP16/QF16 内核也能达到同样结果。系数舍入、输入量化、最终舍入以及边界映射都可能改变最大误差点。CPU 工具的价值不是替代 DSP,而是提前复现 DSP 的数值语义,把候选方案筛选成本降下来。
5. 真正该优化的是模型热点
单独的 Unary SiLU 并不是 LLM 中唯一、也不一定是最重要的调用形式。Qwen 的门控 MLP 主要命中融合 MulSiLU,因此让 Binary 和 Loop 路径共享同一个 learned8 内核,才把 PWL 的收益带进真实模型 profile。
八、总结
这次 Hexagon 激活函数优化的核心,是把 Sigmoid、Tanh、GELU、SiLU 和 MulSiLU 从“实时计算指数、Tanh、倒数及其组合”改造成 HVX 友好的“选段、查系数、一次乘加”数据流:Sigmoid/Tanh/GELU 使用非均匀 companded16,SiLU/MulSiLU 使用面向 HVX 位编码和精度约束搜索得到的 learned8。
最终在满足 FP16 推理精度要求的前提下,MulSiLU 单算子加速 2.10 倍,Qwen3-0.6B 的激活类算子耗时降低 27%~43%。更重要的是,这个过程说明端侧算子优化不能只从公式出发:近似算法必须和数据类型、查表指令、寄存器布局、融合位置以及真实模型热点一起设计。