企业官网建设流程全解析

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 在 MoE（Mixture 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目前支持：

算子按VecCore编译（见 CMakeLists.txt 中--cce-soc-version=Ascend910B1 --cce-soc-core-type=VecCore --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 入口getRoutingConfigurationsSBufV2（get_routing_configurations.cpp）可确认每个参数的真实类型与语义：

关键参数实现细节

• 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_token：README 标注固定为 1024，host 侧入口同样以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三个参数共同描述“当前 EP（Expert Parallel）rank 负责的专家区间”，end_expert_id是开区间右端点（即“最后专家的 id 的后一位 id”）。

约束说明

README 给出的约束为：

输入输出仅支持 BFLOAT16 类型。

从实现细节看，还需要注意以下约束（来自源码与测试的印证）：

• scores 必须为 bfloat16：Kernel 内gmScores_声明为bfloat16_t，测试脚本dtype = torch.bfloat16并注明“只测试 BF16”。
• indices 必须为 int32，token_idx_intra_expert、final_token_table为 int32，init_token_list为 int64（测试脚本与 Kernel GM buffer 类型一致）。
• topk 支持范围：测试脚本注明“只测试 topk=8”，实际使用时应按token_num * topk的连续内存布局传入。
• 所有 Tensor 必须位于 NPU 设备：host 入口对 7 个张量逐一TORCH_CHECK(torch_npu::utils::is_npu(...))校验。

设计方案

Tilling 策略（分核策略）

README 将整个算子按数据处理阶段拆成四个子任务，分别采用不同的分核方式：

对应到源码，这三个 Kernel 分别实现为getRoutingConfigurationsSBufV2_InitTable_kernel、getRoutingConfigurationsSBufV2_InitOther_kernel、getRoutingConfigurationsSBufV2_Final_kernel（README 中的init_list统计逻辑被合并进InitTable阶段的结果中，InitOther阶段再基于init_token_table做掩码统计）。

分核的负载均衡细节

三个 Kernel 都采用“平均分配 + 余数摊派”的经典负载均衡写法，例如InitTable/Final：

bkTokenNum_ = 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_count（ubLoopCount_ = (bkTokenNum_ + ubMaxToken_ - 1) / ubMaxToken_，即按每轮 UB 内最多ub_max_token个 token 划分迭代轮数）；
• 建立 GM Tensor 映射（SetGlobalBuffer绑定 indices/scores/各输出表）；
• 初始化队列缓冲区（pipe_.InitBuffer(...)，对inQueIndices_、inQueScores_、outQueInitTokenTable_、outQueFinalScoreTable_等按对齐后的字节数申请 UB）；
• 初始化临时缓冲区（如InitOther中的tmpMask_、tmpBuffer_）。

计算流程（Compute）

README 中描述的“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）与搬出（CopyOut）

• CopyIn：将输入数据从inGM搬入到inQueue。例如InitTable阶段用DataCopyPad(local_indices, gmIndices_[...], params, pad_params)按 token 块搬入 indices 与 scores，DataCopyExtParams中显式给出搬运长度，并对末块做对齐处理（AlignUp(ubTokenNum, 32 / sizeof(int32_t))）。
• CopyOut：将输出结果数据从outQueue搬出到outGM。InitTable阶段写回init_token_table（按 expert 行、带行间距gbTokenNum_ - ubTokenNum的DataCopyPad）与final_score_table；InitOther阶段写回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等）做流水同步，或用PipeBarrier<PIPE_V/PIPE_ALL>保证向量计算与搬入搬出之间的数据依赖（见 get_routing_configurations.cpp 的Process()实现）。

调用方式与验证

编译集成

get_routing_conf通过 CMakeLists.txt 独立成库：以file(GLOB ... *.cpp)收集源码，使用Ascend910B1+VecCore的 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 算子输出逐项对比：

1. 构造测试数据：TOKEN_NUM=16、TOPK=8、LOCAL_EXPERT_NUM=4、BLOCK_DIM=8，固定随机种子保证可复现；indices用torch.randint生成、scores用torch.randn(..., dtype=torch.bfloat16)生成。
2. CPU 参考实现：get_routing_conf_cpu按 README 语义复刻逻辑——init_token_table用-1初始化、init_token_list统计每个 expert 的 token 数、final_score_table记录每个 token-expert 对的得分，并用seen_mask保证每个 token 对同一 expert 只登记一次。
3. NPU 调用：将全部张量.npu()后调用torch.ops.ascend_ops.get_routing_conf(...)，返回码记入日志。
4. 结果对比：把 NPU 输出的 flat 张量view成 2D（init_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_token=1024为粒度切分循环，所有 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