摘要
随着基于大型语言模型的代理在操作系统工作流中的广泛应用，笔记本级系统芯片（SoC）上能效推理的重要性日益凸显。尽管云端卸载仍然常见，但它引发的可靠性和隐私问题在代理工作负载中尤为棘手。因此，近期的笔记本SoC集成了针对能效优化的神经处理引擎（NPU），然而将注意力机制有效映射到NPU上依然面临挑战，这主要源于架构多样性和显式数据移动编程模型。
在本研究中，我们提出了STEEL，这是首个针对XDNA类NPU的FlashAttention开源实现。STEEL引入了一种数据流形式的预填充注意力，使得空间并行性和片上内存的高效利用成为可能。此外，STEEL通过利用稀疏感知的管道放置解决了因因果掩码引起的负载不平衡问题，减少了同步开销，提高了利用率。
我们在AMD Ryzen AI 9 HX 370 SoC上评估了STEEL，并将其性能与优化的CPU和GPU实现进行了比较。实验结果表明，STEEL相较于CPU和GPU基线平均减少了9.17倍和1.75倍的能耗。在XDNA 1上，STEEL实现了平均9.6倍的延迟减少，相较于XDNA 2上的逐层注意力实现，平均实现了22.8倍的加速。
博主点评： STEEL的提出不仅展示了在能效优化方面的巨大潜力，也为长序列推理提供了有效解决方案。通过稀疏感知机制的应用，STEEL显著降低了能耗和延迟，为未来的NPU架构设计提供了新思路。