【开源源码】从零实现一个玩具版LLM推理引擎（一）：用最简单代码彻底搞懂 Pre-fill、Decode 和 KV Cache

Superjet
10 阅读15分钟

目标读者：

  • 了解Transformer基础，但不清楚推理流程细节的工程师
  • 想入门LLM推理系统的学生/爱好者

引言

最近在学习 AI Infrastructure，我毕业多年，从事服务器BMC行业，只用过C/C++语言，被迫零基础从Transformer 架构学起。网上相关介绍文章很多，但是大多都是《Attention is all you need》的工程讲解文，理论居多，编码实操讲解较少。我死记硬背之后自我感觉良好，以为已经熟练掌握Transformer的推理流程，便开始下载vLLM 源码，企图开始推理引擎的源码学习。 但打开代码瞬间破防：我对vLLM的实现思想和需要解决的问题完全没有概念，对推理流程也只是停留在一个模糊的理论认知上。完全不知道从何开始。 痛定思痛，我决定从零开始实现一个MiniVLLM推理引擎，通过实践来理解：

  • LLM的推理核心流程（Prefill/Decode）
  • KV Cache 是什么，为什么需要它
  • Continuous Batching 为什么能提升吞吐
  • 为什么naive batching 会遇到性能瓶颈 *代码已经开源，GitHub 项目

本系列预计共 3~4篇，以记录学习，分享结论为主。

  • 第一篇：Prefill/Decode/KV Cache 基础（本文）
  • 第二篇：Continuous Batching 实现挑战
  • 第三篇：PagedAttention 原理与实现
  • 第四篇：生产级优化（待定）

一 LLM推理的两个阶段

1.Transformer架构的推理流程简单回顾:

image.png

LLM推理流程概括为三步骤：

  1. 输入编码：将prompt通过分词器（Tokenize），向量化（Embedding），后注入位置信息（位置编码），完成模型计算的初始输入序列。
  2. 自回归生成循环：将输入序列进行Attention计算，经过FFN网络非线性变换后输出词表各token的概率分布，通过采样策略(Sampling)预测序列下一个Token。若未生成EOS Token，则将新Token拼回序列。反复进行此步骤，直到生成EOS为止。
  3. 结果输出：将最终token序列通过反编码为输出结果，并返回给用户。

2.KV Cache 的缘由

众所周知，Attention计算公式为 Attention(Q,K,V)=softmax(Q⋅KTdk)⋅VAttention(Q, K, V) = softmax\left(\frac{Q \cdot K^T}{\sqrt{d_k}}\right) \cdot V , 其中，Q，K，V分别为Query，Key，Value 矩阵，矩阵值为上述输入序列（token通过embedding后得到的向量表示）与模型预训练的权重矩阵(Wq，Wk，Wv)相乘(线性变换/Linear层)所得。

从自回归生成循环流程可以直观发现，如果不做任何优化，每次新生成Token时，都需要将已生成的历史文本+新生成Token重新输入给模型。而在Attention计算中，历史文本Token的K，V向量实际上是固定不变的。

为了消除这种巨大的计算浪费，推理引擎引入了KVCache机制。有了KVCache后，推理过程发生了质变：每次新生成Token时，模型不再输入完整历史文本+新Token长序列，而仅仅把最新生成的那1个Token作为输入。模型只需要计算这1个新token的Query，Key，Value向量。接着，用这个新Query直接与缓存在显存中的K，V向量进行attention计算，同时将当前新token的K，V向量追加到缓存中。

以单头注意力，层数为1的LLM举例。输入prompt为“今天很”三字，经分词器，embedding和位置编码后，我们得到 I 矩阵，shape为(3,hiddim)

一.进行QKV投影，I矩阵分别于Wq，Wk，Wv相乘，得到QKV。

image.png

二.计算相似度矩阵：S = Q @ K^T

image.png

三，Softmax + V 加权计算，得到Attention矩阵

