第二十四课：Transformer 中的 MLP、GEMM 与算子融合

导语：这一课讲 Transformer 中计算量最大的 MLP 模块，以及 GEMM Epilogue、GELU/SiLU/SwiGLU 融合和张量并行的数据流。

Transformer 中，除了 Attention，另一个主要计算模块是 MLP，也常称为：

• FFN：Feed-Forward Network；
• 前馈网络；
• 通道混合层。

一个典型 Transformer Block 可以简化为：

输入
  ↓
Attention
  ↓
Residual + Norm
  ↓
MLP / FFN
  ↓
Residual + Norm

在很多大模型中，MLP 的参数量和计算量甚至高于 Attention。

---

一、最基本的 MLP 结构

经典 Transformer FFN 通常是：

$$H=\phi(XW_1+b_1)$$ $$Y=HW_2+b_2$$

其中：

• $X$：输入；
• $W_1$：升维权重；
• $W_2$：降维权重；
• $\phi$：激活函数；
• $H$：中间隐藏特征。

数据流为：

X
↓
Linear 1
↓
Activation
↓
Linear 2
↓
Y

如果输入形状是：

$$X\in\mathbb{R}^{M\times H}$$

第一层权重为：

$$W_1\in\mathbb{R}^{H\times F}$$

第二层权重为：

$$W_2\in\mathbb{R}^{F\times H}$$

则：

$$XW_1\in\mathbb{R}^{M\times F}$$

这里：

• $M$：本次处理的 Token 数；
• $H$：模型隐藏维度；
• $F$：FFN 中间维度。

通常：

$$F>H$$

也就是先升维，再降维。

---

二、为什么 MLP 要先升维

升维后，每个 Token 可以在更大的特征空间中进行非线性变换。

例如：

$$H=4096$$

中间维度可能是：

$$F=11008$$

或者其他接近 $3H$、$4H$ 的配置。

MLP 的作用不是在不同 Token 之间交换信息，而是：

对每个 Token 的特征维度进行非线性变换。

因此：

Attention：
混合不同 Token 的信息

MLP：
混合同一个 Token 内不同通道的信息

---

三、Linear 层为什么等价于 GEMM

假设有 $M$ 个 Token，每个 Token 有 $H$ 个特征。

把所有 Token 组成矩阵：

$$X\in\mathbb{R}^{M\times H}$$

线性层为：

$$Y=XW$$

其中：

$$W\in\mathbb{R}^{H\times F}$$

这就是标准矩阵乘法：

$$[M\times H]\times[H\times F] = [M\times F]$$

因此 Transformer 中的大量 Linear 层，最终都映射为 GEMM。

包括：

• Q、K、V 投影；
• Attention 输出投影；
• MLP 升维层；
• MLP 降维层；
• 分类头；
• 部分专家层。

---

四、MLP 的计算量有多大

第一层：

$$XW_1$$

计算量约为：

$$2MHF$$

第二层：

$$HW_2$$

计算量也是：

$$2MFH$$

总计算量约为：

$$4MHF$$

这里乘加通常按2个 FLOPs 计算：

一次乘法
+
一次加法

如果：

$$F=4H$$

则总计算量约为：

$$16MH^2$$

所以 MLP 是 Transformer 中非常大的计算来源。

---

五、Attention 与 MLP 的计算量对比

Attention 的主要矩阵计算约为：

$$QK^T:\ 2N^2d$$ $$PV:\ 2N^2d$$

总计约：

$$4N^2d$$

MLP 约为：

$$4NHF$$

如果：

$$H=n_hd$$

并且：

$$F\approx4H$$

MLP 计算约为：

$$16NH^2$$

因此：

• 序列较短时，MLP 往往占很大比例；
• 序列非常长时，Attention 的 $N^2$ 项会快速增长；
• 自回归解码时，每次只处理少量新 Token，MLP 权重读取又可能成为主要成本。

---

六、MLP 的两个 GEMM 有什么区别

第一个 GEMM：升维

