DeepSeek开源周 Day01：从FlashMLA背后原理回顾KV Cache

FlashMLA

今天DeepSeek开源周第一天，开放了FlashMLA仓库，1小时内星标1.6k！

FlashMLA 是一个高效的 MLA 解码内核，专为 Hopper GPU 优化，适用于可变长度序列。该项目目前发布了 BF16 和具有 64 块大小分页 kvcache 的功能。在 H800 SXM5 上，使用 CUDA 12.6，内存受限配置下可达 3000 GB/s，计算受限配置下可达 580 TFLOPS。

Github仓库地址：https://github.com/deepseek-ai/FlashMLA

这里提到两个比较关键的功能就是BF16精度计算以及Paged kvcache缓存技术

好巧不巧，近期DeepSeek 发布了一篇新论文，提出了一种改进版的注意力机制 NSA，即Native Sparse Attention，可以直译为「原生稀疏注意力」；但其实就在同一天，月之暗面也发布了一篇主题类似的论文，提出了一种名为 MoBA 的注意力机制，即 Mixture of Block Attention，可以直译为「块注意力混合」。注意机制最近这么火爆的背景下，不妨我们趁机复习下Kv Cache相关概念以及相关注意力机制模型。

KV Cache简介

这部分主要参考LLM推理算法简述,可以快速回顾下KV Cache概念，关于更多LLM推理算法讲解大家可以阅读https://zhuanlan.zhihu.com/p/685794495。

LLM 推理服务的吞吐量指标主要受制于显存限制。研究团队发现现有系统由于缺乏精细的显存管理方法而浪费了 60% 至 80% 的显存，浪费的显存主要来自 KV Cache。因此，有效管理 KV Cache 是一个重大挑战。

什么是KV Cache？

Transformer 模型具有自回归推理的特点，即每次推理只会预测输出一个 token，当前轮输出 token 与历史输入 tokens 拼接，作为下一轮的输入 tokens，反复执行多次。该过程中，前后两轮的输入只相差一个 token，存在重复计算。KV Cache 技术实现了将可复用的键值向量结果保存下来，从而避免了重复计算。如下图所示，展示了有无 kv cache 的流程：

• Without KV Cache 每次需要计算全 Wq(X), Wk(X), Wv(X)，每次需要计算全量 Attn
• With KV Cache，第一步计算完整 $Attn$，将 $KV$ 保存成 $KV\_cache$
• With KV Cache，第二步，取第二步的 Next Token 计算 $QKV$，将 $[KV\_cache,KV]$ 联合，计算出$QKV$

图片来源：陆淳，https://zhuanlan.zhihu.com/p/685794495

如何优化KV cahce？

KV cache的峰值显存占用大小计算公式：2 x Length x batch_size x [d x n_kv_heads] x Layers x k-bits，由此我们可以看出影响KV cache的具体因素：

• k-bits: 数据类型，FP16 占2个bytes。（量化）
• 2: 代表 Key/Value 两个向量现
• Length: 输入输出序列的长度（循环队列管理窗口KV，减少长度kv）
• Layers：模型层数
• d x n_kv_heads：kv维度（MQA/GQA通过减少KV的头数减少显存占用）
• batch_size : KV Cache 与 batchsize 度呈线性关系，随着 batch size 的增大，KV cache 占用的显存开销快速增大，甚至会超过模型本身。
• 操作系统管理：现GPU的KV Cache的有效存储率较低低 （page-attention）

在bf16格式下的13B模型中，我们只有大约10G的空间来存储kv cache。

KV Cache 的引入也使得推理过程分为如下两个不同阶段，进而影响到后续的其他优化方法。

• 预填充阶段 (Prefill)：发生在计算第一个输出 token 过程中，计算时需要为每个 Transformer layer 计算并保存 key cache 和 value cache；FLOPs 同 KV Cache 一致，存在大量 GEMM (GEneral Matrix-Matrix multiply) 操作，属于 Compute-bound 类型计算。