image.png

后续经过系列处理后，生成新Token "高" ，在没有KVCache情况下，我们需要把这个Token 经过embedding和位置编码后的向量拼接到原 I 序列矩阵，得到 I1（4，hiddim）矩阵。开始重复自回归生成循环：

一，进行QKV投影，I1矩阵分别与Wq，Wk，Wv相乘，得到Q1, K1, V1。

image.png

从计算结果可直观发现，Q1, K1, V1 相对上次计算，历史序列"今天很"对应的向量值是一样的，仅新增了"高"这一向量

二，计算相似度矩阵：S = Q1 @ K1^T

image.png hidden数据，不深究，继续

三，Softmax + V 加权计算，得到Attention矩阵

image.png 经过计算后的Attention矩阵所有向量的值都发生了变化。但对于文本生成任务，大模型只关心最后一个位置的Attention输出（对应token “高”），虽然attention矩阵计算了所有位置的表示，但只有最后一个位置回被送入FFN 和 LM Head 来预测下一个Token。

结论：

可以看出，Q，K，V生成，历史序列对应的向量都是没有变化的，每次重复计算都是无用消耗。 我们将K，V值保存下来，计算新token的时候直接使用。只需计算该新token的Q，K，V 投影，并将KV添加在历史KV矩阵之后，后续使用该Token的Q向量进行计算即可。直接把attention计算的时间复杂度从O(n²) ->O(n)

为什么不缓存Q值呢？- 因为不需要。

attention计算完成后选取最后一行向量进行下一步计算，在有KVCache优化的系统中，仅计算新token的Q矩阵即可获得同样结果。 推导一下： 假设当前序列长度为 N （前 N−1 个是历史，第 N 个是新 token "高"）。

1. 无优化方案

我们输入整个序列 X1:N。 计算全量 Q,K,V： Q=X1:NWqQ=X_{1:N}W_q, K=X1:NWkK=X_{1:N}W_k, V=X1:NWvV=X_{1:N}W_v

  • 注意： Q是一个 (N,d) 的矩阵，其中第 N行是 qNq_N 。 K,V 也是 (N,d) 的矩阵。
  • 计算 Attention 矩阵 A=Softmax(QKT/d)A=Softmax(QK^T/\sqrt{d}) 这是一个 (N,N) 的矩阵。
  • 关键点：我们只关心第 N 行的输出（对应新 token "高" 的表示）。 第 N 行的计算公式是： OutputN=∑j=1NSoftmax(qNkjTd)⋅vj\text{Output}_N = \sum_{j=1}^{N} \text{Softmax}\left(\frac{q_N k_j^T}{\sqrt{d}}\right) \cdot v_j

这里， qNq_N 是第 N个 query，kjk_j , vjv_j 是第 j 个 key/value 。

2. 有 KV Cache 方案
  • 历史部分：我们从显存中读取之前算好的 K1:N−1K_{1:N-1} 和 V1:N−1V_{1:N−1} 。
  • 当前部分：我们只输入新 token: XNX_N向量 ，计算：

qnewq_{new}=xNWqx_{N}W_{q}, knew=xNWkk_{new} =x_{N}W_{k} , vnew=xNWvv_{new} =x_{N}W_{v}

（注意：因为线性变换 W 是逐 token 独立的，所以这里的 qnewq_{new} 数学上完全等于上面的 qNq_N ，同理 knew=kNk_{new}=k_N, vnew=vNv_{new}=v_N ）。

  • 拼接： Kfull=[K1:N−1,knew],Vfull=[V1:N−1,vnew]K_{full}=[K_{1:N−1} ,k_{new}], V_{full}=[V_{1:N−1}, v_{new}]
  • 计算 Attention： 此时 Query 只有 qnewq_{new} (形状 1×d )。

