CANN ops-nn 算子解读：AIGC 多模态模型中的 Softmax 与 LayerNorm 实现

本文基于 CANN ops-nn 仓库中的 Softmax 和 LayerNorm 算子，解析其在 AIGC 多模态模型（如 CLIP、BLIP）中的关键作用与优化实现。

---

一、多模态 AIGC 与归一化算子

1.1 多模态 AIGC：让 AI 同时理解图像和文字

"看图说话"、"以文生图"------这些曾经只有人类才能做到的事情，如今 AIGC 多模态模型也能轻松完成。

多模态 AIGC 正在改变我们与 AI 交互的方式：

• GPT-4V：能够理解图像内容并回答问题
• CLIP：将图像和文字映射到同一空间，支持文生图
• BLIP-2：图文理解与生成的统一模型
• LLaVA：开源的多模态对话模型

这些模型的共同特点是：需要处理图像和文本两种模态的信息，并让它们相互理解、相互影响。
CLIP 架构
图像编码器

ViT
对比学习
文本编码器

Transformer
LayerNorm
Softmax

在多模态模型中，Softmax 和 LayerNorm 的调用频率极高，是除 MatMul 外最关键的算子。CANN ops-nn 仓库提供了这些算子的高效实现，支持 AIGC 多模态应用的部署。

1.2 ops-nn 归一化算子的价值

算子	多模态场景	ops-nn 优化点
Softmax	Attention 权重、对比损失	数值稳定、向量化
LayerNorm	每层输出归一化	融合优化、精度保证
RMSNorm	LLaMA 系列模型	简化计算、高效实现

---

二、ops-nn Softmax 实现解析

2.1 数值稳定的 Softmax

ops-nn 实现了数值稳定的 Softmax，避免溢出：
输入 x
减去最大值

x - max
指数运算

exp
求和
归一化

除以 sum
输出

对 AIGC 的意义：

• CLIP 的对比学习使用大规模 Softmax（batch_size × batch_size）
• 数值稳定性确保训练和推理的一致性

2.2 Flash Attention 中的 Online Softmax

ops-nn 支持 Flash Attention 所需的在线 Softmax 计算：
Online Softmax
分块计算 QK^T
局部 Softmax
在线更新统计量
修正输出
标准 Softmax
计算全部 QK^T
全局 Softmax

内存优化：从 O(N²) 降至 O(N)，支持更长的上下文。

---

三、ops-nn LayerNorm 实现解析

3.1 LayerNorm 计算流程

渲染错误: Mermaid 渲染失败: Parse error on line 4: ... B --> D归一化 (x-μ)/√(σ²+ε) C -----------------------^ Expecting 'SQE', 'DOUBLECIRCLEEND', 'PE', '-)', 'STADIUMEND', 'SUBROUTINEEND', 'PIPE', 'CYLINDEREND', 'DIAMOND_STOP', 'TAGEND', 'TRAPEND', 'INVTRAPEND', 'UNICODE_TEXT', 'TEXT', 'TAGSTART', got 'PS'

3.2 融合优化

ops-nn 支持 LayerNorm 与相邻算子的融合：
融合后
Add_LayerNorm_Linear
融合前
Add
LayerNorm
Linear

AIGC 收益：

• ViT 每层都有 Add + LayerNorm 模式
• 融合后减少 2 次全局内存访问

3.3 RMSNorm 支持

LLaMA 等模型使用 RMSNorm 替代 LayerNorm：
LayerNorm

均值 + 方差
计算量: 2N
RMSNorm

仅方差
计算量: N

ops-nn 提供专门的 RMSNorm 实现，计算量减少约 40%。

---

四、多模态场景优化

4.1 CLIP 对比学习 Softmax

CLIP 的对比损失需要计算图像-文本相似度矩阵的 Softmax：

# 相似度矩阵 [batch, batch]
# ops-nn 优化：
# 1. 分块计算避免内存溢出
# 2. 温度参数融合到 Softmax 中
logits = image_features @ text_features.T / temperature
probs = aclnn_softmax(logits, dim=-1)

