高效Attention引擎是怎样炼成的？陈天奇团队FlashInfer打响新年第一枪！

原标题：高效Attention引擎是怎样炼成的？陈天奇团队FlashInfer打响新年第一枪！
文章来源：新智元
内容字数：6055字

FlashInfer：高效LLM推理引擎

新年伊始，陈天奇团队发布了FlashInfer论文，旨在提升大语言模型（LLM）的推理效率。该论文详细介绍了FlashInfer如何通过块稀疏、可组合、可定制以及负载均衡等技术，在各种推理场景下显著提升性能。

1. LLM推理的挑战

   传统的Transformer模型在实际应用中面临诸多工程挑战，例如多样化的工作负载、个性化的Attention实现以及对延迟和吞吐量的严格要求。不同的推理阶段（Prefill、Decode、多轮对话）和技术（Speculative Decoding）都带来了不同的计算访存比，影响GPU利用率。此外，各种LLM框架（如vLLM、SGLang）在KV cache存储方面也存在差异，增加了框架的复杂性。

2. 块稀疏格式：高效存储KV Cache

   FlashInfer提出使用块稀疏（BSR）矩阵来统一存储各种形式的KV cache。BSR类似于CSR，但操作单元由单个元素变为数据块，从而提高GPU的利用率。对于Paged Attention、Radix Tree和Token Tree等不同形式的KV cache，BSR都能有效地减少冗余，并最大化Tensor Core的利用率。通过将稀疏块数据加载到共享内存并排布成密集格式，FlashInfer消除了硬件浪费。

3. 块并行与可组合性

   FlashInfer采用与BPT相同的块并行方法分解KV cache，实现并行计算，并通过可组合特性支持不同大小块的组合，在共享和全局内存之间取得平衡，进一步优化性能。

4. 可定制的注意力机制

   FlashInfer为FlashAttention开发了CUDA/CUTLASS模板，支持不同架构的Nvidia GPU，并提供可自定义的Attention机制。开发者可以通过插入自定义函数来实现个性化的Attention计算，例如ALiBi、Logits SoftCap、RoPE等，提高了框架的灵活性。

5. 高效的数据移动和负载均衡

   FlashInfer采用高效的数据加载方法，将tiles从全局内存转移到共享内存，并利用异步复制指令最大化内存带宽。同时，它根据硬件资源和工作负载选择合适的块大小，并通过负载均衡调度算法，根据用户请求动态调整资源分配，并兼容CUDAGraph。

6. 性能提升

   实测结果表明，FlashInfer显著提升了LLM推理性能，token间延迟降低了29%-69%，长上下文推理延迟降低了28%-30%，并行生成速度提高了13%-17%。

总而言之，FlashInfer通过巧妙地结合块稀疏格式、块并行计算、可组合性和负载均衡等技术，为构建高效且可定制的LLM推理引擎提供了一种新的解决方案，显著提升了LLM的推理效率。

联系作者

文章来源：新智元
作者微信：
作者简介：智能+中国主平台，致力于推动中国从互联网+迈向智能+新纪元。重点关注人工智能、机器人等前沿领域发展，关注人机融合、人工智能和机器人对人类社会与文明进化的影响，领航中国新智能时代。