https://github.com/flashinfer-ai/flashinfer

FlashInfer 0.2.5 documentation

FlashInfer 是一个用于大型语言模型的库和内核生成器，可提供 LLM GPU 内核（如 FlashAttention、SparseAttention、PageAttention、Sampling 等）的高性能实现。FlashInfer 专注于 LLM 服务和推理，并在各种场景中提供一流的性能。

FlashInfer的核心功能包括：

  1. 高效的稀疏/密集注意核：CUDA 核心和 Tensor 核心（FA2 和 FA3）模板上稀疏（分页）/密集 KV 存储的高效单/批量注意。向量稀疏注意可以实现相同问题规模的密集核带宽的 90%。
  2. 负载平衡调度：FlashInfer 解耦注意力计算 plan/run阶段，我们在 plan 阶段中安排可变长度输入的计算，以缓解负载不平衡问题。
  3. 内存效率：FlashInfer 为分层 KV-Cache 提供级联注意力，并实现 Head-Query 融合以加速分组查询注意力，以及为低精度注意力提供高效内核和为压缩 KV-Cache 提供融合 RoPE 注意力。
  4. 可定制的注意力：通过 JIT 编译带来您自己的注意力变体。
  5. CUDAGraph 和 torch.compile 兼容性：FlashInfer 内核可以被 CUDAGraphs 和 torch.compile 捕获，以实现低延迟推理。
  6. 高效的 LLM 特定运算符：用于 Top-P、Top-K/Min-P 采样的高性能融合内核，无需排序。