🪴 Quartz 4.0

7-Flash-Infer

11 min read

概述

FlashInfer 是一个专为大语言模型（LLM）设计的库和内核生成器，提供高性能的 LLM GPU 内核实现，包括 FlashAttention、PageAttention 和 LoRA。FlashInfer 专注于 LLM 服务和推理，在各种场景下提供最先进的性能。


graph TB

A[FlashInfer] --> B[FlashAttention]

A --> C[PageAttention]

A --> D[LoRA]

A --> E[LLM 服务]

A --> F[LLM 推理]

B --> G[高性能 GPU 内核]

C --> G

D --> G

E --> H[多样化场景]

F --> H

H --> I[最先进性能]

Python 包

FlashInfer 作为 Python 包提供，基于 PyTorch 构建，可轻松集成到您的 Python 应用程序中。

系统要求

组件要求
操作系统仅限 Linux
Python 版本3.8, 3.9, 3.10, 3.11, 3.12
PyTorch 版本2.2/2.3/2.4/2.5/2.6
CUDA 版本11.8/12.1/12.4/12.6
GPU 架构支持sm75, sm80, sm86, sm89, sm90

重要: FlashInfer 当前使用的包名是 flashinfer-python，而不是 flashinfer

检查 PyTorch CUDA 版本：

 
python -c "import torch; print(torch.version.cuda)"

完整版本兼容性矩阵


graph TB

subgraph "PyTorch 版本"

PT22[PyTorch 2.2]

PT23[PyTorch 2.3]

PT24[PyTorch 2.4]

PT25[PyTorch 2.5]

PT26[PyTorch 2.6]

end

subgraph "CUDA 版本"

C118[CUDA 11.8]

C121[CUDA 12.1]

C124[CUDA 12.4]

C126[CUDA 12.6]

end

PT22 --> C118

PT22 --> C121

PT23 --> C118

PT23 --> C121

PT24 --> C121

PT24 --> C124

PT25 --> C124

PT25 --> C126

PT26 --> C126

快速开始

通过 pip 安装 FlashInfer 是最简单的方式。我们为不同的 PyTorch 版本和 CUDA 版本提供带索引 URL 的 wheels。

完整安装命令矩阵

PyTorch 2.6:

 
# CUDA 12.6
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu126/torch2.6/

PyTorch 2.5:

 
# CUDA 12.4
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu124/torch2.5/
 
# CUDA 12.6 (需要 PyTorch 2.5+)
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu126/torch2.5/

PyTorch 2.4:

 
# CUDA 12.1
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu121/torch2.4/
 
# CUDA 12.4 (PyTorch 2.4+ 支持)
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu124/torch2.4/

PyTorch 2.3:

 
# CUDA 11.8
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu118/torch2.3/
 
# CUDA 12.1
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu121/torch2.3/

PyTorch 2.2:

 
# CUDA 11.8
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu118/torch2.2/
 
# CUDA 12.1
 
pip install flashinfer-python -i https://flashinfer.ai/whl/cu121/torch2.2/

版本选择指南


flowchart TD

A[开始安装] --> B{检查 CUDA 版本}

B --> C[CUDA 12.6]

B --> D[CUDA 12.4]

B --> E[CUDA 12.1]

B --> F[CUDA 11.8]

C --> G{PyTorch 版本}

D --> H{PyTorch 版本}

E --> I{PyTorch 版本}

F --> J{PyTorch 版本}

G --> G1[2.6 ✓]

G --> G2[2.5 ✓]

H --> H1[2.5 ✓]

H --> H2[2.4 ✓]

I --> I1[2.4 ✓]

I --> I2[2.3 ✓]

I --> I3[2.2 ✓]

J --> J1[2.3 ✓]

J --> J2[2.2 ✓]

G1 --> K[安装对应 wheel]

G2 --> K

H1 --> K

H2 --> K

I1 --> K

I2 --> K

