DeepSeek中的MLA无疑是一种具有开创性的注意力机制。它巧妙地运用低秩联合压缩技术，将多个注意力头所对应的键（K）和值（V）精准地映射到一个低维潜在空间之中。这一精妙的设计，如同为数据洪流搭建了一座高效的桥梁，使得原本繁杂的KV缓存存储需求得到了显著削减，计算复杂度也大幅降低。在模型性能丝毫不打折扣的前提下，推理效率得到了质的飞跃，显存占用更是被有效控制，为深度学习模型的轻量化与高效运行开辟了新的路径。

值得一提的是，DeepSeek还巧妙地融合了MoE架构，即DeepSeekMoE技术。这一结合不仅将计算量降至最低，而且极大地提升了模型的总体性能，实现了质的飞跃。DeepSeek总体架构如图9-1所示。

图9-1 DeepSeek总体架构

在本章中，我们将以MLA为起点，开启一场探索DeepSeek技术奥秘的奇妙之旅。从MLA这一核心技术的基础原理出发，逐步深入剖析其背后的数学逻辑和实现细节，让读者能够清晰地理解它是如何在模型中发挥关键作用的。随后，我们会沿着技术的脉络，逐步展开对DeepSeek其他核心技术的介绍，带领读者全方位、多层次地领略DeepSeek这一深度学习领域杰出成果的独特魅力与卓越实力，一同探寻其在人工智能发展道路上的无限可能。

9.1 从推理角度详解MLA注意力模型与代码实现

MLA采用了一种极具巧思的压缩形式来对缓存进行优化处理。直白来讲，相较于直接存储键（K）和值（V），或是采用共享键值对的方式，MLA另辟蹊径，它先将键和值压缩成一个更为紧凑的潜在向量。在模型实际运行过程中，当需要使用这些键和值时，再精准地将其恢复出来。这种独特的处理方式，犹如为数据存储与调用打造了一个高效的“压缩包”，极大地提升了缓存的利用效率。

具体来看，MLA的核心创新之处主要体现在以下两个方面：

低秩键值联合压缩：MLA突破了传统方法的局限，不再对键和值进行单独处理，而是将键和值所蕴含的信息进行深度融合，共同压缩到一个更小的潜在空间之中。这一创新举措，使得原本分散、冗余的数据信息得以高度整合，在减少数据存储空间的同时，也保留了数据的关键特征，为后续的高效计算奠定了坚实基础。
矩阵乘法优化：在计算过程中，MLA通过巧妙地重排计算顺序，对矩阵乘法进行了深度优化。这一优化策略，如同为计算引擎注入了一剂高效催化剂，使得计算效率得到了显著提升，大大缩短了模型的推理时间。

在本节中，我们将从缓存优化的独特视角出发，深入剖析注意力模型在推理阶段的缓存占用情况，并通过严谨的量化计算，为读者呈现清晰的数据对比和分析。通过这一环节，读者将能够直观地了解到缓存优化对于模型性能提升的重要性。

之后，本节将详细演示MLA计算公式的推导过程，带领读者逐步揭开MLA背后的数学奥秘。同时，我们还将结合推导出的公式，实现对应的MLA注意力模型，让读者在实践中进一步加深对MLA技术的理解和掌握。

9.1.1 大模型的推理过程

我们在前面章节实现了经典生成模型，从推理输出上来看，在推理阶段，由于模型由多层Transformers堆叠而成，因此主要的计算负担落在了注意力模型内部，涉及MHA和前馈神经网络（或MoE）等核心操作。在MHA中，需要计算查询（Q）、键（K）和值（V）矩阵，以进行多头注意力的相关计算。

在大语言模型（LLM）的生成过程中，模型是基于前面的词序列来预测下一个词的。在这个过程中，每个词（token）仅与其前面的词进行交互以计算注意力，这种特定的注意力机制被称为因果注意力（Causal Attention）。在矩阵计算层面，通过一个下三角形状的因果注意力掩码（Causal Attention Mask）来确保每个词仅感知其前面的词序列。

KV Cache缓存的使用如图9-2所示。

图9-2 KV Cache缓存的使用

同样地，我们在介绍KV Cache缓存的章节中，已通过图示向读者清晰展示了输入的查询（Q）只与其对应位置前的键（K）和值（V）进行计算。因此，为了提升训练和推理的效率，在token-by-token的生成过程中，我们需要避免对前序token的重复计算。

为此，研究者们提出了一个有效的解决方案：将前序已经计算好的键（K）和值（V）缓存起来。这就是目前被广泛采用的KV Cache机制。KV Cache本质上是通过牺牲存储空间来换取时间效率的方法。然而，考虑到当前大语言模型（LLM）的规模都相当庞大，而GPU的显存空间又十分有限，使用显存来保存KV Cache势必会引发访问瓶颈。

