模型推理框架——vllm原理及整体框架 - Big-Yellow

PageAttention原理分析

Page Attention也是一种优化方法（区别于MLA，page attention是对内存进行分配管理）。参考论文^[1]中描述，对于KV-cache存在3个问题：

1、预留浪费 (Reserved)：为将来可能的 token 预留的空间，这些空间被保留但暂未使用，其他请求无法使用这些预留空间；
2、内部内存碎片化问题（internal memory fragmentation）：系统会为每个请求预先分配一块连续的内存空间，大小基于最大可能长度(比如2048个token)，但实际请求长度往往远小于最大长度，这导致预分配的内存有大量空间被浪费。
3、外部内存碎片化问题（external memory fragmentation）：不同内存块之间的零散空闲空间，虽然总空闲空间足够，但因不连续而难以使用。

只有 20.4%-38.2% 的token是被使用的，大部分都被浪费掉了。Page Attention允许在非连续的内存空间中存储连续的 key 和 value 。具体来说，Page Attention将每个序列的 KV-cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，Page Attention内核可以有效地识别和获取这些块。如何理解上面描述呢？还是借用论文中的描述：

比如说按照上面Prompt要输出（假设只输出这些内容）：“fathers brought a car”，一般的套路可能是：比如说：“Four score and seven years ago our xxxxx”（xxx代表预留空间）因为实际不知道到底要输出多少文本，因此会提前预留很长的一部分空间（但是如果只输出4个字符，这预留空间就被浪费了），因此在page attention里面就到用一种“分块”的思想处理，以上图为例，分为8个Block每个Block只能存储4个内容，因此就可以通过一个Block Table来建立一个表格告诉那些Block存储了多少，存储满了就去其他Blobk继续存储。整个过程如下：

上述过程描述如下：具体而言，Page Attention 首先将 Key/Value 的连续显存空间划分为固定大小的 Block（页），每个 Block 作为最小的内存分配与调度单元。随后，引入一个 Block Table（页表） 来维护逻辑序列位置与物理 Block 之间的映射关系，用于记录每个 Block 当前的存储状态与可用容量。
一个小问题：分块之后注意力计算过程，因为我的KV被存储在不同的block中，由于Block table存在可以直接去索引不同Blcok中KV值，这样一来对于Q、K、V三者计算不成问题，不过关键问题就是：Softmax 的分母需要全局信息，Block （不管是Flash Attn还是Page Attn都需要面对这个问题）是分开的，怎么办？
在softmax计算过程中：\(\sigma= \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}\) 由于分块可能导致值过小进而导致数值溢出问题，除此之外计算需要所有token的分数一起归一化，因此首先会对上面的公式改进为：\(\sigma= \frac{e^{z_i-m}}{\sum_{j=1}^K e^{z_j-m}}\)也就是将每块都去减去当前的最大值（避免溢出问题）。在处理全局问题上：只需要考虑两个值的更新：1、当前最大值；2、归一化因子（\(\sum_{j=1}^K e^{z_j-m}\)）因此这个过程就可以处理为：

\[l_{t+1}=\sum_{i\in B_{≤t+1}} e^{z_j- m_{t+1}}=\sum_{i\in B_{≤t}}e^{z_j- m_{t+1}}+ \sum_{i\in B_{t+1}}e^{z_j- m_{t+1}}\\ =\sum_{i\in B_{≤t}}e^{z_j- m_{t}} e^{m_t-m_{t+1}}+ \sum_{i\in B_{t+1}}e^{z_j- m_{t+1}} \]

这样一来就可以转化为：\(l_{t+1}=l_t e^{m_t-m_{t+1}}+ \sum_{i\in B_{t+1}}e^{z_j- m_{t+1}}\)

基本使用方式

在使用vllm上有两种方式：1、离线使用；2、在线使用（直接将使用过程转化为调用API方式）：

