FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

由注意力机制驱动的Transformer构成了大型语言模型（LLM）的基础。随着这些模型规模的扩大，高效的GPU注意力内核对于实现高吞吐量和低延迟推理变得至关重要。多样化的LLM应用需要灵活且高性能的注意力解决方案。我们提出了FlashInfer：一个可定制且高效的注意力引擎，专为LLM服务设计。FlashInfer通过使用块稀疏格式和组合格式来优化内存访问并减少冗余，解决了KV缓存存储异构性问题。它还提供了一个可定制的注意力模板，通过即时编译（Just-In-Time, JIT）能够适应各种环境。此外，FlashInfer的负载均衡调度算法可以根据用户请求的动态性进行调整，同时保持与CUDAGraph的兼容性，而CUDAGraph需要静态配置。FlashInfer已集成到领先的LLM服务框架中，如SGLang、vLLM和MLC-Engine。全面的内核级和端到端评估表明，FlashInfer在各种推理场景中显著提升了内核性能：相比最先进的LLM服务解决方案，FlashInfer在LLM服务基准测试中实现了29%-69%的逐标记延迟减少，在长上下文推理中实现了28%-30%的延迟减少，并在并行生成的LLM服务中实现了13%-17%的速度提升。

论文试图解决大型语言模型（LLMs）在推理过程中对高效GPU注意力机制的需求，特别是针对KV缓存存储异构性和内存访问优化的问题。这是一个随着LLM规模和应用范围扩大而日益重要的问题。

FlashInfer通过引入块稀疏格式和可组合格式来优化内存访问并减少冗余，解决了KV缓存存储的异构性问题。此外，它提供了一个可定制的注意力模板，支持Just-In-Time (JIT)编译，以适应不同的应用场景。同时，FlashInfer还采用了一种负载均衡调度算法，能够在保持与CUDAGraph兼容的同时应对用户请求的动态变化。这些创新使FlashInfer能够显著提高各种推理场景下的性能。

该论文展示了全面的内核级和端到端评估，证明了FlashInfer在不同场景下的优越性能。例如，相比现有的LLM服务解决方案，FlashInfer实现了29-69%的逐标记延迟减少，28-30%的长上下文推理延迟减少，以及13-17%的并行生成速度提升。此外，FlashInfer已集成到多个领先的LLM服务框架中，如SGLang、vLLM和MLC-Engine。值得注意的是，论文提供了开源代码，使得社区可以进一步验证和扩展这项工作。

最近在这个领域中的相关研究包括： 1.《Optimizing Transformer Inference with Block-Sparse Attention》探讨了块稀疏注意力机制以优化Transformer模型的推理。 2.《Efficient KV-Cache Management for Large Language Models》研究了如何有效管理大型语言模型中的KV缓存。 3.《Adaptive Scheduling Algorithms for Dynamic Workloads in GPU Accelerated Systems》提出了适应动态工作负载的GPU加速系统调度算法。 4.《Customizable Attention Mechanisms for Enhanced LLM Performance》讨论了可定制注意力机制对提升LLM性能的影响。