cann-bench TopK算子API描述

TopK 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介返回 k 个最大或最小的元素及其索引。主要应用场景推荐系统中选取得分最高的 k 个候选项分类任务中获取 Top-K 预测类别及其置信度搜索与排序场景中的部分排序加速MoEMixture of Experts路由中选取 Top-K 专家算子特征难度等级L3SortSelect单输入双输出值和索引支持 1-8 维输入支持沿指定维度选取最大或最小的 k 个元素2. 算子定义数学公式$$ y, idx \text{topk}(x, k, dim) $$沿指定维度 dim 对输入张量 x 进行部分排序返回前 k 个最大值当 largesttrue或前 k 个最小值当 largestfalse及其对应的索引。3. 接口规范算子原型cann_bench.top_k(Tensor x, int k, int dim, bool largest) - (Tensor y, Tensor idx)输入参数说明参数类型默认值描述xTensor必选输入张量支持 1-8 维kint必选返回的 topk 数量取值范围1 k dim_sizedimint必选排序维度取值范围-ndim ~ ndim-1largestbooltrue是否返回最大值false 时返回最小值输出参数Shapedtype描述y与输入相同但 dim 维大小变为 k与输入 x 相同输出张量topk 值idx与 y 相同int64输出索引张量始终为 int64数据类型输入 dtype输出 dtypey输出 dtypeidxint8int8int64uint8uint8int64int32int32int64int64int64int64float16float16int64float32float32int64bfloat16bfloat16int64规则与约束输入支持 1-8 维张量k 的取值范围为 1 k 指定维度的大小dim 支持负数索引取值范围为 -ndim ~ ndim-1当 largesttrue 时返回最大的 k 个元素largestfalse 时返回最小的 k 个元素输出 shape 与输入相同仅 dim 维度大小变为 k4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch TopK算子Torch Golden参考实现 返回k个最大或最小的元素及其索引 公式: y, idx topk(x, k, dim) def top_k( x: torch.Tensor, k: int, dim: int, largest: bool True ) - tuple[torch.Tensor, torch.Tensor]: 返回k个最大或最小的元素及其索引 公式: y, idx topk(x, k, dim) Args: x: 输入张量 k: 返回的topk数量 (取值范围: 1 k dim_size) dim: 排序维度 (取值范围: -ndim ~ ndim-1) largest: 是否返回最大值 (false时返回最小值) Returns: y, idx values, indices torch.topk(x, kk, dimdim, largestlargest) return values, indices6. 额外信息算子调用示例import torch import cann_bench x torch.randn(1024, 1024, dtypetorch.float16, devicenpu) y, idx cann_bench.top_k(x, 10, -1, True) # 每行取最大的10个元素 x torch.randn(2, 8, 256, 256, dtypetorch.float32, devicenpu) y, idx cann_bench.top_k(x, 10, -1, False) # 每行取最小的10个元素【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考