from vllm import LLM, SamplingParams
prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)
for output in outputs:prompt = output.promptgenerated_text = output.outputs[0].textprint(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

vllm整体框架分析

基于：Version: 0.11.0

在vllm中主要是两种调用方式：1、离线调用；2、在线调用（这个就类似在本地启动一个服务，而后其他及其直接访问ip端口等进行访问处理）

上图中在线调用方式（Asy）和离线调用（Syn）

对于具体的LLMEngine的结构描述见后面的描述

以离线调用方式进行解释，直接使用官方代码为例：

from vllm import LLM, SamplingParams
prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)
for output in outputs:prompt = output.promptgenerated_text = output.outputs[0].textprint(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

从上面代码分析发现感觉和平时使用Transformer框架和相似：加载模型-->编码输入-->输入模型-->模型输出并且解码。差异在于使用vllm首先会使用一个LLM去处理你的模型，而后你其他的方式都是在这个LLM中，因此了解一下在模型接受到我的prompt之前模型都在做什么。

vllm初始化过程

按照PPT^[2]中对于模型加载的描述：

在模型进行输出之前主要是进行3步：1、初始化并且加载模型；2、预分配显存过程；3、将预分配的KV Cache加载到gpu上。

模型初始化过程

在vllm中定义一个llm过程为：

# vllm/entrypoints/llm.py
class LLM:...self.llm_engine = LLMEngine.from_engine_args(...)# vllm/v1/engine/llm_engine.py
class LLMEngine:def __init__(...):self.engine_core = EngineCoreClient.make_client(...)def generate(...):...def add_request(...):...
# /vllm/v1/engine/core_client.py 中 EngineCoreClient通过多种（异步/多进程，这也就意味这在linux有些可能需要使用`multiprocessing.set_start_method('spawn', force=True)`）方式进行加载模型

在LLMEngine代码中定义了基本所有函数功能，如生成等（后续解释具体过程）。

预分配显存过程

这个给过程的的话首先是去计算预分配的KV Cache大小，而后将预分配的KV Cache加载（一般就是初始化为0的向量）到gpu上

• 计算预分配的KV Cache

计算预分配的KV Cache^[3]：可用显存大小×预分配vllm比率- 非kv cache占用大小得到kv cache的可用（字节）大小，而后通过总共可用大小计算可用分多少个block：可分配大小//KV cache block 的字节大小//所有 kv_cache_groups中层数的最大值

在调用代码LLM(model="facebook/opt-125m")实际过程中会使用load_model进行模型加载（代码：vllm/v1/worker/gpu_model_runner.py）在加载模型之后，模型会进行一个显存的预分配处理，这个过程（代码：vllm/v1/core/kv_cache_utils.py）描述如下：
1、计算需要分配多少显存给vllm：可用显存大小*初始化分配大小（self.requested_memory=self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization，比说24G（实际可能比24G要小，因为还有模型占用）显卡那么的第一项结果就是：24*1024^3，后面一下就是最开始的参数）
2、计算分配给kv cache的显存占用字节大小：可以显存大小-除去KV cache显存外其他大小（self.available_kv_cache_memory_bytes = self.requested_memory - profile_result.non_kv_cache_memory）
在计算完毕之后（以上面模型加载为例，得到KV cache大小为：20.44GiB）接下来就是计算GPU上 KV Cache 内总token数量：num_tokens = num_blocks // len(kv_cache_groups) * min_block_size
1、num_blocks计算过程：int(available_memory // page_size // num_layers)，其中page_size代表是一个 KV cache block 的字节大小（page_size = 2(K+V) * 16(block_size) * 12(num_kv_heads) * 64(head_size) * 2(dtype_bytes 其中fp16对应2)=49152，里面num_kv_heads对应你的模型结构使用数量）；num_layers：所有 kv_cache_groups中层数的最大值，比如说在模型facebook/opt-125m中总共有12层decode（即 12 层进行注意力计算）并且这些attn计算方式完全相同那么就是1个group分组（如果还有其他attn那么可能就是多个group但是最后还是取最大值：group_size = max(len(group.layer_names) for group in kv_cache_groups)。最后计算得到结果为：num_blocks = 21946158284//49152// 12=37207
2、min_block_size = min([group.kv_cache_spec.block_size for group in kv_cache_groups]) 计算得到:16。

实际调试过程中（直接在需要调试位置使用logger.info），输出kv_cache_groups看到的比如（其实这个参数也就是记录cache需要发生位置，一般就是attn计算，不过可能对于不同的attn存在差异，有些是常规有些有可能是window-attn等）：[KVCacheGroupSpec(layer_names=['model.decoder.layers.0.self_attn.attn', ..., 'model.decoder.layers.11.self_attn.attn'], kv_cache_spec=FullAttentionSpec(block_size=16, num_kv_heads=12, head_size=64, dtype=torch.float16, sliding_window=None, attention_chunk_size=None))] 除此之外这部分结果会直接存入KVCacheConfig中。在后续代码（vllm/v1/worker/gpu_model_runner.py）中对于initialize_kv_cache（具体解释见下面）还会为每一块model.decoder.layers.0.self_attn.attn取分配一个初始化（具体函数：initialize_kv_cache_tensors）为0的向量大小为：[2, num_blocks, block_size, num_kv_heads, head_size]

因此最后就可以直接得到：num_tokens = 37207// 1*16 = 595,312。

• 将预分配的KV Cache加载到gpu上

在上述步骤中计算得到了预分配的KV cache大小以及num blocks，接下来就是直接将其先放置到gpu上，实现显存的预分配，以后这块显存就是专门用来做KV Cache。具体过程中还是使用上面得到的kv_cache_groups这个参数

# vllm/v1/worker/gpu_model_runner.py
def initialize_kv_cache_tensors(self, kv_cache_config: KVCacheConfig):# Initialize the memory buffer for KV cachekv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)# Change the memory buffer to the desired shapekv_caches = self._reshape_kv_cache_tensors(kv_cache_config, kv_cache_raw_tensors)...num_attn_module = 2 if self.model_config.hf_config.model_type == "longcat_flash" else 1bind_kv_cache(kv_caches,self.compilation_config.static_forward_context,self.kv_caches, num_attn_module)return kv_cachesdef _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig):kv_cache_raw_tensors: dict[str, torch.Tensor] = {}logger.info(kv_cache_config)for kv_cache_tensor in kv_cache_config.kv_cache_tensors:tensor = torch.zeros(kv_cache_tensor.size,dtype=torch.int8,device=self.device)for layer_name in kv_cache_tensor.shared_by:kv_cache_raw_tensors[layer_name] = tensor...return kv_cache_raw_tensors

对于参数KVCacheConfig就是上面的kv_cache_groups结果，只不过还会取计算每层的大小也就是会更新为：KVCacheConfig(num_blocks=37207, kv_cache_tensors=[KVCacheTensor(size=1828798464, shared_by=['model.decoder.layers.0.self_attn.attn']), ..., KVCacheTensor(size=1828798464, shared_by=['model.decoder.layers.11.self_attn.attn'])], kv_cache_groups=[KVCacheGroupSpec(layer_names=['model.decoder.layers.0.self_attn.attn', ...,'model.decoder.layers.10.self_attn.attn', 'model.decoder.layers.11.self_attn.attn'], kv_cache_spec=FullAttentionSpec(block_size=16, num_kv_heads=12, head_size=64, dtype=torch.float16, sliding_window=None, attention_chunk_size=None))])

在函数self._allocate_kv_cache_tensors中很容易理解直接初始化一个全部为0的张量，而后再去通过函数_reshape_kv_cache_tensors将张量的形状改为[num_blocks, block_size, num_kv_heads, head_size]

总结

本文主要是简单介绍了一些vllm的显存分配过程中，预分配显存比较简单:可用显存大小×预分配vllm比率- 非kv cache占用大小，而后就是直接去计算KV Cache中总token数量这部分计算过程是：num_tokens = num_blocks // len(kv_cache_groups) * min_block_size。

参考

---

1. https://dl.acm.org/doi/pdf/10.1145/3600006.3613165 ↩︎

2. vllm-ppt ↩︎

3. https://zhuanlan.zhihu.com/p/691045737 ↩︎