$$[M\times H]\times[H\times F]$$

输出形状：

$$[M\times F]$$

特点：

• 输出特征维度较大；
• 后面通常接激活函数；
• 可能同时生成多个分支，例如 SwiGLU。

第二个 GEMM：降维

$$[M\times F]\times[F\times H]$$

输出回到：

$$[M\times H]$$

特点：

• 输出通常接 Residual Add；
• 可能融合 Bias；
• 可能需要跨 GPU All-Reduce。

---

七、激活函数为什么必要

如果两个 Linear 层之间没有激活函数：

$$Y=(XW_1)W_2$$

可以合并为：

$$Y=X(W_1W_2)$$

两个线性变换仍然只是一个线性变换。

加入非线性：

$$Y=\phi(XW_1)W_2$$

后，模型才能表达更复杂的函数。

常见激活函数包括：

• ReLU；
• GELU；
• SiLU；
• SwiGLU。

---

八、ReLU

ReLU 为：

$$\operatorname{ReLU}(x)=\max(0,x)$$

优点：

• 计算简单；
• 不需要指数或误差函数；
• 硬件实现容易。

缺点：

• 负半轴全部变为0；
• 函数不平滑；
• 现代大语言模型中已不如 GELU、SiLU、SwiGLU 常见。

GPU 上可能直接用比较和最大值指令实现。

---

九、GELU

GELU 为：

$$\operatorname{GELU}(x) = x\Phi(x)$$

其中 $\Phi(x)$ 是标准正态分布的累积分布函数。

常见近似公式为：

$$\operatorname{GELU}(x) \approx \frac{x}{2} \left[ 1+ \tanh \left( \sqrt{\frac{2}{\pi}} \left( x+0.044715x^3 \right) \right) \right]$$

与 ReLU 的硬截断不同，GELU 会对负数进行平滑衰减。

示意：

很大的正数：
GELU(x) ≈ x

接近 0：
平滑变化

很大的负数：
GELU(x) ≈ 0

---

十、GELU 为什么比 ReLU 更昂贵

GELU 可能需要：

• 乘法；
• 三次方；
• 加法；
• tanh；
• 或 erf；
• 其他近似指令。

所以单元素计算比 ReLU 复杂。

但与大型 GEMM 相比，GELU 的算术量通常仍然较小。

真正的问题通常不是 GELU 计算本身，而是：

是否需要把 GEMM 输出写到显存，再启动一个单独 Kernel 读取并执行 GELU。

---

十一、SiLU

SiLU，也称 Swish，定义为：

$$\operatorname{SiLU}(x) = x\cdot\sigma(x)$$

其中：

$$\sigma(x)=\frac{1}{1+e^{-x}}$$

因此：

$$\operatorname{SiLU}(x) = \frac{x}{1+e^{-x}}$$

它同样是一种平滑非线性函数。

执行中通常涉及：

• 指数；
• 加法；
• 倒数；
• 乘法。

---

十二、什么是 GLU

GLU 是 Gated Linear Unit，门控线性单元。

它把输入投影成两个分支：

$$A=XW_a$$ $$B=XW_b$$

然后输出：

$$H=A\odot\sigma(B)$$

其中：

$$\odot$$

表示逐元素乘法。

一个分支提供主要特征，另一个分支控制门的开启程度。

示意：

        ┌→ Linear A → 特征 A ─┐
X ──────┤                     ├→ 逐元素乘法
        └→ Linear B → Gate ───┘

---

十三、什么是 SwiGLU

SwiGLU 使用 SiLU 作为门控函数。

常见形式为：

$$H= \operatorname{SiLU}(XW_g) \odot (XW_u)$$

也可以看到分支名称：

Gate Projection
Up Projection

即：

$$G=XW_g$$ $$U=XW_u$$ $$H=\operatorname{SiLU}(G)\odot U$$

然后：

$$Y=HW_d$$

其中 $W_d$ 是 Down Projection。

完整结构：

                ┌→ Gate Projection → SiLU ─┐
