flashinfer

<picture> <source media="(prefers-color-scheme: dark)" srcset="https://github.com/flashinfer-ai/web-data/blob/main/logo/FlashInfer-black-background.png?raw=true"> <img alt="FlashInfer" src="https://yellow-cdn.veclightyear.com/835a84d5/dcb2322b-8cbc-4e3b-9bab-033066ecc65c.png?raw=true" width=55%> </picture> <h1 align="center"> 大语言模型服务的内核库 </h1> | <a href="https://flashinfer.ai">博客</a> | <a href="https://docs.flashinfer.ai">文档</a> | <a href="https://github.com/orgs/flashinfer-ai/discussions">讨论论坛</a> |

FlashInfer 是一个大语言模型库，提供高性能的 LLM GPU 内核实现，如 FlashAttention、SparseAttention、PageAttention、采样等。FlashInfer 专注于 LLM 服务和推理，在各种场景下都能提供最先进的性能。

FlashInfer 的独特特性包括：

全面的注意力内核：涵盖 LLM 服务所有常见用例的注意力内核，包括单请求和批处理版本的预填充、解码和追加内核，适用于不同格式的 KV 缓存（填充张量、不规则张量和页表）。
优化的共享前缀批量解码：FlashInfer 通过级联增强共享前缀批量解码性能，与基准 vLLM PageAttention 实现相比，可实现高达 31 倍的加速（针对 32768 个 token 的长提示和 256 的大批量大小）。
加速压缩/量化 KV 缓存的注意力：现代 LLM 通常使用量化/压缩的 KV 缓存来减少内存流量。FlashInfer 通过优化分组查询注意力、融合 RoPE 注意力和量化注意力的性能来加速这些场景。

FlashInfer 支持 PyTorch、TVM 和 C++（仅头文件）API，可以轻松集成到现有项目中。

新闻

[2024 年 1 月 31 日] 博客文章级联推理：内存高效的共享前缀批量解码
[2024 年 1 月 31 日] 博客文章使用 FlashInfer 加速 LLM 服务的自注意力

开始使用

使用我们的 PyTorch API 是最简单的入门方式：

安装

我们为 Linux 提供预构建的 wheel 包，您可以使用以下命令尝试 FlashInfer：

# 适用于 CUDA 12.4 和 torch 2.4
pip install flashinfer -i https://flashinfer.ai/whl/cu124/torch2.4
# 其他 CUDA 和 torch 版本，请查看 https://docs.flashinfer.ai/installation.html

或者您可以从源代码构建：

git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer/python
pip install -e .

为减少构建和测试时的二进制大小：

git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer/python
# 参考 https://pytorch.org/docs/stable/generated/torch.cuda.get_device_capability.html#torch.cuda.get_device_capability
export TORCH_CUDA_ARCH_LIST=8.0
pip install -e .

尝试使用

以下是使用 FlashInfer 单请求解码/追加/预填充注意力内核的最小示例：

import torch
import flashinfer

kv_len = 2048
num_kv_heads = 32
head_dim = 128

k = torch.randn(kv_len, num_kv_heads, head_dim).half().to(0) 
v = torch.randn(kv_len, num_kv_heads, head_dim).half().to(0) 

# 解码注意力

num_qo_heads = 32
q = torch.randn(num_qo_heads, head_dim).half().to(0)

o = flashinfer.single_decode_with_kv_cache(q, k, v) # 不使用即时 RoPE 的解码注意力
o_rope_on_the_fly = flashinfer.single_decode_with_kv_cache(q, k, v, pos_encoding_mode="ROPE_LLAMA") # 使用 LLaMA 风格即时 RoPE 的解码

# 追加注意力
append_qo_len = 128
q = torch.randn(append_qo_len, num_qo_heads, head_dim).half().to(0) # 追加注意力，KV 缓存中的最后 128 个 token 是新 token
o = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True) # 不使用即时 RoPE 的追加注意力，应用因果掩码
o_rope_on_the_fly = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True, pos_encoding_mode="ROPE_LLAMA") # 使用 LLaMA 风格即时 RoPE 的追加注意力，应用因果掩码

# 预填充注意力
qo_len = 2048
q = torch.randn(qo_len, num_qo_heads, head_dim).half().to(0) # 预填充注意力
o = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=False) # 不使用即时 RoPE 的预填充注意力，不应用因果掩码

查看文档了解批量解码/追加/预填充内核和共享前缀级联内核的用法。

运行基准测试

我们使用 nvbench 对 FlashInfer 内核性能进行分析，您可以使用以下命令编译和运行基准测试：

mkdir build
cp cmake/config.cmake build # 您可以修改 config.cmake 以启用/禁用基准测试并更改 CUDA 架构
cd build
cmake ..
make -j12

您可以运行 ./bench_{single/batch}_{prefill/decode} 来测试性能（例如，./bench_single_prefill 用于单请求预填充注意力）。./bench_{single/batch}_{prefill/decode} --help 将显示可用选项。