换句话说，如果不使用KV Cache，模型将直接进行计算（包括重复计算前序token），这是一个计算密集型任务；而增加了KV Cache之后，键（K）和值（V）不再是通过实时计算获得，而是直接从存储介质中读取，这意味着GPT内核与存储介质之间需要频繁地进行数据读写，从而使任务转变为访存密集型。因此，虽然KV Cache机制解决了重复计算的问题，但存储访问的速度却直接影响了训练和推理的速度。

为了进一步优化这一流程，未来的研究可以探索更高效的存储解决方案，例如利用更快的存储介质或者优化数据访问模式，以减少访存延迟，从而进一步提高KV Cache机制的整体性能。此外，也可以考虑设计更加智能的缓存替换策略，以适应不同大小和特性的语言模型，确保在有限的显存资源下实现最优的性能提升。

9.1.2 通用大模型的显存占用量化计算

对于通用的大模型在推理阶段时，主要有三部分数据会放到显存中，如下所示：

模型参数：包括Transformers、Embedding等模型参数会存到显存里。模型大小固定后，这个存储空间是固定的。
KV Cache：如上一节所述，前序token序列计算的KV结果，会随着后面token推理过程逐步存到显存中。存储的量随着Batch和sequence_len长度而动态变化。
运行时中间数据：推理过程中产出的一些中间数据会临时存到显存中，即用即释放，一般占用空间比较小。

完整的模型推理时显存空间的占用情况如图9-3所示。

图9-3 完整的模型推理时显存空间的占用

可以看到推理阶段主要存储消耗是两部分：模型参数和KV Cache。那么模型参数占多少，KV Cache又占多少？首先我们先以一个token的计算过程为例，看下一个token计算要存储多少KV？计算公式如下：

存储的KV数目=2×注意力头数×层数

下面我们举一个简单的例子，当我们设置的Config如下：

class Config:

n_embd = hidden_size = 384

num_attention_heads = 6

vocab_size = 4000

vocab_size = 7000 #在进行SFT时的长度为7000

assert hidden_size % num_attention_heads == 0, 'hidden_size must be divisible by num_head'

intermediate_size = hidden_size * 4

dropout = 0.1

layer_norm_epsilon = layer_norm_eps = 1e-12

num_hidden_layers = n_layers = 6

根据上面的公式计算，我们的一个KV占用的 2×6×6＝72，这表示推理一个token就要缓存72个KV个数，而统计占据了多少缓存，计算公式如下：

每个token_占据空间=每个参数占据空间×参数个数×存储的KV数目

假设推理阶段是半精度（BF16，128位）参数，每个参数占2B。最终一个token的存储占用如下：

我们现在知道了一个Token计算后需要缓存的KV数量和存储量。那么对于一个实际的推理场景，还要考虑批量batch_size和序列长度seq_len两个维度，来确认整体KV Cache的存储消耗。下面是序列长度消耗公式：

单条缓存占用=单token存储×batch_size×seq_len

通过计算我们可以很轻易地得到模型在推断阶段在不同batch_size下进行推理时的显存占用。

9.1.3 手把手MLA注意力公式的总体推导

前面我们在讲解MHA多头注意力时，在讲解对于不同的输入查询（Q）、键（K）和值（V）的计算如下：

而对于MLA中的注意力维度参数矩阵的计算，我们则修改了这里可训练参数矩阵的大小，即将原有的统一维度大小的可训练参数矩阵，修改如下：

1. 查询（Q）

2. 键（K）与值（V）

可以看到，通过这种方式MLA实现了参数量的节省，并在进行注意力计算时完成了速度优化。

9.2 从缓存角度详解MLA注意力模型与代码实现

在深度学习模型日益庞大、复杂的当下，推理空间占用的优化显得尤为关键，而MLA凭借其独特的设计，在这方面表现得极为出色。其通过先进的压缩算法和巧妙的数据处理方式，将原本需要大量空间存储的键（K）和值（V）信息进行高效压缩，从而大幅降低了模型在推理过程中对内存和显存的占用。

节省推理空间占用只是MLA带来的基础益处，由此还引发了一系列积极的连锁反应。由于占用空间的减少，模型在推理时的数据读取和传输速度得到了显著提升。这就好比在一条原本拥堵的道路上，车辆（数据）能够更加顺畅地行驶，减少了等待和延误，从而加快了整个推理流程。推理速度的提升意味着模型能够更快地给出预测结果，这对于一些对实时性要求极高的应用场景，如自动驾驶、在线实时翻译等，具有至关重要的意义。

此外，MLA节省推理空间占用的特性还为模型的部署提供了更大的灵活性。在资源有限的设备上，如移动设备、嵌入式系统等，传统的注意力模型可能因空间占用过大而无法顺利运行。而MLA则打破了这一限制，使得这些设备也能够轻松承载和运行复杂的深度学习模型，极大地拓展了模型的应用范围。

在接下来的内容中，我们将深入探究MLA是如何实现节省推理空间占用的具体机制，包括其背后的算法原理、关键技术的运用等。同时，我们还将通过实际案例和实验数据，直观地展示MLA在节省推理空间占用以及提升推理效率方面的卓越表现，让读者对这一创新技术有更全面、更深入的认识。

9.2.1 优化的MLA模型实现1：压缩低秩空间

MLA的创新之一就是低秩空间压缩，我们通过将输入的向量压缩到低秩维度从而完成了维度变换。而在具体实现上，我们则可以通过定义和组成对应的可训练参数完成模型的结构。

1. 变换参数的定义

#首先我们需要定义MLA中的参数，在这里我们定义参数如下：

self.q_proj_up = d_model * 2 #先把Query的维度升高

self.qk_proj_down = d_model #再降低Query的维度进行计算

self.query_head_dim = self.qk_proj_down // n_heads

#这里是进行压缩kv的维度，可以设置为原有d_model大小的2/3为好

self.kv_lora_rank = int((d_model * 3)//2) #变为hidden压缩了维度

在上面代码中，我们分别设置了MLA中的参数，q_proj_up与qk_proj_down是对输入的查询（Q）进行计算，我们采用双参数期望能够在输入的参数较少（在推理时只有1个token）的情况下，也能获得一个较为合理的特征向量。

kv_lora_rank对整体的Key和Value整合计算维度，其作用是建立一个过度的低维处理向量，用于后期的整合。

2. Query的维度变换处理

下面我们首先看Query的维度变换，代码如下所示：

self.W_dq = torch.nn.Parameter(0.01 * torch.randn(d_model, self.q_proj_dim))

self.q_norm = torch.nn.LayerNorm(self.q_proj_dim)

self.W_Uq = torch.nn.Parameter(0.01 * torch.randn(self.q_proj_dim, self.qk_proj_dim))

在这里，我们首先依据公式定义了多个映射参数，旨在对输入的特征向量x进行变换。其中，qk_proj_dim与qk_head_num的设定是为了调整维度的尺寸，以便我们人为地修正输入和输出的维度。在进行qk计算时，我们仅需确保最后一个维度保持一致即可。

而Query的计算代码如下所示：

compressed_q = x @ self.W_dq

compressed_q = self.q_norm(compressed_q)

query = compressed_q @ self.W_Uq

q_nope = self._split_heads(query, self.num_heads, self.query_head_dim) # 将query分割为多头 shape = [-1,6,48,64]

3. Key与Value的维度变换处理

对于Key和Value的维度变换与处理，我们同样按推理公式中首先对其进行变换，之后将其变换为输入维度，代码如下所示：

q_absorb = torch.einsum('bqd,hdc->bhqc', compressed_kv, self.kv_b_proj[:,:,
:self.query_head_dim])

这里我们首先完成了Key与Value的总体维度变换，将其拆分后我们所需要的Key和Value维度矩阵，计算过程如下：

v_nope = torch.einsum('bqc,hcd->bqhd', compressed_kv, self.kv_b_proj[:,:,self.query_head_dim:])

attn_output = torch.einsum('bhql,blhd->bhqd', attn_weights,v_nope)

顺便说一下，MLA中的低秩压缩的核心思想是：我们可以用一个较小的矩阵来近似表示一个大矩阵。而这个小矩阵正常来说是通过矩阵分解得到的，但在MLA中是直接设计小矩阵的维度，让它们通过训练自动学习到一个好的低秩表示。

秩就是信息压缩的维度，我们通过一个具体的例子来理解这个过程：想象一个5120维的向量，它可能代表了一幅图像的所有像素值。当我们用一个秩为512的变换去处理它时，本质上我们是在说：这5120个数字中，实际上可以用512个独立的特征来表达主要信息。

而这种压缩之所以有效，是因为在实际应用中：

真实数据通常存在大量冗余。
并非所有维度都同等重要。

因此在实际中保留最重要的维度往往足以表达数据的主要特征。

9.2.2 优化的MLA模型实现2：核心注意力矩阵计算

注意力模型MLA核心优化之一就是对核心计算使用矩阵计算的方法，在这里我们按公式完成了新的注意力计算优化，代码如下所示：

def _attn(self, q_nope, compressed_kv, attention_mask=None):

q_cope = q_nope.clone()

# 为了与compressed_kv结合计算attention_score

if True:

q_absorb = torch.einsum('bqd,hdc->bhqc', compressed_kv, self.kv_b_proj[:,:,:self.query_head_dim])

attn_weights = torch.einsum('bhqc,bhlc->bhlq', q_absorb, q_nope)

else:

pass

# 缩放注意力权重

attn_weights = attn_weights / torch.full([], q_nope.size(-1) ** 0.5, dtype=attn_weights.dtype,device=attn_weights.device)

query_length, key_length = q_nope.size(-2), compressed_kv.size(-2)

causal_mask = self.bias[:, :, key_length - query_length: key_length, :key_length].to(q_nope.device)

attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), self.mask_value)

