NVIDIA、CMU 和华盛顿大学发布了“FlashInfer”：一个为 LLM 推理和服务提供最先进内核实现的内核库 原创

01、概述

在人工智能的浩瀚星河中，大型语言模型（LLMs）犹如一颗耀眼的明星，为现代AI应用注入了强劲动力。它们是聊天机器人、代码生成器等工具背后的“大脑”，凭借海量数据和复杂算法，能够理解、生成自然语言，甚至编写代码，展现出惊人的智能。然而，随着对这些大模型依赖程度的加深，其推理过程中的诸多问题也逐渐浮出水面，成为制约其发展的瓶颈。

02、大模型推理的困境

注意力机制的局限

注意力机制是大模型的核心组件之一，它帮助模型聚焦于输入信息中最重要的部分，从而提高处理效率和准确性。但现有的注意力机制，如FlashAttention和SparseAttention，在面对多样化的任务负载、动态变化的输入模式以及有限的GPU资源时，常常显得力不从心。

• 多样化任务负载：不同的应用场景对模型的要求各异。例如，在聊天机器人中，用户的问题可能涉及多个话题，需要模型快速切换注意力焦点；而在代码生成中，模型则需要关注代码的逻辑结构和语法规范。现有的注意力机制往往难以灵活应对这种多变的任务需求。
• 动态输入模式：现实世界中的输入信息是不断变化的。以在线客服为例，用户在咨询过程中可能会随时插入新的问题或改变话题，导致输入序列的长度和内容不断变化。这种动态性使得模型难以提前规划注意力分配，影响推理效率。
• GPU资源限制：大模型通常需要大量的计算资源，尤其是GPU。但在实际应用中，GPU资源往往是有限的，需要在多个任务或模型之间共享。现有的注意力机制在资源受限的情况下，容易出现计算瓶颈，导致推理速度大幅下降。

高延迟与内存瓶颈

• 高延迟：在一些实时性要求较高的场景，如自动驾驶辅助系统，模型需要在极短的时间内做出决策。然而，现有的大模型推理过程往往存在较高的延迟，无法满足这种实时性需求。延迟过高会导致系统反应迟缓，影响用户体验和系统性能。
• 内存瓶颈：大模型通常包含数十亿甚至数千亿个参数，存储这些参数需要大量的内存空间。在推理过程中，模型还需要存储大量的中间计算结果，进一步增加了内存需求。现有的内存架构和管理方式难以高效地支持这种大规模的内存需求，导致推理速度受限，甚至出现内存溢出的情况。

03、FlashInfer的破局之道

为了解决上述问题，来自华盛顿大学、NVIDIA、Perplexity AI和卡内基梅隆大学的研究人员联合开发了FlashInfer，这是一个专为大模型推理量身定制的AI库和内核生成器。FlashInfer通过一系列创新的技术手段，为大模型推理带来了显著的性能提升和更高的灵活性。

技术特点与优势

1）全面的注意力内核支持

FlashInfer支持多种注意力机制，包括预填充、解码和追加注意力等，能够与各种KV缓存格式兼容。这种广泛的适应性使得FlashInfer在单请求和批量服务场景下都能发挥出色的性能。例如，在处理长文本生成任务时，FlashInfer能够根据不同的注意力机制灵活调整计算策略，提高生成效率和质量。

2）优化的共享前缀解码

FlashInfer采用了分组查询注意力（GQA）和融合旋转位置嵌入（RoPE）注意力等技术，实现了显著的速度提升。以长提示解码为例，FlashInfer在与vLLM的Page Attention实现相比，能够达到31倍的速度提升。这种优化不仅加快了推理速度，还使得模型能够更高效地处理复杂的长文本任务。

3）动态负载均衡调度

FlashInfer的调度器能够根据输入的变化动态调整，减少GPU的空闲时间，确保资源的高效利用。其与CUDA图的兼容性进一步增强了其在生产环境中的适用性。例如，在处理多个并发请求时，FlashInfer能够根据每个请求的计算需求和优先级，动态分配GPU资源，避免资源浪费和任务阻塞。

4）可定制的JIT编译

FlashInfer允许用户自定义并编译特定的注意力变体，生成高性能的内核。这一功能为满足特殊用例需求提供了极大的便利，如滑动窗口注意力或RoPE变换等。用户可以根据自己的应用场景和需求，灵活地定制注意力机制，从而进一步提升推理性能。

性能提升实例

FlashInfer在多个基准测试中都展现出了显著的性能提升：

• 延迟降低：与现有的解决方案如Triton相比，FlashInfer将令牌间延迟降低了29%至69%。在涉及长上下文推理和平行生成的场景中，这种延迟降低尤为明显。例如，在长文本生成任务中，FlashInfer能够更快地生成每个令牌，从而缩短整体推理时间。
• 吞吐量提升：在NVIDIA H100 GPU上，FlashInfer在平行生成任务中实现了13%至17%的速度提升。这表明FlashInfer在处理高需求应用时具有更高的效率，能够更快地完成大规模的推理任务。
• GPU利用率增强：FlashInfer的动态调度器和优化内核提高了带宽和FLOP利用率，尤其在处理倾斜或均匀序列长度的场景中表现突出。这意味着在有限的GPU资源下，FlashInfer能够更充分地发挥其计算能力，提升整体性能。
• 并行解码任务优势：FlashInfer在并行解码任务中也表现出色，其可组合的格式使得首次令牌时间（TTFT）大幅降低。以Llama 3.1模型（70B参数）为例，在特定配置下，FlashInfer将TTFT降低了22.86%。这使得在需要快速响应的应用场景中，如实时翻译或语音识别，FlashInfer能够更快地提供结果。

04、结语

FlashInfer为大模型推理提供了一个实用且高效的解决方案，显著提升了性能和资源利用率。其灵活的设计和强大的集成能力，使其成为推动大模型服务框架发展的重要工具。通过解决关键的效率问题并提供强大的技术解决方案，FlashInfer为更易用、更可扩展的AI应用铺平了道路。作为一个开源项目，FlashInfer还期待与研究社区的进一步合作与创新，以应对AI基础设施中不断出现的新挑战，持续推动人工智能技术的进步。

参考：

1. ​​https://arxiv.org/abs/2501.01005​​
2. ​​https://github.com/flashinfer-ai/flashinfer​​

本文转载自公众号Halo咯咯 作者：基咯咯

原文链接：​​https://mp.weixin.qq.com/s/LPKvpSs83691Ih_7TNNNMQ​​