"""MLA：Multi-head Latent Attention（多头潜在注意力）。 这是 DeepSeek-V2 起最核心的省显存设计。 标准 MHA 的 KV cache 每 token 要存 2 * n_heads * head_dim 个数； MLA 先把每个 token 压成一个低维潜在向量 c_kv（kv_lora_rank 维）， 用的时候再用 W_UK / W_UV 升维还原出 K 和 V。于是 cache 每 token 只需要存 kv_lora_rank + qk_rope_head_dim 个数，nano 档就是 64+16=80，而同规模 MHA 需要 2*8*32=512，省了 6.4 倍。 一个麻烦：RoPE 是位置相关的旋转，如果 K 是从 c_kv 现算出来的， 旋转矩阵没法和 W_UK 交换位置，"矩阵吸收"技巧就失效了。 DeepSeek 的解法是「解耦 RoPE」：额外切出 qk_rope_head_dim 维专门承载位置信息， 这部分所有 head 共享、直接缓存；剩下的 qk_nope 部分完全不带位置信息。 两条实现路径： naive —— 把 K/V 显式还原出来，走 F.scaled_dot_product_attention（能吃到 flash 内核）， 训练时用它最快。 absorb —— 把 W_UK 吸收进 Q、W_UV 吸收进输出投影，全程在潜在空间里算注意力， 长上下文解码时中间激活小得多。 两条路径数学上完全等价，tests.py 里有对拍。 """ import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from .layers import RMSNorm, apply_rope class MLA(nn.Module): def __init__(self, cfg): super().__init__() self.dim = cfg.dim self.n_heads = cfg.n_heads self.q_lora_rank = cfg.q_lora_rank self.kv_lora_rank = cfg.kv_lora_rank self.qk_nope_head_dim = cfg.qk_nope_head_dim self.qk_rope_head_dim = cfg.qk_rope_head_dim self.qk_head_dim = cfg.qk_head_dim self.v_head_dim = cfg.v_head_dim self.attn_impl = cfg.attn_impl self.softmax_scale = self.qk_head_dim ** -0.5 # ---- Query 投影：可选低秩分解（大模型上能省不少参数）---- if self.q_lora_rank == 0: self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False) else: self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False) self.q_norm = RMSNorm(self.q_lora_rank, cfg.norm_eps) self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False) # ---- KV 联合压缩：一次投影同时产出潜在向量 c_kv 和共享的 k_pe ---- self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False) self.kv_norm = RMSNorm(self.kv_lora_rank, cfg.norm_eps) self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False) self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False) self.dropout_p = cfg.dropout # 推理缓存（训练时为 None） self.kv_cache: Optional[torch.Tensor] = None self.pe_cache: Optional[torch.Tensor] = None # ------------------------------------------------------------------ def setup_cache(self, max_batch: int, max_seq_len: int, device, dtype): self.kv_cache = torch.zeros(max_batch, max_seq_len, self.kv_lora_rank, device=device, dtype=dtype) self.pe_cache = torch.zeros(max_batch, max_seq_len, self.qk_rope_head_dim, device=device, dtype=dtype) def clear_cache(self): self.kv_cache = None self.pe_cache = None def cache_bytes(self, seq_len: int) -> int: if self.kv_cache is None: return 0 per_token = self.kv_lora_rank + self.qk_rope_head_dim return per_token * seq_len * self.kv_cache.element_size() # ------------------------------------------------------------------ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, start_pos: int = 0, mask: Optional[torch.Tensor] = None) -> torch.Tensor: B, T, _ = x.shape end = start_pos + T # ---------------- Query ---------------- if self.q_lora_rank == 0: q = self.wq(x) else: q = self.wq_b(self.q_norm(self.wq_a(x))) q = q.view(B, T, self.n_heads, self.qk_head_dim).transpose(1, 2) # (B,H,T,qk) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_pe = apply_rope(q_pe, cos, sin) # ---------------- KV 压缩 ---------------- kv = self.wkv_a(x) # (B,T,c+rope) c_kv, k_pe = kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) c_kv = self.kv_norm(c_kv) # 缓存的就是它 k_pe = apply_rope(k_pe.unsqueeze(1), cos, sin).squeeze(1) # (B,T,rope)，所有 head 共享 if self.kv_cache is not None: self.kv_cache[:B, start_pos:end] = c_kv.to(self.kv_cache.dtype) self.pe_cache[:B, start_pos:end] = k_pe.to(self.pe_cache.dtype) c_kv_all = self.kv_cache[:B, :end].to(x.dtype) k_pe_all = self.pe_cache[:B, :end].to(x.dtype) else: c_kv_all, k_pe_all = c_kv, k_pe if self.attn_impl == "naive": out = self._attn_naive(q_nope, q_pe, c_kv_all, k_pe_all, mask) else: out = self._attn_absorb(q_nope, q_pe, c_kv_all, k_pe_all, mask) out = out.transpose(1, 2).reshape(B, T, self.n_heads * self.v_head_dim) return self.wo(out) # ------------------------------------------------------------------ def _attn_naive(self, q_nope, q_pe, c_kv_all, k_pe_all, mask): """显式还原 K/V，交给 SDPA（可命中 flash-attention 内核）。""" B, H, T, _ = q_nope.shape S = c_kv_all.shape[1] kv = self.wkv_b(c_kv_all).view(B, S, H, self.qk_nope_head_dim + self.v_head_dim) kv = kv.transpose(1, 2) # (B,H,S,nope+v) k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) k = torch.cat([k_nope, k_pe_all.unsqueeze(1).expand(B, H, S, -1)], dim=-1) q = torch.cat([q_nope, q_pe], dim=-1) return F.scaled_dot_product_attention( q, k, v, attn_mask=mask, is_causal=(mask is None and T > 1), dropout_p=self.dropout_p if self.training else 0.0, scale=self.softmax_scale, ) def _attn_absorb(self, q_nope, q_pe, c_kv_all, k_pe_all, mask): """矩阵吸收：全程在 kv_lora_rank 维的潜在空间里做注意力。""" B, H, T, _ = q_nope.shape W = self.wkv_b.weight.view(H, self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank) W_uk = W[:, : self.qk_nope_head_dim] # (H, nope, c) W_uv = W[:, self.qk_nope_head_dim:] # (H, v, c) # q_nope @ W_UK —— 把升维矩阵吸收进 Q，K 就不用还原了 q_absorb = torch.einsum("bhtd,hdc->bhtc", q_nope, W_uk.to(q_nope.dtype)) scores = (torch.einsum("bhtc,bsc->bhts", q_absorb, c_kv_all) + torch.einsum("bhtd,bsd->bhts", q_pe, k_pe_all)) * self.softmax_scale S = c_kv_all.shape[1] if mask is not None: scores = scores + _to_additive(mask, scores.dtype) elif T > 1: causal = torch.ones(T, S, dtype=torch.bool, device=scores.device).tril(S - T) scores = scores.masked_fill(~causal, torch.finfo(scores.dtype).min) attn = scores.softmax(dim=-1, dtype=torch.float32).type_as(scores) if self.training and self.dropout_p > 0: attn = F.dropout(attn, self.dropout_p) x_lat = torch.einsum("bhts,bsc->bhtc", attn, c_kv_all) # 仍在潜在空间 return torch.einsum("bhtc,hdc->bhtd", x_lat, W_uv.to(x_lat.dtype)) # 输出时才用 W_UV 还原 def _to_additive(mask: torch.Tensor, dtype) -> torch.Tensor: if mask.dtype == torch.bool: return torch.zeros_like(mask, dtype=dtype).masked_fill(~mask, torch.finfo(dtype).min) return mask.to(dtype)