if attention_mask is not None:

# 如果有额外的注意力掩码，应用它

attn_weights = attn_weights + attention_mask

attn_weights += self.cope(q_cope, attn_weights)

attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) # 计算Softmax

attn_weights = attn_weights.type(compressed_kv.dtype)

attn_weights = self.attn_dropout(attn_weights) # 应用dropout

if True:

v_nope = torch.einsum('bqc,hcd->bqhd', compressed_kv, self.kv_b_proj[:,:,self.query_head_dim:])

attn_output = torch.einsum('bhql,blhd->bhqd', attn_weights,v_nope)

else:

pass

return attn_output, attn_weights

在上面代码中，我们分别对两个部分进行优化，首先是attn_weights的计算，而同样的attn_output部分我们也进行优化，将维度变换过程整合到一个完整过程中实现。另外，读者在具体操作时，可能注意到代码的最后部分使用了if…else条件语句进行判断。这是由于在这个位置进行维度变换时可以将if条件下的两个变换过程整合成一个。这样的计算过程虽然会节省时间，但是这种变换对于硬件资源消耗较大，部分GPU并不适合这种直接维度变换的操作，因此我们建议有兴趣的读者可以自行尝试。

9.2.3 优化的MLA模型实现3：对显存KV Cache部分的压缩

与传统的我们分别对生成的Key和Value值进行缓存不同，MLA中首先对Key和Value进行一个整体变换，将其压缩后通过维度计算重新获取对应的Key和Value值。在这个过程中我们可以通过缓存中间过程，也就是压缩的整体值从而完成对值的缓存。代码如下所示：

if layer_past is not None:

current_kv = x @ self.W_dkv

compressed_kv = torch.cat([layer_past, current_kv], dim=1)

else:

compressed_kv = x @ self.W_dkv

present = compressed_kv #在这里进行了键值对的压缩

compressed_kv = compressed_kv@self.W_duv

compressed_kv = self.kv_norm(compressed_kv)

在这里我们完成了对键值对的联合压缩过程，这也是MLA的创新，我们可以将其理解为：

x：是输入的原始向量。
W_dkv：是一个特殊的下投影矩阵，它同时服务于键和值的压缩。
compressed_kv：是压缩后的潜在向量，它包含了键和值共享的信息。

compressed_kv作为一个存储关键信息供模型在计算时快速访问它包含了键和值的压缩信息，这个向量的维度一般认为比原始的键值对小得多，因此存储这个压缩向量比存储完整的键和值更节省空间。

9.2.4 带有缓存的MLA注意力模型完整实现

我们已经对MLA模型的各个模块进行了详细的讲解，而MLA注意力机制的一个显著优点，就是其在推理阶段能够出色地完成缓存的优化与利用。这一特性不仅提升了模型的运行效率，还为实际应用带来了更多的便利性。

接下来，本小节将着手构建一个带有缓存功能的MLA完整模型。通过整合缓存机制，我们将进一步展示MLA模型在实际应用中的优势。以下是实现这一完整模型的代码示例：

class CoPEMLA(torch.nn.Module):

def __init__(self, config, layer_idx=None):

super().__init__()

self.config = config

self.max_position_embeddings = max_positions = config.max_position_embeddings

# 创建一个下三角矩阵，用于因果掩码（causal mask）

self.bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(1, 1, max_positions,max_positions)

self.mask_value = torch.tensor(-1E+9) # 用于掩码的极大负值

# 维度参数

d_model = config.hidden_size

self.d_model = torch.tensor(d_model)

self.num_heads = n_heads = config.num_attention_heads

# 投影维度

self.q_proj_up = d_model * 2 # 先把Query的维度升高

self.qk_proj_down = d_model # 再降低Query的维度进行计算

