文件最后提交记录最后更新时间
9 天前
1 个月前
1 个月前
1 个月前
README

算子名称:InitRouing

产品支持情况

产品 是否支持
Atlas A2 训练系列产品

功能说明

  • 算子功能:初始化路由,将输入token按expert分组并连续排列。

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
blockDim 输入 AI CORE的数量,比如:Ascend910B是40。 int64_t -
stream 输入 Device端的stream AclrtStream -
in 输入 公式中的输入张量x,shape为(token_num, hidden_size) BFLOAT16 ND
token_table 输入 token到expert的映射表, shape为(expert_num, token_num) int32_t ND
token_list 输入 每个expert累计token数量,shape为(expert_num) int64_t ND
out 输出 同一expert的token连续排列,shape为(expert_num*topk, hidden_size) BFLOAT16 ND
expert_num 输入 expert数量 int64_t ND
copy_byte 输入 每个token拷贝字节数,等于hidden_size * element_size int64_t -

约束说明

  • token_list最后一个元素的值为总token数
  • copy_byte按64bytes对齐

调用说明

torch.ops.npu_ops_transformer_ext.init_routing(block_dim, input, token_table,token_list, output)