• 解码阶段 (Decoder)：发生在计算第二个输出 token 至最后一个 token 过程中，这时 KV Cache 已存有历史键值结果，每轮推理只需读取 KV Cache，同时将当前轮计算出的新 Key、Value 追加写入至 Cache；GEMM 变为 GEMV (GEneral Matrix-Vector multiply) 操作，FLOPs 降低，推理速度相对预填充阶段变快，这时属于 Memory-bound 类型计算。

解码中的KV Cache

我们下面用一个例子更加详细的解释什么是KV Cache，了解一些背景的计算问题，以及KV Cache的概念。

无论是encoder-decoder结构，还是现在我们最接近AGI的decoder-only的LLM，解码生成时都是自回归auto-regressive的方式。也就是说，解码的时候，先根据当前输入 $inpu{t}_{i-1}$，生成下一个token，然后把生成的token拼接在 $inpu{t}_{i-1}$ 后面，获得新的输入 $inpu{t}_i$，再用 $inpu{t}_i$ 生成 $toke{n}_{i+1}$，依此选择，直到生成结果。

比如我们输入“窗前明月光下一句是”，那么模型每次生成一个token，输入输出会是这样（方便起见，默认每个token都是一个字符）

step0: 输入=[BOS]窗前明月光下一句是；输出=疑
step1: 输入=[BOS]窗前明月光下一句是疑；输出=是
step2: 输入=[BOS]窗前明月光下一句是疑是；输出=地
step3: 输入=[BOS]窗前明月光下一句是疑是地；输出=上
step4: 输入=[BOS]窗前明月光下一句是疑是地上；输出=霜
step5: 输入=[BOS]窗前明月光下一句是疑是地上霜；输出=[EOS]

(其中[BOS]和[EOS]分别是开始和结束的标记字符)

我们看一下在计算的过程中，如何输入的token “是” 的最后是hidden state如何传递到后面的类Token预测模型，以及后面每一个token，使用新的输入列中最后一个时刻的输出。

我们可以看到，在每一个step的计算中，主要包含了上一轮step的内容，而且只在最后一步使用（一个token）。那么每一个计算也就包含了上一轮step的计算内容。

从公式来看是这样的，回想一下我们attention的计算：

$${\alpha }_{i,j}=\text{softmax}(q_i{k}_j^T)$$

$$o_i=\sum _{j=0}^i{\alpha }_{i,j}{v}_{ij}$$

注意对于decoder的时候，由于mask attention的存在，每个输入只能看到自己和前面的内容，而看不到后面的内容。

假设我们当前输入的长度是3，预测第4个字，那么每层attention所做的计算有：

$$o_0={\alpha }_0,0v_0$$

$$o_1={\alpha }_{1,0}{v}_0+{\alpha }_{1,1}{v}_1$$

$$o_2={\alpha }_{2,0}{v}_0+{\alpha }_{2,1}{v}_1+{\alpha }_{2,2}{v}_2$$

预测完第4个字，放到输入里，继续预测第5个字，每层attention所做的计算有：

$$o_0={\alpha }_{0,0}{v}_0$$

$$o_1={\alpha }_{1,0}{v}_0+{\alpha }_{1,1}{v}_1$$

$$o_2={\alpha }_{2,0}{v}_0+{\alpha }_{2,1}{v}_1+{\alpha }_{2,2}{v}_2$$

可以看到，在预测第5个字时，只有最后一步引入了新的计算，而 $o_0$ 到 $o_2$ 的计算部分是完全重复的。

但是模型在推理的时候可不管这些，无论你是否只是要最后一个字的输出，它都会把所有输入计算一遍，给出所有输出结果。

也就是说中间有很多我们不需要的计算，这样就造成了浪费。

而且随着生成的结果越来越多，输入的长度也越来越长，上面这个例子里，输入长度是step0的10个， 每步骤，直接step5到15个。如果输入的instruction是规范型任务，那么可能有800个step。这个情况下，step0就变得有800次，step1被重复了799次——这样浪费的计算资源显然不可忍受。

有没有什么方法可以重利用上一个step里已经计算过的结果，减少浪费呢？

