与TopK算子实战)
训练营简介 2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。
报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro
前言
在通用计算机科学中，排序算法（如 QuickSort, MergeSort）通常依赖复杂的分支跳转（Branching）。
CPU：有强大的分支预测器，适合处理逻辑判断。
NPU (AI Core)：是极其强悍的“向量吞吐机器”，最讨厌逻辑判断。
这就导致了一个悖论：AI 模型越来越依赖 TopK（如 LLM 的采样、搜索推荐），但 AI 芯片却不擅长排序。
为了解决这个问题，我们需要换一种思维：不要让硬件去“判断”该不该交换，而是设计一个固定的“交换网络”，让数据流过这个网络，出来时自然就是有序的。
这就是双调排序（Bitonic Sort）的核心思想。本期文章，我们将利用 Ascend C 的 Vector 单元，手写一个并行排序算子。
一、 核心图解：排序网格的魔法
双调排序不依赖数据的值来决定控制流，它的比较路径是固定的。
二、 算法原理：从双调序列到有序序列
2.1 什么是双调序列 (Bitonic Sequence)？
一个序列如果是“先单调递增，再单调递减”（或者通过循环移位能变成这样），它就是双调的。 例如：[1, 3, 5, 9, 8, 6, 2]。
2.2 排序步骤
Bitonic Sort 的神奇之处在于：
构造：把任意序列变成多个小的双调序列。
合并：把两个小的双调序列合并成一个大的双调序列。
排序：把一个双调序列变成有序序列。
对于 Ascend C 开发者，我们只需要关注最核心的操作：Compare-and-Swap (CAS)。 对于长度为 $N$ 的向量，我们需要进行 $\log N$ 个阶段的比较，每个阶段包含若干次跨步比较。
三、 实战：Ascend C 实现 TopK
虽然 Ascend C 提供了内置的Sort指令（基于硬件加速），但它通常有长度限制（如 UB 大小或特定元素个数）。对于超长序列（如 32K 词表）的 TopK，我们需要分治 + 归并。
3.1 Kernel 类定义
假设我们要实现一个行级 TopK：从[Batch, N]中选出前 K 个。
class KernelTopK { public: __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t K, uint32_t N) { // ... Init ... // Tiling: 每次处理一行 this->K = K; this->N = N; } __aicore__ inline void Process() { for (int i = 0; i < batchSize; i++) { Compute(i); } } };3.2 核心排序逻辑 (Bitonic Merge)
假设 N 很大，无法一次性 Sort。我们采用分块排序 + 归并策略。
__aicore__ inline void Compute(int32_t i) { // 1. CopyIn // 假设 N=4096，我们可以一次性搬进 UB (如果 UB 够大) // 如果放不下，需要多轮归并（类似归并排序） DataCopy(dataLoc, xGm[offset], N); // 2. Local Sort (利用硬件指令) // Ascend C 的 Sort 指令通常支持对较短序列的全排序 // Sort(dst, src, len, repeat) // 假设硬件支持最大 128 长度的 Sort // 我们先把数据切成 32 个 128 的块，分别排序 // 伪代码：Block Sort // for (int j=0; j<N/128; j++) Sort(dataLoc[j*128], dataLoc[j*128], 128); // 3. Bitonic Merge (手写归并网络) // 现在我们有 32 个有序小块。需要两两合并。 // Merge 两个有序序列 A (升序) 和 B (降序) 成为一个双调序列，再排序 // 这里演示最核心的 Compare-and-Swap 步骤 // 假设我们要合并两个长度为 L 的有序块 // 步长 stride 从 L/2 递减到 1 for (uint32_t stride = L/2; stride > 0; stride /= 2) { // 构造比较对：data[i] 和 data[i+stride] // 技巧：利用 Vector 指令的地址偏移 // vec1 = data[0 ... N-stride] // vec2 = data[stride ... N] // maxVec = Max(vec1, vec2) // minVec = Min(vec1, vec2) // 此时需要根据排序方向，将 min/max 写回正确的位置 // 这通常需要 Muls (mask) 或者 Select 指令 // Ascend C 提供了 MrgSort (MergeSort) 高阶指令来加速这一步 // MrgSort(dst, src, mrg_list) } // 4. Extract TopK // 排序完成后，取前 K 个即可 DataCopy(yGm[offset], dataLoc, K); }3.3 索引追踪 (ArgMax 问题)
TopK 最大的麻烦在于：我们不仅要 Value，还要 Index。 但在 Vector 排序时，Index 会丢失。
黑魔法：Value-Index Packing我们可以把 Value 和 Index 打包成一个 struct，或者更 trick 的做法：利用高位存 Value，低位存 Index。
假设 Value 是 FP16，Index 是 Uint16。 我们可以把它们拼成一个Uint32。
高 16 位：Value (如果不全是正数，需要做 Flip 符号位处理，保证整数比较逻辑与浮点一致)。
低 16 位：Index。
这样直接对 Uint32 数组进行排序，自然就是按 Value 排序，且 Index 紧随其后！
__aicore__ inline void PackAndSort() { // 1. Pack // value_int32 = (reinterpret_cast<uint16>(value_fp16) << 16) | index_uint16 // 需要用到 Vector 的 Shift 和 Or 指令 // 2. Sort (Uint32) Sort(packedLoc, packedLoc, N); // 3. Unpack // topk_val = packed >> 16 // topk_idx = packed & 0xFFFF }四、 性能优化的“胜负手”
排序算子是Compute Bound的，比较次数极多。
利用
Sort指令：尽可能利用硬件内置的排序指令处理小块，而不是手写比较网络。利用
MrgSort指令：硬件加速的归并指令比手动 Compare-Swap 快得多。只排必要的数据：
如果 $K=1$，直接
ReduceMax。如果 $K$ 很小（如 Top5），且 $N$ 很大，可以先分块求 TopK，最后再汇总求 TopK（Map-Reduce 思想），避免对全量数据排序。
五、 总结
在 AI Core 上实现排序，是对 Ascend C 编程能力的极致考验。
算法选择：放弃快排，拥抱双调（Bitonic）。
数据技巧：利用 Packing 技术解决 ArgMax 索引追踪问题。
硬件协同：混合使用
Sort（小块）和MrgSort（合并）指令。
掌握了 TopK，你就打通了 LLM 解码阶段的最后一道关卡，让模型生成的每一个 Token 都既快又准。
 算子)