Score=qnew⋅KfullT=[qnew⋅k1T,…,qnew⋅kN−1T,qnew⋅knewT]Score=q_{new} \cdot K_{full}^T=[q_{new} \cdot k_1^T, … ,q_{new} \cdot k_{N−1}^T, q_{new} \cdot k_{new}^T]

这与无优化方案中 QKT 矩阵的第 N 行完全一致

  • 加权求和：
  • Outputcache=Softmax(Score)⋅VfullOutput_{cache} =Softmax(Score)⋅V_{full} 展开后：

Outputcache=∑j=1NSoftmax(qNkjTd)⋅vjOutput_{cache}= \sum_{j=1}^{N} \text{Softmax}\left(\frac{q_N k_j^T}{\sqrt{d}}\right) \cdot v_j

结果一模一样。有无KVCache 优化最终生成的概率分布于不使用Cache的结果在数学上是严格等价的，所以KVCache是一种无损优化。它并没有改变模型的数学逻辑，也没有近似计算。它仅仅是利用了线性层的逐点独立性和Attention的因果掩码特性将原本需要重复计算的O(N)次矩阵乘法变成了O(1)次针对新token，通过显存换时间，复用了历史计算结果。

3.Prefill 和 Decode 两阶段

从以上两段不难总结LLM推理的两个场景：

  • prompt处理，一次生成prompt所有token序列的KV Cache，供后续自回归逐个生成token
  • 自回归生成流程每次生成一个token，计算Query，刷新 KV Cache

于是，推理引擎将LLM推理分成两个本质不同的阶段：Prefill 和 Decode。

image.png

推理流程如上，两个阶段：

Prefill阶段（预填充）：

  • 输入：完整的prompt序列 （如“今天很” 3个token）
  • 计算：并行计算，所有token并行处理。并行计算是GPU的长项。
  • 输出：
    1. 所有prompt tokens 的KV Cache（存如显存）
    2. 第一个生成的token（如“高”）
  • 特征：计算密集模式（Compute-bound）：因为是一次性把整个prompt塞进来，一次性计算生成KV Cache 和首token，GPU算力利用率很高。

Decode（解码）：

  • 输入：上一步生成的token（如“高”）
  • 计算：串行计算 1. 计算上一步新生成token的 Q，K，V 2. 从显存读取KV Cache 3. 将新Token 的 K，V追加到KV Cache中
  • 输出：下一个新token （如“兴”）
  • 特征：内存访问密集性（Memory-bound）：因为自回归每次仅计算一个token，GPU大部分时间消耗在显存访问读取KV Cache中，GPU大多时间处于摸鱼状态,等待数据进行计算。

二 代码实现

理论回顾完成，开始编码。我们避开推理引擎复杂细节优化（其实是不懂~），专注核心推理流程，基于pytorch，HF transformers 框架实现MiniVLLM的轻量级（玩具）引擎。

Kick off~

1. 核心依赖

PyTorch: 深度学习框架，提供张量运算和 GPU 加速 Transformers: HuggingFace 库，用于加载预训练模型和分词器 Accelerate：简化模型在不同设备(CPU/GPU)上的加载

2. 项目结构

项目采用模块化设计，清晰划分职责：

  • model.py: 模型封装层。负责加载 HuggingFace 模型，并封装 forward 方法，使其支持 past_key_values 的显式传递。
  • tokenizer.py: 文本处理层。封装分词与解码逻辑。
  • mini_vllm.py: 核心引擎层。实现推理编排逻辑，包含 prefill 和 decode 两个关键阶段。
  • main.py: 入口脚本。用于快速测试模型生成效果。
  • benchmark_mini_vllm.py: 性能分析工具。通过采集 GenerationState 数据，量化分析 Prefill 与 Decode 的耗时差异。
3. 模块设计

当前我们的MiniVLLM推理引擎包含主要三个模块：

  • MiniVLLM：引擎类，统筹全局，控制推理流程
  • Tokenizer：分词器类，负责将输入文本转换成token，将token转成文本
  • Model：LLM管理类，负责HuggingFace 模型加载，并封装forward方法实现模型的前向传播计算
  • GenerationState：学习用的推理观测数据