X ──────────────┤                          ├→ Elementwise Multiply
                └→ Up Projection ──────────┘
                                               ↓
                                        Down Projection
                                               ↓
                                               Y

---

十四、SwiGLU 为什么有三个权重矩阵

普通两层 MLP：

W1：Up Projection
W2：Down Projection

SwiGLU：

Wg：Gate Projection
Wu：Up Projection
Wd：Down Projection

因为第一阶段需要生成两个中间张量：

$$G=XW_g$$

和：

$$U=XW_u$$

然后做：

$$H=\operatorname{SiLU}(G)\odot U$$

---

十五、SwiGLU 的中间维度为什么常小于 $4H$

如果两个分支都使用 $4H$ 中间维度，参数量会明显增大。

普通 FFN 参数量约为：

$$H\times4H+4H\times H = 8H^2$$

SwiGLU 参数量约为：

$$H\times F + H\times F + F\times H = 3HF$$

为了与普通 $4H$ FFN 参数量接近，可以令：

$$3HF\approx8H^2$$

得到：

$$F\approx\frac{8}{3}H$$

因此实际 SwiGLU 中间维度常接近：

$$2.67H$$

再根据 Tensor Core 对齐要求调整为合适整数。

---

十六、朴素 SwiGLU 的数据流

最直接实现可能分为四个阶段：

Kernel/GEMM 1：
G = XWg

Kernel/GEMM 2：
U = XWu

Kernel 3：
H = SiLU(G) × U

Kernel/GEMM 4：
Y = HWd

中间需要保存：

G
U
H

这些张量通常都很大。

例如：

$$G,U,H\in\mathbb{R}^{M\times F}$$

如果分别写入和读取 HBM，会产生大量带宽开销。

---

十七、什么是 GEMM Epilogue

GEMM 的主体计算为：

$$C=AB$$

矩阵乘法完成后，结果通常先位于：

• Tensor Core 累加器；
• 寄存器；
• 片上临时状态。

在写入 Global Memory 前执行的后处理，称为：

Epilogue，尾处理阶段。

常见 Epilogue 操作包括：

$$D=\alpha AB+\beta C$$

进一步还可以执行：

• Bias；
• ReLU；
• GELU；
• SiLU；
• Scale；
• Residual Add；
• Quantization；
• Clamp；
• 类型转换。

---

十八、为什么 Epilogue 融合很重要

不融合时：

GEMM
  ↓
C 写入 HBM
  ↓
新 Kernel 读取 C
  ↓
Bias + Activation
  ↓
结果再次写回 HBM

融合后：

GEMM 累加结果仍在寄存器
          ↓
直接执行 Bias + Activation
          ↓
只写最终结果

减少：

• 一次中间写入；
• 一次中间读取；
• 一次额外 Kernel 启动；
• 中间张量显存占用。

---

十九、Bias + GELU 融合

假设：

$$Z=XW+b$$ $$Y=\operatorname{GELU}(Z)$$

可以在 GEMM Epilogue 中直接执行：

Tensor Core 累加器
      ↓
加 Bias
      ↓
转换为 GELU 计算格式
      ↓
执行 GELU
      ↓
转换到输出精度
      ↓
写入 HBM

这样中间的 $Z$ 不需要进入显存。

---

二十、SwiGLU 的融合方式

SwiGLU 需要两个 GEMM 结果：

$$G=XW_g$$ $$U=XW_u$$

一种方式是把两个权重矩阵拼接：

$$W_{\text{fused}} = \begin{bmatrix} W_g & W_u \end{bmatrix}$$

一次 GEMM：

$$Z=XW_{\text{fused}}$$

输出：

$$Z= \begin{bmatrix} G & U \end{bmatrix}$$

然后在 Epilogue 或紧接的融合逻辑中计算：

$$H=\operatorname{SiLU}(G)\odot U$$

这样两个投影可以使用一次较大的 GEMM 完成。

---

二十一、合并 Gate 和 Up Projection 的优势

原本：

