vLLM中paged attention算子分析

简要分析vLLM中PA的代码架构和v1与v2的区别

vLLM版本：0.8.4

整体结构分析

首先从torch_bindings.cpp入手分析：

这里可以看到vLLM向pytorch中注册了两个PA算子：v1和v2
其中paged_attention_v1和paged_attention_v2分别实现在csrc/attention/paged_attention_v1.cu以及csrc/attention/paged_attention_v2.cu中：


paged_attention_v1和paged_attention_v2都是PA实现的顶层接口，是Python接口调用的入口点，接收来自PyTorch的张量和各种配置参数。
根据csrc/attention/paged_attention_v1.cu中的代码逻辑，我们可以知道vLLM使用了多层分发机制来实现灵活的算子调用，以csrc/attention/paged_attention_v1.cu中的paged_attention_v1为例，分发层级为：
paged_attention_v1：
-> DISPATCH_BY_KV_CACHE_DTYPE
-> CALL_V1_LAUNCHER_BLOCK_SIZE
-> CALL_V1_LAUNCHER_SPARSITY
-> CALL_V1_LAUNCHER
-> paged_attention_v1_launcher
-> LAUNCH_PAGED_ATTENTION_V1
-> paged_attention_v1_kernel （csrc/attention/attention_kernels.cuh实现）
-> paged_attention_kernel（csrc/attention/attention_kernels.cuh实现）

可知paged_attention_kernel是具体的PA算子，csrc/attention/attention_kernels.cuh的代码主要包含以下几个函数：

paged_attention_v1_kernel和paged_attention_v2_kernel分别是v1和v2对paged_attention_kernel的封装，功能是类似的。

v1与v2的具体区别

总体上：

• v1版本提供了直接、无分区的注意力计算，适用于中短序列，实现简单，内核启动开销低。
• v2版本引入了分区计算和归约合并的两阶段处理，解决了长序列的共享内存限制问题，提高了GPU利用率。

函数参数

paged_attention_v1和paged_attention_v2参数对比可以发现，v2版本新增了三个参数：exp_sums，max_logits，tmp_out，他们是用于分区注意力计算。
这三个参数的维度包含了max_num_partitions，表示序列被划分的最大分区数。

分区计算策略

v2版本引入了固定大小的分区（默认512个token），将长序列分割成多个分区并行处理。

 // v1版本 - 整体处理序列
dim3 grid(num_heads, num_seqs, 1);  // 注意z维度为1// v2版本 - 分区处理序列
int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);// 默认分区大小PARTITION_SIZE = 512
dim3 grid(num_heads, num_seqs, max_num_partitions);  // z维度为分区数

v2两阶段计算

#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \PARTITION_SIZE>                              \<<<grid, block, shared_mem_size, stream>>>(                              \exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank,  \blocksparse_local_blocks, blocksparse_vert_stride,                   \blocksparse_block_size, blocksparse_head_sliding_step);              \vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \PARTITION_SIZE>                       \<<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \max_num_partitions);

相对于v1，v2包含两阶段计算，分别是：

// 第一阶段：分区计算
vllm::paged_attention_v2_kernel<<< ... >>>(/* 参数 */);// 第二阶段：合并结果
vllm::paged_attention_v2_reduce_kernel<<< ... >>>(/* 参数 */);

而v1版本只需一个核函数：

vllm::paged_attention_v1_kernel<<< ... >>>(/* 参数 */);

共享内存的使用上

// v1版本 - 需要为整个序列分配共享内存
int padded_max_seq_len = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
int logits_size = padded_max_seq_len * sizeof(float);// v2版本 - 只需为一个分区分配共享内存
int logits_size = PARTITION_SIZE * sizeof(float);

v2版本显著减少了每个线程块需要的共享内存大小.

v2版本的规约

v2版本在paged_attention_v2_reduce_kernel中实现了复杂的跨分区归约。

// 计算全局最大logit
max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));// 计算重新缩放的指数和
float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);// 聚合分区结果
float acc = 0.0f;
for (int j = 0; j < num_partitions; ++j) {acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] * inv_global_exp_sum;
}

若有错误，欢迎指正！