【CANN训练营学习笔记】Transformer算子专场——大模型时代的算子优化之道

训练营简介

2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro

📖 本文思维导图

mindmap
  root((Transformer算子))
    OpenTransformer
      覆盖场景
        FNN
        Attention
        MOE
        Window Attention
      设计原则
        生态兼容
        高内聚低耦合
        隔离原则
      算子分类
        基础算子
        MHA算子
        Gated MHA
        FNN算子
        MOE算子
        TOPK算子
    Dispatch和Combine
      MOE通信挑战
      共享内存方案
        IB+UBM
      关键优化
        双分区
        量化处理
        双缓冲
    FIA算子
      Flash Attention
      通用优化
        负载均衡
        核间流水
        核内流水
        数据驻留
      MLA适配
        Amable算法
        NZ格式
        访存优化
    GMM算子
      分组矩阵乘法
      低精度优化
        INT8/INT4
        量化方案
        MSC优化
      深度融合
        生产者消费者
        软同步机制

一、踏入大模型算子优化的殿堂

作为CANN算子开源周系列直播的最后一场，Transformer算子专场可以说是压轴大戏！今天有幸听到四位技术大牛的分享：刘丹老师介绍OpenTransformer算子仓、陈建军老师讲解Dispatch和Combine优化、黄伟老师分享FIA算子在MLA中的应用、余阳老师剖析低精度GMM算子设计。作为一个大模型研发人员，这次学习让我对Transformer算子的底层实现有了全新的认识。

二、OpenTransformer算子仓：大模型算子的集大成者

2.1 算子仓的使命与定位

刘丹老师开场就明确了OpenTransformer的目标：为Transformer类大模型提供高性能算子库

这个定位太重要了！现在AI领域基本上被Transformer统治：

• NLP领域：BERT、GPT、ChatGPT
• CV领域：ViT、CLIP、DALL-E
• 多模态：Flamingo、GPT-4

所有这些模型的核心都是Transformer架构，算子优化直接决定了模型的推理速度。

2.2 覆盖场景全面

OpenTransformer覆盖了四大核心场景：

1. FNN（Feed Forward Network）

• Transformer的前馈网络部分
• 通常包含两个线性层和一个激活函数

2. Attention（注意力机制）

• Transformer的灵魂
• Self-Attention、Multi-Head Attention等

3. MOE（Mixture of Experts）

• 大模型的新宠
• 通过专家网络提升模型容量

4. Window Attention

• 用于处理长序列
• 降低注意力计算复杂度

2.3 设计三原则

刘老师特别强调了三个设计原则，我觉得非常有见地：

原则1：生态兼容

• 尊重社区主流使用习惯
• 接口设计参考PyTorch等框架
• 降低迁移成本

原则2：高内聚低耦合

• 功能相似的算子集中管理
• 模块间依赖关系清晰
• 便于维护和扩展

原则3：隔离原则

• 稳定算子与活跃算子隔离管理
• 保证核心算子的稳定性
• 新算子快速迭代

2.4 算子分类体系

OpenTransformer将算子分为六大类：

1. 基础算子库
2. MHA（Multi-Head Attention）算子库
3. Gated MHA算子库
4. FNN相关算子
5. MOE类计算算子
6. TOPK选择算子

每个算子都包含五大交付模块：

• 头文件：接口定义
• Python绑定：调用方式
• 图模式接口：类型和形状推导
• Host端代码：CPU侧逻辑
• Device端代码：NPU侧实现

三、Dispatch和Combine：MOE架构的通信优化

3.1 MOE架构的挑战

陈建军老师的分享让我对MOE有了更深的理解。MOE（Mixture of Experts）架构的核心是：

• 动态数据分配
• 多专家并行计算
• 输出整合

问题：当专家并行规模扩大时，通信开销急剧增加！

传统的All-to-All通信方式效率不高，成为性能瓶颈。

3.2 共享内存通信方案

陈老师提出了一个创新的解决方案：使用IB+UBM进行共享内存通信

通信四步骤：

1. 发数据：内存拷贝到远端数据区
2. 写状态：标识数据发送完成
3. 等状态：等待其他卡数据发送完成
4. 读数据：从共享内存读取数据

