JUNO++：通过光线追踪核心优化ANNS并实现大规模语言模型（LLM）中的高效稀疏注意力机制

《ACM Transactions on Architecture and Code Optimization》：JUNO++: Optimizing ANNS and Enabling Efficient Sparse Attention in LLM via Ray Tracing Core

基于光线追踪的近似最近邻搜索优化与注意力计算集成方法研究。提出JUNO++系统，包含光线追踪核心的ANNS搜索管道和注意力集成模块，解决产品量化中的冗余计算问题，实验表明吞吐量提升2.2-8.5倍，延迟降低46%。

　　

摘要

近似最近邻搜索（ANNS）是现代智能应用中的关键技术，包括推荐系统和向量数据库。随着大型语言模型（LLMs）的出现，ANNS在实现注意力剪枝机制方面发挥了关键作用，这些机制利用了注意力的稀疏性，例如top-K注意力机制和检索注意力机制。因此，ANNS的效率变得日益重要。在本文中，我们发现了基于产品量化的现有ANNS方法中的一个主要效率问题：即与代码本进行成对距离计算时的冗余计算和数据积累。为了解决这个问题，我们提出了JUNO++系统，该系统包括：i) 一个基于光线追踪核心的端到端ANNS搜索流程，该流程采用了考虑稀疏性的算法；ii) 将基于光线追踪的ANNS搜索流程集成到注意力计算中。在四个数据集上的评估表明，搜索吞吐量提高了2.2倍到8.5倍。对于基于ANNS的稀疏注意力机制，JUNO++将q × k^?计算的延迟降低了46%，同时保持了几乎相同的准确率。这不仅是基于检索的稀疏注意力机制的关键组成部分，也是在长上下文场景中的主导因素，意味着整体性能有了显著提升。