self.query_head_dim = self.qk_proj_down // n_heads

# 这里是进行压缩kv的维度，可以设置为原有d_model大小的2/3为好

self.kv_lora_rank = int((d_model * 3)//2) # 变为hidden压缩了维度

self.qk_nope_head_dim = self.query_head_dim * 2

self.kv_attn_dim = (self.query_head_dim + self.qk_nope_head_dim)

# Q投影

self.W_dq = torch.nn.Parameter(0.01 * torch.randn(d_model, self.q_proj_up))

self.q_norm = torch.nn.RMSNorm(self.q_proj_up)

self.W_Uq = torch.nn.Parameter(0.01 * torch.randn(self.q_proj_up, self.qk_proj_down))

# KV投影

self.W_dkv = torch.nn.Parameter(0.01 * torch.randn((d_model), (self.kv_lora_rank)))

self.W_duv = torch.nn.Parameter(0.01 * torch.randn((self.kv_lora_rank), (self.kv_lora_rank)))

self.kv_norm = torch.nn.RMSNorm((self.kv_lora_rank))

self.kv_b_proj = torch.nn.Parameter(0.01 * torch.randn(size=(self.num_heads, (self.kv_attn_dim), self.kv_lora_rank)))

# 输出

self.W_o = torch.nn.Parameter(0.01 * torch.randn(d_model, d_model))

self.attn_dropout = torch.nn.Dropout(config.dropout)

#kq_weight要参与多头后的query与key计算

self.uk_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

#v_proj_weight要参与value计算，将压缩的内容重新映射回value

self.uv_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

self.cope = updat_moudle.CoPE(config.max_position_embeddings, self.query_head_dim)

def forward(self, x, layer_past=None, attention_mask=None):

# Q投影

compressed_q = x @ self.W_dq

compressed_q = self.q_norm(compressed_q)

query = compressed_q @ self.W_Uq

q_nope = self._split_heads(query, self.num_heads, self.query_head_dim)

# KV投影

if layer_past is not None:

current_kv = x @ self.W_dkv

compressed_kv = torch.cat([layer_past, current_kv], dim=1)

else:

compressed_kv = x @ self.W_dkv

present = compressed_kv #在这里进行了键值对的压缩

compressed_kv = compressed_kv@self.W_duv

compressed_kv = self.kv_norm(compressed_kv)

# 计算注意力输出和注意力权重

attn_output, attn_weights = self._attn(q_nope, compressed_kv,attention_mask)

attn_output = self._merge_heads(attn_output, self.num_heads, self.query_head_dim) # 合并多头

attn_output = attn_output @ self.W_o

attn_output = self.attn_dropout(attn_output) # 应用dropout

outputs = (attn_output, present) # 返回注意力输出和当前的key、value

return outputs

def _split_heads(self, tensor, num_heads, attn_head_size):

"""

将隐藏层维度分割为多头注意力的头和头的大小。

"""

new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)

tensor = tensor.view(new_shape)

return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)

def _attn(self, q_nope, compressed_kv, attention_mask=None):

q_cope = q_nope.clone()

#为了与compressed_kv结合计算attention_score

if True:

q_absorb = torch.einsum('bqd,hdc->bhqc', compressed_kv, self.uk_weight)

attn_weights = torch.einsum('bhqc,bhlc->bhlq', q_absorb, q_nope)

else:

pass

# 缩放注意力权重

attn_weights = attn_weights / torch.full([], q_nope.size(-1) ** 0.5, dtype=attn_weights.dtype,device=attn_weights.device)

query_length, key_length = q_nope.size(-2), compressed_kv.size(-2)

causal_mask = self.bias[:, :, key_length - query_length: key_length, :key_length].to(q_nope.device)

attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), self.mask_value)

if attention_mask is not None:

# 如果有额外的注意力掩码，应用它

attn_weights = attn_weights + attention_mask

attn_weights += self.cope(q_cope, attn_weights)

attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) # 计算softmax

attn_weights = attn_weights.type(compressed_kv.dtype)

attn_weights = self.attn_dropout(attn_weights) # 应用dropout

if True:

v_nope = torch.einsum('bqc,hcd->bqhd', compressed_kv, self.uv_weight)

attn_output = torch.einsum('bhql,blhd->bhqd', attn_weights,v_nope)

else:

#下面这个压缩代码在训练时会报错

pass

return attn_output, attn_weights

def _merge_heads(self, tensor, num_heads, attn_head_size):

"""

将多头注意力的头和头的大小合并回隐藏层维度。

"""

tensor = tensor.permute(0, 2, 1, 3).contiguous()

new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)

return tensor.view(new_shape)