I3 --> K

J1 --> K

J2 --> K

核心功能

1. 递归注意力 (Recursive Attention)

递归注意力是 FlashInfer 的核心特性之一，允许将注意力计算分解为更小的组件，提高内存效率和计算性能。


flowchart TD

A[输入序列] --> B[分解注意力计算]

B --> C[递归处理]

C --> D[合并结果]

D --> E[输出]

F[内存效率优化] --> B

G[计算性能提升] --> C

subgraph "递归注意力特性"

H[状态管理]

I[层次化计算]

J[动态调度]

end

C --> H

C --> I

C --> J

主要优势:

  • 内存效率: 通过分解计算减少内存占用

  • 性能优化: 递归处理提高计算效率

  • 可扩展性: 支持长序列处理

  • 状态保持: 维护注意力状态以支持增量计算

2. KV-Cache 布局

FlashInfer 提供灵活的 KV-Cache 布局管理，支持多种缓存策略以优化不同场景下的性能。


graph TB

A[KV-Cache 管理] --> B[Paged KV-Cache]

A --> C[连续布局]

A --> D[分层缓存]

A --> E[动态调整]

B --> F[页表管理]

B --> G[内存碎片优化]

C --> H[顺序访问优化]

C --> I[批处理友好]

D --> J[Cascade Attention]

D --> K[层次化存储]

E --> L[运行时优化]

E --> M[负载均衡]

F --> N[高效内存利用]

G --> N

H --> N

I --> N

J --> N

K --> N

L --> N

M --> N

KV-Cache 特性:

  • 分页管理: 支持 Paged KV-Cache 减少内存碎片
  • 布局优化: 针对不同访问模式优化数据布局
  • 动态管理: 运行时动态调整缓存策略
  • 页表支持: 完整的页表布局文档和 API

技术架构

graph TB
    subgraph api["Python API Layer"]
        flashinfer["FlashInfer Core<br/>PyTorch"]
        kernels["GPU Kernels<br/>C"]
        memory["Memory Manager"]
        scheduler["Task Scheduler"]
    end
    
    subgraph gpu["GPU 计算层"]
        cuda["CUDA Runtime"]
        sm["SM 架构"]
        tensor["Tensor Cores"]
    end
    
    subgraph optimization["优化层"]
        flash["FlashAttention"]
        page["PageAttention"]
        lora["LoRA"]
    end
    
    flashinfer --> kernels
    kernels --> memory
    kernels --> scheduler
    memory --> cuda
    scheduler --> cuda
    cuda --> sm
    cuda --> tensor
    flash --> kernels
    page --> kernels
    lora --> kernels

性能特性

内存优化技术

特性描述适用场景性能提升
Cascade Attention分层 KV-Cache 管理长序列推理30-50% 内存节省
Head-Query 优化头查询级别优化多头注意力20-30% 缓存效率提升
动态内存分配运行时内存管理变长序列减少内存碎片
Paged KV-Cache分页缓存机制批处理推理显著减少内存占用

计算优化策略


pie title 性能提升来源分布

"FlashAttention 内核优化" : 35

"PageAttention 加速" : 25

"内存访问模式优化" : 20

"并行计算调度" : 15

"数据布局优化" : 5

GPU 架构支持


graph LR

A[支持的 GPU 架构] --> B[sm75 - Turing]

A --> C[sm80 - Ampere A100]

A --> D[sm86 - Ampere RTX 30xx]

A --> E[sm89 - Ada Lovelace]

A --> F[sm90 - Hopper H100]

B --> G[GTX 16xx, RTX 20xx]

C --> H[A100, A40]

D --> I[RTX 3060-3090]

E --> J[RTX 40xx 系列]

F --> K[H100, H800]

使用场景

1. LLM 推理服务


sequenceDiagram

participant Client

participant LoadBalancer

participant Server

participant FlashInfer

participant GPU