答案就是KV Cache，利用一个缓存，把需要重复利用的时序计算结果保存下来，减少重复计算。

而 $K$ 和 $V$ 就是需要保存的对象。

想一想，下图就是缓存的过程，假设我们第一次输入的输入长度是3个，我们第一次预测输出预测第4个字，那么由于下图给你看的是每个输入步骤的缓存，每个时序步骤都需要存储一次，而我们依旧会有些重复计算的情况。则有：

kv_cache下标l表示模型层数。在进行第二次预测时，也就是预测第5个字的时候，在第l层的时候，由于前面我们缓存了每层的$k$, $v$值，那层就不需要算新的$o_3$，而不再算$o_1$, $o_2$。因为第l层的$o_0$, $o_1$ 本来经过FFN层之后进到$l+1$层，再经过新的投影变换，成为$l+1$层的$k$, $v$值，但是是$l+1$层的$k$, $v$值就已经保留了!

然后我们把本次新算出来的$k$, $v$值也存储起来。

$$cach{e}_1=\left[(k_0,v_0),(k_1,v_1)\right]$$

$$cach{e}_2=\left[(k_0,v_0),(k_1,v_1),(k_2,v_2)\right]$$

然后我们再做下一次计算出的结果:

$$cach{e}_3=\left[(k_0,v_0),(k_1,v_1),(k_2,v_2),(k_3,v_3)\right]$$

这样就节省了attention和FFN的很多重复计算。

transformers中，生成的时候传入use_cache=True就会开启KV Cache。

也可以简单看下GPT2中的实现，中文注释的部分就是使用缓存结果和更新缓存结果

Class GPT2Attention(nn.Module):......def forward(self,hidden_states: Optional[Tuple[torch.FloatTensor]],layer_past: Optional[Tuple[torch.Tensor]] = None,attention_mask: Optional[torch.FloatTensor] = None,head_mask: Optional[torch.FloatTensor] = None,encoder_hidden_states: Optional[torch.Tensor] = None,encoder_attention_mask: Optional[torch.FloatTensor] = None,use_cache: Optional[bool] = False,output_attentions: Optional[bool] = False,) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:if encoder_hidden_states is not None:if not hasattr(self, "q_attn"):raise ValueError("If class is used as cross attention, the weights `q_attn` have to be defined. ""Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`.")query = self.q_attn(hidden_states)key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)attention_mask = encoder_attention_maskelse:query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)query = self._split_heads(query, self.num_heads, self.head_dim)key = self._split_heads(key, self.num_heads, self.head_dim)value = self._split_heads(value, self.num_heads, self.head_dim)# 过去所存的值if layer_past is not None:past_key, past_value = layer_pastkey = torch.cat((past_key, key), dim=-2)  # 把当前新的key加入value = torch.cat((past_value, value), dim=-2)  # 把当前新的value加入if use_cache is True:present = (key, value)  # 输出用于保存else:present = Noneif self.reorder_and_upcast_attn:attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)else:attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)attn_output = self.c_proj(attn_output)attn_output = self.resid_dropout(attn_output)outputs = (attn_output, present)if output_attentions:outputs += (attn_weights,)return outputs  # a, present, (attentions)

总的来说，KV Cache是以空间换时间的做法，通过使用快速的缓存存储，减少了重复计算。（注意，只能在decoder结构的模型可用，因为有mask attention的存在，使得前面的token可以不用关照后面的token）

但是，用了KV Cache之后也不是立刻万事大吉。

我们简单计算一下，对于输入长度为$L$，层数为$L$，hidden size为$d$的模型，需要缓存的参数量为

$$2\times L\times s\times d$$

如果使用的是半精度浮点数，那么每个值所需要的空间就是

$$2\times 2\times L\times s\times d$$

以Llama2 7B为例，有$L=32$，$L=4096$，那么每个token所需的缓存空间就是524,288 bytes，约524k，假设$s=1024$，则需要占用536,870,912 bytes，超过500M的空间。

这些参数的大小是batch size=1的情况，如果batch size增大，这个值是很容易就超过1G。

