加速Firefox本地AI运行时
去年我们推出了Firefox AI运行时，这个引擎默默支撑着PDF.js生成替代文本等特性，以及最近推出的智能标签分组功能。系统虽然能工作，但速度未达我们预期。
本文阐述了我们如何通过替换Transformers.js默认使用的onnxruntime-web（改用现在内置于Firefox的原生C++版本）来加速推理过程。
起点
Transformers.js是Hugging Face Python库的JavaScript对应版本。其底层依赖onnxruntime-web，这是ONNX Runtime的WebAssembly（WASM）构建版本。
典型的推理周期包括：
- JavaScript中的预处理（标记化、张量整形）
- WASM中的模型执行
- 回到JavaScript的后处理
即使有预热缓存，这个流程仍需跨越多个层级。真正的热点是矩阵乘法，在CPU上运行时使用通用SIMD实现。
为什么纯WASM不够
WASM SIMD很棒，但无法超越硬件特定指令，比如Apple Silicon上的NEON或现代Intel芯片上的AVX-512。
Firefox Translations（使用Bergamot）已经证明深入原生代码可以加速：它使用WASM内置函数，这些小型钩子让WASM能够调用使用这些内部函数编译的C++代码。这个昵称为gemmology的项目效果显著。
我们尝试将此技巧移植到ONNX，但大量运算符使得逐个重写难以维护。而且每次冷启动仍需支付JS/WASM预热开销。
切换到ONNX C++
Transformers.js通过一个极小的接口与ONNX Runtime通信。它创建一个会话，推送一个张量，然后拉取结果。这使得在不触及功能代码的情况下交换后端变得简单。
我们实现此目标的步骤是：
- 将ONNX Runtime C++直接引入Firefox代码树
- 通过薄WebIDL层将其暴露给JavaScript
- 将Transformers.js连接到新后端
从PDF替代文本等功能的视角来看，没有任何变化，它仍然调用await pipeline(...)。但在底层，张量现在直接进入原生代码。
将ONNX Runtime集成到构建系统
上游ONNX运行时不支持我们所有的构建配置，且代码量庞大。因此我们选择不将其加入代码树。相反，可以使用配置标志提供已编译的ONNX运行时版本。它最终会自动从Taskcluster下载（我们在那里为一系列支持的配置构建它）或由下游开发人员提供。这提供了灵活性，同时不拖慢我们通常的构建过程且维护需求低。
在Taskcluster上构建ONNX需要一些配置更改和上游补丁。目标是在速度和二进制大小之间找到平衡，同时与Firefox仓库的原生代码要求兼容。
最值得注意的是：
- 在没有异常和RTTI支持的情况下构建需要一些上游补丁
- 默认构建配置设置为MinSizeRel，编译使用LTO
成效
由于原生后端是即插即用的替代品，我们可以逐个功能启用它并收集实际数据。早期基准测试显示推理速度提升2到10倍，且零WASM预热开销。
例如，在首次运行时可能滞后的智能标签分组主题建议，现在相当迅捷，这是我们为Firefox 142逐步迁移到此后端的第一个功能。
用于PDF.js替代文本功能的图像到文本模型也受益于此更改。在相同硬件上，延迟从3.5秒降至350毫秒。
下一步计划
我们正在整个夏季逐步将此新后端推广到更多功能，因此所有基于Transformers.js构建的功能都能利用它。
而且手头有了C++ API，我们计划解决一些长期存在的痛点，并启用GPU支持。
这些更改将随我们定制的ONNX Runtime一起发布，并在未来为基于Transformers.js的功能在我们的运行时中提供最佳性能。
1. DequantizeLinear实现多线程
DequantizeLinear操作是单线程的，常常主导推理时间。虽然上游工作最近合并了一项改进（PR #24818），但我们构建了一个补丁来将工作分散到多个核心，让编译器自动向量化内部循环。结果几乎是线性的加速，尤其是在多核机器上。
2. 矩阵转置实现多线程
类似地，在执行推理任务时，通常需要转置非常大的矩阵（数十兆字节）。此操作原本使用嵌套for循环简单完成。切换到多线程缓存感知的平铺转置方案，并利用SIMD，使得能够充分利用现代硬件，将此操作加速超线性因子，通常是分配给此任务的线程数的两倍，例如使用4个线程实现8倍加速。
这可以解释为：朴素的for循环虽然自动向量化，但对CPU缓存的使用效果不佳。
3. 缓存编译图
在推理运行之前，ONNX Runtime会为当前平台编译模型图。在大型模型（如Qwen 2.5 0.5B）上，每次启动可能耗时长达五秒。
我们可以动态地将编译图与权重分开缓存，节省从几毫秒到整整五秒的时间。
4. 使用GPU
目前，我们仅集成了基于CPU的提供程序。下一步是支持GPU加速的ONNX后端，这将需要更多努力。这是因为GPU支持需要额外的沙盒化，以安全地与底层硬件交互。
结论
这次迁移的有趣之处在于，我们能够在逐步迁移功能的同时如此大幅度地提升性能，而且所有这些都是在完全隔离的情况下完成的，无需更改任何功能代码。
虽然从用户体验的角度来看，加速已经可见，但我们相信未来可以并且将会发生许多改进，进一步提高基于ML的功能的效率，并使它们对更广泛的受众更加易用。
有想法、问题或错误报告？在Discord的firefox-ai频道（https://discord.gg/TBZXDKnz）联系我们，或在Bugzilla上提交问题，我们洗耳恭听。