4.2 Cross Attention 中的 Softmax

Stable Diffusion 的 Cross Attention 连接图像和文本：
文本 Value 文本 Key 图像 Query 文本 Value 文本 Key 图像 Query MatMul (Q × K^T) Softmax (ops-nn) MatMul (Attn × V)

---

五、性能数据

5.1 Softmax 性能

Shape	数据类型	耗时
1, 12, 1024, 1024	FP16	0.45ms
1, 12, 4096, 4096	FP16	6.2ms
256, 256 (CLIP)	FP32	0.08ms

5.2 LayerNorm 性能

Shape	数据类型	耗时
1, 1024, 768	FP16	0.05ms
1, 4096, 4096	FP16	0.35ms

5.3 多模态模型推理

模型	任务	优化前	优化后
CLIP ViT-L	图像编码	12ms	5ms
BLIP-2	图文理解	85ms	38ms

---

六、开发者实践

6.1 调用 ops-nn Softmax

// 标准 Softmax
aclnnSoftmax(workspace, workspaceSize,
             input, dim, output, stream);

// 带 Scale 的 Softmax（用于 Attention）
aclnnScaledSoftmax(workspace, workspaceSize,
                   input, scale, output, stream);

6.2 调用 ops-nn LayerNorm

aclnnLayerNorm(workspace, workspaceSize,
               input, normalizedShape,
               weight, bias, eps,
               output, meanOut, rstdOut, stream);

---

七、多模态模型架构

7.1 CLIP 架构详解

文本编码器
图像编码器
图像
Patch Embedding
ViT Transformer
图像嵌入
文本
Token Embedding
Transformer
文本嵌入
对比学习

7.2 Softmax 和 LayerNorm 的调用频率

模型	Softmax 调用	LayerNorm 调用
CLIP ViT-L	24×12 = 288	24×2 = 48
BLIP-2	500+	100+
LLaVA	1000+	200+

---

八、ops-nn Softmax 优化技术

8.1 数值稳定实现

输入 x
找最大值 max
x - max
exp
求和 sum
除以 sum
输出

为什么要减最大值：

• 避免 exp(x) 溢出
• 保持数值稳定性

8.2 Flash Attention 中的 Online Softmax

分块计算 QK^T
局部 Softmax
记录局部 max 和 sum
下一块
更新全局统计量
修正之前的输出

收益：内存从 O(N²) 降至 O(N)

---

九、ops-nn LayerNorm 优化技术

9.1 融合实现

融合后
LayerNorm 融合 Kernel
融合前
计算均值
计算方差
归一化
缩放平移

9.2 RMSNorm 对比

特性	LayerNorm	RMSNorm
计算	均值 + 方差	仅方差
参数	γ, β	仅 γ
速度	基准	快 40%
使用模型	BERT, CLIP	LLaMA

---

十、多模态场景优化

10.1 CLIP 对比学习 Softmax

# 图像-文本相似度矩阵
# [batch, batch] 大小
logits = image_features @ text_features.T / temperature

# 需要对两个方向做 Softmax
image_probs = softmax(logits, dim=1)  # 图像到文本
text_probs = softmax(logits, dim=0)   # 文本到图像

10.2 Cross Attention 优化

优化技术	方法	收益
KV 缓存	缓存文本 K/V	减少重复计算
融合 Softmax	与 MatMul 融合	减少访存
Flash Attention	分块计算	支持长序列

---

十一、开发者实践指南

11.1 完整调用示例

#include "aclnn/acl_nn.h"

// Softmax 调用
aclnnStatus softmaxStatus = aclnnSoftmax(
    workspace, workspaceSize,
    input,              // [B, Seq, Vocab] 或 [B, Heads, Seq, Seq]
    -1,                 // dim (最后一维)
    output,
    stream
);

// 带 Scale 的 Softmax (用于 Attention)
aclnnStatus scaledSoftmaxStatus = aclnnScaledSoftmax(
    workspace, workspaceSize,
    input,              // Attention scores
    1.0 / sqrt(dim),    // scale
    output,
    stream
);