整体框架如下：

image.png

4.流程实现

image.png

generate实现

        def generate(self, prompt: str, max_new_tokens: int = 100, 
	                 temperature: float = 1.0, 
	                 top_p: float = 1.0, 
	                 return_generation_state: bool = False):
	  
	        # prompt -> tokens
	        input_ids = self.tokenizer.encode(prompt)
	        input_ids = input_ids.to(self.model.model.device)
	
	        # 处理 prefill阶段，获取初始的past_key_values和第一个生成的token id
	        # past_key_values 就是KV Cache，会被decode阶段不断更新, 保存在outputs.past_key_values里
	        outputs = self.prefill(input_ids, temperature=temperature, top_p=top_p, generation_state=generation_state)
	        
	        #处理 decode阶段，逐步生成新token     
	        generated_ids = self.decode(outputs.next_token_id, 
	                                    outputs.past_key_values, 
	                                    max_new_tokens=max_new_tokens - 1, 
	                                    eos_token_id=self.model.eos_token_id, 
	                                    temperature=temperature, 
	                                    top_p=top_p,
	                                    generation_state=generation_state)
	        
	        # tokens -> 文本
	        generated_text = self.tokenizer.decode(generated_ids[0])
	
        return generated_text

Prefill 实现

prefill的目标是吃掉prompt，并生成第一批KV Cache

def prefill(self, input_ids: torch.Tensor, 
	                temperature: float = 1.0, 
	                top_p: float = 1.0, 
	                generation_state=None)->PrefillOutput:
	   
	        prefill_output = PrefillOutput(outputs=None, generated_ids=None, next_token_id=None, past_key_values=None)
	        
	        # 性能观测：计算生成第一个token的时间
	        if generation_state is not None and generation_state.return_generation_state:
	            torch.cuda.synchronize()
	            generation_state.time_to_first_token = time.time()
	        #1. 一次性将完整的 prompt输入给模型
	        prefill_output.outputs = self.model.forward(input_ids)
	        # 2. 获取初始 KV Cache（包含了 Prompt 中所有 token 的信息）
	        # past_key_values 的形状通常是: [层数, 2 (K和V), Batch, Head数, 序列长, Head维度]
	        prefill_output.past_key_values = prefill_output.outputs.past_key_values
	        
		      # 3. 采样出第一个生成的 token
		   # 这里的 -1 是因为在 Prefill 阶段，我们其实计算了所有输入 token 的预测结果，但为了生成第一个新 token，我们只取最后一个位置的预测。
	        logits = prefill_output.outputs.logits[:, -1, :] 
	
	
	        prefill_output.next_token_id = self._sample(logits, temperature, top_p)
	        
	        # 性能观测：记录prefill阶段的耗时
	        if generation_state is not None and generation_state.return_generation_state:
	            torch.cuda.synchronize()
	            generation_state.time_to_first_token = time.time() - generation_state.time_to_first_token
	        # 4 新token拼接到生成序列中
	        prefill_output.generated_ids = torch.cat([input_ids, prefill_output.next_token_id], dim=-1)
	        return prefill_output

Decode 实现 一旦有了初始 Cache，后续的生成就不再需要重新读 Prompt 了。

	    def decode(self, next_token_id, past_key_values, 
	               max_new_tokens: int = 100, 
	               eos_token_id: int = None, 
	               temperature: float = 1.0, 
	               top_p: float = 1.0,
	               generation_state: GenerationState = None)->torch.Tensor:
	        
			...
			
	        # 逐步生成新 token
	        for _ in range(max_new_tokens):
	            # ！这里是核心关键：每次只传入【一个】token_id 和 之前的past_key_values
	            outputs = self.model.forward(next_token_id, past_key_values=past_key_values)
	            
	            # 刷新 KV cache：模型会将当前token的kv 追加到历史缓存中
	            past_key_values = outputs.past_key_values
	            # 获取最后一个位置的 logits，采样下一个 token
	            logits = outputs.logits[:, -1, :]
	            next_token_id = self._sample(logits, temperature, top_p)
	            
	            # 拼接到生成序列中
	            new_tokens_list.append(next_token_id)
	            # 如果生成了结束符，则停止生成
	            if eos_token_id is not None and next_token_id.item() == eos_token_id:
	                break
            ...
