走向专精：我的NLP特化算子开发之旅

训练营简介

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

报名链接：https://www.hiascend.com/developer/activities/cann20252

前言

从通用型的图像分类模型优化项目结束后，我开始思考：是否应该在某个特定领域深化？直到有一天，我被邀请参加某互联网公司的大模型推理优化项目，我才意识到——NLP模型的算子特性和CV模型完全不同。

这次经历让我从"全能的算子开发者"转变为"NLP领域的专家"。这篇文章记录了我如何在一个领域走向专精的过程。

从CV到NLP的思维转变

第一次遭遇的困惑

接到NLP项目时，我满信心地拿出自己优化好的Conv2D、MatMul算子套件，以为可以直接用。

结果却发现，大模型推理的关键瓶颈完全不是Conv2D！Profiling结果显示：

总耗时：200ms
├─ MatMul（线性层）：80ms (40%)
├─ Attention（自注意力）：70ms (35%)  ← 这是啥？
├─ LayerNorm（归一化）：30ms (15%)
└─ 其他算子：20ms (10%)

我之前从未专门优化过Attention和LayerNorm！而这两个算子合起来占了50%的时间。在开发者说专题中听到的一句话浮现在脑子里：

“不同领域的模型有完全不同的算子热点。深入一个领域，比什么都通学得更深有价值。”

这时我才明白专精的重要性。 💭

LayerNorm的性能陷阱

看似简单的算子

LayerNorm是一个看起来很简单的操作：

y = (x - mean) / sqrt(var + eps) * gamma + beta

只是四则运算，应该很快啊？实际测试后发现性能只有35%，效率低得惊人。

为什么这么慢？

深入分析后，我发现问题出在两个地方：

问题1：多次扫描数据
标准实现需要扫描两遍数据——第一遍计算均值和方差，第二遍做归一化。但我的Tiling方案导致同一块数据被重复加载。

问题2：同步点太多
均值计算完后需要等待，方差计算完后又需要等待。这些同步点限制了硬件的并行度。

算法级别的优化

在社区讨论中，一位资深贡献者分享了一个关键思路：Fused Kernel。

我重新设计了算法，把原来的两遍扫描融合成一遍。关键trick是使用Welford算法计算在线方差，这样只需一遍数据扫描就能同时得到均值和方差。

优化后性能从35%提升到72%！🚀

进一步的优化空间

通过Profiling发现向量化还不充分。我利用CANN提供的2D访问指令，把多个元素的操作并行化。

最终LayerNorm算子达到了78%的峰值算力，时间从30ms降到了16ms。

Attention机制的复杂性

核心瓶颈的识别

Attention是大模型的灵魂，但也是性能杀手。标准的Self-Attention计算流程是：

Q = X @ W_q
K = X @ W_k  
V = X @ W_v
Attention = softmax(Q @ K^T / sqrt(d_k)) @ V

看起来就是几个MatMul，但实际上Attention有独特的特性：内部维度很小，但序列长度很长。

比如，一个序列长度为2048的文本，embedding维度才768。这个特性导致标准的矩阵优化策略不适用。

Multi-Head的额外挑战

BERT、GPT这样的模型都用Multi-Head Attention。8个或12个头并行计算，这给并行度带来了新的维度。

我在码力全开特辑中学到的融合技巧在这里派上用场。关键想法是：把多个head的计算融合在一起，减少内存往返。

从零开始实现高效Attention

经过研究和反复试验，我实现了以下优化：

优化1：QKV融合计算
原来Q、K、V是分别计算的。我把三个MatMul融合成一个矩阵乘法，减少数据搬运。

优化2：Online Softmax
原来需要先算Q@K，再整体做softmax，最后和V相乘。改为在计算过程中边算边归一化，避免中间结果的显存爆炸。

优化3：Multi-Head并行
充分利用Cube单元和Vector单元的并行能力，让多个头真正并行执行而不是串行。

优化4：Flash Attention思想
借鉴Flash Attention的思想，减少访存次数。虽然完全实现Flash Attention比较复杂，但核心思想——减少HBM访问次数——可以借鉴。

经过这轮优化，Attention从原来的70ms降到了28ms，峰值算力从42%提升到74%！✨

Softmax的数值稳定性与性能的平衡

陷入的困境

优化Attention时，我需要频繁计算Softmax。但标准的Softmax实现e^x / sum(e^x)容易溢出。

我实现了数值稳定的版本：e^(x-max) / sum(e^(x-max))。这个版本更稳定，但多了一次数据扫描，性能反而下降了。

在精度和性能间的平衡

