CANN ops-transformer 中 GetRoutingConfigV2 算子详解MoE Routing 数据预处理的 Tiling 与 Kernel 设计【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读GetRoutingConfigV2仓库路径experimental/moe/get_routing_conf/README.md是 CANN ops-transformer 在 MoEMixture of Experts场景下用于MoeInitRouting计算前的数据准备算子它负责把 token 级的indices/scores路由信息整理为按 expert 组织、可直接喂给后续路由/置换算子的多种表格token 表、score 表、专家内编号等。本文将以该 README 为主线结合其 Kernel 实现 get_routing_configurations.cpp 与测试脚本 test_get_routing_conf.py完整讲解算子的输入输出语义、分核Tiling策略、三阶段 Kernel 流水设计以及调用与验证方式帮助你理解并复用这套 MoE 路由数据预处理方案。产品支持与算子定位支持的产品形态根据 README 的“产品支持情况”表GetRoutingConfigV2目前支持产品是否支持Atlas A2 训练系列产品是算子按VecCore编译见 CMakeLists.txt 中--cce-soc-versionAscend910B1 --cce-soc-core-typeVecCore --cce-auto-sync -xcce编译选项依赖向量核完成计算README 中blockDim以 Ascend910B 的 40 个 AI Core 为例说明。功能定位MoE 路由的数据整理员算子功能moeinitrouting计算前的数据准备工作README“功能说明”。价值/作用为 MoERouting 计算前准备必要的数据。计算公式README 中未给出具体公式留空从实现看该算子并非做数学变换而是做路由信息的重排、去重与编号分配为后续 init_routing 这类“将输入 token 按 expert 分组并连续排列”的算子准备可直接索引的映射表。可以推断在完整的 MoE 计算链中GetRoutingConfigV2负责产出路由配置下游算子如InitRouting、MoeTokenPermute等基于这些配置表完成 token 的搬运与重排因此它的正确性直接决定整条 MoE 数据通路的质量。参数说明README 以表格形式给出了算子的全部参数。结合 Kernel 入口getRoutingConfigurationsSBufV2get_routing_configurations.cpp可确认每个参数的真实类型与语义参数名输入/输出/属性描述数据类型数据格式blockDim输入AI CORE 的数量比如Ascend910B 是 40int64_t-stream输入Device 端的 streamAclrtStream-indices输入专家索引shape(token, topk)int32_tNDscores输入专家得分shape(token, topk)bfloat16NDinit_token_table输出当前卡上每个 token 对应的专家的 token 位置shape 为 ((expert, token))int64NDfinal_token_table输出每条 routing 信息token-expert的编号shape 为 (token, expert)int32NDfinal_score_table输出每条 routing 信息的对应得分shape (token, expert)BFLOAT16NDtoken_idx_intra_expert输出中间结果每条 routing 信息在其对应专家内部的编号是 initTokenTable 的逆映射shape(expert, token)int32NDstart_expert_id输入当前 eprank 上专家的起始 idint64-end_expert_id输入当前 eprank 上最后专家的 id 的后一位 idint64-local_expert_num输入每个 eprank 上的专家数int64-token_num输入所有 eprank 总 token 数int64-ub_max_token输入ub 能存放的最大 token, 1024int64_t-关键参数实现细节blockDim 与 block 数裁剪在 launch 函数中三个子 kernel 的 block 数并非直接取block_dim而是做min(可用量, block_dim)裁剪inittable_block_dim token_num block_dim ? block_dim : token_num、initother_block_dim local_expert_num block_dim ? block_dim : local_expert_num见 get_routing_configurations.cpp 的getRoutingConfigurationsSBufV2_launch避免在 token 数或专家数少于 AI Core 数时浪费核资源。ub_max_tokenREADME 标注固定为 1024host 侧入口同样以constexpr int64_t UB_MAX_TOKEN 1024;硬编码传入 kernel。int64 与 int32 混用README 表格中init_token_table标注为int64但 Kernel 实际将其当作int32_t的 GM buffer 使用gmInitTokenTable_.SetGlobalBuffer((__gm__ int32_t*)(init_token_table), ...)测试脚本也以torch.int32创建该张量这一点在移植该算子到其他产品时需注意对齐。eprank 语义start_expert_id/end_expert_id/local_expert_num三个参数共同描述“当前 EPExpert Parallelrank 负责的专家区间”end_expert_id是开区间右端点即“最后专家的 id 的后一位 id”。约束说明README 给出的约束为输入输出仅支持 BFLOAT16 类型。从实现细节看还需要注意以下约束来自源码与测试的印证scores 必须为 bfloat16Kernel 内gmScores_声明为bfloat16_t测试脚本dtype torch.bfloat16并注明“只测试 BF16”。indices 必须为 int32token_idx_intra_expert、final_token_table为 int32init_token_list为 int64测试脚本与 Kernel GM buffer 类型一致。topk 支持范围测试脚本注明“只测试 topk8”实际使用时应按token_num * topk的连续内存布局传入。所有 Tensor 必须位于 NPU 设备host 入口对 7 个张量逐一TORCH_CHECK(torch_npu::utils::is_npu(...))校验。设计方案Tilling 策略分核策略README 将整个算子按数据处理阶段拆成四个子任务分别采用不同的分核方式子任务分核方式说明init_table分割 tokenNum每个核处理一段连续的 token负责把本段 token 的路由写入 initTokenTable / initScoreTableinit_list单核计算计算每个 expert 的 token 计数表initTokenList由单核串行完成前缀统计init_other分割 localExpertNum按专家维度分核得到每个 expert 的专属 token 编号tokenIdxIntraExpertfinal分割 tokenNum再次按 token 分核组合前缀偏移生成最终编号finalTokenTable对应到源码这三个 Kernel 分别实现为getRoutingConfigurationsSBufV2_InitTable_kernel、getRoutingConfigurationsSBufV2_InitOther_kernel、getRoutingConfigurationsSBufV2_Final_kernelREADME 中的init_list统计逻辑被合并进InitTable阶段的结果中InitOther阶段再基于init_token_table做掩码统计。分核的负载均衡细节三个 Kernel 都采用“平均分配 余数摊派”的经典负载均衡写法例如InitTable/FinalbkTokenNum_ gbTokenNum_ / blockNum_; bkTokenNumStartIdx_ bkTokenNum_ * blockIdx_; if (blockIdx_ gbTokenNum_ % blockNum_) { bkTokenNum_ 1; // 前 (tokenNum % blockNum) 个核多处理 1 个 token bkTokenNumStartIdx_ blockIdx_; } else { bkTokenNumStartIdx_ gbTokenNum_ % blockNum_; }InitOther阶段则按专家数local_expert_num做同样的余数摊派并进一步按ub_max_token切分内层循环保证每个核的 UB 占用可控。Kernel 侧设计Init Process 三阶段流水README 明确指出 Kernel 侧分为Init和Process两个阶段其中 Process 包括数据搬入CopyIn、计算Compute、数据搬出CopyOut实现 RoutingConfig 算子计算。初始化Init计算loop_countubLoopCount_ (bkTokenNum_ ubMaxToken_ - 1) / ubMaxToken_即按每轮 UB 内最多ub_max_token个 token 划分迭代轮数建立 GM Tensor 映射SetGlobalBuffer绑定 indices/scores/各输出表初始化队列缓冲区pipe_.InitBuffer(...)对inQueIndices_、inQueScores_、outQueInitTokenTable_、outQueFinalScoreTable_等按对齐后的字节数申请 UB初始化临时缓冲区如InitOther中的tmpMask_、tmpBuffer_。计算流程ComputeREADME 中描述的“CAST 转 float → AscendC 命令计算 → CAST 转回 bfloat16/half”在InitOther阶段有非常清晰的对应实现Cast(local_in_table_fp32, local_in_table, RoundMode::CAST_NONE, ubTokenNum); // int32 - float CompareScalar(local_mask, local_in_table_fp32, 0.0f, CMPMODE::GE, ubTokenNumAlign); // 生成有效位掩码 GatherMask(local_out_table, local_in_table, local_mask_uint32, true, ubTokenNumAlign, {1, 1, 0, 0}, rsvdCnt); // 紧凑收集这里先用Cast把 int32 的 token 表转成 float再用CompareScalar与 0 比较生成掩码因为init_token_table初始化为 -1非负值即有效 token 位置最后用GatherMask把有效 token 紧凑收集并计数rsvdCnt从而得到每个 expert 的专属 token 列表及其内部编号。在InitTable阶段计算逻辑则是纯标量循环向量核上用GetValue/SetValue遍历indices中每个(token, topk)条目命中当前 EP 专家区间expert_id gbStartExpertId_ expert_id gbEndExpertId_时换算本地专家号local_expert_id expert_id % gbLocalExpertNum_通过“只在init_token_table首次为 -1 时写入”实现每个 (token, expert) 只登记一次的去重语义README 中“每条 routing 信息token-expert的编号”的由来同步写入final_score_table中该 token 对应专家的得分。Final阶段则先对init_token_list做串行前缀和for iExpert: sum val把“每个专家 token 数”转成“每个专家起始偏移”再对每个 token-expert 对计算全局编号final_token_table[token, expert] token_idx_intra_expert[expert, token] 前缀偏移。数据搬入CopyIn与搬出CopyOutCopyIn将输入数据从inGM搬入到inQueue。例如InitTable阶段用DataCopyPad(local_indices, gmIndices_[...], params, pad_params)按 token 块搬入 indices 与 scoresDataCopyExtParams中显式给出搬运长度并对末块做对齐处理AlignUp(ubTokenNum, 32 / sizeof(int32_t))。CopyOut将输出结果数据从outQueue搬出到outGM。InitTable阶段写回init_token_table按 expert 行、带行间距gbTokenNum_ - ubTokenNum的DataCopyPad与final_score_tableInitOther阶段写回token_idx_intra_expert与紧凑后的init_token_table。流程调度Process遍历loop_count按CopyIn → Compute → CopyOut的顺序调度各阶段之间通过SetFlag/WaitFlag硬件事件如HardEvent::MTE2_S、HardEvent::V_S、HardEvent::S_MTE3、HardEvent::MTE2_V、HardEvent::V_MTE3等做流水同步或用PipeBarrierPIPE_V/PIPE_ALL保证向量计算与搬入搬出之间的数据依赖见 get_routing_configurations.cpp 的Process()实现。调用方式与验证编译集成get_routing_conf通过 CMakeLists.txt 独立成库以file(GLOB ... *.cpp)收集源码使用Ascend910B1VecCore的 CCE 编译选项生成对象库get_routing_conf_objects再经 experimental/moe/CMakeLists.txt 的目录遍历机制被上层统一add_subdirectory集成。Host 侧算子注册Kernel 文件末尾通过 Torch 库注册机制暴露为自定义算子TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m) { m.impl(get_routing_conf, getRoutingConfigurationsSBufV2); }即 PyTorch 侧可通过torch.ops.ascend_ops.get_routing_conf(...)调用测试脚本的调用方式。Host 侧函数签名如下int64_t getRoutingConfigurationsSBufV2( int64_t block_dim, torch::Tensor indices, torch::Tensor scores, torch::Tensor init_token_table, torch::Tensor init_token_list, torch::Tensor final_token_table, torch::Tensor final_score_table, torch::Tensor token_idx_intra_expert, const int64_t start_expert_id, const int64_t end_expert_id, const int64_t local_expert_num, const int64_t token_num, const int64_t topk);注意init_token_list每个 expert 的 token 计数虽然不在 README 参数表中但它是 host 侧必需的中间/输出张量是连接InitTable与Final两个 Kernel 的关键桥梁。端到端验证方式仓库提供了可直接运行的验证脚本 test_get_routing_conf.py其验证思路是CPU 参考实现 vs NPU 算子输出逐项对比构造测试数据TOKEN_NUM16、TOPK8、LOCAL_EXPERT_NUM4、BLOCK_DIM8固定随机种子保证可复现indices用torch.randint生成、scores用torch.randn(..., dtypetorch.bfloat16)生成。CPU 参考实现get_routing_conf_cpu按 README 语义复刻逻辑——init_token_table用-1初始化、init_token_list统计每个 expert 的 token 数、final_score_table记录每个 token-expert 对的得分并用seen_mask保证每个 token 对同一 expert 只登记一次。NPU 调用将全部张量.npu()后调用torch.ops.ascend_ops.get_routing_conf(...)返回码记入日志。结果对比把 NPU 输出的 flat 张量view成 2Dinit_token_table→(local_expert_num, token_num)final_score_table→(token_num, local_expert_num)分别输出init_token_list、init_token_table、final_score_table的 CPU/NPU 结果及最大/平均绝对误差、相对误差。这套“flat layout 输出 2D view 对齐 CPU 参考对照”的测试结构同样适用于二次开发该算子时的正确性回归。小结GetRoutingConfigV2是一个典型的“面向后续计算的数据整理型”MoE 算子它本身不做数值变换而是把(token, topk)的路由决策整理成专家视角的 token 表、score 表与编号表。其设计要点可以总结为三阶段子任务流水InitTable按 token 分核登记专家 token 与得分→InitOther按专家分核生成专家内编号→Final前缀和 全局编号各阶段通过硬件事件同步实现 CopyIn → Compute → CopyOut 的流水调度双维度负载均衡token 维度和 expert 维度分别做“均分 余数摊派”并在 host 侧对 block 数做裁剪UB 受限的分块处理以ub_max_token1024为粒度切分循环所有 GM 访问按 512 字节对齐MIN_ALIGN_BYTES保证 UB 队列可容纳去重语义借助-1哨兵值 GatherMask掩码实现“一个 token 对同一 expert 只产生一条 routing 记录”的约束。若你需要将该算子移植到其他 SoC 或接入新的 MoE 训练链路建议以 get_routing_configurations.cpp 的三个 kernel 为起点对照 test_get_routing_conf.py 的 CPU 参考实现逐步验证各输出表的语义一致性。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考