5.深度解析：KVCache 到底占了多少显存？

在MiniVLLM类中，我们实现了一个有趣的方法 kv_cache_size_mb 专门用于计算KV Cache显存占用

  • 计算原理： KVCache Size = BatchSize x LayerNum x 2(K+V) x KV_HeadNum x HeadDim x SeqLen x 精度（bytes）

举例，对于一个Llama-7B 模型：

  • 层数(LayerNum):32
  • KV_Head数(HeadNum)：32 （MHA 时等于 Q_HeadNum）
  • Head维度(HeadDim)：128
  • 序列长度（Seqlen）：5（假设当前序列为“今天很高兴”）
  • BatchSize : 单序列，batch size = 1
  • 精度：FP16（2Bytes）

计算结果：KV Cache Size = 322321285*2 = 2,621,440 (bytes)

*注：本计算基于 MHA 全量注意力，若模型使用 GQA 架构，KV Head 数需按实际比例缩减

  • 对于 GQA 架构（如 Qwen2.5），KV_HeadNum < Q_HeadNum
  • 对于 MQA 架构，KV_HeadNum = 1*

kv_cache_size_mb 代码实现：

	    #--------------------------------------------------------------
	    # 计算KV Cache的大小，单位MB
	    # kv cache 大小的计算规则: size_bm = batch_size * head_num * head_dim * seq_len * 2(K&V) * 2(FP16) / (1024*1024)
	    #--------------------------------------------------------------
	    def kv_cache_size_mb(self, past_key_values):
	        total_size_mb = 0.0
	        for layer_kv in past_key_values:
	            # layer_kv 是一个 tuple，包含 (key, value)
	            # 常见时(k,v)
	            for kv in layer_kv:
	                if torch.is_tensor(kv):
	                    total_size_mb += kv.element_size() * kv.numel()
        return total_size_mb / (1024 ** 2)
6.执行结果

运行环境： GPU： RTX 4060 8GB HOST RAM ： 16GM Model：Qwen2.5-1.5B-Instruct

执行代码：

	if __name__ == "__main__":
	    model_id = "Qwen/Qwen2.5-1.5B-Instruct"
	    mllm = MiniVLLM(model_id)
	    prompt = "你好，请写一首关于码农写代码的五言绝句,谢谢。表现出码农面临家庭，工作，社会的经济压力和情感压力，还要努力debug，加班加点地工作。要现实主义风格但又带有一点浪漫主义的豪放。回复控制在100字以内。"
	    generated_text = mllm.generate(prompt, max_new_tokens=100, temperature=0.7, top_p=0.9)
	    
    print(f"\n模型回答: {generated_text}")

执行结果

模型回答: 题目：《码农之歌》 在代码的海洋，日夜奋战， 家的呼唤与社会的重担并肩。 debug的火焰，燃烧在键盘边， 加班加点，豪放的歌声在耳边。

这首诗表达了码农在工作与生活压力中的挣扎和坚持，以及对生活的热爱。它通过现实主义的描写，同时融入了一些浪漫主义的元素，展现了码农的奋斗精神。希望您喜欢。

结果分析

….. 这好像不是五言绝句，不过基于我的要求比较高，它好歹也正常生成了文本…… 嗯…..先这样吧…..

