在 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 路径大致需要:

  1. 把 FP16 输入拆分并转换为 FP32 向量;
  2. 用 exp2 近似计算指数项;
  3. 计算分母并通过倒数初值和 Newton 迭代求逆;
  4. 完成 Sigmoid、SiLU 和上投影分支的乘法;
  5. 再把 FP32 结果转换回 FP16。

Sigmoid、Tanh、GELU 和 SiLU 的 HVX 路径虽然已经是向量化实现,但它们仍然包含指数、倒数迭代、多次乘加或复合函数展开。对于本身只有逐元素计算的算子,这些数学步骤和 FP16/FP32 往返就是主要成本。

原始非线性路径与 PWL 路径对比

图中的核心变化可以概括为:不再实时计算完整的非线性函数,而是先判断输入属于哪一段,再查出该段的斜率和偏置,最后执行一次向量乘加。


二、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\]

运行时只需要完成三件事:

  1. 根据 x 确定分段编号 i;
  2. 从查找表取出 a[i] 和 b[i];
  3. 用 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 耗时中位数。

Hexagon PWL 优化前后 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%。更重要的是,这个过程说明端侧算子优化不能只从公式出发:近似算法必须和数据类型、查表指令、寄存器布局、融合位置以及真实模型热点一起设计。