读取 X → GEMM Gate
再次读取 X → GEMM Up

合并后：

读取 X
    ↓
一次更宽的 GEMM
    ↓
同时产生 Gate 和 Up

优势包括：

• X 的加载和 Tile 复用更充分；
• 减少 Kernel 启动；
• 更容易提高 GEMM 尺寸；
• 减少调度开销；
• 更适合 Tensor Core。

但输出宽度变成：

$$2F$$

输出累加器和写回压力也会增加。

---

二十二、能否把 SwiGLU 和 Down Projection 全部融合

理论上希望：

X
↓
Gate/Up GEMM
↓
SiLU × Multiply
↓
Down GEMM
↓
Y

全部在一个 Kernel 中完成，不保存中间 $H$。

但实现非常困难。

原因是：

• 第一 GEMM 输出为 $M\times F$；
• 第二 GEMM 需要把 $H$ 与 $W_d$ 相乘；
• $H$ 是第二个 GEMM 的输入；
• 第二 GEMM 需要在多个输出 Tile 间复用 $H$；
• 片上无法保存完整的大型 $H$；
• 两个 GEMM 的 Tile 映射可能不同；
• 寄存器和 Shared Memory 压力极高。

所以常见融合范围通常是：

Gate GEMM + Up GEMM + SiLU + Multiply

然后把 $H$ 写出，再执行 Down Projection。

---

二十三、为什么完全融合不一定划算

完全融合可能减少 HBM 流量，但也可能造成：

• Kernel 极其复杂；
• 寄存器 Spill；
• Occupancy 很低；
• Tile 选择受到双重限制；
• Tensor Core 利用率下降；
• 指令缓存压力增加；
• 无法使用高度优化的独立 GEMM Kernel。

因此优化目标不是“融合最多”，而是：

总执行时间最短。

有时两个高效 GEMM 加一个中间写回，比一个低效的超大型融合 Kernel 更快。

---

二十四、Residual Add Epilogue

Down Projection 后通常接残差：

$$Y=X+\operatorname{MLP}(X)$$

朴素实现：

Down GEMM 输出 M
        ↓
写入 HBM
        ↓
新 Kernel 读取 M 和 X
        ↓
执行 Residual Add
        ↓
写出 Y

可以在 Down GEMM Epilogue 中直接执行：

$$Y=M+X$$

数据流：

GEMM 累加结果
     ↓
读取残差 X
     ↓
相加
     ↓
写最终 Y

避免保存中间的 MLP 输出。

---

二十五、Residual + Norm 能否继续融合

如果网络结构是 Post-Norm：

$$Y=\operatorname{Norm}(X+\operatorname{MLP}(X))$$

Residual Add 后还需要 LayerNorm 或 RMSNorm。

但 Norm 需要对整行做 Reduction。

普通 GEMM Epilogue 通常按输出 Tile 独立工作，未必掌握完整一行。

因此融合需要解决：

• 跨 Tile Reduction；
• 全行统计；
• 输出回放；
• Tile 间同步。

这比单纯 Residual Add 融合复杂得多。

---

二十六、量化 GEMM 的 Epilogue

在 INT8 或 FP8 GEMM 中，累加器精度通常高于输出精度。

例如 INT8：

$$\text{INT8}\times\text{INT8} \rightarrow \text{INT32 Accumulate}$$

Epilogue 需要执行：

INT32 Accumulator
      ↓
乘 Scale
      ↓
加 Bias
      ↓
Activation
      ↓
Round
      ↓
Saturation
      ↓
INT8 输出

公式可能类似：

$$y= \operatorname{sat}_{8} \left( \operatorname{round} \left( s\cdot acc+b \right) \right)$$

这些操作非常适合与 GEMM 写回阶段融合。

---

二十七、为什么布局转换也可以融合进 Epilogue

下一个算子可能需要不同数据布局。

例如当前 GEMM 输出为：

Row-Major

后续 Kernel 希望：

某种 Interleaved 或 Tensor Core 友好布局

