FlashInfer项目中的KV-Cache布局详解

FlashInfer项目中的KV-Cache布局详解

引言

在大型语言模型(LLM)的推理过程中，KV-Cache(键值缓存)是提高推理效率的关键技术。FlashInfer作为一个高性能的推理加速库，提供了多种KV-Cache的存储布局方案。本文将深入解析FlashInfer中KV-Cache的各种布局方式及其适用场景。

KV-Cache基础布局：NHD与HND

FlashInfer提供了两种基础的KV-Cache三维布局方案：

1. NHD布局：维度顺序为(序列长度, 头数, 头维度)
2. HND布局：维度顺序为(头数, 序列长度, 头维度)

NHD布局更符合自然的数据组织方式，因为它与注意力机制中计算xW_k和xW_v的输出形状一致，无需转置操作。而HND布局在KV-Cache使用低精度数据类型(如fp8)时对GPU实现更友好。

实际测试表明，在fp16精度的KV-Cache上，这两种布局的性能差异不大。FlashInfer默认使用NHD布局以提高代码可读性，但同时也实现了两种布局的注意力计算内核，用户可根据需求选择。

不规则张量(Ragged Tensor)布局

在处理批量预填充(batch prefill)自注意力时，不同请求的序列长度可能各不相同。FlashInfer采用不规则张量来存储这些变长的Q/K/V张量。

数据结构设计

不规则张量将所有请求的Q/K/V数据打包到一个连续的data张量中，不进行填充。同时使用一个indptr数组(长度为请求数+1，首元素固定为0)来记录每个请求的序列长度信息：

• indptr[i+1]-indptr[i]表示第i个请求的序列长度
• data张量的形状为(indptr[-1], num_heads, head_dim)(NHD布局时)
• 可通过data[indptr[i]:indptr[i+1]]切片获取第i个请求的键或值

重要提示：FlashInfer库中的所有indptr数组必须使用int32类型，使用int64类型可能导致索引错误。

应用场景

这种布局特别适合处理批量请求中序列长度差异较大的场景，避免了传统填充方式造成的内存浪费和计算冗余。

掩码布局(2D不规则张量)

注意力掩码可以看作是二维不规则张量。当批量大小大于1时，不同请求可能有不同的查询长度和KV长度。

数据结构设计

FlashInfer使用两个数组来存储变长序列信息：

• qo_indptr：记录每个请求的查询长度
• kv_indptr：记录每个请求的KV长度

所有请求的掩码数据被展平并连接成一个一维数组mask_data。FlashInfer会隐式创建一个mask_indptr数组来记录每个请求掩码的起始偏移量。

内存优化

为节省内存，FlashInfer支持将布尔掩码数组打包为位压缩数组(每个元素1bit，8个元素打包为1个uint8)。FlashInfer既接受原始布尔掩码，也接受位压缩掩码，如果提供布尔掩码，库会在内部自动进行位压缩。

页表布局(Page Table Layout)

在动态KV-Cache场景(如追加或解码阶段)，将所有键值打包存储效率不高。FlashInfer借鉴了vLLM的思想，采用页表方式组织KV-Cache。

数据结构设计

页表可视为块稀疏矩阵，使用CSR格式索引KV-Cache中的页面。每个请求维护以下信息：

• page_indices：记录该请求使用的页面
• last_page_len：记录最后一个页面中的有效条目数

请求i的KV序列长度为：page_size * (len(page_indices[i]) - 1) + last_page_length[i]

重要约束：每个请求的last_page_len必须大于0且不超过page_size。

存储方式

KV数据可以存储为单个5D张量或键值分离的4D张量元组。在NHD布局下：

• 单张量存储：(max_num_pages, 2, page_size, num_heads, head_dim)
• 分离存储：k_data和v_data形状均为(max_num_pages, page_size, num_heads, head_dim)

多级级联推理数据布局

在多级级联推理场景中，查询和输出存储在Ragged Tensor中，而所有级别的KV-Cache存储在统一的页式KV-Cache中。

设计特点

• 每个级别有独立的qo_indptr数组
• 所有级别共享相同的基础数据布局
• 通过不同的索引数组实现多级视图

这种设计允许在不同级别重用前缀，显著提高了推理效率。

常见问题解答

Q: FlashInfer如何管理KV-Cache？

A: FlashInfer本身不负责页表管理(如页面分配/回收等)，这些策略由用户决定。FlashInfer专注于计算查询与KV-Cache中键值之间的注意力。

总结

FlashInfer提供了灵活多样的KV-Cache布局方案，从基础的NHD/HND布局到高级的页表布局和级联推理布局，满足了不同场景下的性能需求。理解这些布局的特点和适用场景，有助于开发者充分发挥FlashInfer的性能潜力，构建高效的推理系统。