FastTPS：AI加速器LLM词元阶段的优化方法

大语言模型 (LLM) 的流行加剧了对有效推理的持续需求。然而，由于在仅解码器 LLM 推理中的词元阶段对词元进行顺序处理，固有的低并行性会导致人工智能 (AI) 加速器上计算单元的吞吐量降低和次优利用率，特别是在处理会带来大量内存开销的长序列输入时。最近，许多报道的方法已被开发为潜在的解决方案，因为它们出现了数值偏差。本文提出了 FastTPS，一种高性能、低精度损失的方法，用于加速通用 AI 加速器上 LLM 推理中的词元阶段，该方法包括三个关键组件：（1）AI 加速器支持的免重载 KV 缓存串联，可减少内存访问开销并实现注意力的完全融合，（2）基于平铺优化 FLAT 的高效率和高精度“RoPE”注意力，以及（3）具有细粒度流水线调度的高度融合 MLP。我们的结果证实，FastTPS 显着缓解了词元阶段的内存瓶颈，在具有 BF16 精度的 AMD Ryzen AI 300 系列 NPU 上实现了 6 倍的速度提升（与无融合相比），同时在 Phi3-mini-4k 指令推理期间维持 93% 的峰值内存带宽利用率。