当前位置：首页 > news >正文

【昇腾CANN训练营·前沿篇】解构DeepSeek：基于Ascend C实现MLA (Multi-Head Latent Attention) 算子

news 2026/8/30 9:27:22

训练营简介 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）节省了一半的带宽。