大语言模型服务的内核库
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,可以轻松集成到现有项目中。
新闻
开始使用
使用我们的 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
将显示可用选项。
C++ API 和 TVM 绑定
FlashInfer 还提供 C++ API 和 TVM 绑定,详情请参阅文档。
采用情况
目前,以下项目采用了 FlashInfer:
致谢
FlashInfer 受到 FlashAttention 1&2、vLLM、stream-K 和 cutlass 项目的启发。