可以在 Epilogue 写回时直接改变：

• 元素排列；
• 向量打包；
• 数据类型；
• Scale 布局；
• Transpose 方向。

这样避免一个独立 Layout Transform Kernel。

---

二十八、GEMM 为什么通常是计算受限

大矩阵 GEMM 具有很高的数据复用。

一个 $A$ Tile 可以与多个 $B$ 数据复用，一个 $B$ Tile也可以被多个线程复用。

计算量约为：

$$2MNK$$

数据量约为：

$$MK+KN+MN$$

当矩阵足够大时，算术强度很高。

因此大规模 GEMM 通常可以接近：

• Tensor Core 峰值；
• 计算 Roofline。

但这需要：

• Tile 合理；
• Tensor Core 利用充分；
• 数据搬运被隐藏；
• 矩阵尺寸足够大。

---

二十九、什么时候 MLP GEMM 可能不是计算受限

以下情况可能变成内存或延迟受限：

• $M$ 很小；
• Batch Size 很小；
• 自回归解码每次只有一个 Token；
• 权重无法在 Cache 中复用；
• GEMM 变成 GEMV；
• 矩阵维度不适合 Tensor Core；
• 量化解包或 Scale 访问开销大。

尤其是在解码阶段：

$$M=1$$

此时：

$$[1\times H]\times[H\times F]$$

更像矩阵向量乘法。

每个权重通常只被使用一次，算术强度很低。

---

三十、Prefill 和 Decode 的区别

Prefill

一次处理整个输入序列。

例如：

$$M=B\times N$$

可能很大。

Linear 层表现为较大的 GEMM：

多个 Token
×
同一组权重

权重可以在大量 Token 间复用。

因此：

• Tensor Core 利用率较高；
• 更容易计算受限；
• 批处理效率较高。

Decode

每一步每个序列只生成一个新 Token。

如果 Batch 较小：

$$M\approx1$$

Linear 层变成接近 GEMV：

读取整个权重矩阵
只为少量 Token 计算

权重带宽成为主要瓶颈。

---

三十一、为什么量化对 Decode 特别重要

解码阶段每生成一个 Token，都要读取大量模型权重。

假设某层权重使用 FP16：

$$2\text{ Bytes/weight}$$

改成 INT8：

$$1\text{ Byte/weight}$$

理论权重流量约减半。

改成4位：

$$0.5\text{ Byte/weight}$$

理论权重流量进一步降低。

因为 Decode 常是权重带宽受限，所以低比特权重可以显著提高 Token 生成吞吐率。

但还需要承担：

• 反量化；
• Scale 读取；
• 数据解包；
• 低精度误差；
• 内核适配。

---

三十二、MLP 中的 Batch Size 为什么重要

更大的 $M$ 表示一次处理更多 Token。

权重 $W$ 可以被更多行输入复用：

一次加载某个 Weight Tile
        ↓
服务多个 Token

因此 Batch 或 Token 数越大，通常：

• GEMM 更规则；
• Tensor Core 利用率更高；
• 权重流量更容易摊薄；
• Kernel 启动开销占比更低。

但 Batch 增大会增加：

• 激活显存；
• 推理延迟；
• KV Cache；
• 调度等待。

---

三十三、张量并行中的列并行 MLP

假设第一层权重：

$$W_1\in\mathbb{R}^{H\times F}$$

使用两块 GPU 沿输出维度切分：

$$W_1= \begin{bmatrix} W_1^{(0)} & W_1^{(1)} \end{bmatrix}$$

其中：

$$W_1^{(0)},W_1^{(1)} \in\mathbb{R}^{H\times F/2}$$

两块 GPU 都拥有完整输入 $X$。

分别计算：

$$H_0=XW_1^{(0)}$$ $$H_1=XW_1^{(1)}$$

因此：

GPU 0：计算前一半中间通道
GPU 1：计算后一半中间通道

输出在特征维度上分片，不需要立即求和。

这叫列并行 Linear。

---

