「随笔小记」大模型随笔

MOE基本类型

MOE全称Mixtrue of Experts (混合专家)

简单来说，就是使用MOE替换Transformer之前的FFN（前馈网络）结构，从而获取更多的信息。

Transformer直接的堆叠，在层数深了之后会导致提取特征的低秩现象。添加残差和前馈神经网络（FFN）可以避免这种现象。

MOE的基本结构如下图：

这里的 $FFN1,...,FFN4$ 就是不同的4个专家

对于输入 $X_1,X_2$ ，首先会通过一个门控网络决定走哪一个FFN，然后将计算结果与对应的门控权重 $p$ 进行乘积，将结果与输入进行残差连接和归一化，最后输出结果。

稀疏MOE与稠密MOE

门控函数（路由函数）：协调专家与其各自输出组合
门控函数的分类
1. 稀疏门控：激活部分专家
2. 稠密门控：激活所有专家
3. 软门控：输入Token合并和专家合并

DeepSeek MOE

Expert共享机制：部分Expert在不同Token或层间共享参数，减少模型冗余。

Expert共享多了一个Expert作为Shared Expert，每个Expert的计算都使用它

内存优化：MLA+KV Cache优化，减少生成任务中的浮点运算量

DeepSeek V2

MLA

multi-head latent attention 利用了一套低秩联合压缩的架构节省了很多KVcache的显存

对比算法MQA（multi-query attention）,GQA（Grouped-Query Attention），MHA（multi-head Attention）

MQA：所有的Q共用一个KV

GQA：对Q进行分组，相同组内共用KV

MHA：对每个Q都使用不同的KV

基本思路

对于输入 $h_t$ ，按照self-attention的思路，他会将其映射成Q,K,V三个embedding，这里在映射之前，提前利用低秩联合压缩将 $h_t$ 压缩成 $c_t^{kv}$ 和 $c_t^Q$ ，首先看Q的生成，使用 $c_t^Q$ 乘一个Q权重，然后生成 $q_{t,i}^C$ ，另一个是在这基础之上应用RoPE位置编码，最后将两个部分拼接起来合成最终的Q。接着看KV的部分， $c_t^{KV}$ 乘上对应的K权重和V权重，生成 $K_{t,i}^C$ 和 $V_{t,i}^C$ 两个向量，其中 $V_{t,i}^c$ 就是最终的V向量，而最终的K向量还需要使用 $h_t$ 计算RoPE得到 $K_t^R$ 最后将 $K_{t,i}^C$ 和 $K_t^R$ 拼接得到最终的K向量。最后把Q,K,V输入到MHA中，完成MLA的计算。

图中有颜色的圈表示在实际推理过程中，会进行缓存的部分，对比于传统的MHA，大大降低了KV Cache

用公式表示如下：（含有D的权重表示低秩压缩（大小降低），含有U的权重表示还原（大小变大）），相比于MHA，就多了 $c_t^{KV}$ 和 $k_t^R$

$\begin{aligned} c_t^Q &= W^{DQ}h_t\\ q_t^C &=W^{UQ}c_t^Q\\ [q_{t,1}^R;q_{t,2}^R;...;q_{t,n_h}^R] &= RoPE(W^{QR}c_t^Q) \\ q_{t,i} &= [q_{t,i}^C;q_{t,i}^R] \\ c_t^{KV}&=W^{DKV}h_t\\ k_t^{C}&=W^{UK}c_t^{KV}\\ k_t^R&=RoPE(W^{KR}h_t)\\ k_{t,i} &= [k^C_{t,i};k_t^R]\\ v_t^{C}&=W^{UV}c_t^{KV}\\ o_{t,i} &= \sum_{j=1}^t Softmax_j(\frac {q_{t,i}^Tk_{j,i}}{\sqrt{d_h+d_h^R}})v^C_{j,i}\\ u_t &= W^o[o_{t,1};o_{t,2};...;o_{t,n_h}] \end{aligned}$

Flash Attention

参考 [Attention优化][2w字]🔥原理篇: 从Online-Softmax到FlashAttention V1/V2/V3 - 知乎 (zhihu.com)

Online-Softmax

save-softmax

$softmax({x_1,x_2,...,x_n})=\frac{e^{x_i-m}}{\sum_{j=1}^n e^{x_j-m}},m=max({x_1,x_2,...,x_n})$

3-pass softmax

先算最大值： $m_i \leftarrow max(m_{i-1}, x_i)$
再算求和： $d_i \leftarrow d_{i-1}+e^{x_i-m_N}$
最后算每个元素的值： $res_i \leftarrow \frac {e^{x_i}-m}{d_N}$

2-pass online-softmax