注意力相关模型

笔者之前回顾了一些注意力机制相关模型：从MHA、MQA、GQA、MLA到NSA、MoBA，写了下面一篇文章来解读这些注意力模型，

注意力机制进化史：从MHA到MoBA，新一代注意力机制的极限突破！

这里为了节约篇幅，只提一下模型简介，具体原理可以阅读上面综述。

MHA:Multi-Head Attention

论文标题：Attention Is All You Need
论文链接：https://arxiv.org/pdf/1706.03762

MHA在2017年就随着《Attention Is All You Need》一起提出，主要干的就是一个事：把原来一个attention计算，拆成多个小份的attention，并行计算，分别得出结果，最后再合回原来的维度。

MQA:Multi-Query Attention

论文标题：Fast Transformer Decoding: One Write-Head is All You Need
论文链接：https://arxiv.org/pdf/1911.02150

MQA就是减少所有所需要的键值缓存内存消耗的。

Google在2019年就提出了《Fast Transformer Decoding: One Write-Head is All You Need》提出了MQA，不过那时候主要是针对的人不多，那是大家主要还是关注在用Bert也开始创新上。

MQA的做法其实很简单。在MHA中，输入分别经过$W_Q,W_K,W_V$的变换之后，都切成$8$份（$n$=头数），维度也从$d_{model}$降到$d_{head}$，分别进行attention计算再拼接。而MQA这一步，在运算过程中，首先对$Q$进行切分（和MHA一样），而$K,V$则直接在在线变换的时候把维度压到$d_{head}$（而不是切分开），然后返回每个Query头分别和一份$K,V$进行attention计算，之后最终结果拼接起来。

简而言之，就是MHA中，每个注意力头的$K,V$是不一样的，而MQA这里，每个注意力头的$K,V$是一样的，值是共享的。而性别效果和MHA一样。

GQA:Grouped Query Attention

论文标题：GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
论文链接：https://arxiv.org/pdf/2305.13245

既然MQA对效果有点影响，MHA存储又有不下，那2023年GQA（Grouped-Query Attention）就提出了一个折中的办法，既能减少MQA效果的损失，又相比MHA需要更少的存储。

GQA是，$Q$ 还是按原来MHA/MQA的做法不变。只使用一套共享的$K,V$就能效果不好吗，那就还是多个头。但是要不要太多，数量还是比$Q$的头数少一些，这样相当于把多个头分成group，同一个group内的$K,V$共享，同不group的$Q$所用的$K,V$不同。

MHA可以认为是$K,V$头数最大时的GQA（有几个Q就有几个K，V），而MQA可以认为是$K,V$头数最少时的GQA(K的头数只有1个)。

MLA:Multi-head Latent Attention

论文标题：DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
论文链接：https://arxiv.org/abs/2405.04434

MLA（Memory-efficient Latent Attention） 的核心思想是将注意力输入 $h_t$ 压缩到一个低维的潜在向量，记作 $c_t^{KV}$，其维度 $d_c$ 远小于原始的 $h_n\cdot d_h$ 维度。这样，在计算注意力时，我们可以通过映射将该潜在向量恢复到高维空间，以重构键（keys）和值（values）。这种方法的优势在于，只需存储低维的潜在向量，从而大幅减少内存占用。

这一过程可以用以下公式描述：

• $c_t^{KV}$ 是低维的潜在向量。
• $W^{DKV}$ 是一个压缩矩阵（down-projection matrix），用于将 $h_t$ 的维度从 $h_n\cdot d_h$ 降维到 $d_c$（其中 D 代表“降维”）。
• $W^{UK}$ 和 $W^{UV}$ 是两个向上投影矩阵（up-projection matrices），分别用于将共享的潜在向量映射回高维空间，以恢复键(K)和值(V)。

多头注意力（MHA）、分组查询注意力（GQA）、多查询注意力（MQA）和多头潜在注意力（MLA）的简化示意图。通过将键（keys）和值（values）联合压缩到一个潜在向量中，MLA在推理过程中显著减少了KV缓存的大小。

