在 Hexagon DSP 上,激活函数通常不是模型的主要耗时,但 exptanh、倒数等非线性运算的实现成本较高:它们往往需要近似计算、迭代求逆以及 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 读取 ab。因此更通用、也更准确的描述是:

选段(位运算或 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.01edgeedge+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 前原始实现、companded16learned8 使用交替顺序运行并取 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%。更重要的是,这个过程说明端侧算子优化不能只从公式出发:近似算法必须和数据类型、查表指令、寄存器布局、融合位置以及真实模型热点一起设计。