合并1-2步

$i==N$ 是刚好满足3-pass的第二步

$\begin{align} d_i &= \sum_{j=1}^{i}e^{x_j-m_i}\\ &= (\sum_{j=1}^{i-1}e^{x_j-m_i}) + e^{x_i-m_i}\\ &= (\sum_{j=1}^{i-1}e^{x_j-m_{i-1}} \times e^{m_{i-1}-m_i}) + e^{x_i-m_i} \\ &= (\sum_{j=1}^{i-1}e^{x_j-m_{i-1}} )e^{m_{i-1}-m_i} + e^{x_i-m_i} \\ &= d_{i-1}e^{m_{i-1}-m_i} + e^{x_i-m_i} \end{align}$

这样只用存储上一步的最小值 $m_{i-1}$ 和上一步的结果 $d_{i-1}$ 就可以计算当前这一步了，规避了计算全局max

Self-Attention

Standard self-Attention

$O = softmax(\frac{QK^T}{\sqrt{d}})V$ ，后续默认 $\sqrt{d}=1$

多阶段计算

$S = QK^T\\ P=softmax(S)\\ O=PV$

Multi-pass Self-Attention

利用2-pass online-softmax的基础上添加QKV计算

计算QK

$\begin{aligned} x_i & \leftarrow Q[k,:] K^T[:, i] \\ m_i & \leftarrow \max \left(m_{i-1}, x_i\right) \\ d_i^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_i}+e^{x_i-m_i} \end{aligned}$
计算结果

$\begin{aligned} & a_i \leftarrow \frac{e^{x_i-m_N}}{d_N^{\prime}} \\ & \boldsymbol{o}_i \leftarrow \boldsymbol{o}_{i-1}+a_i V[i,:] \end{aligned}$

可以变化为

$o_i \leftarrow \boldsymbol{o}_{i-1}+\frac{e^{x_j-m_N}}{d_N^{\prime}} V[i,:]$

1-pass FlashAttention v1(核心)

令 $\boldsymbol{o}_i=\left(\sum_{j=1}^i \frac{e^{x_j-m_i}}{d_i^{\prime}} V[j,:]\right)$

当 $i==N$ 时，刚好和multi-pass self-attention的第二步相同

同样的，可以推导出 $o_i$ 与 $o_{i-1}$ 的关系

$\begin{align} o_i &=(\sum_{j=1}^i\frac{e^{x_j-m_i}}{d^{\prime}_i} V[j,:])\\ &=(\sum_{j=1}^{i-1}\frac{e^{x_j-m_i}}{d^{\prime}_i} V[j,:]) + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:]\\ &=(\sum_{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d^{\prime}_{i-1}} \frac{e^{x_j-m_i}}{e^{x_j-m_{i-1}}} \frac{d^{\prime}_{i-1}}{d^{\prime}_i} V[j,:]) + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:] \\ &=(\sum_{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d^{\prime}_{i-1}} \frac{d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i} V[j,:]) + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:] \\ &=(\sum_{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d^{\prime}_{i-1}} V[j,:])\frac{d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i} + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:] \\ &=o_{i-1}\frac{d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i} + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:] \\ \end{align}$

这样就可以看到 $o_i$ 只依赖于上一次的 $d_{i-1},m_{i-1},o_{i-1}$ 与本次的 $d_i,m_i$ ，可以在一个循环中全部计算完成，就不用两阶段了

所以得到1-pass的计算为

$\begin{aligned} x_i & \leftarrow Q[k,:] K^T[:, i] \\ m_i & \leftarrow \max \left(m_{i-1}, x_i\right) \\ d_i^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_i}+e^{x_i-m_i}\\ o_i & \leftarrow o_{i-1}\frac{d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i} + \frac{e^{x_i-m_i}}{d^{\prime}_i} V[i,:] \end{aligned}$

上面的伪代码是按列进行计算的，外层循环是要遍历列的。当然这个步骤可以进行分块计算（tiling）这样可以减少外层循环的次数，还可以增强数据的访问效率。

这里取K时多取了几列，计算 $m_i$ 时先计算每一个分块每一行的最大值，最后是 $d_i$ 后面加 $e^{x_i-m_i}$ 的部分修改成了对当前块所有列的求和

$\begin{aligned} x_i & \leftarrow Q[k,:] K^T[:, (i-1)b:ib] \\ m_i^{local} & =max_{j=1}^{b}(x_i[j]),每一行的最大值\\ m_i & \leftarrow \max \left(m_{i-1}, m_i^{local}\right) \\ d_i^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_i}+ \sum_{j=1}^{b} e^{x_i[j]-m_i}\\ o_i & \leftarrow o_{i-1}\frac{d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i} + \sum_{j=1}^{b}\frac{e^{x_i[j]-m_i}}{d^{\prime}_i} V[j+(i-1)b,:] \end{aligned}$

