 算子)
训练营简介 2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。
报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro
前言
在 Transformer 的军备竞赛中，KV Cache一直是显存杀手。
MHA (Multi-Head Attention)：显存占用巨大，性能好。
GQA (Grouped-Query Attention)：LLaMA 采用的方案，显存减半，但仍有瓶颈。
MLA (Multi-Head Latent Attention)：DeepSeek 提出的方案，通过低秩投影将 KV 压缩为一个极小的Latent Vector。
MLA 的神奇之处在于：推理时，它看起来像是 MHA（算力强），但显存占用比 GQA 还低。这是通过数学上的“矩阵吸收”技巧实现的。但在算子层面，这意味着我们不能简单地套用 FlashAttention。我们需要处理两路 Query（一路用于内容，一路用于 RoPE）和压缩的 KV。
本期文章，我们将深入 DeepSeek 的心脏，用 Ascend C 复现这一精妙设计。
一、 核心图解：把大象装进火柴盒
MLA 的核心思想是：不要直接存储巨大的 $K$ 和 $V$ 矩阵，而是存储它们“压缩”后的形态 $c_{KV}$。
二、 算法原理：解耦 RoPE 与 矩阵吸收
2.1 压缩 (Compression)
在 MLA 中，Key 和 Value 共享一个压缩的隐向量 $c_{KV}$。
$$c_{KV} = X \cdot W_{DKV}$$
这是我们在显存中实际存储的东西（KV Cache）。
2.2 解耦 RoPE (Decoupled RoPE)
由于 RoPE 对位置敏感，不能直接压缩。DeepSeek 将 Query 和 Key 拆分为两部分：
Content Part (内容部分)：携带语义信息，参与压缩。
RoPE Part (位置部分)：携带位置信息，不参与压缩，单独计算。
2.3 矩阵吸收 (The Magic)
在推理阶段，我们需要计算 $Q^T K$。 原始公式：$q = [q_{content}, q_{rope}], k = [UP(c_{KV}), k_{rope}]$ 其中 $UP$ 是升维矩阵 $W_{UK}$。
如果不优化，我们需要先把 $c_{KV}$ 升维回 $k_{content}$，这会浪费算力。 MLA 的技巧是：将升维矩阵$W_{UK}$吸收到 Query 的投影矩阵中。
$$Score = (q_{content} W_{UQ}) \cdot c_{KV}^T + (q_{rope} \cdot k_{rope}^T)$$
结论：在算子层面，我们需要同时进行两个矩阵乘法（一个针对 Latent，一个针对 RoPE），然后相加。
三、 实战：Ascend C 实现 Fused MLA Kernel
我们需要实现一个融合算子，输入是 Query 的两个部分和压缩后的 KV Cache。
3.1 Kernel 类定义
输入：
q_content_absorb: 吸收了 $W_{UK}$ 的 Query，Shape[B, 1, H, LatentDim]。q_rope: 原始的 RoPE Query，Shape[B, 1, H, RopeDim]。kv_latent: 压缩的 KV Cache，Shape[B, SeqLen, LatentDim]。k_rope: 缓存的 RoPE Key，Shape[B, SeqLen, RopeDim]。
class KernelMLA { public: __aicore__ inline void Init(...) { // Init... // Tiling 策略： // 由于 LatentDim 通常较小 (e.g. 512)，而 SeqLen 很长 // 我们依然采用 FlashDecoding 的 Split-K 策略 } __aicore__ inline void Process() { // 并行处理 SeqLen 分块 } };3.2 Compute 核心逻辑：双路 Attention
这是 MLA 与标准 FlashAttention 最大的不同：Score 是两部分之和。
__aicore__ inline void Compute(int32_t i) { // 1. Load Data // Latent Stream DataCopy(qContentLoc, qContentGm, ...); DataCopy(kvLatentLoc, kvLatentGm[offset], blockSize * latentDim); // RoPE Stream DataCopy(qRopeLoc, qRopeGm, ...); DataCopy(kRopeLoc, kRopeGm[offset], blockSize * ropeDim); // 2. Compute Score Part 1: Content (Latent) // S_content = Q_absorbed * C_kv^T // 这是一个 [1, Latent] * [Block, Latent]^T 的 GEMV Matmul(sContent, qContentLoc, kvLatentLoc); // 3. Compute Score Part 2: RoPE (Position) // S_rope = Q_rope * K_rope^T // 这是一个 [1, Rope] * [Block, Rope]^T 的 GEMV // 注意：kRopeLoc 需要在 Host 侧预先做过 RoPE 旋转，或者在这里做 // DeepSeek 通常缓存的是旋转后的 K_rope Matmul(sRope, qRopeLoc, kRopeLoc); // 4. Fuse Scores // S = S_content + S_rope // Ascend C 向量加法 Add(scoresLoc, sContent, sRope, blockSize); // 5. Softmax & Update // 后续逻辑与标准 FlashDecoding 一致 (Online Softmax) // ... Softmax ... // 6. Compute Output // O = P * V // 注意：这里的 V 也是压缩的 Latent Vector (c_KV)！ // 也就是说，我们不需要读两遍内存，c_KV 既充当 K 也充当 V (部分共享) // 或者 DeepSeek 可能有独立的 c_V，视具体配置而定 // 假设 V = c_KV (KV 解耦不完全时) 或者 V = c_V Matmul(outputLoc, probsLoc, kvLatentLoc); // [1, Block] * [Block, Latent] // 7. Write Back // ... }四、 性能优化的“胜负手”
MLA 算子的性能瓶颈在于Vector (Add) 与 Cube (Matmul) 的频繁切换。
4.1 流水线掩盖
我们有两路 Matmul（Content 和 RoPE）。优化策略：
启动
Matmul(Content)。在等待 Content 结果时，启动
DataCopy(RoPE)。启动
Matmul(RoPE)。 利用多级流水线，掩盖小矩阵计算的 Latency。
4.2 显存复用 (Cache Locality)
MLA 的精髓在于 $c_{KV}$ 非常小。 在计算 $QK^T$ 和 $PV$ 时，如果 $V$ 也是基于 $c_{KV}$ 投影的（或者直接复用），那么 $c_{KV}$ 只需要加载一次到 L1，就可以被两个 Matmul 阶段复用！Ascend C 实现：确保kvLatentLoc在 UB/L1 中常驻，直到 $QK^T$ 和 $PV$ 都算完再释放。这比标准 FlashAttention（读 K 再读 V）节省了一半的带宽。
4.3 吸收矩阵的预计算
虽然这不属于 Kernel 内部，但作为算子开发者，必须告诉算法同事： $W_{UQ}$和$W_{UK}$的合并必须在 Host 侧或模型初始化时完成。如果在 Kernel 里现场做 $Q \cdot W_{UK}$，MLA 的性能优势将荡然无存。
五、 总结
DeepSeek 的 MLA 架构是算法与算子协同设计 (Co-design)的典范。
算法层：通过低秩分解减少存储。
算子层：通过矩阵吸收减少计算，通过双路 Attention 保持精度。
Ascend C：利用高带宽的 UB 复用 Latent Vector，完美契合 MLA 的“小数据、高计算”特性。
掌握了 MLA 算子，你不仅能看懂 DeepSeek 的论文，更能亲手部署这个当前最强的开源模型。