Client->>LoadBalancer: 推理请求

LoadBalancer->>Server: 路由请求

Server->>FlashInfer: 处理请求

FlashInfer->>GPU: 执行优化内核

GPU-->>FlashInfer: 计算结果

FlashInfer-->>Server: 返回结果

Server-->>LoadBalancer: 响应

LoadBalancer-->>Client: 最终响应

Note over FlashInfer,GPU: 使用 Paged KV-Cache<br/>和 FlashAttention

2. 批处理优化场景

特性优势:

  • 负载均衡: 计划阶段缓解负载不平衡问题
  • 批处理效率: 优化批量请求处理
  • 资源利用: 最大化 GPU 资源利用率
  • 动态调度: 根据请求特征动态调整处理策略

3. 长序列处理


flowchart TD

A[长序列输入] --> B{序列长度判断}

B -->|短序列| C[标准 FlashAttention]

B -->|中等序列| D[递归注意力]

B -->|超长序列| E[Cascade Attention]

C --> F[直接计算]

D --> G[分块递归处理]

E --> H[分层缓存处理]

F --> I[输出结果]

G --> I

H --> I

J[内存监控] --> B

K[性能分析] --> B

高级特性

CUDAGraph 兼容性

FlashInfer 内核支持 CUDAGraph 捕获，可以进一步提升推理性能：

# CUDAGraph 使用示例
import torch
from flashinfer import single_decode_with_kv_cache
 
# 启用 CUDAGraph
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
	output = single_decode_with_kv_cache(...)

编译时优化

FlashInfer 支持编译时优化，可以根据具体硬件和使用场景生成最优内核：


graph TB

A[源代码] --> B[编译时分析]

B --> C[硬件特性检测]

B --> D[使用模式分析]

C --> E[架构优化]

D --> F[内存布局优化]

E --> G[生成优化内核]

F --> G

G --> H[运行时执行]

安装故障排除

常见问题及解决方案

问题原因解决方案
CUDA 版本不匹配PyTorch 和 CUDA 版本不兼容参考兼容性矩阵选择正确版本
GPU 架构不支持GPU 计算能力低于 sm75升级到支持的 GPU 架构
包名错误使用了 flashinfer 而非 flashinfer-python使用正确的包名安装
wheel 不存在版本组合不在支持矩阵中选择支持的 PyTorch/CUDA 组合

验证安装

import torch
import flashinfer
# 检查版本
print(f"PyTorch: {torch.__version__}")
print(f"CUDA: {torch.version.cuda}")
print(f"FlashInfer: {flashinfer.__version__}")
 
# 检查 GPU 可用性
print(f"GPU 可用: {torch.cuda.is_available()}")
print(f"GPU 数量: {torch.cuda.device_count()}")
 
# 简单功能测试
if torch.cuda.is_available():
	device = torch.cuda.current_device()
	print(f"当前设备: {torch.cuda.get_device_name(device)}")

总结

FlashInfer 作为专业的 LLM 内核库，通过以下核心技术提供卓越性能：

  1. 完整的版本支持: 支持 PyTorch 2.2-2.6 和 CUDA 11.8-12.6 的完整组合矩阵

  2. 高性能内核: FlashAttention、PageAttention、LoRA 的优化实现

  3. 先进的内存管理: Cascade Attention、Paged KV-Cache 和分层缓存

  4. 灵活的架构支持: 从 Turing 到 Hopper 的全系列 GPU 架构

  5. 易于集成: Python 包形式，与现有工作流程无缝集成

  6. 企业级特性: CUDAGraph 支持、编译时优化、负载均衡

FlashInfer 为 LLM 服务和推理场景提供了完整且高性能的解决方案，在保证最佳性能的同时简化了开发和部署流程。通过精确的版本兼容性管理和丰富的优化特性，FlashInfer 已成为 LLM 推理加速的首选解决方案。

Graph View

Backlinks

  • No backlinks found