这是我第一次真正面对"鱼和熊掌"的选择：

• 标准Softmax：快，但不稳定
• 稳定Softmax：慢，但精准

在某个社区讨论中，我看到一位工程师的建议：根据input range动态选择。

我实现了这个想法：

• 如果input数据范围小（max-min < 50），用快速版本
• 如果范围大，用稳定版本

这样既保证了稳定性，在大多数情况下也保证了性能。

Token Embedding的特殊优化

NLP独有的需求

NLP模型的一个独特之处是Token Embedding。原始的embedding lookup看起来很简单，但在大词表（几万到几百万）的情况下，有独特的性能特性。

从Embedding到优化

标准的embedding lookup是一个索引操作：

output = embedding_table[token_ids]

这个操作的特点是：随机访问、数据量小、内存带宽不是主要瓶颈。

基于这个特性，我设计了以下优化：

优化1：批量预取
提前预测可能的token序列（基于概率或上下文），预先加载对应的embedding向量。

优化2：量化存储
使用INT8存储embedding table，大幅减少显存占用，加快加载速度。只在计算时转为FP32。

优化3：局部缓存
热点token的embedding保持在L1缓存中，避免反复加载。

经过这些优化，embedding部分的延迟从原来的15ms降到了8ms。

建立NLP算子的完整工具链

从单算子到完整解决方案

经过几个月的优化，我意识到不能把这些优化分散开，需要建立一个完整的工具链：

工具1：NLP算子库
把所有优化好的NLP相关算子（LayerNorm、Softmax、Attention、Embedding等）整理成一个模块化的库。

工具2：模型编译器
自动识别模型中的算子，选择最优的实现版本，生成高效的执行计划。

工具3：性能分析器
提供NLP特定的性能分析视图，帮助开发者快速定位瓶颈。

这套工具链的开发花了一个月，但效果显著——新的NLP模型只需要简单配置就能获得优化。

从独行到团队协作

意外的邀请

工具链完成后，我把它贡献给了社区。意外的是，这引起了某大厂算法负责人的关注。他邀请我参加他们的NLP优化项目——作为技术负责人。

这是我第一次从单独贡献者变成项目领导者。😲

带领团队的挑战

从一个人写代码到带领5个人的团队，挑战远比想象大：

• 怎样清晰地传达技术方案？
• 怎样让团队成员快速上手CANN？
• 怎样在紧凑的时间表下保证质量？

我意识到，分享和教学是比单独优化更高层次的技能。

建立知识体系

我开始系统地整理CANN在NLP领域的最佳实践：

文档体系
详细的算子设计文档、性能优化指南、常见问题解答。

培训体系
内部培训课程，让新人能快速理解NLP算子的特性和优化思路。

工程规范
代码规范、测试规范、性能评测规范，确保整个团队的输出质量一致。

专精领域的职业转变

从通才到专家

通过深化NLP领域的优化，我的职业路径发生了清晰的转变。

面试时，不再是笼统地讲"我优化过各种算子"，而是能深入讲解NLP模型的性能特性、关键的优化思路、团队的工程实践。

这种专精得到了高度认可。我最终加入的公司正是因为看重我在NLP优化领域的深度。

学术与工程的结合

专精的另一个好处是能与学术结合。我开始研究最新的NLP优化论文（如Flash Attention、PagedAttention等），把学术思想应用到实践中。

这种学术与工程的结合，让我的工作有了更高的层次。

给想走向专精的建议

建议1：选择有潜力的领域
选择一个：问题空间大、有广阔应用前景、技术热点多的领域。NLP、推荐系统、多模态都是不错的选择。

建议2：深度优于广度
一开始可以尝试多个领域，但要坚定地选择一个来深化。深度的专业知识往往比浅尝辄止更有价值。

建议3：构建知识体系
不仅优化代码，更要整理方法论。写文章、做分享、建立最佳实践库。

建议4：参与开源建设
贡献你的优化成果到社区。不仅能帮助他人，也能获得反馈和认可。

建议5：从实践到理论
在深入优化的过程中，学习相关的学术理论。理论能指导实践走得更远。

总结

从CV通才到NLP专家，这段转变让我体会到了专精的力量。不仅技术能力提升了，职业前景也豁然开朗。

CANN训练营的系统化课程为这段深化之路打下了基础。更重要的是社区提供的开放平台，让我能自由地选择专精方向，不受限制地深化。

如果你已经掌握了CANN基础，我的建议是：选择一个感兴趣的领域，坚定地走下去。从通才到专家，往往就差这一步决定。💪

期待在CANN社区看到更多领域专家的出现！

---