NSA:Native Sparse Attention

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

论文地址：https://arxiv.org/abs/2502.11089

NSA的技术方法涵盖算法设计与内核优化。其整体框架基于对注意力机制的重新定义，通过设计不同的映射策略构建更紧凑、信息更密集的键值对表示，以减少计算量。同时，针对硬件特性进行内核优化，提升实际运行效率。

MoBA: Mixture of Block Attention for Long-Context LLMs

论文标题：Mixture of Block Attention for Long-Context LLMs
论文地址：https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf

它基于MoE原理，应用于Transformer模型的注意力机制，通过将上下文划分为块，并采用门控机制选择性地将查询令牌路由到最相关的块，提高LLMs效率，使模型能处理更长更复杂的提示，同时降低资源消耗。

Page Attention

论文标题：Efficient Memory Management for Large Language
Model Serving with PagedAttention
论文地址：https://arxiv.org/pdf/2309.06180

在 vLLM 库中 LLM 服务的性能受到内存瓶颈的影响。在自回归 decoder 中，所有输入到 LLM 的 token 会产生注意力 key 和 value 的张量，这些张量保存在 GPU 显存中以生成下一个 token。这些缓存 key 和 value 的张量通常被称为 KV cache，其具有以下特点：

• 显存占用大：在 LLaMA-13B 中，缓存单个序列最多需要 1.7GB 显存；
• 动态变化：KV 缓存的大小取决于序列长度，这是高度可变和不可预测的。因此，这对有效管理 KV cache 挑战较大。该研究发现，由于碎片化和过度保留，现有系统浪费了 60% - 80% 的显存。

为了解决这个问题，该研究引入了 PagedAttention，这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在非连续的内存空间中存储连续的 key 和 value 。具体来说，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。

https://arxiv.org/pdf/2309.06180

具体来讲，Paged Attention 将每个序列的 KV Cache 分成若干块，每个块包含固定数量token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理 key 和 value ，就像在操作系统的虚拟内存中一样：可以将块视为页面，将 token 视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块中。物理块在生成新 token 时按需分配。在 PagedAttention 中，内存浪费只会发生在序列的最后一个块中。这使得在实践中可以实现接近最佳的内存使用，仅浪费不到 4%。

FlashAttention

论文标题：FlashAttention: Fast and Memory-Efficient Exact Attention
with IO-Awareness
论文地址：https://arxiv.org/pdf/2205.14135

此前针对Transformer类模型的优化工作致力于减少FLOP从而优化计算速度，忽略了显存IO访问上的优化。

attention计算所带来的O(n^2)复杂度的矩阵对HBM的重复读写是制约模型推理的一个主要瓶颈。

如下图所示GPU SRAM的读写（I/O）的速度为19TB/s 和GPU HBM 的读写（I/O）速度 1.5TB/s 相差十几倍，而对比存储容量也相差了好几个数量级。从图中可以看出Flash Attention的计算是从HBM中读取块，在SRAM中计算之后再写到HBM中。

要解决这些问题，需要做两件主要的事情：

• 在不访问整个输入的情况下计算 softmax
• 不为反向传播存储大的中间 attention 矩阵

为此 FlashAttention 提出了两种方法来分布解决上述问题：tiling 和 recomputation。

• tiling - 注意力计算被重新构造，将输入分割成块，并通过在输入块上进行多次传递来递增地执行softmax操作。
• recomputation - 存储来自前向的 softmax 归一化因子，以便在反向中快速重新计算芯片上的 attention，这比从HBM读取中间矩阵的标准注意力方法更快。

由于重新计算，这虽然导致FLOPS增加，但是由于大量减少HBM访问，FlashAttention运行速度更快。该算法背后的主要思想是分割输入，将它们从慢速HBM加载到快速SRAM，然后计算这些块的 attention 输出。在将每个块的输出相加之前，将其按正确的归一化因子进行缩放，从而得到正确的结果。

参考资料

• LLM（十七）：从 FlashAttention 到 PagedAttention, 如何进一步优化 Attention 性能
• LLM推理算法简述