已关闭
[Feature]: MoeInitRoutingV3 新增支持quantmode为16、17 #235
wuxiyuan创建于  7月2日关闭于  7月4日
wuxiyuan
7月2日 创建

提交提案之前,请先检索仓库内是否已有相同的提案,如已有请在同一提案中进行讨论。

💻 需求背景、当前现状、期望实现的功能内容、具体的设计方案、以及测试方案

【修改方案】

MoeInitRoutingV3 新增支持 quant_mode 为 16、17(MXFP8 RoundScale + Amax 钳位量化)

修改概述

为 npu_moe_init_routing_v2 算子新增两种量化模式支持:

  • quant_mode=16:MXFP8 RoundScale + Amax 钳位量化,输出类型为 float8_e5m2
  • quant_mode=17:MXFP8 RoundScale + Amax 钳位量化,输出类型为 float8_e4m3fn

修改文件

文件 修改内容
docs/zh/custom_APIs/torch_npu/torch_npu-npu_moe_init_routing_v2.md 更新 API 文档,补充所有 quant_mode 取值说明及对应输出数据类型、shape 描述
op_plugin/ops/opapi/MoeInitRoutingV2KernelNpuOpApi.cpp 新增 QUANT_MODE_MXFP8_ROUNDSCALE_AMAX_E5M2(16) 和 QUANT_MODE_MXFP8_ROUNDSCALE_AMAX_E4M3FN(17) 常量定义;扩展 IsQuantModeMXFP8 判断逻辑;在 expanded_x 构建 switch 中复用 MXFP8 分支
op_plugin/python/meta/_meta_registrations.py meta 函数中 quant_mode 合法值列表新增 16、17;dtype 和 scale shape 推导逻辑中将 16/17 与 2/3 合并处理
test/test_custom_ops/test_npu_moe_init_routing_v2.py 新增 mxfp8_roundscale_quantize 量化参考实现(含 amax 钳位);golden 计算逻辑适配 quant_mode 16/17;测试参数列表新增 16、17

替代方案

补充说明

欢迎加入社区,感谢您对社区的贡献 🎉!

likedislike
ascend-robotascend-robot成员
7月2日 添加了label:feature
ascend-robotascend-robot成员
7月4日 关闭了 issue
ascend-robotascend-robot成员
7月13日 添加了label:resolved