已开启
[Feature]: 【feat】MoeInitRoutingV3算子支持mxfp4透传&Mxfp4量化 #82
cpy_123456创建于  5月13日
cpy_123456
5月13日 创建

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

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

MoeInitRoutingV3算子在非量化模式(quant_mode=-1)下
第0输入x新增支持类型float4_e2m1。按原逻辑搬运到第0输出expanded_x,类型仍为float4_e2m1
第2输入scale作为可选传参,新增支持shape为(N, CeilDiv(H, 64), 2)&&Dtype为float8_e8m0的数据。按原逻辑搬运到第3输出expanded_scale,类型仍为float8_e8m0
新增mxfp4量化模式(quant_mode=9)下
第0输入x支持类型bf16,fp16。量化后搬运到第0输出expanded_x,类型为float4_e2m1
量化的scale输出到到第3输出expanded_scale,shape为(N, CeilDiv(H, 64), 2)&&Dtype为float8_e8m0的数据

替代方案

输入Tensor
x:第0输入,表示MoE要处理的N行tokens,每个token长度为H
必须输入
Shape: (N, H)
DType: float16、bfloat16、float32、int8、hifloat8、float8_e4m3fn、float8_e5m2、float4_e2m1
Format: ND
约束:
quant_mode为9时:DType仅支持float16、bfloat16
expert_idx:第1输入,表示x的每一行对应的K个专家ID
必须输入
Shape: (N, K)
DType: int32
Format: ND
约束: 每个元素id,要求0<=id<expert_num
scale:第2输入,量化模式时用作smoothScale,非量化模式时携带与x的每一行一起排序
可选输入。quant_mode为9时,不使用、不校验该输入
Shape:
quant_mode为-1时: (N,)或空
quant_mode为1时: (expertEnd-expertStart, H),其中activateExpertRangeOptional=[expertStart, expertEnd]
quant_mode为7时:(1,) //每个tensor使用同一个scale

DType: float32
Format: ND
offset:第3输入,用于量化。目前不使用、不校验
可选输入。目前不使用、不校验,可以传任意值
1.2.4 输入属性
quant_mode:第6属性,表示是否对expanded_x进行量化,以及量化的模式
可选输入(默认值为-1)
数据类型:int32
约束(取值):
-1:不量化
1:动态量化,输出expanded_x量化到int8
2:MXFP8量化,输出expanded_x量化到float8_e5m2
3: MXFP8量化,输出expanded_x量化到float8_e4m3fn
​6: HIF8直传量化,输出expanded_x量化到HIF8
​7: HIF8静态量化​(pertensor),输出expanded_x量化到HIF8
​8: HIF8动态量化​(pertoken),输出expanded_x量化到HIF8
9: MXFP4量化,输出expanded_x量化到float4_e2m1
1.2.5 输出Tensor
expanded_x:第0输出,表示x根据expert_idx和active_expert_range扩增、排序、筛选后的tokens
Shape:(NK, H)
有效Shape:(availIdxNum, H),即前availIdxNum行有效,后面行无意义(随机值)
DType:
quant_mode为-1时:与x一致
quant_mode为1时:int8
quant_mode为2时:float8_e5m2
quant_mode为3时:float8_e4m3fn
quant_mode为6时:对应直转方式,输出为hif8
quant_mode为7时:对应外部传入scale的静态量化(pertensor),输出为hif8
quant_mode为8时:对应无scale的hif8动态量化,scale动态生成(pertoken),输出为hif8
quant_mode为9时:float4_e2m1
Format: ND
expanded_row_idx:第1输出,表示x与expanded_x的映射关系,具体由属性rowIdxType决定
Shape:(NK,)
有效Shape:
rowIdxType为0时:(NK,)整个Shape有效,其中有效元素与无效元素随机出现,无效元素填充-1
rowIdxType为1时:(availIdxNum,),即前availIdxNum个元素有效,后面元素无意义(随机值)
DType:int32
Format: ND
expert_tokens_count_or_cumsum:第2输出,表示每个有效专家ID对应的tokens的数目,具体由属性expert_tokens_num_type决定
Shape:
expert_tokens_num_type为1时,即Count模式:(expertEnd-expertStart,)
expert_tokens_num_type为2时,即KeyValue模式:(expert_num, 2)。每行为(id, count),表示有效专家id对应count个tokens。
有效Shape:
expert_tokens_num_type为1时:(expertEnd-expertStart,)整个Shape有效
expert_tokens_num_type为2时:(nonZeroCount, 2),其中nonZeroCount为去除了count值为0的(id,count)的行数。若nonZeroCount<expert_num,则有效行末尾会填充(0,0)表示有效行的结束。
expanded_scale:第3输出,表示量化时计算的scale,或是非量化时随x携带的scale,具体由属性quant_mode和输入scale决定
Shape:
quant_mode为-1或1时:(NK,)
quant_mode为2或3时:(NK, M),其中M=CeilAlign(CeilDiv(H,32),2),为token长度H向上整除32(一个block大小)后,再向上对齐到2的倍数
quant_mode为6时:(NK,)
quant_mode为7时:无输出
quant_mode为8时:(N*K, 1 )
quant_mode为9时:(N, CeilDiv(H, 64), 2)
有效Shape:
quant_mode为-1,未输入scale时:整个Shape的元素均无意义(随机值)
quant_mode为-1,有输入scale时:(availIdxNum,),即前availIdxNum个元素有效,后面元素无意义(随机值)
quant_mode为1时:(availIdxNum,)
quant_mode为2或3时:(availIdxNum, M),M值与上述Shape描述一致
quant_mode为6时,未输入scale时:整个Shape的元素均无意义(随机值)
quant_mode为6时,有输入scale时:(availIdxNum,),即前availIdxNum个元素有效,后面元素无意义(随机值)
quant_mode为7时:无输出
quant_mode为8时:(availIdxNum, 1)
quant_mode为9时:(N, CeilDiv(H, 64), 2)
DType:
quant_mode为-1、1时:float32
quant_mode为2、3时:float8_e8m0
quant_mode为6时:float32
quant_mode为7时:无输出
quant_mode为8时:float32
quant_mode为9时:float8_e8m0
Format:ND

补充说明

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

likedislike
ascend-robotascend-robot成员
5月13日 添加了label:feature
ascend-robotascend-robot成员
7月8日 关联了看板:MindStudio ISSUE管理
TorchNPU-BotTorchNPU-Bot
6 天前 添加了label:bot-triaged
TorchNPU-Bot
TorchNPU-Bot
6 天前 评论:

检测到当前 issue 已关联 PR,自动添加标签:bot-triaged

likedislike