三十四、激活函数可以在本地执行

因为第一层输出已经按中间维度切分：

GPU 0：H₀
GPU 1：H₁

激活函数是逐元素操作：

$$\phi(H_0)$$ $$\phi(H_1)$$

各 GPU 可以独立执行，不需要通信。

SwiGLU 也可以让 Gate 和 Up 的对应中间通道放在同一 GPU：

GPU 0：
Gate 分片 0
Up 分片 0
本地执行 SiLU × Up

GPU 1：
Gate 分片 1
Up 分片 1
本地执行 SiLU × Up

这样门控操作不需要跨 GPU 交换数据。

---

三十五、张量并行中的行并行 Down Projection

Down Projection：

$$Y=HW_d$$

其中：

$$H= \begin{bmatrix} H_0 & H_1 \end{bmatrix}$$

相应地沿输入维度切分权重：

$$W_d= \begin{bmatrix} W_d^{(0)}\\ W_d^{(1)} \end{bmatrix}$$

各 GPU 计算部分输出：

$$Y_0=H_0W_d^{(0)}$$ $$Y_1=H_1W_d^{(1)}$$

最终：

$$Y=Y_0+Y_1$$

因此需要：

All-Reduce

或者：

Reduce-Scatter

把各 GPU 的部分结果合并。

---

三十六、完整的 Tensor Parallel MLP

两块 GPU 的数据流为：

完整输入 X
   ├─────────────┐
   ↓             ↓
GPU 0          GPU 1
列并行 W1      列并行 W1
   ↓             ↓
H0             H1
   ↓             ↓
本地激活       本地激活
   ↓             ↓
行并行 W2      行并行 W2
   ↓             ↓
Y0             Y1
   └──── All-Reduce ────┘
              ↓
              Y

MLP 中通常只需要在 Down Projection 后进行一次关键集合通信。

---

三十七、为什么第一层不立即 All-Gather

第一层列并行后，每块 GPU 只有一部分中间通道。

如果下一步激活和第二层权重切分方式匹配，就没有必要把完整中间张量聚合到每块 GPU。

这样可以避免：

大型 M×F 中间激活

的 All-Gather。

设计张量并行时的重要原则是：

让相邻算子的分片方式相互匹配，尽量延迟或消除通信。

---

三十八、SwiGLU 张量并行如何切分

Gate 和 Up 两个投影都沿中间维度切分。

例如两块 GPU：

GPU 0：
Wg 的前 F/2 列
Wu 的前 F/2 列

GPU 1：
Wg 的后 F/2 列
Wu 的后 F/2 列

各自计算：

$$G_0,U_0$$

以及：

$$G_1,U_1$$

然后本地：

$$H_0=\operatorname{SiLU}(G_0)\odot U_0$$ $$H_1=\operatorname{SiLU}(G_1)\odot U_1$$

最后分别进入行并行 Down Projection。

---

三十九、通信能否与 GEMM 重叠

Down Projection 产生的部分输出需要归约。

可以把输出矩阵按 Tile 或 Chunk 分段。

当某一部分计算完成后，立即开始 Reduce-Scatter：

时间 →
计算 Chunk 0 | 计算 Chunk 1 | 计算 Chunk 2
通信：         Reduce 0      | Reduce 1

这样通信可以被后续计算部分覆盖。

但需要：

• 合适的 Chunk 大小；
• 通信和计算使用不同硬件资源；
• 软件框架支持异步操作；
• 网络带宽足够；
• 避免过多小通信调用。

---

四十、为什么小消息通信很难隐藏

如果每个 Chunk 很小：

• 集合通信启动延迟占比高；
• GPU/NCCL 调度开销高；
• 带宽利用率低。

如果 Chunk 太大：

• 通信启动太晚；
• 可与计算重叠的时间不足。

因此需要在：

尽早开始通信
与
让每次通信足够大

之间平衡。

---

四十一、MoE 中的 MLP

Mixture-of-Experts 中，每个 Expert 通常就是一个独立 MLP。

