vLLM Paged Attention 内核设计详解分页 KV Cache 内存布局与线程级计算流程【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm本文以 vLLM 官方设计文档 Paged Attention 为核心系统讲解 vLLM 单查询多头注意力内核multi-head query attention kernel的设计从 KV 分页缓存paged KV cache的内存布局到 Query/Key/Value 的线程数据分派再到 QK 点积、Softmax 归一化、LV 归约的完整计算流程。读完本文后你可以理解 vLLM 如何将一个 batch 中每个序列的 decode 注意力分解到线程组、warp、thread block 三个层级上并行计算以及其访存合并memory coalescing策略背后的原理。需要先说明适用前提该文档自述为基于 vLLM 原始论文PagedAttention, SOSP 2023的历史文档描述的是早期csrc/attention/attention_kernels.cu内核实现并不代表当前 vLLM 正在运行的全部注意力代码。但在当前仓库中这一设计仍然留有清晰痕迹——例如csrc/attention/目录下仍保留着内核使用的向量类型头文件 attention_dtypes.h、各数据类型的实现头 dtype_float16.cuh、dtype_bfloat16.cuh、dtype_float32.cuh、dtype_fp8.cuh以及向量运算的通用抽象 attention_generic.cuh。从源码结构看attention_generic.cuh 中的VecT, VEC_SIZE、dot、mul、sum等模板设施正是文档中q_vecs/k_vecs/v_vec计算所依赖的底层原语。而 vLLM 当前实际调度的注意力实现已迁移至 vllm/v1/attention/backends/ 目录下的多后端体系flash_attn、triton_attn、flashinfer、mla 等KV 分页缓存的block_table机制则在 cache_kernels.cu 等缓存内核中继续发挥作用。因此本文将严格按文档脉络展开末尾再补一节说明文档设计与当前仓库代码的对应关系。内核输入指针参数与模板参数内核函数接收一系列参数用于当前线程执行其分配到的工作。其中四个最重要的指针参数分别指向全局内存上的查询、键、值数据和输出位置templatetypename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS, int PARTITION_SIZE 0 __device__ void paged_attention_kernel( ... // Other side args. const scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] const scalar_t* __restrict__ q, // // [num_seqs, num_heads, head_size] const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] ... // Other side args. )这四个指针实际上指向多维数组但每个线程只访问分配给它的那部分数据。这里省略了其余运行时参数以保持简洁。函数签名上方还有一组编译期确定的模板参数各参数的含义如下模板参数含义scalar_t查询、键、值数据元素的数据类型例如 FP16HEAD_SIZE每个注意力头的元素个数head sizeBLOCK_SIZE每个 block 中存储的 token 个数NUM_THREADS每个 thread block线程块中的线程数PARTITION_SIZE张量并行 GPU 数量文档为简化假设其为 0即禁用张量并行拿到这些参数后内核需要做一系列准备工作包括计算当前 head 索引、block 索引及其他必要变量。文档建议先跳过这些准备直接进入实际计算掌握整体流程后再回看这些索引推导会更清晰。核心概念Sequence、Context、Vec、线程组、Block、Warp、Grid在进入计算流程之前需要先明确后续章节会用到的术语体系。注意此处的 blockvLLM 分页注意力块与 GPU 的 thread block线程块是两个不同概念Sequence序列一个序列代表一个客户端请求。例如q指向的数据形状为[num_seqs, num_heads, head_size]表示共有num_seqs个查询序列。由于该内核是单查询注意力single query attention内核每个序列只有一个 query token所以num_seqs等于本 batch 中处理的 token 总数。Context上下文上下文由该序列已生成的 token 组成。例如[What, is, your]是 context token而输入查询 token 是name模型可能接着生成 token?。Vec向量一次取出并一起计算的元素列表。对 query 和 key 数据vec 大小VEC_SIZE被确定得让每个线程组thread group一次取 16 字节对 value 数据vec 大小V_VEC_SIZE被确定得让每个线程一次取 16 字节。例如当scalar_t为 FP162 字节、THREAD_GROUP_SIZE为 2 时VEC_SIZE为 4V_VEC_SIZE为 8。Thread group线程组由THREAD_GROUP_SIZE个线程组成的小组一次取并计算一个 query token 与一个 key token每个线程只处理该 token 数据的一部分。一个线程组处理的元素总数记为x。例如线程组含 2 个线程、head size 为 8 时线程 0 处理索引 0、2、4、6 处的 query 和 key 元素线程 1 处理 1、3、5、7。Block分页块vLLM 中 key/value 缓存数据被切分为 block每个 block 在单个 head 上存储固定数量BLOCK_SIZEtoken 的数据每个 block 可能只包含整个 context 的一部分 token。例如 block size 为 16、head size 为 128 时一个 head 的一个 block 可存 16 × 128 2048 个元素。Warp32 个线程WARP_SIZE组成在流多处理器SM上同时执行。在本内核中每个 warp 一次处理一个 query token 与一整个 block 的 key token 之间的计算可通过多次迭代处理多个 block。例如 4 个 warp、6 个 block 时warp 0 处理第 0、4 块warp 1 处理第 1、5 块warp 2 处理第 2 块warp 3 处理第 3 块。Thread block线程块NUM_THREADS个线程组成的、可访问同一块共享内存的组。每个 thread block 含多个 warpNUM_WARPS在本内核中每个 thread block 处理一个 query token 与整个 context所有 key token 之间的计算。Gridthread block 的集合定义集合形状。本内核中形状为(num_heads, num_seqs, max_num_partitions)因此每个 thread block 只负责一个 head、一个 sequence、一个 partition 的计算。Query 的存储与读取每个线程组取一个 query token 数据而线程本身只处理该 query token 的一部分。warp 内部的每个线程组取的是同一份query token 数据但会与不同的 key token 数据相乘。线程首先计算自己指向全局内存中已分配 query token 数据的指针const scalar_t* q_ptr q seq_idx * q_stride head_idx * HEAD_SIZE;每个线程定义自己的q_ptr。例如当VEC_SIZE为 4、HEAD_SIZE为 128 时q_ptr指向的数据共 128 个元素被划分为 128 / 4 32 个 vec。接下来把q_ptr指向的全局内存数据读入共享内存q_vecs__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];关键在于每个 vec 被分配到不同的行例如THREAD_GROUP_SIZE为 2 时线程 0 处理第 0 行的 vecs线程 1 处理第 1 行的 vecs。这样读 query 数据可以让线程 0 与线程 1 等相邻线程读取相邻内存实现**访存合并memory coalescing**以提升性能。Key 的存储与读取与 Query 一节类似这里介绍 key 的内存布局与分派方式。每个线程组一次内核运行只处理一个 query token但可能跨多次迭代处理多个 key token每个 warp 会跨多次迭代处理多个 block 的 key token从而保证整个线程组在内核运行结束后处理完所有 context token。这里的 handle 指对 query 数据与 key 数据执行点积。与q_ptr不同每个线程的k_ptr在不同迭代中指向不同的 key tokenconst scalar_t* k_ptr k_cache physical_block_number * kv_block_stride kv_head_idx * kv_head_stride physical_block_offset * x;上式基于k_cache上的已分配 block、已分配 head 和已分配 token 定位 key token 数据。上图展示了 key 数据的内存布局。图中假设BLOCK_SIZE为 16、HEAD_SIZE为 128、x为 8、THREAD_GROUP_SIZE为 2共 4 个 warp。每个矩形代表一个 head 上一个 key token 的全部元素将由一个线程组处理左半部分是 warp 0 对应的共 16 个 block 的 key token 数据右半部分表示留给其他 warp 或其他迭代处理的 key token 数据。每个矩形内共有 32 个 vec一个 token 的 128 个元素由 2 个线程一个线程组分别处理。接下来把k_ptr处的 key token 数据读入寄存器内存k_vecsK_vec k_vecs[NUM_VECS_PER_THREAD]之所以用寄存器存放k_vecs是因为它只会被一个线程访问一次而q_vecs会被多个线程多次访问所以放共享内存。每个k_vecs包含后续计算要用的多个向量每个 vec 在内层迭代的每轮被设置一次。vec 的分派方式让 warp 内的相邻线程一起读取相邻内存同样促进访存合并例如线程 0 读 vec 0线程 1 读 vec 1下一轮内层循环中线程 0 读 vec 2线程 1 读 vec 3依此类推。QKQuery 与 Key 的点积按下面伪代码所示在整个 for 循环块之前先取一个 token 的 query 数据存入q_vecs然后在外层 for 循环中遍历指向不同 token 的不同k_ptr在内层 for 循环中准备k_vecs最后执行q_vecs与每个k_vecs之间的点积q_vecs ... for ... { k_ptr ... for ... { k_vecs[i] ... } ... float qk scale * Qk_dotscalar_t, THREAD_GROUP_SIZE::dot(q_vecs[thread_group_offset], k_vecs); }如前所述每个线程一次只取 query 和 key token 数据的一部分。但Qk_dot::dot内部会发生跨线程组归约cross thread group reduction所以返回的qk不只是一部分 query 与 key 的点积而是完整 query token 与完整 key token 之间的点积结果。例如HEAD_SIZE为 128、THREAD_GROUP_SIZE为 2 时每个线程的k_vecs只包含 64 个元素但返回的qk实际上是 128 个 query 元素与 128 个 key 元素点积的结果。点积与归约的细节可参考Qk_dot::dot的实现文档为简洁起见不展开。Softmax全线程块范围的归一化得到所有qk后需要对其计算归一化 softmax每个 $x$ 代表一个qk。为此必须得到所有qk的归约值qk_max$m(x)$与exp_sum$\ell(x)$归约应覆盖整个 thread block即涵盖 query token 与全部 context key token 之间的计算结果$$ \begin{gather*} m(x):\max _i \quad x_i \ \quad f(x):\left[\begin{array}{lll}e^{x_1-m(x)} \ldots e^{x_B-m(x)}\end{array}\right]\ \quad \ell(x):\sum_i f(x)_i \ \quad \operatorname{softmax}(x):\frac{f(x)}{\ell(x)} \end{gather*} $$qk_max与logits拿到qk结果后立即做两件事用qk设置临时logits结果最终logits应存储归一化 softmax 结果并对当前线程组算出的所有qk比较收集qk_maxif (thread_group_offset 0) { const bool mask token_idx context_len; logits[token_idx - start_token_idx] mask ? 0.f : qk; qk_max mask ? qk_max : fmaxf(qk_max, qk); }注意这里的logits位于共享内存每个线程组只设置其已分配 context token 对应的字段logits的整体大小应等于 context token 数量。随后通过 warp 内 shuffle 得到 warp 级的归约qk_max。核心思路是让 warp 内线程互相通信以获得最终最大的qkfor (int mask WARP_SIZE / 2; mask THREAD_GROUP_SIZE; mask / 2) { qk_max fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); } if (lane 0) { red_smem[warp_idx] qk_max; }最后把 thread block 内所有 warp 的qk_max比较合并得到全线程块归约后的qk_max并广播到每个线程for (int mask NUM_WARPS / 2; mask 1; mask / 2) { qk_max fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); } qk_max VLLM_SHFL_SYNC(qk_max, 0);exp_sum与归一化与qk_max类似也需要得到整个 thread block 的归约求和值for (int i thread_idx; i num_tokens; i NUM_THREADS) { float val __expf(logits[i] - qk_max); logits[i] val; exp_sum val; } ... exp_sum block_sumNUM_WARPS(red_smem[NUM_WARPS], exp_sum);首先把每个线程组负责范围内的 exp 值求和同时把logits的每个条目从qk转换为exp(qk - qk_max)此时qk_max已是全 thread block 的最大值。然后对exp_sum做全线程块归约方式与qk_max相同。const float inv_sum __fdividef(1.f, exp_sum 1e-6f); for (int i thread_idx; i num_tokens; i NUM_THREADS) { logits[i] * inv_sum; }至此得到归一化 softmax 结果并写回logits。该变量将在后续步骤与 value 数据做点积——此时它应保存所有已分配 context token 的归一化 softmax 权重。Value 的存储与读取与 query 和 key 不同value 数据没有线程组的概念。value token 的内存布局与 key token 不同同一列的元素对应同一个 value token。对一个 block 的 value 数据共有HEAD_SIZE行、BLOCK_SIZE列被划分为多个v_vec。现在需要取出 value 数据并与logits做点积。每个线程一次总是取相同V_VEC_SIZE个 token 上的V_VEC_SIZE个元素因此单个线程通过多次内层迭代从不同行、相同列取出多个v_vec。每个v_vec都要与对应的logits_vec点积——logits_vec同样是来自logits的V_VEC_SIZE个元素。通过多次内层迭代每个 warp 处理一整个 block 的 value token通过多次外层迭代整个 context 的 value token 全部处理完毕float accs[NUM_ROWS_PER_THREAD]; for ... { // Iteration over different blocks. logits_vec ... for ... { // Iteration over different rows. v_vec ... ... accs[i] dot(logits_vec, v_vec); } }如伪代码所示外层循环中类似k_ptrlogits_vec遍历不同 block 并从logits读取V_VEC_SIZE个元素内层循环中每个线程把相同 token 上的V_VEC_SIZE个元素作为v_vec读取并执行点积。注意每次内层迭代中线程取的是相同 token的不同 head 位置元素。点积结果累加进accs因此accs的每个条目映射到当前线程负责的一个 head 位置。举例说明BLOCK_SIZE为 16、V_VEC_SIZE为 8 时每个线程一次取 8 个 token 上的 8 个 value 元素每个元素来自不同 token 的同一 head 位置。若HEAD_SIZE为 128、WARP_SIZE为 32则每次内层循环一个 warp 需取WARP_SIZE * V_VEC_SIZE 256个元素即一个 warp 处理一整个 block 的 value token 共需 128 × 16 / 256 8 次内层迭代。每个线程的accs含 8 个元素分别在 8 个不同 head 位置累加。对线程 0 而言accs的 8 个元素分别是对应 value head 的第 0、32、…、224 个元素由分配给它的 8 个 token 累加而来。LVaccs 的 warp 内与跨 warp 归约接下来在每个 warp 内对accs做归约使每个线程累加该 block 内所有 token 在其已分配 head 位置上的accsfor (int i 0; i NUM_ROWS_PER_THREAD; i) { float acc accs[i]; for (int mask NUM_V_VECS_PER_ROW / 2; mask 1; mask / 2) { acc VLLM_SHFL_XOR_SYNC(acc, mask); } accs[i] acc; }然后跨所有 warp 归约accs使每个线程拥有整个 context 所有 token 在其已分配 head 位置上accs的累加值。注意每个线程的accs只保存了整个 head 的一部分元素在所有 context token 上的累加但从整体看输出的全部结果此时都已算出只是分散存放在各线程的寄存器内存中float* out_smem reinterpret_castfloat*(shared_mem); for (int i NUM_WARPS; i 1; i / 2) { // Upper warps write to shared memory. ... float* dst out_smem[(warp_idx - mid) * HEAD_SIZE]; for (int i 0; i NUM_ROWS_PER_THREAD; i) { ... dst[row_idx] accs[i]; } // Lower warps update the output. const float* src out_smem[warp_idx * HEAD_SIZE]; for (int i 0; i NUM_ROWS_PER_THREAD; i) { ... accs[i] src[row_idx]; } // Write out the accs. }这段跨 warp 归约采用树形合并轮次从NUM_WARPS减半到 2上半部分 warp 把部分和写入共享内存下半部分 warp 从共享内存读回并累加逐轮折叠直到结果收敛到少数 warp 中。Output写回全局内存最后把寄存器中的计算结果写入最终输出的全局内存。先定义out_ptr指向已分配 sequence 与已分配 head 的起始地址scalar_t* out_ptr out seq_idx * num_heads * max_num_partitions * HEAD_SIZE head_idx * max_num_partitions * HEAD_SIZE partition_idx * HEAD_SIZE;然后遍历各已分配的 head 位置基于out_ptr写出对应的累加结果for (int i 0; i NUM_ROWS_PER_THREAD; i) { const int row_idx lane / NUM_V_VECS_PER_ROW i * NUM_ROWS_PER_ITER; if (row_idx HEAD_SIZE lane % NUM_V_VECS_PER_ROW 0) { from_float(*(out_ptr row_idx), accs[i]); } }至此整个内核的计算流程闭环分页 KV 缓存 → 按 warp/block 分片的 QK 点积 → 全 thread block 的 softmax 归约 → 按 warp/block 分片的 LV 累加 → warp 间归约 → 写回输出。文档设计与当前仓库代码的对应关系从源码结构看当前仓库对该历史设计有两处承接内核基础设施仍然保留csrc/attention/目录保留了内核使用的类型与向量设施。attention_dtypes.h 定义Q_vec/K_vec/V_vec等向量类型的类型别名attention_generic.cuh 提供VecT, VEC_SIZE模板结构、dot、mul、sum、zero等模板运算各 dtype 的实例化分别位于 dtype_float16.cuh、dtype_bfloat16.cuh、dtype_float32.cuh、dtype_fp8.cuh可印证文档中q_vecs/k_vecs/v_vec点积所依赖的底层机制。分页 KV 机制延续到新的注意力体系分页块表block table与 KV 缓存重排逻辑在当前实现中体现于 cache_kernels.cu 等缓存内核而具体 decode 注意力的计算调度已按硬件与算法拆分为多个后端统一入口见 vllm/v1/attention/backends/init.py后端注册见 vllm/v1/attention/backends/registry.py。因此本文所述的一个 thread block 处理一个 (head, seq) 完整 context的划分思想是理解 vLLM 注意力性能优化的良好基础但实际 kernel 细节以各后端当前实现为准。参考资料文档原文给出 PagedAttention 论文的引用信息SOSP 2023inproceedings{kwon2023efficient, title{Efficient Memory Management for Large Language Model Serving with PagedAttention}, author{Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica}, booktitle{Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles}, year{2023} }【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 vLLM Paged Attention 内核设计详解：分页 KV Cache 内存布局与线程级计算流程
vLLM Paged Attention 内核设计详解分页 KV Cache 内存布局与线程级计算流程【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm本文以 vLLM 官方设计文档 Paged Attention 为核心系统讲解 vLLM 单查询多头注意力内核multi-head query attention kernel的设计从 KV 分页缓存paged KV cache的内存布局到 Query/Key/Value 的线程数据分派再到 QK 点积、Softmax 归一化、LV 归约的完整计算流程。读完本文后你可以理解 vLLM 如何将一个 batch 中每个序列的 decode 注意力分解到线程组、warp、thread block 三个层级上并行计算以及其访存合并memory coalescing策略背后的原理。需要先说明适用前提该文档自述为基于 vLLM 原始论文PagedAttention, SOSP 2023的历史文档描述的是早期csrc/attention/attention_kernels.cu内核实现并不代表当前 vLLM 正在运行的全部注意力代码。但在当前仓库中这一设计仍然留有清晰痕迹——例如csrc/attention/目录下仍保留着内核使用的向量类型头文件 attention_dtypes.h、各数据类型的实现头 dtype_float16.cuh、dtype_bfloat16.cuh、dtype_float32.cuh、dtype_fp8.cuh以及向量运算的通用抽象 attention_generic.cuh。从源码结构看attention_generic.cuh 中的VecT, VEC_SIZE、dot、mul、sum等模板设施正是文档中q_vecs/k_vecs/v_vec计算所依赖的底层原语。而 vLLM 当前实际调度的注意力实现已迁移至 vllm/v1/attention/backends/ 目录下的多后端体系flash_attn、triton_attn、flashinfer、mla 等KV 分页缓存的block_table机制则在 cache_kernels.cu 等缓存内核中继续发挥作用。因此本文将严格按文档脉络展开末尾再补一节说明文档设计与当前仓库代码的对应关系。内核输入指针参数与模板参数内核函数接收一系列参数用于当前线程执行其分配到的工作。其中四个最重要的指针参数分别指向全局内存上的查询、键、值数据和输出位置templatetypename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS, int PARTITION_SIZE 0 __device__ void paged_attention_kernel( ... // Other side args. const scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] const scalar_t* __restrict__ q, // // [num_seqs, num_heads, head_size] const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] ... // Other side args. )这四个指针实际上指向多维数组但每个线程只访问分配给它的那部分数据。这里省略了其余运行时参数以保持简洁。函数签名上方还有一组编译期确定的模板参数各参数的含义如下模板参数含义scalar_t查询、键、值数据元素的数据类型例如 FP16HEAD_SIZE每个注意力头的元素个数head sizeBLOCK_SIZE每个 block 中存储的 token 个数NUM_THREADS每个 thread block线程块中的线程数PARTITION_SIZE张量并行 GPU 数量文档为简化假设其为 0即禁用张量并行拿到这些参数后内核需要做一系列准备工作包括计算当前 head 索引、block 索引及其他必要变量。文档建议先跳过这些准备直接进入实际计算掌握整体流程后再回看这些索引推导会更清晰。核心概念Sequence、Context、Vec、线程组、Block、Warp、Grid在进入计算流程之前需要先明确后续章节会用到的术语体系。注意此处的 blockvLLM 分页注意力块与 GPU 的 thread block线程块是两个不同概念Sequence序列一个序列代表一个客户端请求。例如q指向的数据形状为[num_seqs, num_heads, head_size]表示共有num_seqs个查询序列。由于该内核是单查询注意力single query attention内核每个序列只有一个 query token所以num_seqs等于本 batch 中处理的 token 总数。Context上下文上下文由该序列已生成的 token 组成。例如[What, is, your]是 context token而输入查询 token 是name模型可能接着生成 token?。Vec向量一次取出并一起计算的元素列表。对 query 和 key 数据vec 大小VEC_SIZE被确定得让每个线程组thread group一次取 16 字节对 value 数据vec 大小V_VEC_SIZE被确定得让每个线程一次取 16 字节。例如当scalar_t为 FP162 字节、THREAD_GROUP_SIZE为 2 时VEC_SIZE为 4V_VEC_SIZE为 8。Thread group线程组由THREAD_GROUP_SIZE个线程组成的小组一次取并计算一个 query token 与一个 key token每个线程只处理该 token 数据的一部分。一个线程组处理的元素总数记为x。例如线程组含 2 个线程、head size 为 8 时线程 0 处理索引 0、2、4、6 处的 query 和 key 元素线程 1 处理 1、3、5、7。Block分页块vLLM 中 key/value 缓存数据被切分为 block每个 block 在单个 head 上存储固定数量BLOCK_SIZEtoken 的数据每个 block 可能只包含整个 context 的一部分 token。例如 block size 为 16、head size 为 128 时一个 head 的一个 block 可存 16 × 128 2048 个元素。Warp32 个线程WARP_SIZE组成在流多处理器SM上同时执行。在本内核中每个 warp 一次处理一个 query token 与一整个 block 的 key token 之间的计算可通过多次迭代处理多个 block。例如 4 个 warp、6 个 block 时warp 0 处理第 0、4 块warp 1 处理第 1、5 块warp 2 处理第 2 块warp 3 处理第 3 块。Thread block线程块NUM_THREADS个线程组成的、可访问同一块共享内存的组。每个 thread block 含多个 warpNUM_WARPS在本内核中每个 thread block 处理一个 query token 与整个 context所有 key token 之间的计算。Gridthread block 的集合定义集合形状。本内核中形状为(num_heads, num_seqs, max_num_partitions)因此每个 thread block 只负责一个 head、一个 sequence、一个 partition 的计算。Query 的存储与读取每个线程组取一个 query token 数据而线程本身只处理该 query token 的一部分。warp 内部的每个线程组取的是同一份query token 数据但会与不同的 key token 数据相乘。线程首先计算自己指向全局内存中已分配 query token 数据的指针const scalar_t* q_ptr q seq_idx * q_stride head_idx * HEAD_SIZE;每个线程定义自己的q_ptr。例如当VEC_SIZE为 4、HEAD_SIZE为 128 时q_ptr指向的数据共 128 个元素被划分为 128 / 4 32 个 vec。接下来把q_ptr指向的全局内存数据读入共享内存q_vecs__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];关键在于每个 vec 被分配到不同的行例如THREAD_GROUP_SIZE为 2 时线程 0 处理第 0 行的 vecs线程 1 处理第 1 行的 vecs。这样读 query 数据可以让线程 0 与线程 1 等相邻线程读取相邻内存实现**访存合并memory coalescing**以提升性能。Key 的存储与读取与 Query 一节类似这里介绍 key 的内存布局与分派方式。每个线程组一次内核运行只处理一个 query token但可能跨多次迭代处理多个 key token每个 warp 会跨多次迭代处理多个 block 的 key token从而保证整个线程组在内核运行结束后处理完所有 context token。这里的 handle 指对 query 数据与 key 数据执行点积。与q_ptr不同每个线程的k_ptr在不同迭代中指向不同的 key tokenconst scalar_t* k_ptr k_cache physical_block_number * kv_block_stride kv_head_idx * kv_head_stride physical_block_offset * x;上式基于k_cache上的已分配 block、已分配 head 和已分配 token 定位 key token 数据。上图展示了 key 数据的内存布局。图中假设BLOCK_SIZE为 16、HEAD_SIZE为 128、x为 8、THREAD_GROUP_SIZE为 2共 4 个 warp。每个矩形代表一个 head 上一个 key token 的全部元素将由一个线程组处理左半部分是 warp 0 对应的共 16 个 block 的 key token 数据右半部分表示留给其他 warp 或其他迭代处理的 key token 数据。每个矩形内共有 32 个 vec一个 token 的 128 个元素由 2 个线程一个线程组分别处理。接下来把k_ptr处的 key token 数据读入寄存器内存k_vecsK_vec k_vecs[NUM_VECS_PER_THREAD]之所以用寄存器存放k_vecs是因为它只会被一个线程访问一次而q_vecs会被多个线程多次访问所以放共享内存。每个k_vecs包含后续计算要用的多个向量每个 vec 在内层迭代的每轮被设置一次。vec 的分派方式让 warp 内的相邻线程一起读取相邻内存同样促进访存合并例如线程 0 读 vec 0线程 1 读 vec 1下一轮内层循环中线程 0 读 vec 2线程 1 读 vec 3依此类推。QKQuery 与 Key 的点积按下面伪代码所示在整个 for 循环块之前先取一个 token 的 query 数据存入q_vecs然后在外层 for 循环中遍历指向不同 token 的不同k_ptr在内层 for 循环中准备k_vecs最后执行q_vecs与每个k_vecs之间的点积q_vecs ... for ... { k_ptr ... for ... { k_vecs[i] ... } ... float qk scale * Qk_dotscalar_t, THREAD_GROUP_SIZE::dot(q_vecs[thread_group_offset], k_vecs); }如前所述每个线程一次只取 query 和 key token 数据的一部分。但Qk_dot::dot内部会发生跨线程组归约cross thread group reduction所以返回的qk不只是一部分 query 与 key 的点积而是完整 query token 与完整 key token 之间的点积结果。例如HEAD_SIZE为 128、THREAD_GROUP_SIZE为 2 时每个线程的k_vecs只包含 64 个元素但返回的qk实际上是 128 个 query 元素与 128 个 key 元素点积的结果。点积与归约的细节可参考Qk_dot::dot的实现文档为简洁起见不展开。Softmax全线程块范围的归一化得到所有qk后需要对其计算归一化 softmax每个 $x$ 代表一个qk。为此必须得到所有qk的归约值qk_max$m(x)$与exp_sum$\ell(x)$归约应覆盖整个 thread block即涵盖 query token 与全部 context key token 之间的计算结果$$ \begin{gather*} m(x):\max _i \quad x_i \ \quad f(x):\left[\begin{array}{lll}e^{x_1-m(x)} \ldots e^{x_B-m(x)}\end{array}\right]\ \quad \ell(x):\sum_i f(x)_i \ \quad \operatorname{softmax}(x):\frac{f(x)}{\ell(x)} \end{gather*} $$qk_max与logits拿到qk结果后立即做两件事用qk设置临时logits结果最终logits应存储归一化 softmax 结果并对当前线程组算出的所有qk比较收集qk_maxif (thread_group_offset 0) { const bool mask token_idx context_len; logits[token_idx - start_token_idx] mask ? 0.f : qk; qk_max mask ? qk_max : fmaxf(qk_max, qk); }注意这里的logits位于共享内存每个线程组只设置其已分配 context token 对应的字段logits的整体大小应等于 context token 数量。随后通过 warp 内 shuffle 得到 warp 级的归约qk_max。核心思路是让 warp 内线程互相通信以获得最终最大的qkfor (int mask WARP_SIZE / 2; mask THREAD_GROUP_SIZE; mask / 2) { qk_max fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); } if (lane 0) { red_smem[warp_idx] qk_max; }最后把 thread block 内所有 warp 的qk_max比较合并得到全线程块归约后的qk_max并广播到每个线程for (int mask NUM_WARPS / 2; mask 1; mask / 2) { qk_max fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); } qk_max VLLM_SHFL_SYNC(qk_max, 0);exp_sum与归一化与qk_max类似也需要得到整个 thread block 的归约求和值for (int i thread_idx; i num_tokens; i NUM_THREADS) { float val __expf(logits[i] - qk_max); logits[i] val; exp_sum val; } ... exp_sum block_sumNUM_WARPS(red_smem[NUM_WARPS], exp_sum);首先把每个线程组负责范围内的 exp 值求和同时把logits的每个条目从qk转换为exp(qk - qk_max)此时qk_max已是全 thread block 的最大值。然后对exp_sum做全线程块归约方式与qk_max相同。const float inv_sum __fdividef(1.f, exp_sum 1e-6f); for (int i thread_idx; i num_tokens; i NUM_THREADS) { logits[i] * inv_sum; }至此得到归一化 softmax 结果并写回logits。该变量将在后续步骤与 value 数据做点积——此时它应保存所有已分配 context token 的归一化 softmax 权重。Value 的存储与读取与 query 和 key 不同value 数据没有线程组的概念。value token 的内存布局与 key token 不同同一列的元素对应同一个 value token。对一个 block 的 value 数据共有HEAD_SIZE行、BLOCK_SIZE列被划分为多个v_vec。现在需要取出 value 数据并与logits做点积。每个线程一次总是取相同V_VEC_SIZE个 token 上的V_VEC_SIZE个元素因此单个线程通过多次内层迭代从不同行、相同列取出多个v_vec。每个v_vec都要与对应的logits_vec点积——logits_vec同样是来自logits的V_VEC_SIZE个元素。通过多次内层迭代每个 warp 处理一整个 block 的 value token通过多次外层迭代整个 context 的 value token 全部处理完毕float accs[NUM_ROWS_PER_THREAD]; for ... { // Iteration over different blocks. logits_vec ... for ... { // Iteration over different rows. v_vec ... ... accs[i] dot(logits_vec, v_vec); } }如伪代码所示外层循环中类似k_ptrlogits_vec遍历不同 block 并从logits读取V_VEC_SIZE个元素内层循环中每个线程把相同 token 上的V_VEC_SIZE个元素作为v_vec读取并执行点积。注意每次内层迭代中线程取的是相同 token的不同 head 位置元素。点积结果累加进accs因此accs的每个条目映射到当前线程负责的一个 head 位置。举例说明BLOCK_SIZE为 16、V_VEC_SIZE为 8 时每个线程一次取 8 个 token 上的 8 个 value 元素每个元素来自不同 token 的同一 head 位置。若HEAD_SIZE为 128、WARP_SIZE为 32则每次内层循环一个 warp 需取WARP_SIZE * V_VEC_SIZE 256个元素即一个 warp 处理一整个 block 的 value token 共需 128 × 16 / 256 8 次内层迭代。每个线程的accs含 8 个元素分别在 8 个不同 head 位置累加。对线程 0 而言accs的 8 个元素分别是对应 value head 的第 0、32、…、224 个元素由分配给它的 8 个 token 累加而来。LVaccs 的 warp 内与跨 warp 归约接下来在每个 warp 内对accs做归约使每个线程累加该 block 内所有 token 在其已分配 head 位置上的accsfor (int i 0; i NUM_ROWS_PER_THREAD; i) { float acc accs[i]; for (int mask NUM_V_VECS_PER_ROW / 2; mask 1; mask / 2) { acc VLLM_SHFL_XOR_SYNC(acc, mask); } accs[i] acc; }然后跨所有 warp 归约accs使每个线程拥有整个 context 所有 token 在其已分配 head 位置上accs的累加值。注意每个线程的accs只保存了整个 head 的一部分元素在所有 context token 上的累加但从整体看输出的全部结果此时都已算出只是分散存放在各线程的寄存器内存中float* out_smem reinterpret_castfloat*(shared_mem); for (int i NUM_WARPS; i 1; i / 2) { // Upper warps write to shared memory. ... float* dst out_smem[(warp_idx - mid) * HEAD_SIZE]; for (int i 0; i NUM_ROWS_PER_THREAD; i) { ... dst[row_idx] accs[i]; } // Lower warps update the output. const float* src out_smem[warp_idx * HEAD_SIZE]; for (int i 0; i NUM_ROWS_PER_THREAD; i) { ... accs[i] src[row_idx]; } // Write out the accs. }这段跨 warp 归约采用树形合并轮次从NUM_WARPS减半到 2上半部分 warp 把部分和写入共享内存下半部分 warp 从共享内存读回并累加逐轮折叠直到结果收敛到少数 warp 中。Output写回全局内存最后把寄存器中的计算结果写入最终输出的全局内存。先定义out_ptr指向已分配 sequence 与已分配 head 的起始地址scalar_t* out_ptr out seq_idx * num_heads * max_num_partitions * HEAD_SIZE head_idx * max_num_partitions * HEAD_SIZE partition_idx * HEAD_SIZE;然后遍历各已分配的 head 位置基于out_ptr写出对应的累加结果for (int i 0; i NUM_ROWS_PER_THREAD; i) { const int row_idx lane / NUM_V_VECS_PER_ROW i * NUM_ROWS_PER_ITER; if (row_idx HEAD_SIZE lane % NUM_V_VECS_PER_ROW 0) { from_float(*(out_ptr row_idx), accs[i]); } }至此整个内核的计算流程闭环分页 KV 缓存 → 按 warp/block 分片的 QK 点积 → 全 thread block 的 softmax 归约 → 按 warp/block 分片的 LV 累加 → warp 间归约 → 写回输出。文档设计与当前仓库代码的对应关系从源码结构看当前仓库对该历史设计有两处承接内核基础设施仍然保留csrc/attention/目录保留了内核使用的类型与向量设施。attention_dtypes.h 定义Q_vec/K_vec/V_vec等向量类型的类型别名attention_generic.cuh 提供VecT, VEC_SIZE模板结构、dot、mul、sum、zero等模板运算各 dtype 的实例化分别位于 dtype_float16.cuh、dtype_bfloat16.cuh、dtype_float32.cuh、dtype_fp8.cuh可印证文档中q_vecs/k_vecs/v_vec点积所依赖的底层机制。分页 KV 机制延续到新的注意力体系分页块表block table与 KV 缓存重排逻辑在当前实现中体现于 cache_kernels.cu 等缓存内核而具体 decode 注意力的计算调度已按硬件与算法拆分为多个后端统一入口见 vllm/v1/attention/backends/init.py后端注册见 vllm/v1/attention/backends/registry.py。因此本文所述的一个 thread block 处理一个 (head, seq) 完整 context的划分思想是理解 vLLM 注意力性能优化的良好基础但实际 kernel 细节以各后端当前实现为准。参考资料文档原文给出 PagedAttention 论文的引用信息SOSP 2023inproceedings{kwon2023efficient, title{Efficient Memory Management for Large Language Model Serving with PagedAttention}, author{Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica}, booktitle{Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles}, year{2023} }【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考