FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

LLM MQIA Long Context
2025年01月02日
由注意力机制驱动的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%的速度提升。
Transformers, driven by attention mechanisms, form the foundation of large language models (LLMs). As these models scale up, efficient GPU attention kernels become essential for high-throughput and low-latency inference. Diverse LLM applications demand flexible and high-performance attention solutions. We present FlashInfer: a customizable and efficient attention engine for LLM serving. FlashInfer tackles KV-cache storage heterogeneity using block-sparse format and composable formats to optimize memory access and reduce redundancy. It also offers a customizable attention template, enabling adaptation to various settings through Just-In-Time (JIT) compilation. Additionally, FlashInfer's load-balanced scheduling algorithm adjusts to dynamism of user requests while maintaining compatibility with CUDAGraph which requires static configuration. FlashInfer have been integrated into leading LLM serving frameworks like SGLang, vLLM and MLC-Engine. Comprehensive kernel-level and end-to-end evaluations demonstrate FlashInfer's ability to significantly boost kernel performance across diverse inference scenarios: compared to state-of-the-art LLM serving solutions, FlashInfer achieve 29-69% inter-token-latency reduction compared to compiler backends for LLM serving benchmark, 28-30% latency reduction for long-context inference, and 13-17% speedup for LLM serving with parallel generation.
许愿