nanoVLM中实现KV缓存：提高38%生成速度的秘诀

我们在nanoVLM代码库中从零实现了KV缓存机制，使得模型在生成过程中得到了38%的速度提升。nanoVLM是一个小型的PyTorch代码库，用于训练自己的视觉语言模型（Vision Language Model）。本文详细介绍了KV缓存的工作原理以及我们的实现经验，这些经验和技巧可以应用到所有自回归语言模型的生成过程中。 自回归语言模型的工作原理 自回归语言模型通过逐个采样生成文本。在推理阶段，模型处理输入序列，预测下一个词元，将其追加到序列末尾，然后再重复这一过程直到满足某个停止条件。这种逐个生成的方式虽然有效，但存在计算冗余的问题，尤其是在生成较长的序列时，每次都会重新计算整个序列的Q、K、V值。 变压器架构重温 为了理解KV缓存的作用，我们先回顾一下变压器模型中的自注意力机制。自注意力机制涉及三个主要步骤： 输入嵌入投影：输入序列通过权重矩阵W_q、W_k、W_v分别投影得到查询（Q）、键（K）和值（V）。 计算注意力分数：通过Q和K的点积计算注意力分数，再应用因果掩码防止未来词元的泄露。 加权求和：用softmax函数处理注意力分数后，与V进行加权求和，得到最终的输出。 计算冗余问题 在自回归生成中，每当产生一个新词元时，模型会重新计算所有之前词元的K和V值，这导致了大量的计算冗余。例如，对于长度为5的输入序列，在生成第6个词元时，前5个词元的K和V值与之前的计算结果完全相同，只有新词元的部分需要更新。 KV缓存的工作方式 为了解决这个问题，我们引入了KV缓存机制： 缓存结构：每个层都有一个缓存字典，包含"key"和"value"两个部分，形状分别为(batch_size, num_heads, seq_len_cached, head_dim)。 增量更新：只计算当前新词元的K和V值，并将其追加到缓存中，而不是重新计算整个序列。 两阶段生成：生成过程分为预填充（Prefill）和解码（Decode）两个阶段。预填充阶段处理输入提示并建立初始缓存；解码阶段利用缓存逐个生成新词元。 KV缓存在nanoVLM中的实现 Attention Block：在LanguageModelGroupedAttention类中，我们修改了forward函数以接受和更新KV缓存。现在只计算新词元的Q、K、V值，并将其追加到缓存中。 Layer-wise Cache Tracking：在LanguageModel类中，我们引入了按层缓存跟踪的机制。通过start_pos参数帮助模型正确计算新生成词元的旋转位置编码。 生成循环：在VisionLanguageModel类的generate方法中，我们将生成过程分为预填充和解码两个阶段。预填充阶段处理输入提示并构建初始缓存，解码阶段则使用缓存逐个生成新词元。 实现后的变化总结 LanguageModelGroupedAttention.forward：从每步重新计算Q、K、V变为使用和更新KV缓存。 LanguageModel.forward：不再记忆之前的态，而改为按层跟踪KV缓存，并处理start_pos参数。 VisionLanguageModel.generate：从单一阶段生成循环变更为预填充和解码两个阶段。 KV缓存的重要性 增量增长：缓存随新词元逐行增长，减少了计算负担。 位置感知解码：start_pos参数确保新词元的位置编码计算正确。 效率提升：将每个词元的推理复杂度从O(seq_len^2)降至O(seq_len)，特别是在生成长序列和实时应用中效果显著。 KV缓存机制通过减少重复计算，提升了自回归语言模型生成的效率和速度，使其能够在资源受限的环境中运行。尽管这种方法可能会使代码变得更复杂，并限制某些高级推理方案（如束搜索）的应用，但在实践中，它带来了显著的性能改进，使其成为现代大型语言模型加速推理的标准技术之一。 业内评价 业界普遍认为，KV缓存是一种高效的技术，能够大幅提高自回归模型的推理速度，尤其适用于长文本生成和实时对话系统。nanoVLM作为一个简洁、自包含的PyTorch代码库，展示了如何在实际项目中应用这项技术，非常适合初学者学习和理解。