这让我想起了操作系统中的生产者-消费者模式！

3.3 Dispatch实现方案

内存分区设计：

共享内存
├── 数据共享区
│   ├── 量化后的token数据
│   └── 动态量化参数
└── 状态共享区
    └── 通信状态标志位

数据处理流程：

第一步：数据拼接

• 量化处理
• 数据信息打包

第二步：地址映射

• 专家编号 → Rank ID
• 确定目标卡位置

第三步：数据发送

• 通过IB+UBM发送
• 写入远端共享内存

第四步：状态同步

• 多核同步
• 确保数据完整性

3.4 Combine实现方案

Combine是Dispatch的逆过程，用于整合多个专家的输出。

内存排布：

按Batch连续排布
├── M个专家数据（前）
└── 共享专家数据（后）

处理流程：

1. 循环发送数据和状态位
2. 等待Batch对应状态位
3. 数据搬出和加权融合计算
4. 共享专家数据累加

3.5 关键优化技术

优化1：双分区方案

问题：算力运行速度不一致导致数据踩踏

什么叫数据踩踏？就是前一轮的数据还没读完，下一轮的数据就覆盖上来了！

解决方案：

• 共享内存分为0区和1区
• 交替使用
• 避免数据冲突

内存分配：

• 数据区：200MB
• 状态区：包含标识区

优化2：量化处理

量化时机：通信前量化

好处：

• 减少数据量
• 降低通信开销
• FP16量化为INT8，数据量减半！

量化方式：Block量化

• 部分数据作为一次量化单元
• 动态量化参数随数据发送
• 平衡精度和压缩率

优化3：双缓冲技术

这个技术在之前的分享中也出现过，在这里同样适用：

缓冲区A: 量化计算 | 缓冲区B: 数据拷贝
缓冲区A: 数据拷贝 | 缓冲区B: 量化计算

实现流水并行，提升性能！

四、FIA算子：Flash Attention的昇腾实现

4.1 FIA算子概述

黄伟老师的分享聚焦在FIA（Flash Attention Inference Algorithm）算子上。

基础：经典Flash Attention算法的融合算子

计算流程：

1. QK MatMul：Query和Key的矩阵乘法
2. SoftMax：注意力权重计算
3. Value1计算：Attention Score和Value的矩阵乘法
4. Value2计算：输出投影的矩阵乘法

传统实现中，这四个步骤是分开的，中间结果需要反复读写内存。Flash Attention的核心思想就是融合这些操作，减少内存访问！

4.2 通用优化手段

优化1：全局负载均衡

切分维度：

• head维度切分
• batch维度切分
• sequence维度切分

并行策略：不同分块并行在不同核上运行

这让我想起了并行计算的基本原理——把大任务分解为小任务，并行执行。

优化2：核间流水优化

问题：Cube和Vector计算单元存在空闲浪费

解决方案：

• 流水线排布
• MHA计算提前发射
• Cube和Vector并行执行

优化3：核内流水优化

技术：双缓冲技术

应用范围：

• Cube计算
• Vector计算

效果：单个核内计算排布优化

优化4：数据驻留优化

原理：Q数据块重复利用

在Attention计算中，Q数据需要和所有的K数据做矩阵乘法，如果每次都重新加载Q数据，会浪费大量带宽。

实现：Q数据驻留在L1内存中

效果：Q数据只加载一次，避免重复搬移

4.3 MLA结构的特殊适配

MLA技术特点：

• Multi-Latent Attention
• 降低KV Cache存储量
• 新增旋转位置编码数据

针对性优化1：Amable算法

背景：MLA结构中head维度较大（512），历史Attention数据内存占用大

算法原理：
$$\text{FP32乘法}\Rightarrow \text{INT32加法}$$

这个转换很巧妙！利用了浮点数的特性进行等价变换。

硬件利用：利用昇腾芯片的原子加能力

效果：减少GM访存，提升计算效率

针对性优化2：数据驻留优化

挑战：Q+QR的tiling块较大（256×1576）

这么大的数据块，Local Memory放不下！

解决方案：