相比于Standard Self-Attention的计算流程节省了S和P矩阵的显存，减少Q,K的HBM IO

FlashAttention V1原论文对于 $O_1$ 的求解如下：

$O_1\leftarrow diag(l_1^{new})^{-1}(diag(l_1)e^{m_1-m_1^{new}}O_1+e^{\hat{m}_{11}-m_1^{new}}\hat{P}_{12}V_2)$

看着和之前推导的公式差距有点大，那转化一下，上面的 $O_1$ 其实指的是实际内存空间的地址，而之前写的 $o_i$ 和 $o_{i-1}$ 中的 $i$ ，表示的是不同迭代次数， $o_i$ 和 $o_{i-1}$ 其实也是相同的内存空间地址，这里先带入成之前的符号，这里的 $\hat{P_{12}}$ 其实是 $e^{x_i[j]-m_i^{local}}$ 和之前推导公式的计算方式不同，他这里是先用局部最大值计算了每个元素的 $e$ ，而上面推导的公式是在最后直接使用全局的最大值进行计算（第一步），然后这里的 $diag$ 是将一个一维数组转化成一个对角矩阵的方式，因为每次 $o_i$ 的计算其实都是只算了 $x_i$ 一个数据,为了方便理解，可以想象成只有一个元素然后进行转换（第二步）。最后就是将两个 $e$ 的指数进行合并，得到之前推导公式相同的结果（第三步）。

$\begin{aligned} o_i & \leftarrow diag(d^{\prime}_i)^{-1}(diag(d^{\prime}_{i-1})e^{m_{i-1}-m_i}o_{i-1}+e^{m_i^{local}-m_i}e^{x_i[j]-m_i^{local}}V[i,:]) \\ o_i & \leftarrow \frac {d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i}o_{i-1}+\frac{e^{m_i^{local}-m_i+x_i[j]-m_i^{local}}}{d^{\prime}_i}V[i,:] \\ o_i & \leftarrow \frac {d^{\prime}_{i-1}e^{m_{i-1}-m_i}}{d^{\prime}_i}o_{i-1}+ \frac{e^{x_i[j]-m_i}}{d_i^\prime}V[i,:]\\ \end{aligned}$

block size的设置

$M$ 是SRAM的大小，也是L1 cache的大小，通过这样的计算方式，控制Q,O,K,V的大小不会超过SRAM的大小，实现高效的访存。

下面可以推断出通过这样设置 $B_c$ 和 $B_r$ 确保Q,O,K,V的中间变量大小都不会超过 $\frac M4$ 。

当然这里也会有一些剩余的部分，其实就是给 $m_i$ 和 $d_i$ 预留使用，基本上都可以把SRAM打满了

$B_c=\left\lceil\frac{M}{4 d}\right\rceil, B_r=\min \left(\left\lceil\frac{M}{4 d}\right\rceil, d\right)\\ \begin{aligned} & S R A M\left(Q_i\right)=B_r \times d=\min \left(\left\lceil\frac{M}{4 d}\right\rceil, d\right) \times d<\left\lceil\frac{M}{4}\right\rceil \\ & S R A M\left(O_i\right)=B_r \times d=\min \left(\left\lceil\frac{M}{4 d}\right\rceil, d\right) \times d<\left\lceil\frac{M}{4}\right\rceil \\ & S R A M\left(K_j, V_j\right)=2 \times B_c \times d=2 \times\left\lceil\frac{M}{4 d}\right\rceil \times d<\left\lceil\frac{M}{2}\right\rceil \end{aligned}$

稀疏矩阵的拓展

简单来说就是在原本的基础上，检测每次分块的稀疏度，如果稀疏度为0就跳过对这个小块的计算

反向计算

反向最重要的技术就是recompute，前向计算中省略了中间结果S和P，但是反向需要用他们计算梯度值。

所以反向计算时也会进行tiling，将Q,K,V分块加载到SRAM，再通过recompute的方式计算出当前块的S和P的值，用于求取梯度。

无论是否有recompute，都要去load对应的数据到SRAM中，如果不用recompute就要从HBM中拉取(load Q,K,V,dO,dS, )+ (load P,dP) +(write dS, dP, dQ, dV, dK)

但使用了tiling+recompute之后，只用从HBM拉取(load Q,K,V,dO) + (write dQ,dV,dK)，节省了dS,dP,P的IO，虽然recompute技术增加了计算量，但计算过程都是在SRAM中进行的，对比与从HBM拉取数据到SRAM中速度能快很多。

