CANN/cann-bench MoeReRouting算子API描述

MoeReRouting 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介MoE 网络中进行 AlltoAll 操作从其他卡上拿到需要算的 token 后将 token 按照专家顺序重新排列。主要应用场景Mixture of Experts (MoE) 模型中 token 到专家的数据分发分布式 MoE 训练中跨 rank 的 token 重新路由MoE 前馈网络前的 token 重排使同一专家处理的 token 连续排列算子特征难度等级L3LayoutTransform多输入多输出根据每个 rank 的专家 token 数量对输入 token 进行重新排列支持可选的 per_token_scales 同步重排2. 算子定义数学公式通过双重求和计算当前 token 在源位置的偏移量$$ \text{SrcOffset} \sum_{i0}^{\text{cur_rank}} \left( \sum_{j0}^{\text{cur_expert}} \text{expert_token_num_per_rank}(i,j) \right) $$通过双重求和计算当前 token 在目标位置的偏移量$$ \text{DstOffset} \sum_{j0}^{\text{cur_expert}} \left( \sum_{i0}^{\text{cur_rank}} \text{expert_token_num_per_rank}(i,j) \right) $$SrcOffset当前需要移动的 token 源偏移根据输入expert_token_num_per_rank的值进行计算DstOffset当前需要移动的 token 目的偏移cur_rankexpert_token_num_per_rank的纵轴索引表示该 token 原本在的卡cur_expertexpert_token_num_per_rank的横轴索引表示该 token 由卡上专家 cur_expert 计算处理流程根据expert_token_num_per_rank矩阵计算每个 token 的源位置和目标位置将 token 从源位置移动到目标位置实现按专家顺序排列若提供per_token_scales同步进行重排输出重排后的 token、scales、索引及每个专家的 token 数量3. 接口规范算子原型cann_bench.moe_re_routing( Tensor tokens, Tensor expert_token_num_per_rank, Tensor? per_token_scales None, int expert_token_num_type 1, int idx_type 0 ) - (Tensor permute_tokens, Tensor permute_per_token_scales, Tensor permute_token_idx, Tensor expert_token_num)输入参数说明参数类型默认值描述ShapetokensTensor必选待重新排布的 token(A, H)expert_token_num_per_rankTensor必选每张卡上各个专家处理的 token 数矩阵元素 [i,j] 表示从卡 i 获取的专家 j 处理的 token 数(N, E)per_token_scalesTensorNone每个 token 对应的 scale需要随 token 同样进行重新排布(A)expert_token_num_typeint1输出 expert_token_num 的模式0cumsum1count。当前只支持为 1-idx_typeint0输出 permute_token_idx 的索引类型0gather索引1scatter索引。当前只支持为 0-输出参数Shapedtype描述permute_tokens(A, H)与 tokens 相同重新排布后的 tokenpermute_per_token_scales(A)float32重新排布后的 per_token_scales输入不携带时无效permute_token_idx(A)int32每个 token 在原排布方式的索引expert_token_num(E)与 expert_token_num_per_rank 相同每个专家处理的 token 数数据类型tokens dtypeexpert_token_num_per_rank dtypeper_token_scales dtypefloat16int32 / int64float32bfloat16int32 / int64float32int8int32 / int64float32Shape 变量说明Atoken 个数取值要求 Sum(expert_token_num_per_rank) AHtoken 长度hidden_dim取值要求 0 H 16384N卡数rank 数取值无限制E卡上的专家数取值无限制规则与约束tokens的形状为 2D (A, H)expert_token_num_per_rank的形状为 2D (N, E)元素必须大于 0所有元素之和必须等于 ASum(expert_token_num_per_rank) Aexpert_token_num_type当前只支持为 1count 模式idx_type当前只支持为 0gather 索引模式per_token_scales为可选参数存在时 shape 必须为 (A)4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差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 import numpy as np def moe_re_routing( tokens: torch.Tensor, expert_token_num_per_rank: torch.Tensor, per_token_scales: torch.Tensor None, expert_token_num_type: int 1, idx_type: int 0 ) - tuple: MoeReRouting 算子 Torch Golden 参考实现 MoE 网络中将 token 按照专家顺序重新排列 Args: tokens: 待重新排布的 tokenshape (A, H) expert_token_num_per_rank: 每张卡上各个专家处理的 token 数shape (N, E) per_token_scales: 每个 token 对应的 scaleshape (A)可选 expert_token_num_type: 输出 expert_token_num 的模式0cumsum, 1count当前只支持 1 idx_type: 输出 permute_token_idx 的索引类型0gather, 1scatter当前只支持 0 Returns: (permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num) # 获取参数 N, E expert_token_num_per_rank.shape A, H tokens.shape # 确保总和匹配 total_tokens expert_token_num_per_rank.sum().item() assert total_tokens A, fSum of expert_token_num_per_rank ({total_tokens}) must equal A ({A}) # 构建 src_offset 和 dst_offset 映射 # 计算每个 (rank, expert) 位置的源偏移和目标偏移 src_offsets {} # (rank, expert) - src_offset dst_offsets {} # (rank, expert) - dst_offset # 计算 SrcOffset按 rank 和 expert 的顺序累加 src_acc 0 for i in range(N): # cur_rank for j in range(E): # cur_expert src_offsets[(i, j)] src_acc src_acc expert_token_num_per_rank[i, j].item() # 计算 DstOffset按 expert 和 rank 的顺序累加 dst_acc 0 for j in range(E): # cur_expert for i in range(N): # cur_rank dst_offsets[(i, j)] dst_acc dst_acc expert_token_num_per_rank[i, j].item() # 构建重排映射src_pos - dst_pos src_to_dst {} for i in range(N): for j in range(E): num_tokens expert_token_num_per_rank[i, j].item() src_start src_offsets[(i, j)] dst_start dst_offsets[(i, j)] for k in range(num_tokens): src_to_dst[src_start k] dst_start k # 构建反向映射用于 gather 索引 dst_to_src {v: k for k, v in src_to_dst.items()} # 生成 permute_token_idx (gather 索引) permute_token_idx torch.zeros(A, dtypetorch.int32) for dst_pos in range(A): permute_token_idx[dst_pos] dst_to_src[dst_pos] # 重排 tokens permute_tokens tokens[permute_token_idx] # 重排 per_token_scales如果存在 if per_token_scales is not None: permute_per_token_scales per_token_scales[permute_token_idx] else: permute_per_token_scales torch.zeros(A, dtypetorch.float32) # 计算 expert_token_num (count 模式) if expert_token_num_type 1: expert_token_num expert_token_num_per_rank.sum(dim0) # 每个专家的总 token 数 else: # cumsum 模式暂不支持 expert_token_num torch.zeros(E, dtypeexpert_token_num_per_rank.dtype) return permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num6. 额外信息算子调用示例import torch import cann_bench import math import random # 基础示例 tokens_num 16384 tokens_length 7168 rank_num 16 expert_num 16 tokens torch.randint(low-10, high20, size(tokens_num, tokens_length), dtypetorch.int8) expert_token_num_per_rank torch.ones(rank_num, expert_num, dtypetorch.int32) # 设置每个位置的 token 数量确保总和等于 tokens_num tokens_sum 0 for i in range(rank_num): for j in range(expert_num): if i rank_num - 1 and j expert_num - 1: expert_token_num_per_rank[i][j] tokens_num - tokens_sum break rand_num 1 expert_token_num_per_rank[i][j] rand_num tokens_sum rand_num per_token_scales torch.randn(tokens_num, dtypetorch.float32) # 调用算子 permute_tokens, permute_per_token_scales, permute_token_idx, expert_token_num cann_bench.moe_re_routing( tokens, expert_token_num_per_rank, per_token_scalesper_token_scales )【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考