上面代码完整实现了标准的MLA注意力模型。其中加粗的部分代码为注意力计算模块，在具体使用上读者可以将MLA替代原有我们的上一章对比计算中的经典多头注意力模型，读者可以自行尝试。

9.3 MLA注意力模型的完整补充讲解

通过上面两节内容的讲解，我们了解了MLA注意力模型并完整实现了MLA注意力。在实际的MLA注意力计算中，对于多头注意力数值的计算，我们可以使用matmul的方式一次性完整地进行。然而，有时候为了更好地利用硬件计算性能，例如分布式多GPU计算单元的情况，我们则可以把整体的多头计算拆分成分布式多头计算，而在具体处理上我们可以参考如下的形式进行：

#伪代码的形式

import torch

all_querys = []

all_keys = []

all_values = []

for head in n_heads:

query = torch.matmul(per_query,q_weight)

key = torch.matmul(per_key,k_weight)

value = torch.matmul(per_value,v_weight)

all_querys.append(query)

all_keys.append(key)

all_values.append(value)

...

all_attention_outputs = []

for q,k,v in zip(all_querys,all_keys,all_values)

attention_score = torch.matmul(q,k.transpose())

attention_weight = torch.nn.functional.softmax(attention_score,dim=-1)

attention_output = torch.matmul(attention_weight,v)

all_attention_outputs.append(attention_output)

上面代码采用了伪代码的形式，清晰展示了分布式多头计算的方法。针对不同的计算硬件资源，我们可以灵活地手动调整，将单一计算任务分配到不同的计算单元上并行处理，最后将各单元的计算结果汇总整合。

除此之外，在相关公式中，我们还特别提及了为简化计算过程，MLA（假设为某算法缩写）运用了一个简洁的乘法变换法则。这一法则旨在有效减轻矩阵计算的复杂度，提升计算效率。

鉴于这两部分内容均涉及分布式硬件计算的深入应用，本书在此不做进一步的详细讨论和实现展示，仅提供伪代码及基本实现算法作为参考。对此感兴趣的读者，请自行深入探索和实践。

9.3.1 调参、记忆力以及矩阵计算优化

上面我们讲解并完成了MLA模型的基本结构，可以看到相对于传统的注意力模型，MLA在设计上采用了较多巧妙的设计，例如参数调整、矩阵计算优化以及加载了记忆力模块。下面我们分别对这些设计进行讲解。

1. 可调参的MLA注意力模型

MLA模型在设计之初，通过可变维度的参数对注意力模型的计算的大小进行修正，相对于经典的注意力模型，MLA可以根据不同的任务强度对注意力计算参数进行配置。代码如下所示：

# 投影维度