• tiling块局部驻留
• K维度切分，Q块驻留L1
• M维和K维切分，小数据块驻留

针对性优化3：访存效率优化

技术：NZ数据格式

什么是NZ格式？

• N（行）和Z（特殊维度）的数据排布
• Cube计算友好
• 避免数据格式转换

实现：直接以NZ格式读取KVCache

好处：

• 减少格式转换开销
• 提高Cube利用率

针对性优化4：负载均衡优化

技术：FlashAttention的分块技术

原理：

• 精细化评估每个tiling块的计算量
• 均衡算法分配计算到不同核
• 对KV数据切分，均匀分配

五、GMM算子：MOE的性能核心

5.1 GMM算子基础

余阳老师的分享聚焦在GMM（Grouped Matrix Multiplication）算子上。

定义：分组矩阵乘法

应用场景：MOE网络专用算子

网络位置：MOE FFN模块

重要性：在训练和推理中热点占比20%-30%！

基本形式：
$$C_{m\times n}=A_{m\times k}\times B_{k\times n}$$

GMM特性：

• 除了M、K、N外，还有group_list参数
• 每个group对应一个专家
• 不同group的M或K维度可能不同

变种类型：

• 7M GMM：M维度变化（不同专家处理的token数不同）
• 7K GMM：K维度变化

5.2 基本算子方案

分核策略：

在M和N两个轴上分核：

专家1: M1×K×N
专家2: M2×K×N
...
专家E: ME×K×N

每个专家的token数（M维度）不同，需要单独切分。

任务分配：每个子任务负责 tile_M × 1K × 1 与 1K × tile_N 的矩阵乘法

优化技术：

1. 动态调整tile_M和tile_N大小优化访存
2. 向量全带（head较小时）
3. 4D格式在特定场景下的应用
4. 低精度算法

算子外优化：

1. R2 Cache预热：推理阶段带宽瓶颈优化
2. 纯量化：抽成纯量化或双流方向
3. CV双流并发：路由专家和共享专家计算并发

5.3 低精度优化方案

量化支持：

数据类型：

• INT8：8位整数
• INT4：4位整数

量化粒度：

• PerTensor：整个Tensor共享量化参数（粒度粗，精度低，计算简单）
• PerToken：每个Token单独量化（粒度细，精度高，计算复杂）

量化组合：

• A8W8：Activation INT8, Weight INT8
• A4W4：Activation INT4, Weight INT4

全量化方案：

流程：

量化数据 → 矩阵乘法(INT运算) → 反量化 → FP16/FP32结果

数据类型支持：

• FP32：利用FixCast随路转化能力
• BF16：需要完成PerChannel反量化

伪量化方案优化（MSC方案）：

问题：A8W4场景性能上限低

为什么？因为INT8和INT4的矩阵乘法在硬件上支持不够好。

MSC优化思路：

1. 将INT8拆分为2个INT4
2. 分别进行矩阵乘法
3. 合并结果

