vLLM 源码分析系列：KVCache 初始化之 V1 版本分析
今天这篇文章是源码分析系列中最难写的文章，KVCache 到底怎么管理的呢？ 只要看过 vLLM 资料的相比都能脱口而出，PageAttention 嘛，Block 嘛，但是到底怎么实现的呢？
V1 中对于内存的分配改为通过 token 的配额来进行调度，例如一次 prefill 请求的 token 配额等于序列的长度或者 chunked size，一次 decode 的 token 配额就是 1. 那么这个 token 的内存配额是如何在内存中分配的呢？为什么大模型启动之后，内存占用立马就飚上去了？
我带着这些问题，希望搞清楚在 vLLM 的源码中是如何实现的？这一看不要紧，真的是相当的费劲，主要有三点：
运算公式复杂，资料匮乏
代码逻辑不清晰，参数引用混乱
大量异步调用，调用链不直观
vLLM 发展过快，因此代码的结构是有些混乱的，这也是 V1 为什么要对代码结构进行重构的原因，因此梳理清楚内存分配的逻辑也是很费了一番功夫。
内存结构
我们先直奔主题，vLLM 的 KV Cache 到底是什么结构，为什么一上来就会占用那么多内存，这里我是以一个 Qwen2-1.5B 的模型作为开发调试模型进行分析的，vLLM 启动时能看到这么一行日志：
可以看到大概占用了 38.48GiB 的内存，这个数字应该跟 GPU 的总内存、vllm 的内存占用比例参数、模型大小这几个变量有关。总之 vLLM 负责分配这 38GiB 的内存了。
我们直接看下图，内存结构分为三层：
第一层是 Tensor: KVCacheManager 并不直接跟内存进行交互，也就是说创建 Block、修改 Block、释放 Block 都不会直接对物理内存进行操作，物理内存的占用也不会发生变化，至于为什么要用 Tensor，我们后面展开论述
第二层是 Block: 一个大的 Tensor 被分配为若干个 Block，一个 Block 对应一个 Page，所以 Block 的大小就是 page_size_bytes 的大小，tensor 的大小和 page_size 的大小是对齐的
第三层是 token： token 表示一个 token 的 k v 缓存所占用的内存空间，一块 Block 可以容纳的 token 数量就是 block_size
1、为什么要用 Tensor
Tensor 是 Pytorch 的核心数据结构，通过 Tensor 进行内存管理，可以利用 Pytorch 已有的对内存进行分配、管理和释放的成熟生态，避免去进行复杂的内存管理。
另外使用 Tensor 也能够很方便的存储 KV Cache，他们的数据存储方式本身都是一致的，Tensor 的数据结构就是很方便的存取张量。
使用 Tensor 的管理性能也是很高的，不需要频繁的操作内存，很多读写操作可以做到零拷贝，例如 reshape Tensor 的形状，将其维度从 1 维数组改变为适合 token 维度的形状，只需要修改 Tensor 的元数据即可。
2、 Tensor 有多少个，有多大
根据代码的分析，一个模型的 Layer 就对应一个 Tensor,除了有些 Layer 可能会共享 tensor 的情况。那么一个 Tensor 的大小就很简单了，它等于 总内存 // layer nums。
3、 Block 的大小如何计算
Block 的大小计算起来比较麻烦，它需要用用 Token的大小 * token 的数量。而 Block 的数量就等于 。
这里引用一段计算的代码：
4、 token 的大小如何计算
这里直接引用一段 Token 大小计算的公式：
总结：
我们结合上面的描述总结一下这几个层的大小：
源码剖析
分享完上面的部分，我们知道了一个 cache 的整体结构，那么源码层面是如何实现的呢，我将其分成四个环节，如下图所示：
初始化配置，所有的配置都在 vllm_config 中
从模型中获取 kv_cache_spec，这里描述了一个 cache 的各个 layer 的属性
构建 kv_cache_config，这里将 layer 的信息转换成了分配 cache 内存的 Tensor 的信息，并计算了 num_blocks 的信息
初始化 Cache，创建 Tensor，占用物理内存
KV Cache 初始化的起始代码在：
1、 首先是配置初始化
vllm_config 中有很多很多的参数，我们重点关注两个部分：
model_config.hf_config
看名字应该能猜到，这是一个获取模型目录下的那份 config.json 的文件的配置，这个配置中有我们关注的内容：
compilation_config 解析
这个参数里面配置了一些编译相关的参数，例如 cuda_graph 相关，本章节内容聚焦在 kv_cache_manager 上，因此重点关注一个参数叫做 compilation_config.static_forward_context.
如下图所示，在 load_model 的逻辑链中，可以找到初始化 compilation_config.static_forward_context 的方法。
这段代码有很多个实现了，这里只选择当前调试过程中在用的视线。
为什么关注这个参数，是因为他返回了 AttentionSpec 的相关内容，也就是每个 layer 的 block_size,num_kv_heads 等信息，是用来构建 kv_cache_spec 对象的数据来源。
这里就是一点不太明白的，compilation_config 看起来是一个编译相关的配置，不明白为什么 layer 相关的信息不放到 model_config 里，而是放在编译配置中。
2、初始化 kv_cache_spec
调用链如下图所示：
core._initialize_kv_caches
gpu_worker.get_kv_cache_spec
gpu_model_runner.get_kv_cache_spec
从这个代码跟下去，可以看到 kv_cache_spec 的构建逻辑，只是代码相对让人比较困惑的就是其中会引用到很多参数，例如 block_size 、attn_layer，这些参数从哪里来的，会比较难以跟踪。
3、 获取 kv_cache_config
代码的位置比较清晰：
标准的模型每一层的注意力都是一样的，会走到_get_kv_cache_config_uniform_type 的逻辑。最终得到的 kv_cache_config 的结构如下：
Tensor 本质上是一个多维数组，但是它的数据在内存中是以一维连续存储的（通常称为“展平”存储，flattened）.Tensor 的形状只是元数据，描述如何将这一维内存块解释为多维数组。
例如，一个形状为 [2, 3] 的 Tensor（6 个元素）在内存中存储为一个连续的 1D 数组，长度为 6。shape 和 strides（步幅，描述每一维的元素间隔）定义了如何访问这些元素
4、初始化 Cache
根据 kv_cache_config 中描述的 Tensor 信息，gpu_model_runner 会执行初始化的动作：
这里创建的 tensor 就是一个一维的长度为 kv_cache_tensor.size 的数组，根据 block，head_size 等信息，tensor 的形状会被 reshape.
总结
综上，就是 kv_cache 的整个初始化过程了，了解其中的源码结构和 cache 的内存结构，对于 vLLM 是如何实现 PageAttention 是非常重要的，对于下一步了解调度器的配额管理也是至关重要的基础。
后续会展开对于调度器部分的代码分析。
版权声明: 本文为 InfoQ 作者【Jason黄】的原创文章。
原文链接:【http://xie.infoq.cn/article/b7da5a9229d5d0032467f301a】。文章转载请联系作者。
还未添加个人签名 2018-10-18 加入
还未添加个人简介
评论