self.q_proj_up = int((d_model * 3)//2) #先把Query的维度升高

self.qk_proj_down = d_model #再降低Query的维度进行计算

self.query_head_dim = self.qk_proj_down // n_heads

self.kv_lora_rank = int((d_model * 3)//2) #变为hidden压缩了维度

self.qk_nope_head_dim = self.query_head_dim * 2

而为了验证我们的投影维度调整后对模型大小的影响，我们可以采用如下的参数对注意力参数进行统计，代码如下所示：

def get_parameter_number(net):

total_num = sum(p.numel() for p in net.parameters())

trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)

print({'Total': total_num, 'Trainable': trainable_num})

读者可以使用参数打印函数打印在不同条件下的MLA注意力参数。

2. MLA注意力中的记忆力模块

相对于经典的注意力机制，MLA在计算Query和Value之间添加了一个记忆力模块，self.kv_b_proj就起到这个作用，在局部缓存并计算Query、Key和Value的向量。记忆模块如下所示：

#kq_weight要参与多头后的query与key计算

#self.uk_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

#v_proj_weight要参与value计算，将压缩的内容重新映射回value

#self.uv_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

#在这里我们把上面的拟合在一起计算

self.ukv_weight = torch.torch.nn.Parameter(0.01 * torch.randn(2,self.num_heads,self.kv_lora_rank, self.query_head_dim))

在具体处理上：

q_absorb = self.kv_b_proj[0]

...

v_absorb = self.kv_b_proj[1]

q_absorb与v_absorb分别作为中间向量参与了注意力计算。

3. MLA注意力中的矩阵计算优化

我们在计算Query向量以及最终对attn_output输出计算时，分别采用了分阶段的计算方法，然而实际上我们在计算时可以根据要求对这部分的矩阵计算进行进一步的优化，代码如下所示：

q_absorb = torch.einsum('bqd,hdc->bhqc', compressed_kv, self.ukv_weight[0])

attn_weights = torch.einsum('bhqc,bhlc->bhlq', q_absorb, q_nope)

这里我们修正了矩阵计算过程，将原有的分段计算整合成一个完整的计算过程，这样带来的好处是可以加速计算过程，但是由于计算的空间需求较大，在使用上还要读者自行斟酌。

注意力矩阵计算的代码如下所示：

if True:

v_nope = torch.einsum('bqc,hcd->bqhd', compressed_kv, self.ukv_weight[1])

attn_output = torch.einsum('bhql,blhd->bhqd', attn_weights,v_nope)

else:

#下面这个压缩代码在训练时会报错

pass

读者可以自行验证比较。

4. 优化后的MLA注意力模型完整实现

对于优化后的MLA模型的完整实现，我们直接给代码，如下所示：

class CoPEMLA(torch.nn.Module):

def __init__(self, config, layer_idx=None):

super().__init__()

self.config = config

self.max_position_embeddings = max_positions = config.max_position_embeddings

# 创建一个下三角矩阵，用于因果掩码（causal mask）

self.bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(1, 1, max_positions,max_positions)

self.mask_value = torch.tensor(-1E+9) # 用于掩码的极大负值

# 维度参数

d_model = config.hidden_size

self.d_model = torch.tensor(d_model)

self.num_heads = n_heads = config.num_attention_heads

# 投影维度

self.q_proj_up = int((d_model * 3)//2) #先把Query的维度升高

self.qk_proj_down = d_model #再降低Query的维度进行计算

self.query_head_dim = self.qk_proj_down // n_heads

self.kv_lora_rank = int((d_model * 3)//2) #变为hidden压缩了维度

self.qk_nope_head_dim = self.query_head_dim * 2

#这里是用作生成一组memory参数，

self.kv_attn_dim = (self.query_head_dim + self.qk_nope_head_dim)

# Q投影

self.W_dq = torch.nn.Parameter(0.01 * torch.randn(d_model, self.q_proj_up))

self.q_norm = torch.nn.RMSNorm(self.q_proj_up)

self.W_Uq = torch.nn.Parameter(0.01 * torch.randn(self.q_proj_up, self.qk_proj_down))

# KV投影

self.W_dkv = torch.nn.Parameter(0.01 * torch.randn((d_model), (self.kv_lora_rank)))

self.W_duv = torch.nn.Parameter(0.01 * torch.randn((self.kv_lora_rank), (self.kv_lora_rank)))

self.kv_norm = torch.nn.RMSNorm((self.kv_lora_rank))

# 输出

self.W_o = torch.nn.Parameter(0.01 * torch.randn(d_model, d_model))

self.attn_dropout = torch.nn.Dropout(config.dropout)

#kq_weight要参与多头后的query与key计算

#self.uk_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

#v_proj_weight要参与value计算，将压缩的内容重新映射回value

#self.uv_weight = torch.torch.nn.Parameter(0.01 * torch.randn(self.num_heads,self.kv_lora_rank, self.query_head_dim))

#在这里把上面的拟合在一起进行优化

self.ukv_weight = torch.torch.nn.Parameter(0.01 * torch.randn(2,self.num_heads,self.kv_lora_rank, self.query_head_dim))

self.cope = updat_moudle.CoPE(config.max_position_embeddings, self.query_head_dim)

def forward(self, x, layer_past=None, attention_mask=None):

# Q投影

compressed_q = x @ self.W_dq

compressed_q = self.q_norm(compressed_q)

query = compressed_q @ self.W_Uq

q_nope = self._split_heads(query, self.num_heads, self.query_head_dim) # 将query分割为多头 shape = [-1,6,48,64]

# KV投影

if layer_past is not None:

current_kv = x @ self.W_dkv

compressed_kv = torch.cat([layer_past, current_kv], dim=1)

else:

compressed_kv = x @ self.W_dkv

present = compressed_kv #在这里进行了键值对的压缩

compressed_kv = compressed_kv@self.W_duv

compressed_kv = self.kv_norm(compressed_kv)

# 计算注意力输出和注意力权重

attn_output, attn_weights = self._attn(q_nope, compressed_kv,attention_mask)

attn_output = self._merge_heads(attn_output, self.num_heads, self.query_head_dim) # 合并多头

attn_output = attn_output @ self.W_o

attn_output = self.attn_dropout(attn_output) # 应用dropout

outputs = (attn_output, present) # 返回注意力输出和当前的key、value

return outputs

def _split_heads(self, tensor, num_heads, attn_head_size):

"""

将隐藏层维度分割为多头注意力的头和头的大小。

"""

new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)

tensor = tensor.view(new_shape)

return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)

def _attn(self, q_nope, compressed_kv, attention_mask=None):

q_cope = q_nope.clone()

#为了与compressed_kv结合计算attention_score

if True:

q_absorb = torch.einsum('bqd,hdc->bhqc', compressed_kv, self.ukv_weight[0])

attn_weights = torch.einsum('bhqc,bhlc->bhlq', q_absorb, q_nope)

else:

pass

# 缩放注意力权重

attn_weights = attn_weights / torch.full([], q_nope.size(-1) ** 0.5, dtype=attn_weights.dtype,device=attn_weights.device)

query_length, key_length = q_nope.size(-2), compressed_kv.size(-2)

causal_mask = self.bias[:, :, key_length - query_length: key_length, :key_length].to(q_nope.device)

attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), self.mask_value)