// LayerNorm 调用
int64_t normalizedShape[] = {768};
aclnnStatus lnStatus = aclnnLayerNorm(
    workspace, workspaceSize,
    input,              // [B, Seq, Hidden]
    normalizedShape, 1,
    weight,             // [Hidden]
    bias,               // [Hidden]
    1e-5,               // eps
    output,
    meanOut,            // 可选，用于反向传播
    rstdOut,            // 可选，用于反向传播
    stream
);

// RMSNorm 调用 (LLaMA 使用)
aclnnStatus rmsNormStatus = aclnnRmsNorm(
    workspace, workspaceSize,
    input,
    weight,
    1e-6,               // eps
    output,
    stream
);

// CLIP 对比学习实现
void clipContrastiveLoss(
    aclTensor* imageFeatures,   // [B, D]
    aclTensor* textFeatures,    // [B, D]
    float temperature,
    aclTensor* loss
) {
    // 1. 计算相似度矩阵
    aclnnMatmul(workspace, workspaceSize,
                imageFeatures, textFeatures_T,
                logits, 0, stream);
    
    // 2. 温度缩放
    aclnnDiv(workspace, workspaceSize,
             logits, temperature, scaledLogits, stream);
    
    // 3. 图像到文本的 Softmax
    aclnnSoftmax(workspace, workspaceSize,
                 scaledLogits, 1, imageProbs, stream);
    
    // 4. 文本到图像的 Softmax
    aclnnSoftmax(workspace, workspaceSize,
                 scaledLogits, 0, textProbs, stream);
    
    // 5. 计算交叉熵损失
    // labels 是对角线 (正样本)
    aclnnCrossEntropyLoss(...);
}

// Multi-Head Attention with LayerNorm
void mhaWithLayerNorm(
    aclTensor* input,
    aclTensor* output
) {
    // 1. Pre-LayerNorm
    aclnnLayerNorm(workspace, workspaceSize,
                   input, normalizedShape, 1,
                   lnWeight, lnBias, 1e-5,
                   normed, mean, rstd, stream);
    
    // 2. QKV 投影
    aclnnLinear(workspace, workspaceSize,
                normed, qkvWeight, qkvBias, qkv, stream);
    
    // 3. 分离 Q, K, V 并 reshape
    // ...
    
    // 4. Attention Score
    aclnnBatchMatMul(workspace, workspaceSize,
                     Q, K_T, scores, stream);
    
    // 5. Scaled Softmax
    aclnnScaledSoftmax(workspace, workspaceSize,
                       scores, scale, attnProbs, stream);
    
    // 6. Attention Output
    aclnnBatchMatMul(workspace, workspaceSize,
                     attnProbs, V, attnOut, stream);
    
    // 7. 输出投影 + 残差
    aclnnLinear(workspace, workspaceSize,
                attnOut, outWeight, outBias, projected, stream);
    aclnnAdd(workspace, workspaceSize,
             input, projected, 1.0, output, stream);
}

11.2 常见问题与解决方案

问题	原因	解决方案
Softmax 输出 NaN	输入过大	使用数值稳定版本
LayerNorm 精度差	eps 过小	增大 eps
长序列 OOM	Attention 矩阵过大	使用 Flash Attention

---

十二、总结与展望

12.1 核心要点

CANN ops-nn 仓库中的 Softmax 和 LayerNorm 实现具有以下特点：

• 数值稳定：减最大值的 Softmax
• 融合优化：LayerNorm 单 Kernel 实现
• Flash Attention：Online Softmax 支持
• AIGC 适配：针对多模态模型优化

12.2 多模态部署建议

场景	推荐配置	理由
CLIP 编码	FP16 + 融合	速度和精度平衡
长序列 Attention	Flash Attention	内存高效
LLaMA 推理	RMSNorm	速度更快

---

相关链接：

• 🏠 CANN 组织主页：https://atomgit.com/cann
• 📦 ops-nn 仓库地址：https://atomgit.com/cann/ops-nn