出色的 LLM 推理需要两个关键要素：速度和开发者速度。速度是指通过使用高度优化的计算内核算法，最大限度地提高底层硬件的效率。开发者速度是指快速采用这些新内核并加速新模型、算法和硬件的能力。最终，这种速度的基础是快速部署在底层 GPU 上运行的新计算内核，以及将这些内核轻松引入框架。
FlashInfer 是一个可定制的高效库，用于构建高效的 LLM 服务引擎。它使用块稀疏和可组合格式优化 KV 缓存存储，以改善内存访问并减少冗余，并具有可定制的注意力模板，可通过即时 (JIT) 编译来适应各种设置。其负载平衡调度算法可根据动态用户请求进行调整，同时保持与 NVIDIA CUDA Graph 静态配置兼容。FlashInfer 已集成到领先的 LLM 服务框架 (例如 MLC Engine、SGLang 和 vLLM) 以及多个自定义引擎中。
FlashInfer 最初是卡内基梅隆大学 Paul G. Allen 计算机科学与工程学院和 OctoAI (已被 NVIDIA 收购) 的协作研究项目。该团队旨在创建一个灵活的 LLM 推理内核库，该库与引擎无关、高度优化且易于扩展，适用于 KV 缓存重用算法等新技术。它现在是一个蓬勃发展的开源项目，生产部署以及整个 AI 系统社区的研发团队所做的贡献。
技术论文《FlashInfer：Efficient and Customizable Attention Engine for LLM Inference Serving》在MLSys 2025上荣获最佳论文奖。
NVIDIA 目前正在 FlashInfer 中积极发布性能超强的 LLM 推理内核，包括来自 NVIDIA TensorRT-LLM 的推理内核，以便轻松集成到 vLLM、SGLang 和自定义推理引擎中。
FlashInfer 架构概述
作为用于 LLM 服务的专用 NVIDIA GPU Operator Stack，FlashInfer 旨在提高最新内核的速度和开发者速度。推理平台可以采用新创意，而无需等待新库或在 CUDA C++ 中重写内核。这些内核将通过 DLPack API 提供给所有框架，并注册为 PyTorch 运算符，以便轻松集成到许多推理引擎中。JIT 功能使用户能够实现目标模型使用的核函数，这意味着 FlashInfer 的依赖项占用空间最小。
FlashInfer 将 LLM 工作负载分为四个运算符系列 ( Attention、GEMM、通信和采样) ，并通过轻量级、高性能的集合公开每个系列的运算符，这些集合只需更改最少的代码即可放入任何服务引擎。
Attention
现代推理请求的序列长度、KV 缓存块大小、掩码规则和位置编码方案各不相同。FlashInfer 通过以下方式吸收这种活力：
- 统一存储：将每个缓存布局表示为block/vector稀疏矩阵。
- 模板和 JIT 内核：CUDA/CUTLASS 代码库，其专用旋钮、logits/key/query、分组、MLA 和未来变体。
- Inspector-executor 接口：一个 PyTorch 友好型 API，可首先检查请求形状和前缀共享模式，然后通过轻量级调度程序启动经过调优的内核，以保持 GPU 饱和。
GEMM 和通信
LLM 块仍然严重依赖矩阵乘法。除了传统的 GEMV/GEMM 计算和全归约通信之外，近期的进展 (例如混合专家层和 LoRA 层) 引入了新的要求，例如分组 GEMM (单次调用中的许多小矩阵乘法) 和多对多通信。FlashInfer 选择最快的开源或 NVIDIA 内核 (包括 fp4/fp8 tensor-core 路径) ，并将它们提供给一个一致的 API，因此服务堆栈可以在不接触应用逻辑的情况下交换 GPU 或内核。
Token 采样
生成下一个 token 通常会成为 Top-K/ Top-P 过滤的瓶颈。传统的实现会对整个词汇进行分类，而当只有少数 logit 很重要时，这会造成浪费。FlashInfer 使用基于拒绝、无排序的采样器取代了全局排序，该采样器可实时修剪不可能出现的 token，从而降低大型词汇表的延迟，并保持数字准确性。
面向未来的推理
有了这些层，服务框架就可以改变 KV 缓存布局，引入新的注意力设计，批量任意长度，或追求更严格的延迟目标，而无需重写内核或回退到 CPU。从第一个查询到最终 token，FlashInfer 在 GPU 上保留关键推理路径，灵活、面向未来且快速。
使用 FlashInfer
PyPI 上提供了 Flashinfer 软件包。您可以通过以下方式进行尝试：
pip install flashinfer-python
FlashInfer 具有 Torch 原生 API，其设计为 plan/run，可解内核编译/ 选择/ 调整和内核运行。为便于注意，API 如下所示：
from flashinfer.attention import BatchAttention
attention = BatchAttention(backend="cutlass") # we provide multiple backend implementations
attention.plan(
qo_offsets, # offsets of each request in variable length query/output
kv_lens, # kv length of each request in page table
kv_block_table, # block table denoting the block indices in page table, could be packed/padded
num_qo_heads, # number of query/output heads
num_kv_heads, # number of key/value heads
head_dim_qk, # head dimension of query/key
head_dim_vo, # head dimension of value/output
dtype_q=torch.bfloat16, # query data type
dtype_kv=torch.bfloat16, # kv data type
dtype_o=torch.bfloat16, # output data type
**variant_kwargs, # other arguments specifying attention variants
)
O, lse = attention.run(q, (k, v)) # return output/lse
在 plan 阶段执行Kernel选择和调整，该阶段收集Kernel所需的元数据。相同的计划信息可重复用于共享相同元数据 (LLM 生成步骤中的所有层) 的后续运行。
用户可以从多个注意力后端中进行选择，为其用例实现最佳性能。所有内核均支持 CUDAGraph，可实现低延迟 LLM 推理服务。
对于 logits 处理，模块化接口由不同的 logits 处理器组成，而 flashinfer 发出基于融合拒绝采样的高效实现。我们最近的博文解释了 flashinfer rejection sampling 算法的工作原理。
import flashinfer
from flashinfer.logits_processor import LogitsPipe, Temperature, Softmax, TopP, Sample
# Create a pipeline
pipe = LogitsPipe([
Temperature(), # Scale logits by temperature
Softmax(), # Convert logits to probabilities
TopP(), # Apply top-p filtering
Sample() # Sample from the distribution
])
# Apply the pipeline
logits = torch.randn(batch_size, vocab_size, device="cuda")
output_ids = pipe(logits, temperature=0.7, top_p=0.9)
要开始使用 FlashInfer，请参阅 GitHub 资源库和文档。
