大模型推理核心算子：矩阵乘法、Softmax、LayerNorm 与 Attention

各位同学，今天我们来聊聊大模型推理中最核心的几个算子。说实话，搞了这么多年芯片架构，我见过太多团队在算子优化上栽跟头。你想想看，一个 Transformer 模型跑起来，90% 以上的计算时间都花在这几个算子上。不理解它们的计算特征，你做的指令集适配就是空中楼阁。

矩阵乘法（GEMM）：算力的基石

GEMM 说白了就是 C = A × B + C。但别小看这个操作，它可是大模型推理的绝对主角。我个人习惯把 GEMM 分成三类：

• M 维大：batch size 大，比如同时处理多个请求
• N 维大：输出通道多，比如全连接层
• K 维大：输入特征多，比如 embedding 层

我在项目中遇到过最头疼的问题，就是矩阵尺寸不规整。比如某个模型的 hidden size 是 1024，但你的硬件只对 128 的倍数友好。嗯，这里要注意，padding 虽然能解决问题，但会浪费算力。

// 一个简单的 GEMM 伪代码
for i in 0..M:
  for j in 0..N:
    sum = 0
    for k in 0..K:
      sum += A[i][k] * B[k][j]
    C[i][j] = sum

我曾经踩过一个坑：在某个 AI 芯片上，GEMM 的 K 维度如果小于 64，性能会断崖式下跌。后来查了手册才知道，硬件对 K 维度的 tile 大小有硬性要求。所以做指令集适配时，一定要搞清楚硬件的 tile 尺寸约束。

Softmax：数值稳定性的噩梦

Softmax 看起来简单，就是 exp(x_i) / sum(exp(x_j))。但实际做硬件实现时，问题就来了。指数函数 exp 的数值范围极大，x=100 时 exp(100) 已经超出 float16 的表示范围了。

我建议的做法是：先减去最大值，再做 exp。这个技巧叫「数值稳定 Softmax」。你想想看，如果最大值是 100，其他值都减去 100，那 exp 的结果就都在 [0, 1] 区间内了。

Softmax 的计算特征：

• 需要全局归约：先求最大值，再求和
• 指数函数是瓶颈：硬件需要支持快速 exp 近似
• 数据依赖性强：每个元素的计算都依赖全局统计量

LayerNorm：归一化的艺术

LayerNorm 的计算公式是：y = (x - mean) / sqrt(var + eps) * gamma + beta。看起来也不复杂，但它的计算特征和 Softmax 有点像——都需要全局统计量。

我个人习惯把 LayerNorm 拆成三步：

1. 计算均值和方差
2. 归一化
3. 缩放和平移

这里有个关键点：方差计算需要先求平方和，再减去均值的平方。但直接这么算会有精度问题。我建议用 var = mean(x^2) - mean(x)^2，或者用更稳定的 Welford 算法。

LayerNorm 的计算特征：

• 需要两次全局归约：一次求均值，一次求方差
• 除法操作多：归一化时需要除以标准差
• 适合融合操作：可以把均值和方差的计算合并

Attention 机制：大模型的灵魂

Attention 的计算公式是：Attention(Q, K, V) = softmax(Q × K^T / sqrt(d_k)) × V。这里面包含了 GEMM、Softmax、矩阵转置等多种操作。

我把 Attention 拆成四个阶段：

1. QKV 投影：三个 GEMM，把输入映射到 Q、K、V
2. Q × K^T：又一个 GEMM，计算注意力分数
3. Softmax：对注意力分数做归一化
4. Score × V：最后一个 GEMM，得到输出

你想想看，一个 Attention 层里就有 4 个 GEMM 和 1 个 Softmax。这就是为什么大模型推理这么吃算力。

我记得在优化某个 7B 模型时，发现 Attention 的 Q×K^T 这一步，序列长度从 2048 变成 4096 后，计算量翻了 4 倍。这就是 Attention 的 O(n²) 复杂度在作祟。所以现在大家都在搞 Flash Attention，就是为了减少中间结果的访存。

知识体系结构图

这张图把四个核心算子的关系梳理清楚了。你看，GEMM 是基础，Softmax 和 LayerNorm 都需要全局归约，而 Attention 则是它们的集大成者。做指令集适配时，我建议先从 GEMM 入手，因为它最吃算力，优化空间也最大。

好了，这一章的内容就到这里。记住，理解算子的计算特征，是做指令集适配的第一步。下一章我们会深入讨论如何把这些算子映射到具体的硬件指令上。

---