if attention_mask is not None:

# 如果有额外的注意力掩码，应用它

attn_weights = attn_weights + attention_mask

attn_weights += self.cope(q_cope, attn_weights)

attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) # 计算softmax

attn_weights = attn_weights.type(compressed_kv.dtype)

attn_weights = self.attn_dropout(attn_weights) # 应用dropout

if True:

v_nope = torch.einsum('bqc,hcd->bqhd', compressed_kv, self.ukv_weight[1])

attn_output = torch.einsum('bhql,blhd->bhqd', attn_weights,v_nope)

else:

#下面这个压缩代码在训练时会报错

pass

return attn_output, attn_weights

def _merge_heads(self, tensor, num_heads, attn_head_size):

"""

将多头注意力的头和头的大小合并回隐藏层维度。

"""

tensor = tensor.permute(0, 2, 1, 3).contiguous()

new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)

return tensor.view(new_shape)

在具体使用上，读者可以替代原有的注意力模型重新完成模型的训练。

9.3.2 MLA、GQA以及MQA差异详解

MLA（多头潜在注意力）、GQA（分组查询注意力）和MQA（多查询注意力）是3种不同类型的注意力机制，它们在信息处理上的差异，我们可以将其类比为一个图书馆的信息检索系统，以更直观地理解它们各自的特点。

1. MLA

在MLA中，每个图书管理员（注意力头）都拥有自己的检索系统，但这些系统通过智能压缩技术进行了优化。具体来说：

特征提取：每个图书管理员首先会提取图书的关键特征，如科技类图书的关键词、技术领域和出版年份，或文学类图书的作者风格、主题和写作手法。
智能压缩：这些关键特征随后被压缩成一个精简的索引（例如，256维）。这种压缩不是简单地删减信息，而是找到最能代表原始信息的特征组合。
信息重建：在需要时，可以从这些关键特征中重建出详细的检索信息。虽然可能会损失一些细节，但保留了最重要的区分性特征。

因此，MLA在信息损失方面相对较小，因为它通过智能压缩保留了最重要的特征，同时每个头仍然保持了一定的独特性。这使得MLA能够在节省存储空间的同时，仍然保持较好的性能。

2. GQA

GQA则采用了一种分组共享编目系统的方式，具体来说：

分组共享：图书管理员被分成不同的组，每组共享一个专业的编目系统。例如，科技组的管理员共享一个科技文献编目系统，文学组的管理员共享一个文学作品编目系统。
专业化能力：这种方式比MQA好，因为它保留了一定的专业化能力，不同领域可以有各自的检索方式。
组内混合：然而，GQA仍然存在组内信息混合的问题，因为同一组内的管理员仍然需要使用相同的编目系统进行检索。

3. MQA

在MQA中，所有图书管理员共享同一个编目系统，具体来说：

共享系统：无论管理员专注于哪个领域，他们都必须使用同样的检索方式。
信息损失：这会导致信息损失最大，因为所有专业化的检索能力都被限制在一个统一的系统中。

例如，科技类图书可能无法按照技术领域精确分类，文学作品可能无法按照文学流派详细归类。

MLA、GQA和MQA在信息处理上的差异主要体现在它们如何处理信息损失和保留专业化能力上。

MQA：信息损失最大，因为完全共享同一套系统，无法保留不同领域的专业化检索能力。
GQA：信息损失较大，但通过分组共享编目系统保留了一定的专业化能力，仍然存在组内信息混合的问题。
MLA：信息损失相对最小，因为通过智能压缩技术保留了最重要的区分性特征，同时每个头仍然保持了一定的独特性。这使得MLA能够在节省存储空间的同时，仍然保持较好的性能。

因此，MLA通过智能压缩技术实现了在节省存储空间的同时保持较好性能的目标，而GQA和MQA则在不同程度上牺牲了专业化能力以换取存储空间的节省。