Token 根据路由结果被送到不同 Expert：

Token 0 → Expert 2
Token 1 → Expert 0
Token 2 → Expert 2
Token 3 → Expert 5

每个 Expert 执行自己的：

Up/Gate Projection
Activation
Down Projection

MoE 的问题不仅是 GEMM，还包括：

• Token 路由；
• All-to-All；
• Token 排序与打包；
• Expert 负载不均衡；
• 大量小 GEMM。

---

四十二、Grouped GEMM

如果每个 Expert 收到的 Token 数不同：

Expert 0：32个 Token
Expert 1：120个 Token
Expert 2：8个 Token
Expert 3：64个 Token

为每个 Expert 单独启动 GEMM，会产生大量小 Kernel。

Grouped GEMM 可以在一次 Kernel 中处理多个不同尺寸的矩阵乘法：

GEMM 0：M0 × K × N
GEMM 1：M1 × K × N
GEMM 2：M2 × K × N
...

这样可以：

• 减少 Kernel 启动；
• 提高 SM 利用率；
• 动态调度不同 Expert；
• 改善小矩阵效率。

---

四十三、Grouped GEMM 的困难

不同 Expert 的 $M$ 不同，因此：

• Block 数不同；
• 执行时间不同；
• 尾部效应明显；
• 部分 Expert 很小；
• 任务调度不规则。

GPU 需要在多个 GEMM Tile 任务之间动态分配 Block。

这比普通单一大 GEMM 更接近任务队列调度问题。

---

四十四、MLP 融合优化的层次

可以从浅到深分为：

第一层：Elementwise Fusion

Bias + GELU
SiLU + Multiply
Residual Add

第二层：GEMM Epilogue Fusion

GEMM + Bias + Activation + Quantization

第三层：Projection Fusion

Gate Projection + Up Projection

第四层：跨算子深度融合

Gate/Up GEMM
+ SwiGLU
+ Down GEMM
+ Residual

融合程度越深，潜在 HBM 流量越少，但实现难度和资源压力也越大。

---

四十五、分析 MLP Kernel 时看什么

对大 GEMM

重点查看：

• Tensor Core 吞吐率；
• Compute Throughput；
• DRAM/L2 Throughput；
• Tile 和矩阵尺寸；
• 寄存器与 Shared Memory；
• Tensor Core 指令；
• Pipeline Stall。

对 Activation Kernel

重点查看：

• 显存带宽；
• SFU 利用率；
• 是否有多余中间读写；
• 是否适合融合进 Epilogue。

对 SwiGLU Kernel

重点查看：

• Gate 和 Up 是否重复读取输入；
• 两个投影是否合并；
• 中间 G/U/H 的 HBM 流量；
• SiLU 是否成为明显瓶颈；
• 输出布局是否适合 Down Projection。

---

四十六、一个典型优化过程

原始版本：

GEMM Gate
GEMM Up
SiLU Kernel
Multiply Kernel
GEMM Down
Residual Kernel

第一步，融合 Elementwise：

GEMM Gate
GEMM Up
SiLU × Multiply Kernel
GEMM Down
Residual Kernel

第二步，合并 Gate 和 Up：

Fused Gate-Up GEMM
SiLU × Multiply
GEMM Down
Residual

第三步，融合 Gate-Up Epilogue：

Fused Gate-Up GEMM
+ SiLU × Multiply Epilogue
GEMM Down
Residual

第四步，融合 Down Projection 的 Residual：

Fused Gate-Up GEMM + SwiGLU
GEMM Down + Residual Epilogue

这通常已经能消除多数容易消除的中间流量。

---

四十七、ASIC 中的 MLP 数据流

一个专用 MLP 加速器可以设计为：

Activation SRAM
      ↓
Matrix Multiply Array
      ↓
Accumulator Buffer
      ↓
Bias / SiLU / GELU Unit
      ↓
Intermediate Buffer
      ↓
Matrix Multiply Array
      ↓
Residual / Requantization
      ↓
Output SRAM

