第一章
引言——为什么做这件事

推荐系统的算力黑洞

现代推荐系统离不开Embedding层，Embedding表动辄千万甚至上亿行，每次训练迭代都要完成海量反向传播计算。这是推荐系统训练的头号算力瓶颈，原因有三：

GM带宽饥渴：高稀疏key重复场景下，同一个key被不同样本共享，GM读写量膨胀数倍。

原子操作风暴：多个样本共享同一个key时，所有线程争抢同一个GM地址的更新权，流水线停滞。

优化器计算碎片化：Adagrad、Adam、SGD等优化器逐行逐元素计算，无法充分利用向量引擎并行能力。

目标

Meta开源的 [fbgemm](https://github.com/pytorch/FBGEMM) 是推荐系统 Embedding训练的事实标准，在Atlas 350 SuperSuite UBL4超节点套件上达到业界 + fbgemm的2倍性能。

Atlas 350 SuperSuite UBL4超节点套件提供了独特的硬件优势——大容量UB、强大的Vector引擎、灵活的SIMT编程模型，使算法层面优化成为可能。我们从“减少GM原子竞争”、“消除冗余GM读写”和“批量向量化计算”三个方向出发，进行了系统性架构重构。

第二章 整体架构

端到端数据流

图1：Embedding 反向传播端到端数据流 Pipeline

Workspace哈希桶原子标记去重机制

同一个batch内多条样本可能引用相同hash key，不做去重会导致冗余计算和写入冲突。我们的方案是基于哈希桶的原子状态标记去重机制：将hash key映射为Workspace数组下标，利用GM AtomicCas进行原子竞争仲裁，实现跨Block的分布式去重。

图2：Workspace 哈希桶原子标记去重机制

Workspace分为两个连续区域：

多个Block同时处理同一key时，发起AtomicCas(CLEAR → NEED_UPDATE)，仅一个Block成功成为“唯一持有者”走Unique路径直接写入梯度，其余走Non-Unique路径通过AtomicAdd累加。这本质是一个以hash key为锁名的全局分布式锁，一次Kernel调用完成千万级去重仲裁，无需CPU介入。

优化点在流水线中的位置

第三章 优化手段

核心场景：高稀疏Key重复率

推荐系统中，热门视频ID、高频标签、地理位置等特征常跨样本共享，key重复率达4-8x甚至更高，导致梯度反复从GM读取、参数更新产生大量原子竞争。这正是优化的核心瞄准方向。

优化一：UB哈希表Block级去重 —— 原子竞争减少8倍

问题：同一block内不同warp可能处理相同key，导致32个warp对同一GM地址发起CAS竞争。

方案：在UB上分配2048条目哈希表（仅8KB），作为block内跨warp的“已见集合”，只有首次出现的key才走GM CAS。

优化前：GM AtomicCas 直接竞争

优化后：UB Hash Table Block 级去重

图3：优化一前后对比 —— 
GM AtomicCas 直接竞争 vs UB Hash Table Block 级去重

收益：key重复率~8x场景下，GM CAS减少约8x。业界的shared memory容量远小于Atlas 350 SuperSuite UBL4超节点套件的UB，难以复制此方案。

优化二：Warp内Tree-based Leader Election —— 32lane → 1次CAS

问题：同一warp内32个lane仍可能持有相同key，每个lane独立发起CAS。

方案：使用asc_shfl_xor蝴蝶交换（5轮），在warp内选举唯一winner，仅winner执行CAS。

图4：Warp 内 Tree-based Leader Election 蝴蝶选举流程

收益：同一key每个warp内最多1次AtomicCas，而非32次。asc_shfl_xor在warp内交换，O(logN)复杂度完成全warp选举。

优化三：统一走Rowwise计算流程 —— 消除87.5%冗余GM读

这是本次重构中收益最大的单项优化。

问题：Adagrad/Adam/SGD 内核使用ClearGrad() + ComputeGrad() 流程，遍历所有原始索引（含重复），key重复8次就从GM读取8次。

方案：统一采用ProcessIndices<IsUniq>(true) + ProcessIndices<IsUniq>(false)模式，分开处理唯一/非唯一key，编译期确定回写模式，零运行时开销。

优化前：ClearGrad + ComputeGrad

优化后：ProcessIndices 分离处理

图5：优化三前后对比—— 
ClearGrad+ComputeGrad vs ProcessIndices 分离处理

收益：每个唯一key仅从GM读取一次梯度，重复率8x场景下GM读减少约87.5%。代码量从215行缩减至42行。

优化四：Rowwise Adagrad 向量化 —— 逐元素变为批量向量

问题：Rowwise Adagrad逐行逐元素计算，6步操作对每行每元素重复执行。

方案：三阶段批量向量化：(1) 逐行ReduceSum规约；(2) 对N个标量做Sqrt/Reciprocal；(3) Broadcast + 单次Mul完成全批量更新。

图6：Rowwise Adagrad 向量化计算重构对比

收益：O(N×D) 次逐元素运算转为批量Mul + O(N) 次标量运算。Atlas 350 SuperSuite UBL4超节点套件Vector引擎对逐元素规约类操作的效率远高于业界。

优化五：UniqCleanGT选择性清理 —— 大表清理量降低99.99%

问题：ClearGT清理整个hash空间的标记位，10M嵌入表意味着40MB全量清零。

方案：SimtUniqCleanGT仅清零实际被当前batch命中的key槽位，Host侧自动决策大小表策略。

图 7：ClearGT 全量清理 vs UniqCleanGT 选择性清理对比

收益：大表场景清理量降低99.99%。

优化六：SGD全批量向量化 —— 一行代码的性能跃迁

问题：SGD逐行调用Muls，cnt次小规模调用。

方案：单次全批量调用Muls<float>(outLt, inputLt, minusLearningRate, cnt * maxD)。

收益：1次大规模调用，向量引擎批量利用率大幅提升，代码从5行缩减至1行。

优化七：扩展索引步长TRIAD → EXTENDED

方案：索引步长从3扩展至5，新增INPUT_OFFSET（避免重复计算GM偏移）和BAG_LENGTH（支持MEAN_POOL直接计算均值）。

图8：索引步长扩展 TRIAD(3) → EXTENDED(5)

第四章 总结

整体性能收益

在高重复率（8x+）、大嵌入表（> 1M）、Rowwise优化器场景下：

指标

优化幅度

总执行时间

降低 40-60%

GM 带宽消耗

降低 70-80%

原子操作争抢

降低 90-95%

图9：整体性能收益柱状图

方法论总结

图10：方法论总结 —— 三条主线汇合实现端到端 2x 性能提升

三条主线覆盖了三个核心瓶颈：原子竞争拖慢更新、冗余访存浪费带宽、逐元素计算效率低。7项优化环环相扣，从算法到实现、从GM到UB、从warp级到block级，形成系统性性能提升。

核心成果：

GM CAS减少 ~8x：UB哈希表 + Warp Leader Election两级级联抑制原子竞争

GM读数减少 ~87.5%：ProcessIndices去重流程，每个唯一key仅读取一次

Rowwise/SGD向量化2-4x加速：充分利用Vector引擎批量计算

大表清理量降低99.99%：UniqCleanGT仅清零命中槽位

Atlas 350 SuperSuite UBL4超节点套件的大容量UB、Vector引擎和SIMT编程模型是达成这些优化的硬件基础，优化手段本身具有通用性，可推广至同类算子优化。

端到端效果

综合7项优化，Atlas 350 SuperSuite UBL4超节点套件上Embedding反向传播算子的端到端性能达到业界的2倍。

代码仓库：

https://gitcode.com/Ascend/fbgemm-ascend