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

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

前言

在“搜广推”领域，模型的参数量往往比 LLM 还要大（TB 级别），因为 Embedding Table 极其巨大。 为了训练加速，我们通常采用 Sparse Update（稀疏更新）：只更新本 Batch 涉及到的那些 ID 的权重。

这里有一个痛点：如果一个 Batch 里有 1000 个样本都点击了“iPhone”，那么 ID iPhone 会出现 1000 次。

• 笨办法：从 GM 读 1000 次权重，算 1000 次梯度，写回 1000 次。血亏！

• 聪明办法：先做 Unique（去重），发现只有 1 个独立 ID。读 1 次，把 1000 个梯度聚合（Reduce）起来，算 1 次，写回 1 次。

本期我们将挑战如何在 Ascend C 上实现并行 Unique 和 Hash Lookup。

一、 核心图解：快递分拣流水线

处理 ID 类数据，就像是快递分拣。

二、 算法原理：Sort-based Unique

在 GPU/NPU 上做 Unique，最经典的方法是 Sort + Adjacent Difference。

1. Sort: 对 ID 数组进行全排序。

   • 输入：[3, 1, 3, 2, 1]

   • 排序后：[1, 1, 2, 3, 3]

2. Compare: 比较 data[i] 和 data[i+1]。如果不相等，说明 i+1 是一个新的 ID 的开始。

   • Mask: [0, 1, 1, 0, 1] (假设 1 代表边界)

3. Scan (Prefix Sum): 对 Mask 做前缀和，得到去重后的新 Index。

   • Index: [0, 1, 2, 2, 3]

4. Scatter: 把数据写到新位置。

我们在第 41 期学过 Bitonic Sort，在第 42 期学过 Parallel Scan。Unique 算子就是这两个算子的组合拳！

三、 实战：Ascend C 实现 Unique

3.1 Kernel 类定义

输入 x (IDs)，输出 y (Unique IDs), idx (Inverse Indices), counts。

class KernelUnique {
public:
    __aicore__ inline void Process() {
        // 1. CopyIn
        DataCopy(idsLoc, xGm, len);
        
        // 2. Sort (Bitonic Sort)
        // 详见第 41 期，这里直接调用高阶 API 或手写
        Sort(idsLoc, idsLoc, len);
        
        // 3. Find Unique
        ComputeUnique();
        
        // 4. CopyOut
        DataCopy(yGm, uniqueLoc, numUnique);
    }
};

3.2 Compute 核心逻辑

__aicore__ inline void ComputeUnique() {
    // idsLoc 已经有序：[1, 1, 2, 3, 3]
    
    // 1. 构造错位数组
    // idsShift: [X, 1, 1, 2, 3] (右移一位)
    // 实际可以通过 DataCopy 或 Muls 偏移实现
    
    // 2. 比较 (Compare)
    // mask = (idsLoc != idsShift)
    // 结果: [1, 0, 1, 1, 0] (1 表示是新元素的开始)
    // 注意：第一个元素永远是 1
    Compare(maskLoc, idsLoc, idsShift, CMP_NE, len);
    
    // 3. 前缀和 (Parallel Scan) - 详见第 42 期
    // index = Scan(maskLoc)
    // 结果: [0, 0, 1, 2, 2] (这就是去重后的下标)
    
    // 4. 收集 (Gather/Scatter)
    // 我们只把 maskLoc[i] == 1 的那些 idsLoc[i] 挑出来
    // 这里需要利用 mask 作为 predicate
    // 但更高效的是利用 Select 或 Compact 指令
    
    // 伪代码：Compact(dst, src, mask)
    // 将 src 中 mask 为 1 的元素紧凑地写到 dst
    // Ascend C 某些版本支持 GatherMask 或类似指令
    // 如果不支持，需要手写 Scan + Scatter
    
    Scatter(uniqueLoc, idsLoc, indexLoc, len); // 逻辑示意
}

3.3 进阶：Hash Lookup (Open Addressing)

如果 ID 范围极大（比如 $10^{18}$），无法排序怎么办？或者需要实现一个动态插入的 Hash Table？

这需要实现 开放寻址法 (Open Addressing)。

1. Hash: idx = id % capacity。

2. Probe: 检查 table[idx] 是否为空或等于 id。

   • 如果占用且不等，idx = (idx + 1) % capacity (线性探测)。

3. Insert/Get: 写入或读取。

Ascend C 实现难点：

• 需要 AtomicCAS (Compare-And-Swap) 指令来实现无锁并发写入。

• 需要 While 循环处理冲突探测。

// 伪代码：线性探测 Hash
__aicore__ inline void HashInsert(int64_t id) {
    uint32_t idx = Hash(id) % CAP;
    while (true) {
        // 原子读取旧值
        int64_t old_id = AtomicGet(tableGm[idx]);
        
        if (old_id == EMPTY) {
            // 尝试抢占
            if (AtomicCAS(tableGm[idx], EMPTY, id)) break; // 成功
        } else if (old_id == id) {
            break; // 已存在
        }
        
        // 冲突，探测下一个
        idx = (idx + 1) % CAP;
    }
}

四、 性能优化：Memory is Everything

推荐系统算子是典型的 IO Bound。

1. 合并请求：尽可能把多次 Lookup 合并成一次 Gather。

2. Hot/Cold 分离：

   • Hot ID（热门商品）：放在 L2 Cache 或片上内存（如果放得下）。

   • Cold ID（冷门商品）：老老实实去 HBM 捞。

   • 这需要 Host 侧配合做频率统计和重排。

五、 总结

处理推荐系统的算子，就像是在处理数据库的逻辑。

1. Sort-based Unique：适合 Batch 内去重，利用了 Vector 的排序和扫描能力。

2. Hash Table：适合全局去重和动态 Embedding，利用了原子操作和标量循环。

掌握了这两个算子，你就有能力去优化那些参数量达到 TB 级 的巨型推荐模型。