主要问题仍然是：

• 权重带宽；
• 中间激活存储；
• 矩阵阵列利用率；
• 激活函数吞吐；
• 数据精度；
• 两个 GEMM 间的数据流。

---

四十八、ASIC 中如何实现 SwiGLU

可以使用两个权重流：

输入 X
  ├→ GEMM Gate → G ─→ SiLU ─┐
  └→ GEMM Up   → U ─────────┤
                             ↓
                         Gated Multiply
                             ↓
                              H

硬件方案包括：

方案一：两个矩阵阵列

Gate 和 Up 并行计算。

优点：

• 延迟低；
• 吞吐率高。

缺点：

• 面积和功耗大。

方案二：一个阵列分时复用

先计算 Gate，再计算 Up。

优点：

• 面积小。

缺点：

• 延迟增加；
• 需要缓存第一分支结果。

方案三：合并权重矩阵

把 Gate 和 Up 权重拼接，一次较宽矩阵乘法生成两者。

这与 GPU 的 Fused Projection 思路一致。

---

四十九、硬件中的激活函数实现

GELU、SiLU 中包含指数、双曲正切或 Sigmoid。

ASIC 中通常不会直接实现高精度通用数学库，而可能使用：

• 查找表 LUT；
• 分段线性近似；
• 多项式近似；
• CORDIC；
• 特定范围近似；
• 混合 LUT + 插值。

需要权衡：

精度
面积
延迟
吞吐率
功耗

AI 推理中通常允许一定近似，但必须验证模型精度。

---

五十、本课核心结论

第一，Transformer MLP 通常由升维 Linear、激活函数和降维 Linear 组成。

第二，Linear 层本质上是 GEMM：

$$Y=XW$$

第三，MLP 通常占据 Transformer 中很大一部分参数和计算量。

第四，GELU 和 SiLU 比 ReLU 更复杂，但单独激活计算通常不如中间 HBM 读写昂贵。

第五，SwiGLU 使用 Gate Projection 和 Up Projection：

$$H=\operatorname{SiLU}(XW_g)\odot(XW_u)$$

第六，SwiGLU 有三个主要权重矩阵：Gate、Up 和 Down。

第七，GEMM Epilogue 是矩阵乘法累加完成后、写入显存前的尾处理阶段。

第八，可以在 Epilogue 中融合：

Bias
Activation
Scale
Residual
Quantization
Layout Transform

第九，融合能够减少中间显存读写和 Kernel 启动，但过度融合可能增加寄存器、降低 Occupancy 和 Tensor Core 效率。

第十，Gate 和 Up Projection 可以通过拼接权重矩阵合并为一次更宽的 GEMM。

第十一，SwiGLU 与 Down Projection 完全融合较困难，因为中间结果需要被第二个 GEMM 大范围复用。

第十二，Prefill 中 MLP 通常表现为大 GEMM，较容易计算受限。

第十三，Decode 中 $M$ 很小，Linear 层接近 GEMV，常受权重显存带宽限制。

第十四，低比特权重量化对 Decode 尤其有价值，因为它能减少每 Token 的权重读取流量。

第十五，张量并行 MLP 通常采用：

第一层列并行
激活本地执行
第二层行并行
最后 All-Reduce 或 Reduce-Scatter

第十六，MoE 中每个 Expert 通常是一个 MLP，需要 Grouped GEMM 处理不同大小的 Token 分组。

第十七，MLP 优化的核心可以概括为：

$$\boxed{ \text{efficient GEMM} + \text{Epilogue fusion} + \text{reduced intermediate HBM traffic} }$$

下一课是倒数第三课：

第二十五课：稀疏计算、不规则访存，以及 GPU 不擅长什么

重点包括：

稀疏矩阵为什么不一定比稠密矩阵快
CSR、CSC、COO
SpMV 与 SpMM
Gather、Scatter
原子冲突
负载不均衡
结构化与非结构化稀疏
图计算为什么难
GPU 最不擅长哪些类型的任务