分页注意力

此页面记录了在连续批处理中使用的分页注意力前向函数。

Varlen 路径

flash_attn_varlen_func 内核用于处理变长批次。对于预填充阶段中包含大量请求的批次，建议采用此路径。

缓存行为

该内核没有直接与分页缓存交互的机制，因此缓存是通过 ~PagedAttentionCache.update 方法手动读取和写入的。当序列长度变得很大时，这可能会成为瓶颈。

索引机制

该内核使用最大序列长度（max_length_q、max_length_k）和累积序列长度（cu_seq_lens_q、cu_seq_lens_k）来计算每个序列的注意力。

示例

考虑一个包含 3 个序列的批次，其查询长度为 [10, 3, 1]，键长度为 [0, 1, 7]：

cu_seq_lens_q = [0, 10, 13, 14]
cu_seq_lens_k = [0, 0, 1, 8]
max_length_q = 10
max_length_k = 7

输入形状：

Q:      [1, 10+3+1, num_heads, head_dim] = [1, 14, num_heads, head_dim]
K or V: [1, 0+1+7, num_kv_heads, head_dim] = [1, 8, num_kv_heads, head_dim]

该内核使用累积序列长度将每个查询和键/值令牌分配到一个序列中：

Q request index: [r0, r0, r0, r0, r0, r0, r0, r0, r0, r0, r1, r1, r1, r2]
cu_seq_lens_q:   0____________________________________10__________13__14

K request index: [r1, r2, r2, r2, r2, r2, r2, r2]  (r0 has 0 K tokens)
cu_seq_lens_k: 0,0_1_______________________8

解码路径

flash_attn_with_kvcache 内核处理仅包含解码的批次，其中每个序列恰好有一个查询 token。这比变长路径更高效，但无法处理包含预填充请求的批次。

缓存行为

该内核通过 block_table 与分页缓存进行交互，利用它来索引缓存并就地更新。块表的形状为 (batch_size, max_blocks_per_seq)，其中每一行包含某个请求在 KV 缓存张量中对应的缓存块的物理位置。

索引机制

该内核使用 cache_seqlens 来获取每个序列的缓存长度。它假设每个查询 token 属于不同的序列（每个序列一个 token）。

示例

考虑一个包含 3 个序列的批次，其查询长度为 [1, 1, 1]，键长度为 [30, 32, 70]。缓存块大小为 32，每个序列的最大块数为 4。

缓存序列长度即为键的长度：

cache_seqlens = [30, 32, 70]

块表的形状为 (3, 4)。使用示例地址：

block_table = [[2, -1, -1, -1],
               [0,  1, -1, -1],
               [3,  5,  6, -1]]

-1 的值表示未分配的块。

• 序列 0（30 个缓存 token）：缓存位于 KV_cache[2]。新 token 可以放入（30 + 1 = 31 < 32）。
• 序列 1（32 个缓存 token）：缓存位于 KV_cache[0] 和 KV_cache[1] 中。由于 32 + 1 > 32，因此需要第二个块。
• 序列 2（70 个缓存 token）：缓存位于 KV_cache[3]、KV_cache[5] 和 KV_cache[6] 中。请注意，这些块不一定连续，这正是分页缓存的关键优势。新 token 恰好放入第三个块中。