公式表示：
KaTeX parse error: Expected 'EOF', got '_' at position 45: …} = \text{(INT4_̲high × 16 + INT…

优势：

• 避免了A8W4的低效实现
• 性能接近A4W4
• 灵活性更好

5.4 深度融合方案

融合机会：

在MOE FFN中，通常有这样的计算流程：

GMM1 → Silu激活 → GMM2

如果GMM1和GMM2都是低精度量化的，中间还需要反量化和量化操作。

融合痛点：

1. 性能不足：Activation完全执行完才能做后续处理
2. 显存浪费：中间变量数据类型不匹配

深度融合方案：

核心思想：生产者-消费者模式

技术实现：

WedgeWait设置：

• 设置64KB大小的块
• 平衡头尾开销和同步开销

软同步机制：

• 标志位轮询
• 避免全局同步

状态压缩：

• 用bit表示每个flag
• Int32表示所有核的状态
• 高效的状态管理

原子操作：

• TME方式写入
• 避免读写冲突

效果：

• 中间数据直接处理，不写回内存
• 显著降低访存开销
• 提升整体性能

六、社区任务与激励

6.1 任务内容

刘老师最后发布了一个社区任务：Transformer自定义算子开发任务

时间要求：一个月内完成

参考标准：PyTorch官方算子接口

交付内容：

• 完整工程代码
• README文档
• 测试套件

6.2 评审标准

验收用例：10个验收组

关注点：

• 精度表现
• 性能表现

提交方式：可重复提交，48小时内反馈结果

6.3 奖项设置

第一名：华为手表（价值5999元）
第二名：平板电脑
第三名：耳机

这个奖励还是很丰厚的！重点是能学到真东西。

七、互动问答精华

问题1：不同硬件场景性能差异如何分析？

回答：

• 影响因素：算力和带宽硬件参数
• 瓶颈分析：
  • 加法算子：带宽瓶颈
  • MatMul：算力瓶颈
• 瓶颈转移：硬件参数变化可能导致瓶颈转移

问题2：Cube和Vector通信如何设计？

回答：

• 硬件架构：Cube和Vector是分离的
• 通信路径：必须通过Global Memory通信
• 带宽优势：Global Memory带宽是直接访问的3倍以上
• 大缓存优势：192MB A2 Cache优势明显

这个回答让我明白了为什么要通过GM通信——虽然绕了一圈，但带宽更高！

问题3：开源算子的范围是什么？

回答：

• 推理场景：FIA算子
• 训练场景：训练FIA算子和反向FIA算子
• 扩展算子：MLA Prologue等预处理算子

问题4：如何获取技术进展信息？

回答：

• SIG组运作：4个组的日常运作
• 邮件订阅：及时获取信息
• 会议参与：活动会议和纪要
• 社区讨论：SIG组内技术讨论

八、我的学习感悟

8.1 Transformer算子优化的系统性

这次学习让我认识到，Transformer算子优化是一个系统工程：

• 不仅要优化单个算子（如MatMul）
• 还要考虑算子间的融合（如FIA）
• 还要考虑通信优化（如Dispatch/Combine）
• 还要考虑量化等压缩技术

8.2 硬件软件协同优化

所有的优化都是基于对硬件的深刻理解：

• Cube和Vector的特性
• 多级缓存的带宽差异
• 原子操作的支持
• 数据格式的影响

只有深入了解硬件，才能做出正确的软件优化决策。

8.3 大模型时代的新挑战

MOE、MLA等新架构带来了新的挑战：

• 动态的数据分配
• 巨大的KV Cache
• 通信开销的增加

这些都需要新的优化技术来解决。

8.4 开源社区的力量

四位老师都强调了开源社区的重要性：

• 集众人之智慧
• 覆盖更多场景
• 共同推进生态发展

我也计划参与到开源社区中，贡献自己的力量。

九、后续学习计划

基于这次学习，我制定了详细的学习计划：

第一阶段：基础掌握（1个月）

1. 深入学习Flash Attention算法原理
2. 实现一个简单的Attention算子
3. 学习MOE架构的基本原理

第二阶段：优化实践（2个月）

1. 尝试实现FIA算子的基本版本
2. 应用数据驻留等优化技术
3. 学习量化算法并实践

第三阶段：社区贡献（持续）

1. 参与社区任务
2. 优化现有算子
3. 分享学习心得

十、总结

这次CANN算子开源周Meetup的Transformer算子专场可以说是整个系列的华彩乐章！四位老师从不同角度展示了大模型时代算子优化的精髓。

核心要点回顾：

1. OpenTransformer：为Transformer大模型提供完整的高性能算子生态
2. Dispatch/Combine：通过共享内存通信优化MOE架构的通信开销
3. FIA算子：Flash Attention在昇腾平台的高性能实现
4. GMM算子：MOE核心算子的低精度优化和深度融合

优化技术汇总：

• 数据驻留
• 双缓冲流水线
• 负载均衡
• 算子融合
• 量化压缩
• 共享内存通信
• NZ格式优化
• MSC伪量化方案

这些技术不仅适用于Transformer算子，也为其他算子的优化提供了宝贵的参考。

对于想要学习大模型算子优化的同学，我的建议是：

1. 深入理解Transformer架构
2. 掌握基础的算子优化技术
3. 学习硬件架构特性
4. 参与开源社区实践
5. 持续跟踪技术前沿