图谱动态｜学苑周刊 NO.266

 2026

详情访问：

https://conferences.sigappfr.org/apweb2026/

FlashInfer

本周推荐的是arxiv 2025.01上的论文：FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

概述

FlashInfer 提出了一个全面的注意力引擎，旨在解决高效大型语言模型（LLM）推理服务中的复杂挑战。现代 LLM 部署涉及多样的注意力模式、动态工作负载和不同的硬件要求，现有解决方案难以统一处理。FlashInfer 通过结合灵活的数据结构、可定制的计算模板和动态运行时优化，采用统一的方法来解决这些挑战。

图1：FlashInfer 的统一架构结合了块稀疏 KV 缓存格式、动态编译和负载均衡调度，以高效处理多样化的 LLM 服务工作负载。

该系统解决了 LLM 服务中的两个主要瓶颈：KV 缓存存储模式的异构性，以及针对不同硬件和模型配置对专用注意力核函数的需求。通过提供一个能够适应各种场景同时保持高性能的统一解决方案，FlashInfer 代表着向更高效、更易维护的 LLM 推理系统迈出了重要一步。

技术方法

统一的 KV 缓存存储格式

FlashInfer 引入了一种块稀疏矩阵格式（Block-Sparse Row 或 BSR），统一了现代 LLM 服务系统中使用的多种 KV 缓存存储模式。传统的方案，如 PagedAttention 和 RadixAttention，使用不同的内存布局，需要专门的核函数。FlashInfer 的 BSR 格式通过可配置的块大小，能够表示连续和非连续的内存布局，从而使单个核函数实现能够处理多种存储方案。

该系统将这一概念扩展为“可组合格式”——即在单个注意力计算中同时使用多种块稀疏格式的能力。这对于前缀共享场景特别有利，其中公共前缀可以使用更大的块大小存储以最大化内存复用，而独特部分则使用更小的块以保持灵活性。数学表示使用索引将逻辑块映射到物理内存位置：

可定制的注意力模板

FlashInfer 实现了一个即时（JIT）编译框架，允许用户通过基于模板的方法指定自定义注意力变体。用户使用函子（functors）如QueryTransform、KeyTransform、LogitsTransform和OutputTransform作为 CUDA 代码字符串来定义注意力修改。然后，系统通过将这些自定义操作插入到优化的 FlashAttention2/3 模板中来生成专用核函数。

这种方法使得像旋转位置嵌入（RoPE）这样的操作可以直接融合到注意力核函数中，从而消除了单独启动核函数的开销。例如，一个融合的 RoPE 操作可以定义为：

template<typename T>

__device__ void apply_rope(T* q, T* k, float* cos, float* sin, int pos) {

    Custom RoPE implementation fused into attention

}

JIT 编译器会自动生成集成了这些自定义操作的优化 CUDA 核函数，同时提供灵活性和性能。

动态负载均衡调度

FlashInfer 通过受 Stream-K 方法启发的动态调度算法解决了可变序列长度和 KV 缓存大小的挑战。调度器将注意力计算分解为工作单元，并动态地将它们分配给 GPU 流式多处理器（SMs），以最大限度地减少负载不平衡。

调度算法分两个阶段运行：

1. 规划阶段：在 CPU 上执行，分析序列长度分布并生成最优的工作分配计划。

2. 执行阶段：GPU 核函数遵循预先计算的计划，以平衡 SM 之间的工作负载。

此方法通过使用具有固定网格大小和预分配工作区缓冲区的持久化内核，保持了与CUDA-Graph优化的兼容性，确保动态调度不会损害图级性能优势。

主要贡献和成果

性能提升

FlashInfer在多个评估场景中展示了显著的性能提升：

端到端LLM服务：与SGLang框架集成后，FlashInfer在Llama 3.1 8B和70B模型上，针对ShareGPT和合成工作负载，将token间延迟相对于基于Triton的实现降低了29-69%。

动态工作负载处理：FlashInfer在非均匀序列长度分布上显著优于FlashAttention内核，通过其负载均衡调度实现了更高的带宽利用率。对于传统方法因负载不平衡而受损的倾斜分布，性能差距最为明显。

定制化优势：通过FlashInfer的JIT框架生成的融合RoPE内核实现了比单独的RoPE和注意力内核高1.6-3.7倍的带宽利用率，这表明了操作融合的有效性。

高级功能

用于前缀共享的可组合格式：FlashInfer的可组合格式方法对于具有共享前缀的并行生成场景（尤其适用于中等并行度 n=4）将token间延迟降低了13-17%，并将首个token时间降低了16-22%。

细粒度稀疏性支持：对于需要向量级稀疏模式的算法，FlashInfer通过高效利用密集张量核心，即使在小块大小下，也比PyTorch SDPA和FlexAttention加速高达20倍。

硬件优化：FlashInfer的CUDA/CUTLASS后端在各种注意力变体中始终优于基于Triton的方法（如FlexAttention），对于更长的序列，由于更好地利用了高级GPU功能，性能差距更大。

论文、讨论和资源链接：https://www.alphaxiv.org/abs/2501.01005，感兴趣的读者可以关注。