训练营简介 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) 的典范。

1. 算法层：通过低秩分解减少存储。

2. 算子层：通过矩阵吸收减少计算，通过双路 Attention 保持精度。

3. Ascend C：利用高带宽的 UB 复用 Latent Vector，完美契合 MLA 的“小数据、高计算”特性。

掌握了 MLA 算子，你不仅能看懂 DeepSeek 的论文，更能亲手部署这个当前最强的开源模型。