虽然模型没有严格遵循“五言绝句”的格式要求（应为4句，每句5字），但能够理解prompt的语义，并生成了相关诗歌内容，说明：

  • 我们的MiniVLLM引擎能正常工作。
  • KV Cache机制正确第保持了上下文的连贯性
  • 模型对复杂prompt的理解能力还有提升空间

这个测试的重点是验证推理流程的正确性，而非模型的文学创作能力，所以执行成功。

MiniVLLM引擎性能测试

我们针对不同的prompt执行推理，以下观察项：

  • Prefill time： Prefill阶段耗时
  • Decode time：Decode阶段耗时
  • TTFT：Time To First Token，首token生成耗时
  • Throughput：生成速度

测试环境：

  • GPU： RTX 4060 8GB
  • HOST RAM ： 16GB
  • Model：Qwen2.5-1.5B-Instruct

推理参数：

  • MaxNewToken： 500
  • Temperature = 0.7

性能测试代码脚本：benchmark_mini_vllm.py 各位有兴趣自取（见文章末尾 github链接）

执行结果：

image.png

执行到第五个推理请求时nvidia-smi 的输出结果：

image.png

测试结果整理

image.png

数据分析：

Prefill 时间稳定：

  • 与prompt长度相关，prompt = 1: 0.040s, prompt = 129: 0.090s(约2.25倍)。
  • GPU并行计算效率很高，可同时处理多个token的Attention计算。
  • 多出来的时间大多数消耗在显存访问上了，prompt越长，显存访问越次数多。
  • TTFT时间几乎完全等同于整个prefill的时间，这在我们目前的单请求串行推理实现是合理的。但是在高并发场景中，TTFT时间还需要包含prefill未开始前的等待时间，以及batching处理的等待时间。

Decode是主要瓶颈

  • 占总时间99%以上，我们设置的推理参数最大生成token数量为500，如果生成数量降低至50，decode时间会大幅缩小。
  • 因此推理优化的重点应该在decode。

KV Cache 按预期增长

  • Kv cache size符合理论公式
  • 内存管理正确

总结与展望

本文要点

理解了LLM 推理分为 Prefill 和 Decode 两阶段

  • Prefill：并行处理 prompt，生成 KV Cache（Compute-bound）
  • Decode：串行生成 token，复用 KV Cache（Memory-bound）

理解了 KV Cache 的数学原理

  • 利用线性层的逐点独立性
  • 利用 Attention 的因果掩码特性
  • 无损优化：数学上严格等价

实现了完整的单请求推理引擎

  • 清晰的模块化设计
  • 可观测的性能数据
  • 验证了理论分析
当前实现的局限

GPU利用率低

  • decode阶段GPU大部分时间在等显存数据进行计算
  • 从测试数据看，decode占总时间99%+
  • 能否同时处理10个用户请求来提升GPU体用率？

无法处理动态批次

  • 当前不支持batching
  • 即使支持batching，不同的用户prompt长度不一样
  • 生成长度也不一样
  • 气泡（Bubble）问题：传统的 Batching 必须互相等待，这多浪费。

下一步：Continuous Batching

为了解决上述问题，我们将在下一篇实现： 1. 动态批处理

  • 请求可以随时加入/完成
  • 单个请求不需要等待整个batch完成，即可返回结果

2. 发现Navie Batching的性能缺陷

  • Padding 浪费
  • KV Cache merge/split 开销
  • Pageattention铺垫

3. 实际测量性能提升

  • 吞吐量对比
  • GPU利用率分析
  • 瓶颈识别

敬请期待下一篇：《从零实现一个玩具版LLM推理引擎（二）：Continuous Batching 的挑战》

代码已经开源，GitHub 项目

水平有限，欢迎指正，如果觉得有帮助，欢迎
  • GitHub标⭐，Star支持项目
  • 留言讨论
  • 分享给更多人

参考资料 vLLM 论文：Efficient Memory Management for Large Language Model Serving with PagedAttention Transformers 官方文档

avatar
文章
阅读
粉丝