FlashAttention V2

主要优化：

减少大量非matmul的冗余计算，增加Tensor Cores运算比例
forward pass/backward pass均增加seqlen维度的并行，forward pass交替Q,K,V循环顺序
更好的Warp Partitioning策略，避免Split-K

减少非matmul的冗余计算

为啥要减少非matmul？因为matmul有专门的硬件(tensor core)，可以算得更快。

哪里能减少matmul?就是V1中每一轮迭代都进行了rescale（就是softmax中的分母部分），这东西在V1版本中每一轮都用了，但其实可以在QKV都算完之后再除以的

FA2的计算过程如下：

$\begin{aligned} & m^{(1)}=\operatorname{rowmax}\left(\mathbf{S}^{(1)}\right) \in \mathbb{R}^{B_r}, 第一块每一行的最大值 \\ & \ell^{(1)}=\operatorname{rowsum}\left(e^{\mathbf{S}^{(1)}-m^{(1)}}\right) \in \mathbb{R}^{B_r}，第一块每一行的求和 \\ & \tilde{\mathbf{O}}^{(1)}=e^{\mathbf{S}^{(1)}-m^{(1)}} \mathbf{V}^{(1)} \in \mathbb{R}^{B_r \times d}，第一块的结果（无rescale） \\ & m^{(2)}=\max \left(m^{(1)}, \operatorname{rowmax}\left(\mathbf{S}^{(2)}\right)\right)=m ，前两块的最大值\\ & \ell^{(2)}=e^{m^{(1)}-m^{(2)}} \ell^{(1)}+\operatorname{rowsum}\left(e^{\mathbf{S}^{(2)}-m^{(2)}}\right)=\operatorname{rowsum}\left(e^{\mathbf{S}^{(1)}-m}\right)+\operatorname{rowsum} \left(e^{\mathbf{S}^{(2)}-m}\right)=\ell，前两块的求和 \\ & \tilde{\mathbf{O}}^{(2)}=e^{s^{(1)}-m} \mathbf{V}^{(1)}+e^{s^{(2)}-m} \mathbf{V}^{(2)} ，前两块的结果（无rescale)\\ & \mathbf{O}^{(2)}=\operatorname{diag}\left(\ell^{(2)}\right)^{-1} \tilde{\mathbf{O}}^{(2)}=\mathbf{O} ，前两块的结果（带rescale） \end{aligned}$

对比与FA1，计算区别在于那几个 $O$ 的计算，在FA1中每一轮的计算都为，多了 $diag(l)$ 的计算步骤

$\mathbf{O}_i\leftarrow\operatorname{diag}\left(\ell_i^{\text {new }}\right)^{-1}\left(\operatorname{diag}\left(\ell_i\right) e^{m_i-m_i^{\text {new }}} \mathbf{O}_i+e^{\tilde{m}_{i j}-m_i^{\text {new }}} \tilde{\mathbf{P}}_{i j} \mathbf{V}_j\right)$

反向计算不再保存 $m^{(j)}$ 和 $\ell^{(j)}$ ，而是保存 $\operatorname{logsumexp} L^{(j)}=m^{(j)}+\log \left(\ell^{(j)}\right)$

从而减少 $P_{ij}$ 重计算的计算量。FA1→FA $\left.\begin{array}{c} \mathbf{P}_{i j}=\operatorname{diag}\left(l_i\right)^{-1} \exp \left(\mathbf{S}_{i j}^{\text {masked }}\right. \end{array}-m_i\right) \in \mathbb{R}^{B_r \times \boldsymbol{B}_c} \rightarrow \mathbf{P}_i^{(j)}=\exp \left(\mathbf{S}_{i j}-L_i\right) \in \mathbb{R}^{B_r \times \boldsymbol{B}_c}$

增加seqlen维度的并行

在FA1中，是先load K,V子块，再load Q子块，这使得内循环每轮迭代都只是计算了Q的子结果，想要计算完所有的Q的每一行则是需要整个计算过程结束，此外每一次内循环，都要将结果写入到全局内存中，访问开销很大。

在FA2中，是先load Q子块，再load K,V字块，这使得只要内循环结束，Q的一部分行就能计算完成，而不只是Q的子结果。如果我们在外循环中使用一个本地内存去存储O，然后内循环计算的结果全部写入本地内存中，内循环结束后再写入全局内存，阁下又该如何应对？这种方式是对的，这样每个Q行就可以并行起来了，独立去计算这一行Q的结果，每一行Q都会在内循环结束后得到完整的结果，此外每次内循环访问的都是本地内存，访存开销也大大减少。

反观FA1，只在batch_size和headnum做并行