CANN/cann-recipes-infer MoE路由分组量化算子

custom-npu_moe_init_routing_group_quant

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer

产品支持情况

功能说明

• 算子功能：MoE的routing计算，根据MoE Gating TopK的计算结果做routing处理，支持不量化和动态量化模式。

• 计算公式：

  1.对输入expertIdx做排序，得出排序后的结果sortedExpertIdx和对应的序号sortedRowIdx：

  $$ sortedExpertIdx, sortedRowIdx=keyValueSort(expertIdx,rowIdx) $$

  2.以sortedRowIdx做位置映射得出expandedRowIdx：

  $$ expandedRowIdx[sortedRowIdx[i]]=i $$

  3.在drop模式下，对sortedExpertIdx的每个专家统计直方图结果，得出expertTokensCountOrCumsum：

  $$ expertTokensCountOrCumsum[i]=Histogram(sortedExpertIdx) $$

  4.计算quant结果：

  • 动态quant：

    • 若不输入scale： $$ dynamicQuantScaleOut = row_max(abs(x)) / 127 $$

      $$ quantResult = round(x / dynamicQuantScaleOut) $$

    • 若输入scale: $$ dynamicQuantScaleOut = row_max(abs(x * scale)) / 127 $$

      $$ quantResult = round(x / dynamicQuantScaleOut) $$

  5.对quantResult取前NUM_ROWS个sortedRowIdx的对应位置的值，得出expandedXOut：

  $$ expandedX[i]=quantResult[sortedRowIdx[i]%NUM_ROWS] $$

  6.expandedRowIdx的有效元素数量availableIdxNum计算方式为，expertIdx中activeExpertRange范围内的元素的个数 $$ availableIdxNum = |{x\in expertIdx| expert_start \le x<expert_end \ }| $$

函数原型

custom.npu_moe_init_routing_group_quant(Tensor x, Tensor expert_idx, Tensor? scale=None, Tensor? offset=None, int active_num=-1, int expert_capacity=-1, int expert_num=-1, int drop_pad_mode=-1, int expert_tokens_num_type=-1, bool expert_tokens_num_flag=False, int quant_mode=-1, SymInt[] active_expert_range, int row_idx_type=-1, int group_size=128) ->(Tensor, Tensor, Tensor, Tensor)

参数说明

约束说明

• 输入值域限制：

  • active_num 当前未使用，校验需等于NUM_ROWS*K。
  • expert_capacity 当前未使用，仅校验非空。
  • offset 当前未使用。
  • drop_pad_mode 当前只支持0，代表 Dropless 场景。
  • expert_tokens_num_type 当前只支持 1 和 2，分别代表 count 模式和 key_value 模式。
  • expert_tokens_num_flag 只支持 true，代表输出 expert_tokens_count_or_cumsum。
  • quant_mode: 支持-1、0、1、2、3、4、5，其中-1表示不量化，0表示静态量化（输出INT8），1表示动态量化（输出INT8），2表示MXFP8量化到FLOAT8_E5M2，3表示MXFP8量化到FLOAT8_E4M3FN。4模式是groupSize固定为128的PerGroup量化，expand_x量化到float8_e5m2类型，scale为float类型。5模式是groupSize固定为128的PerGroup量化，expand_x量化到float8_e4m3类型，scale为float类型。

• 输入shape约束：

  • x ： shape为 (N, H)
  • expert_id : shape为(N, topK)
  • scale : shape为(N, )或者(expert, H)

• 输出shape约束：

  • expanded_x : shape为 (N*K, H)
  • expanded_row_idx : shape为 (N*K)
  • expert_tokens_count_cumsum : expert_tokens_num_type为1时，shape为 (expert, ) expert_tokens_num_type为2时，shape为 (expert, 2)
  • expanded_scale : quant_mode = 4或者quant_mode = 5时，shape为 (N*K, M)，其中M=Ceil(H, 128)

• 该接口支持推理场景下使用。

• 该接口支持图模式。

• 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。

调用示例

• 详见 test_npu_moe_init_routing_group_quant.py

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer