CANN ops-nn 算子实战ApplyTopKTopPWithSorted 的 top-k/top-p 采样过滤原理与 aclnn 调用指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnApplyTopKTopPWithSorted 是 CANN ops-nn 神经网络算子库中用于大模型解码阶段采样过滤的算子它接收已按升序排序的值与索引sorted_value / sorted_indices通过 top-k 与 top-p 联合或单独过滤将被裁剪的 token 置为 -Inf再按原始顺序还原供后续 softmax 采样使用。本文以 算子 README 为主体结合 aclnn API 文档、op_api 实现、tiling 实现 与 kernel 实现完整讲解其计算原理、参数约束、两段式 aclnn 接口调用流程、底层实现与测试验证帮助读者在 NPU 上快速落地该采样算子。一、功能概述解码采样中的 top-k / top-p 过滤在自回归语言模型解码中从 logits 采样下一个 token 时通常会做两步裁剪top-k 过滤只保留 logits 中排名前 k 的候选其余置为 -Inftop-p核采样nucleus sampling按概率从高到低累计只保留累计概率不超过 p 的最小候选集合其余置为 -Inf。ApplyTopKTopPWithSorted 算子把这两步合并为一次设备端计算它假设sorted_value升序排列与sorted_indices对应原始索引已经由外部排序步骤准备好直接基于排序结果完成过滤并通过logits可选输入或sorted_indices将过滤结果还原回原始顺序。p与k均为可选输入二者至少传入一个因此支持三种模式仅 top-k 过滤、仅 top-p 过滤、top-k 与 top-p 联合过滤。从源码结构看该算子同时提供底层图算子ApplyTopKTopPWithSorted输入为已排序数据与高层 aclnn 接口aclnnApplyTopKTopP输入为原始 logits内部自动完成排序。本文先讲算子本身的过滤语义再讲 aclnn 接口的完整调用。二、计算流程与公式对每个 batch b$0 \le b \lt logits.size(0)$、每个位置 v$0 \le v \lt logits.size(1)$算子按以下步骤计算1. 计算 top-k 保留阈值取每个 batch 中第 k 大的值作为阈值$$topKValue[b][v] sortedValue[b][sortedValue.size(1) - k[b]]$$注意sorted_value按最后一维升序排列因此第 k 大的值位于sortedValue.size(1) - k[b]处。2. 生成 top-k 过滤 mask$$topKMask sortedValue topKValue$$即严格小于第 k 大值的元素需要被过滤。3. top-k 置 -Inf$$ sortedValue[b][v] \begin{cases} -Inf \text{topKMask[b][v]true}\ sortedValue[b][v] \text{topKMask[b][v]false} \end{cases} $$4. 计算概率分布对过滤后的数据沿最后一轴做 softmax得到概率分布$$probsValue softmax(sortedValue, dim-1)$$5. 累计概率从最小的概率开始沿最后一轴计算累计概率cumsum$$probsSum cumsum(probsValue, dim-1)$$6. 生成 top-p 过滤 mask累计概率小于等于 $1-p$ 的位置需要被过滤同时保证每个 batch 至少保留一个元素即最后一个位置永不过滤$$topPMask[b][v] probsSum[b][v] \le 1-p[b]$$ $$topPMask[b][-1] false$$7. top-p 置 -Inf$$ sortedValue[b][v] \begin{cases} -Inf \text{topPMask[b][v]true}\ sortedValue[b][v] \text{topPMask[b][v]false} \end{cases} $$8. 还原到原始顺序按sorted_indices将过滤后的结果还原到原始位置$$out[b][v] sortedValue[b][sortedIndices[b][v]]$$ST系统测试脚本 executor_aclnnApplyTopKTopP.py 用 PyTorch 逐行实现了上述全部公式作为期望结果golden可作为理解语义的参考对照logits_sort, logits_idx logits.sort(dim-1, descendingFalse, stableTrue) kth_idx logits_sort.size(1) - k.to(torch.long) kth_value logits_sort.gather(1, kth_idx.unsqueeze(dim1)) top_k_mask logits_sort kth_value logits_sort.masked_fill_(top_k_mask, -float(inf)) softmax_res logits_sort.to(torch.float32).softmax(dim-1) cumsum_res softmax_res.cumsum(dim-1) top_p_mask cumsum_res 1 - p.unsqueeze(dim1) top_p_mask[:, -1] False logits_sort.masked_fill_(top_p_mask, -float(inf)) logits torch.empty_like(logits_sort).scatter_(dim-1, indexlogits_idx, srclogits_sort)三、参数说明算子ApplyTopKTopPWithSorted各参数定义如下该表格同时对应 算子定义源码 中Input/Output注册的数据类型与格式参数名输入/输出/属性描述数据类型数据格式sorted_value输入表示需要处理的数据的值公式中的sortedValue需要升序顺序FLOAT、FLOAT16、BFLOAT16NDsorted_indices输入表示需要处理的数据的索引公式中的sortedIndicesINT32NDp输入表示 top-p 的阈值公式中的 p。可选输入与 k 至少传入一个。数据类型需要与sorted_value一致shape 需要与sorted_value.size(0)一致FLOAT、FLOAT16、BFLOAT16NDk输入表示 top-k 的阈值公式中的 k。可选输入与 p 至少传入一个。shape 需要与sorted_value.size(0)一致INT32NDlogits输入表示排序前的原始数据公式中的 logits。可选输入数据类型需要与sorted_value一致shape 需要与sorted_value一致。传入时算子将基于该输入完成过滤还原计算FLOAT、FLOAT16、BFLOAT16NDout输出表示过滤后的数据公式中的 out。数据类型、shape 需要与sorted_value一致FLOAT、FLOAT16、BFLOAT16ND针对具体产品还有数据类型限制Atlas 推理系列产品参数sorted_value、p、out的数据类型不支持 BFLOAT16。Kirin X90 / Kirin 9030 处理器系列产品不支持 BFLOAT16。产品支持矩阵来自 算子 README产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√上述支持情况在 op_host/config 目录下按产品分目录的算子二进制配置如 ascend910b 配置 中的 fp32/fp16/bf16 三种 bin以及 算子定义源码 中为ascend910b、ascend910_93、ascend950、kirinx90、kirin9030、ascend310p分别注册的 AICore 配置中得到印证。四、约束说明p和k为可选输入二者不能同时为空至少传入一个。当仅传入p时执行纯 top-p 过滤仅传入k时执行纯 top-k 过滤两者都传入时执行联合过滤。aclnn 接口层aclnnApplyTopKTopP额外约束p的值域为 [0, 1]k小于等于 0 或大于等于logits.size(1)时视为不进行 top-k 过滤logits、p、k、out均支持非连续 tensor源码中第一段接口会调用l0op::Contiguous进行连续化处理参见 aclnn_apply_top_k_top_p.cpp。五、aclnn 接口调用方式两段式 API算子支持通过aclnnApplyTopKTopP接口调用完整样例见 test_aclnn_apply_top_k_top_p.cpp。与 CANN 其它算子一样该接口采用两段式调用模型详见 两段式接口先调用GetWorkspaceSize接口获取 workspace 大小与执行器再调用执行接口完成计算。1. 函数原型第一段接口完成参数校验、构图并返回所需 workspace 大小和执行器aclnnStatus aclnnApplyTopKTopPGetWorkspaceSize( const aclTensor* logits, const aclTensor* p, const aclTensor* k, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口在指定 stream 上执行算子计算aclnnStatus aclnnApplyTopKTopP( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)注意第二段接口aclnnApplyTopKTopP不能重复调用同一个 executor 只允许执行一次否则会出现异常参见 两段式接口。2. 第一段接口参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 tensorlogits输入表示需要处理的数据公式中的 logits-BFLOAT16、FLOAT16、FLOAT32ND2√p输入表示 top-p 的阈值公式中的 p值域为 [0, 1]数据类型需要与 logits 一致shape 需要与 logits.size(0) 一致BFLOAT16、FLOAT16、FLOAT32ND1√k输入表示 top-k 的阈值公式中的 kk 小于等于 0 或大于等于 logits.size(1) 时视为不进行 top-k 过滤shape 需要与 logits.size(0) 一致INT32ND1√out输出表示过滤后的数据公式中的 out数据类型需要与 logits 一致shape 需要与 logits 一致BFLOAT16、FLOAT16、FLOAT32ND2√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----注意aclnn 接口的logits是排序前的原始数据2 维而算子本身的sorted_value/sorted_indices是已排序数据——接口内部会自动完成排序见下文源码分析。3. 返回码与异常场景第一段接口返回aclnnStatus状态码具体取值参见 aclnn 返回码。其中与本算子强相关的校验错误如下返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 logits、out 是空指针或 p 与 k 同为空指针ACLNN_ERR_PARAM_INVALID161002logits、p、k 或 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002logits、p 或 out 的数据类型不匹配ACLNN_ERR_PARAM_INVALID161002logits、p、k 或 out 的 shape 不匹配这些校验在 aclnn_apply_top_k_top_p.cpp 的CheckParams中依次执行CheckNotNull空指针且 p、k 不能同时为空→CheckDtypeValid支持数据类型列表{DT_FLOAT, DT_FLOAT16, DT_BF16}与k的{DT_INT32}并校验 p、out 与 logits 类型一致→CheckShapeValidlogits 必须 2 维、p/k 必须 1 维且size(0)与 logits 一致、out 与 logits 同 shape→CheckFormatValid全部要求 ND 格式。六、源码级原理剖析1. aclnn 接口内部调用链aclnn_apply_top_k_top_p.cpp 的第一段接口实现揭示了完整调用链将输入 logits 连续化l0op::Contiguous支持非连续 tensor若最后一维长度为 1直接l0op::ViewCopy拷贝到 out无需过滤否则调用l0op::Sort(logitsContiguous, -1, false, true, DT_INT32)得到升序、稳定的sortedValue与sortedIndices这正好对应接口文档中的公式sortedValue, sortedIndices sort(logits, dim-1, descendingfalse, stabletrue)将排序结果交给底层算子l0op::ApplyTopKTopPWithSorted(sortedValue, sortedIndices, p, k, logits)执行过滤最后l0op::ViewCopy将结果写入 out兼容非连续输出并返回GetWorkspaceSize()。2. Tiling 与多核调度tiling 实现 通过tilingKey区分不同计算模式tilingKey模式0默认通用路径scatter 还原1仅 top-k 过滤ONLY_TOP_K2仅 top-p 过滤ONLY_TOP_P3优化路径top-k top-p 联合OPT_TOP_K_TOP_P要求传入 logits4优化路径仅 top-kOPT_TOP_K5优化路径仅 top-pOPT_TOP_P关键设计点tiling 时按 batch 维度分配核数usedCoreNum_ min(coreNum_, batchSize_)每个核处理batchPerCore_个 batch余数作为tailBatch_GetUsedCore单次搬入 UB 的数据量以K_VALUE_MAX 1024为上限进行分块ubFactorElement_大词表场景按迭代次数循环处理workspace 大小按模式区分优化路径tilingKey 3/4/5无需额外 workspace仅 top-p 模式需要16MB batch* vocab* 4B其余需要16MB batch* 4BRunKernelTiling310PDAV_2002平台因 DataCopyPad 不可用强制走 tilingKey0 的 scatter 旧路径。3. Kernel 分发与优化路径kernel 入口 根据TILING_KEY在设备侧分发到不同实现通用路径ApplyTopKTopPWithSorted、仅 top-p 的ApplyTopPWithSorted以及基于原始 logits 直接做选择的优化实现ApplyTopKTopPOpt定义于 apply_top_k_top_p_opt.h。优化路径tilingKey 3/4/5的核心思路是利用已排序数组的单调性在排序后的序列上直接定位 top-k 边界索引FindKthStartIdx用二分/分块累计CumsumKoggleStone与FindFirstIndex找到 top-p 边界最后直接在原始 logits 上做比较选择CompareSelectOnLogits、ScatterBoundary并写回 out省去了对整条序列逐元素置 -Inf 与 scatter 还原的访存开销。kernel 中还以常量形式预置了各数据类型的负无穷表示如FLOAT16_NEG_INF 0xFC00、BF16_NEG_INF 0xFF80、FLOAT32_NEG_INF 0xFF800000见 apply_top_k_top_p_opt.h。接口文档亦注明aclnnApplyTopKTopP默认确定性实现相同输入得到相同结果这与排序采用 stabletrue 的语义一致。七、完整调用示例以下示例来自 test_aclnn_apply_top_k_top_p.cpp输入 logits 形状为{3, 4}p 与 k 均为长度为 3 的 batch 级阈值#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_apply_top_k_top_p.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t logitsShape {3, 4}; std::vectorint64_t pShape {3}; std::vectorint64_t kShape {3}; std::vectorint64_t outShape {3, 4}; void* logitsDeviceAddr nullptr; void* pDeviceAddr nullptr; void* kDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* logits nullptr; aclTensor* p nullptr; aclTensor* k nullptr; aclTensor* out nullptr; std::vectorfloat logitsHostData {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; std::vectorfloat pHostData {0.2, 0.4, 0.6}; std::vectorint32_t kHostData {1, 2, 3}; std::vectorfloat outHostData(12, 0); ret CreateAclTensor(logitsHostData, logitsShape, logitsDeviceAddr, aclDataType::ACL_FLOAT, logits); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(pHostData, pShape, pDeviceAddr, aclDataType::ACL_FLOAT, p); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(kHostData, kShape, kDeviceAddr, aclDataType::ACL_INT32, k); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 两段式调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口获取workspace大小与执行器 ret aclnnApplyTopKTopPGetWorkspaceSize(logits, p, k, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnApplyTopKTopPGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 第二段接口执行计算 ret aclnnApplyTopKTopP(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnApplyTopKTopP failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor aclDestroyTensor(logits); aclDestroyTensor(p); aclDestroyTensor(k); aclDestroyTensor(out); // 7. 释放Device资源 aclrtFree(logitsDeviceAddr); aclrtFree(pDeviceAddr); aclrtFree(kDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中值得注意的细节p、k同时传入分别取值 0.2/0.4/0.6 与 1/2/3因此走 top-k 与 top-p 联合过滤而p与k的 shape 均为{3}与logits.size(0) 3严格对应。运行后 out 中未被过滤的位置保留原值被过滤的位置为 -Inf可直接作为 softmax 采样的输入。八、编译与运行编译运行该示例前需保证驱动、固件、CANN 软件包与 ops 包已安装。具体步骤如下详细流程参见 编译与运行样例将 test_aclnn_apply_top_k_top_p.cpp 保存为test_apply_top_k_top_p.cpp并按需编写 CMakeLists.txt设置-stdc11包含${ASCEND_PATH}/include与${ASCEND_PATH}/include/aclnn头文件目录链接libascendcl.so、libnnopbase.so、libopapi_math.so、libopapi_nn.so配置环境变量source ${INSTALL_DIR}/set_env.sh # ${INSTALL_DIR} 为CANN软件安装路径编译并运行mkdir -p build cd build cmake ../ -DCMAKE_CXX_COMPILERg -DCMAKE_SKIP_RPATHTRUE make cd bin ./opapi_test若运行报错可使用aclGetRecentErrMsg获取具体异常信息例如空指针入参时会得到错误码 161001 及对应的参数名提示参考 编译与运行样例 中的示例。九、测试与验证仓库为该算子提供了完整的 UT 与 ST 测试aclnn 接口 UTtest_aclnn_apply_top_k_top_p.cpp覆盖空 tensor、p/k 单空、fp32/fp16/bf16 正常路径、最后一维为 1 的边界场景以及空指针、数据类型不匹配、shape 不匹配、format 非法等全部异常分支逐一断言返回ACLNN_SUCCESS或ACLNN_ERR_PARAM_NULLPTR/ACLNN_ERR_PARAM_INVALIDtiling UTtest_apply_top_k_top_p_with_sorted_tiling.cpp验证不同 shape 下核数分配、UB 分块参数与 workspace 计算的正确性kernel UT 数据生成gen_data.py随机生成 batch×vocab 的 logits、p[0,1) 均匀分布与 k10~30 随机整数并用 PyTorch 的 sort/softmax/cumsum 生成 golden 数据ST 端到端测试executor_aclnnApplyTopKTopP.py以 PyTorch 复现公式见第二节对照代码作为期望输出与 NPU 实际计算结果比对。十、小结ApplyTopKTopPWithSorted 将 top-k 与 top-p 两种采样过滤合并为单个 NPU 算子升序排序输入保证了阈值定位与累计概率计算的单调性p/k 可选输入的组合覆盖了三种过滤模式aclnn 接口则通过内部 sort 过滤 还原的调用链让上层只须传入原始 logits 即可完成完整采样裁剪。若上层框架已具备排序结果亦可直接调用底层ApplyTopKTopPWithSorted算子省去重复排序。理解其公式语义与两段式接口约束即可在解码推理场景中稳定、高效地接入该算子。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn 算子实战：ApplyTopKTopPWithSorted 的 top-k/top-p 采样过滤原理与 aclnn 调用指南
CANN ops-nn 算子实战ApplyTopKTopPWithSorted 的 top-k/top-p 采样过滤原理与 aclnn 调用指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnApplyTopKTopPWithSorted 是 CANN ops-nn 神经网络算子库中用于大模型解码阶段采样过滤的算子它接收已按升序排序的值与索引sorted_value / sorted_indices通过 top-k 与 top-p 联合或单独过滤将被裁剪的 token 置为 -Inf再按原始顺序还原供后续 softmax 采样使用。本文以 算子 README 为主体结合 aclnn API 文档、op_api 实现、tiling 实现 与 kernel 实现完整讲解其计算原理、参数约束、两段式 aclnn 接口调用流程、底层实现与测试验证帮助读者在 NPU 上快速落地该采样算子。一、功能概述解码采样中的 top-k / top-p 过滤在自回归语言模型解码中从 logits 采样下一个 token 时通常会做两步裁剪top-k 过滤只保留 logits 中排名前 k 的候选其余置为 -Inftop-p核采样nucleus sampling按概率从高到低累计只保留累计概率不超过 p 的最小候选集合其余置为 -Inf。ApplyTopKTopPWithSorted 算子把这两步合并为一次设备端计算它假设sorted_value升序排列与sorted_indices对应原始索引已经由外部排序步骤准备好直接基于排序结果完成过滤并通过logits可选输入或sorted_indices将过滤结果还原回原始顺序。p与k均为可选输入二者至少传入一个因此支持三种模式仅 top-k 过滤、仅 top-p 过滤、top-k 与 top-p 联合过滤。从源码结构看该算子同时提供底层图算子ApplyTopKTopPWithSorted输入为已排序数据与高层 aclnn 接口aclnnApplyTopKTopP输入为原始 logits内部自动完成排序。本文先讲算子本身的过滤语义再讲 aclnn 接口的完整调用。二、计算流程与公式对每个 batch b$0 \le b \lt logits.size(0)$、每个位置 v$0 \le v \lt logits.size(1)$算子按以下步骤计算1. 计算 top-k 保留阈值取每个 batch 中第 k 大的值作为阈值$$topKValue[b][v] sortedValue[b][sortedValue.size(1) - k[b]]$$注意sorted_value按最后一维升序排列因此第 k 大的值位于sortedValue.size(1) - k[b]处。2. 生成 top-k 过滤 mask$$topKMask sortedValue topKValue$$即严格小于第 k 大值的元素需要被过滤。3. top-k 置 -Inf$$ sortedValue[b][v] \begin{cases} -Inf \text{topKMask[b][v]true}\ sortedValue[b][v] \text{topKMask[b][v]false} \end{cases} $$4. 计算概率分布对过滤后的数据沿最后一轴做 softmax得到概率分布$$probsValue softmax(sortedValue, dim-1)$$5. 累计概率从最小的概率开始沿最后一轴计算累计概率cumsum$$probsSum cumsum(probsValue, dim-1)$$6. 生成 top-p 过滤 mask累计概率小于等于 $1-p$ 的位置需要被过滤同时保证每个 batch 至少保留一个元素即最后一个位置永不过滤$$topPMask[b][v] probsSum[b][v] \le 1-p[b]$$ $$topPMask[b][-1] false$$7. top-p 置 -Inf$$ sortedValue[b][v] \begin{cases} -Inf \text{topPMask[b][v]true}\ sortedValue[b][v] \text{topPMask[b][v]false} \end{cases} $$8. 还原到原始顺序按sorted_indices将过滤后的结果还原到原始位置$$out[b][v] sortedValue[b][sortedIndices[b][v]]$$ST系统测试脚本 executor_aclnnApplyTopKTopP.py 用 PyTorch 逐行实现了上述全部公式作为期望结果golden可作为理解语义的参考对照logits_sort, logits_idx logits.sort(dim-1, descendingFalse, stableTrue) kth_idx logits_sort.size(1) - k.to(torch.long) kth_value logits_sort.gather(1, kth_idx.unsqueeze(dim1)) top_k_mask logits_sort kth_value logits_sort.masked_fill_(top_k_mask, -float(inf)) softmax_res logits_sort.to(torch.float32).softmax(dim-1) cumsum_res softmax_res.cumsum(dim-1) top_p_mask cumsum_res 1 - p.unsqueeze(dim1) top_p_mask[:, -1] False logits_sort.masked_fill_(top_p_mask, -float(inf)) logits torch.empty_like(logits_sort).scatter_(dim-1, indexlogits_idx, srclogits_sort)三、参数说明算子ApplyTopKTopPWithSorted各参数定义如下该表格同时对应 算子定义源码 中Input/Output注册的数据类型与格式参数名输入/输出/属性描述数据类型数据格式sorted_value输入表示需要处理的数据的值公式中的sortedValue需要升序顺序FLOAT、FLOAT16、BFLOAT16NDsorted_indices输入表示需要处理的数据的索引公式中的sortedIndicesINT32NDp输入表示 top-p 的阈值公式中的 p。可选输入与 k 至少传入一个。数据类型需要与sorted_value一致shape 需要与sorted_value.size(0)一致FLOAT、FLOAT16、BFLOAT16NDk输入表示 top-k 的阈值公式中的 k。可选输入与 p 至少传入一个。shape 需要与sorted_value.size(0)一致INT32NDlogits输入表示排序前的原始数据公式中的 logits。可选输入数据类型需要与sorted_value一致shape 需要与sorted_value一致。传入时算子将基于该输入完成过滤还原计算FLOAT、FLOAT16、BFLOAT16NDout输出表示过滤后的数据公式中的 out。数据类型、shape 需要与sorted_value一致FLOAT、FLOAT16、BFLOAT16ND针对具体产品还有数据类型限制Atlas 推理系列产品参数sorted_value、p、out的数据类型不支持 BFLOAT16。Kirin X90 / Kirin 9030 处理器系列产品不支持 BFLOAT16。产品支持矩阵来自 算子 README产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√上述支持情况在 op_host/config 目录下按产品分目录的算子二进制配置如 ascend910b 配置 中的 fp32/fp16/bf16 三种 bin以及 算子定义源码 中为ascend910b、ascend910_93、ascend950、kirinx90、kirin9030、ascend310p分别注册的 AICore 配置中得到印证。四、约束说明p和k为可选输入二者不能同时为空至少传入一个。当仅传入p时执行纯 top-p 过滤仅传入k时执行纯 top-k 过滤两者都传入时执行联合过滤。aclnn 接口层aclnnApplyTopKTopP额外约束p的值域为 [0, 1]k小于等于 0 或大于等于logits.size(1)时视为不进行 top-k 过滤logits、p、k、out均支持非连续 tensor源码中第一段接口会调用l0op::Contiguous进行连续化处理参见 aclnn_apply_top_k_top_p.cpp。五、aclnn 接口调用方式两段式 API算子支持通过aclnnApplyTopKTopP接口调用完整样例见 test_aclnn_apply_top_k_top_p.cpp。与 CANN 其它算子一样该接口采用两段式调用模型详见 两段式接口先调用GetWorkspaceSize接口获取 workspace 大小与执行器再调用执行接口完成计算。1. 函数原型第一段接口完成参数校验、构图并返回所需 workspace 大小和执行器aclnnStatus aclnnApplyTopKTopPGetWorkspaceSize( const aclTensor* logits, const aclTensor* p, const aclTensor* k, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口在指定 stream 上执行算子计算aclnnStatus aclnnApplyTopKTopP( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)注意第二段接口aclnnApplyTopKTopP不能重复调用同一个 executor 只允许执行一次否则会出现异常参见 两段式接口。2. 第一段接口参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 tensorlogits输入表示需要处理的数据公式中的 logits-BFLOAT16、FLOAT16、FLOAT32ND2√p输入表示 top-p 的阈值公式中的 p值域为 [0, 1]数据类型需要与 logits 一致shape 需要与 logits.size(0) 一致BFLOAT16、FLOAT16、FLOAT32ND1√k输入表示 top-k 的阈值公式中的 kk 小于等于 0 或大于等于 logits.size(1) 时视为不进行 top-k 过滤shape 需要与 logits.size(0) 一致INT32ND1√out输出表示过滤后的数据公式中的 out数据类型需要与 logits 一致shape 需要与 logits 一致BFLOAT16、FLOAT16、FLOAT32ND2√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----注意aclnn 接口的logits是排序前的原始数据2 维而算子本身的sorted_value/sorted_indices是已排序数据——接口内部会自动完成排序见下文源码分析。3. 返回码与异常场景第一段接口返回aclnnStatus状态码具体取值参见 aclnn 返回码。其中与本算子强相关的校验错误如下返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 logits、out 是空指针或 p 与 k 同为空指针ACLNN_ERR_PARAM_INVALID161002logits、p、k 或 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002logits、p 或 out 的数据类型不匹配ACLNN_ERR_PARAM_INVALID161002logits、p、k 或 out 的 shape 不匹配这些校验在 aclnn_apply_top_k_top_p.cpp 的CheckParams中依次执行CheckNotNull空指针且 p、k 不能同时为空→CheckDtypeValid支持数据类型列表{DT_FLOAT, DT_FLOAT16, DT_BF16}与k的{DT_INT32}并校验 p、out 与 logits 类型一致→CheckShapeValidlogits 必须 2 维、p/k 必须 1 维且size(0)与 logits 一致、out 与 logits 同 shape→CheckFormatValid全部要求 ND 格式。六、源码级原理剖析1. aclnn 接口内部调用链aclnn_apply_top_k_top_p.cpp 的第一段接口实现揭示了完整调用链将输入 logits 连续化l0op::Contiguous支持非连续 tensor若最后一维长度为 1直接l0op::ViewCopy拷贝到 out无需过滤否则调用l0op::Sort(logitsContiguous, -1, false, true, DT_INT32)得到升序、稳定的sortedValue与sortedIndices这正好对应接口文档中的公式sortedValue, sortedIndices sort(logits, dim-1, descendingfalse, stabletrue)将排序结果交给底层算子l0op::ApplyTopKTopPWithSorted(sortedValue, sortedIndices, p, k, logits)执行过滤最后l0op::ViewCopy将结果写入 out兼容非连续输出并返回GetWorkspaceSize()。2. Tiling 与多核调度tiling 实现 通过tilingKey区分不同计算模式tilingKey模式0默认通用路径scatter 还原1仅 top-k 过滤ONLY_TOP_K2仅 top-p 过滤ONLY_TOP_P3优化路径top-k top-p 联合OPT_TOP_K_TOP_P要求传入 logits4优化路径仅 top-kOPT_TOP_K5优化路径仅 top-pOPT_TOP_P关键设计点tiling 时按 batch 维度分配核数usedCoreNum_ min(coreNum_, batchSize_)每个核处理batchPerCore_个 batch余数作为tailBatch_GetUsedCore单次搬入 UB 的数据量以K_VALUE_MAX 1024为上限进行分块ubFactorElement_大词表场景按迭代次数循环处理workspace 大小按模式区分优化路径tilingKey 3/4/5无需额外 workspace仅 top-p 模式需要16MB batch* vocab* 4B其余需要16MB batch* 4BRunKernelTiling310PDAV_2002平台因 DataCopyPad 不可用强制走 tilingKey0 的 scatter 旧路径。3. Kernel 分发与优化路径kernel 入口 根据TILING_KEY在设备侧分发到不同实现通用路径ApplyTopKTopPWithSorted、仅 top-p 的ApplyTopPWithSorted以及基于原始 logits 直接做选择的优化实现ApplyTopKTopPOpt定义于 apply_top_k_top_p_opt.h。优化路径tilingKey 3/4/5的核心思路是利用已排序数组的单调性在排序后的序列上直接定位 top-k 边界索引FindKthStartIdx用二分/分块累计CumsumKoggleStone与FindFirstIndex找到 top-p 边界最后直接在原始 logits 上做比较选择CompareSelectOnLogits、ScatterBoundary并写回 out省去了对整条序列逐元素置 -Inf 与 scatter 还原的访存开销。kernel 中还以常量形式预置了各数据类型的负无穷表示如FLOAT16_NEG_INF 0xFC00、BF16_NEG_INF 0xFF80、FLOAT32_NEG_INF 0xFF800000见 apply_top_k_top_p_opt.h。接口文档亦注明aclnnApplyTopKTopP默认确定性实现相同输入得到相同结果这与排序采用 stabletrue 的语义一致。七、完整调用示例以下示例来自 test_aclnn_apply_top_k_top_p.cpp输入 logits 形状为{3, 4}p 与 k 均为长度为 3 的 batch 级阈值#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_apply_top_k_top_p.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t logitsShape {3, 4}; std::vectorint64_t pShape {3}; std::vectorint64_t kShape {3}; std::vectorint64_t outShape {3, 4}; void* logitsDeviceAddr nullptr; void* pDeviceAddr nullptr; void* kDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* logits nullptr; aclTensor* p nullptr; aclTensor* k nullptr; aclTensor* out nullptr; std::vectorfloat logitsHostData {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; std::vectorfloat pHostData {0.2, 0.4, 0.6}; std::vectorint32_t kHostData {1, 2, 3}; std::vectorfloat outHostData(12, 0); ret CreateAclTensor(logitsHostData, logitsShape, logitsDeviceAddr, aclDataType::ACL_FLOAT, logits); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(pHostData, pShape, pDeviceAddr, aclDataType::ACL_FLOAT, p); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(kHostData, kShape, kDeviceAddr, aclDataType::ACL_INT32, k); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 两段式调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口获取workspace大小与执行器 ret aclnnApplyTopKTopPGetWorkspaceSize(logits, p, k, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnApplyTopKTopPGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 第二段接口执行计算 ret aclnnApplyTopKTopP(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnApplyTopKTopP failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor aclDestroyTensor(logits); aclDestroyTensor(p); aclDestroyTensor(k); aclDestroyTensor(out); // 7. 释放Device资源 aclrtFree(logitsDeviceAddr); aclrtFree(pDeviceAddr); aclrtFree(kDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中值得注意的细节p、k同时传入分别取值 0.2/0.4/0.6 与 1/2/3因此走 top-k 与 top-p 联合过滤而p与k的 shape 均为{3}与logits.size(0) 3严格对应。运行后 out 中未被过滤的位置保留原值被过滤的位置为 -Inf可直接作为 softmax 采样的输入。八、编译与运行编译运行该示例前需保证驱动、固件、CANN 软件包与 ops 包已安装。具体步骤如下详细流程参见 编译与运行样例将 test_aclnn_apply_top_k_top_p.cpp 保存为test_apply_top_k_top_p.cpp并按需编写 CMakeLists.txt设置-stdc11包含${ASCEND_PATH}/include与${ASCEND_PATH}/include/aclnn头文件目录链接libascendcl.so、libnnopbase.so、libopapi_math.so、libopapi_nn.so配置环境变量source ${INSTALL_DIR}/set_env.sh # ${INSTALL_DIR} 为CANN软件安装路径编译并运行mkdir -p build cd build cmake ../ -DCMAKE_CXX_COMPILERg -DCMAKE_SKIP_RPATHTRUE make cd bin ./opapi_test若运行报错可使用aclGetRecentErrMsg获取具体异常信息例如空指针入参时会得到错误码 161001 及对应的参数名提示参考 编译与运行样例 中的示例。九、测试与验证仓库为该算子提供了完整的 UT 与 ST 测试aclnn 接口 UTtest_aclnn_apply_top_k_top_p.cpp覆盖空 tensor、p/k 单空、fp32/fp16/bf16 正常路径、最后一维为 1 的边界场景以及空指针、数据类型不匹配、shape 不匹配、format 非法等全部异常分支逐一断言返回ACLNN_SUCCESS或ACLNN_ERR_PARAM_NULLPTR/ACLNN_ERR_PARAM_INVALIDtiling UTtest_apply_top_k_top_p_with_sorted_tiling.cpp验证不同 shape 下核数分配、UB 分块参数与 workspace 计算的正确性kernel UT 数据生成gen_data.py随机生成 batch×vocab 的 logits、p[0,1) 均匀分布与 k10~30 随机整数并用 PyTorch 的 sort/softmax/cumsum 生成 golden 数据ST 端到端测试executor_aclnnApplyTopKTopP.py以 PyTorch 复现公式见第二节对照代码作为期望输出与 NPU 实际计算结果比对。十、小结ApplyTopKTopPWithSorted 将 top-k 与 top-p 两种采样过滤合并为单个 NPU 算子升序排序输入保证了阈值定位与累计概率计算的单调性p/k 可选输入的组合覆盖了三种过滤模式aclnn 接口则通过内部 sort 过滤 还原的调用链让上层只须传入原始 logits 即可完成完整采样裁剪。若上层框架已具备排序结果亦可直接调用底层ApplyTopKTopPWithSorted算子省去重复排序。理解其公式语义与两段式接口约束即可在解码推理场景中稳定、高效地接入该算子。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考