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