Closed
[Requirement|需求建议]: weight_quant_preprocess 算子新增 A16W8 数据流支持 #3655
zhangquanxinAdd on  25 days agoClosed at  22 days ago
zhangquanxin
zhangquanxinCollaborators
25 days ago Create

一、背景信息 (必填)

weight_quant_preprocess 是伪量化 Matmul(weight_quant_batch_matmul_v2,wqbmmv2)的前置参数预处理算子,负责 weight 格式处理与 weightScale/weightOffset/bias 的参数同步,当前已支持 A8W4、A16S4、A16F4、A16MXF4 数据流。为支撑 x 为 FP16/BF16、weight 为 INT8 或 FP8(FLOAT8_E4M3FN/HIFLOAT8)的 A16W8 伪量化推理场景,需为该算子新增 A16W8 数据流支持。

二、价值/作用 (必填)

  1. 覆盖 A16W8 全部 10 个子场景:A16INT8 per-tensor/per-channel/per-group × 转置/非转置(6 个)+ A16FP8HIF8 per-channel(E4M3FN/HIFLOAT8)× 转置/非转置(4 个)
  2. 打通 aclnn 直调与 PyTorch 单算子(torch_npu.npu_weight_quant_preprocess)两条通路,支撑 A16W8 量化模型在 Ascend 950 系列部署
  3. 全场景 ND 直拷透传,Host 侧改动为注册表纯增量插入,存量数据流(A8W4/A16S4/A16F4/A16MXF4)零改动

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)

Aclnn 直调(aclnnWeightQuantPreprocess);PyTorch 单算子(torch_npu.npu_weight_quant_preprocess,op-plugin 侧配套 judge/prepare 注册表同步新增,另仓交付)。

3.2 总体设计

Host 侧沿用"数据流注册表"声明式架构(每个数据流 = 一个 DataFlowEntry{judge, checks, processes},按序匹配),在 NpuArch::DAV_3510 注册表追加 4 个条目;process 全部为 ViewCopy ND 物理直拷 + scale/offset/bias 透传,无独立 kernel。

数据流枚举新增 4 个:MM_A16INT8_PERTENSOR(8)、MM_A16INT8_PERCHANNEL(9)、MM_A16INT8_PERGROUP(10)、MM_A16FP8HIF8_PERCHANNEL(11)。

3.2.1 算子支持的数据类型
  • x:FP16 / BF16
  • weight:INT8 / FLOAT8_E4M3FN / HIFLOAT8(2D)
  • weightScale:FP16 / BF16;weightOffset:FP16 / BF16(可选);bias:可选
  • 输出 outWeight/outWeightScale/outWeightOffset/outBias:与对应输入 ND 直拷透传(格式/形状/数据类型一致)
3.2.2 host侧设计
  • judge 按序匹配:weight dtype + x dtype ∈ {FP16,BF16} + weight 为 2D + weightScale dtype ∈ {FP16,BF16} + xScaleDtype == ACL_DT_UNDEFINED + 量化粒度(由 weightScale 形状隐式判定:{1}/{1,1}=per-tensor、{N}/{1,N}=per-channel、{G,N}(G>1)=per-group)
  • checks 仅保留输出一致性检查(CheckOutWeightSameAsInput / CheckOutWeightScaleSameAsInput / CheckOutWeightOffsetOptionalSameAsInput / CheckOutBiasOptionalSameAsInput);weight/scale/offset/bias 输入侧的 dtype/format/shape/kGroupSize 约束本算子不消费不处理,统一下放下游 wqbmmv2 拦截(与 A16S4 基线口径一致)
  • processes:ProcessWeightDirectCopy / ProcessWeightScaleDirectCopy / ProcessWeightOffsetDirectCopy / ProcessBiasDirectCopy(FP8 流无 offset process);转置 weight(strides [1,K] 视图)直拷天然保持视图语义,无需按转置拆分
3.2.3 kernel侧设计

无新增 kernel:复用 ViewCopy 做 ND 物理字节级直拷;INT8/FP8 均为 1 字节元素,无 A16W4 的 4-bit 打包语义。

3.3 支持硬件

Ascend 950PR / 950DT(NpuArch::DAV_3510)

3.4 算子约束限制

  • 仅支持 outWeight 为 ND 格式(全 ND 直拷);outWeight 传 NZ 等 judge 未命中场景报 161002(dataflow is UNKNOWN),不允许静默落入其它数据流
  • A16W8 无 xScale 语义,xScaleDtype 必须为 ACL_DT_UNDEFINED
  • FP8(E4M3FN/HIFLOAT8)流仅支持 per-channel 量化,且 weightOffset 无透传语义(约束下放下游 wqbmmv2 拦截)
  • kGroupSize:per-tensor/per-channel 场景为 0,per-group 场景 >0 且等于 ceildiv(K, g),粒度匹配校验下放下游
  • 存量数据流(A8W4/A16S4/A16F4/A16MXF4)行为零改动

💡 备注(选填)

配套交付:docs/aclnnWeightQuantPreprocess.md(约束块与调用示例)、examples/arch35/test_aclnn_weight_quant_preprocess_a16w8.cpp(10 场景示例)、tests/ut/op_host/test_aclnn_weight_quant_preprocess_l2.cpp(A16W8 UT 用例)。实现 PR:!5301

likedislike
zhangquanxin
zhangquanxinCollaborators
25 days ago Comments:

/assign @zhangquanxin

likedislike
CANN-robotCANN-robotCollaborators
25 days ago Assign zhangquanxin as responsible person
zhangquanxinzhangquanxinCollaborators
25 days ago Added label: requirement
陈思陈思Collaborators
25 days ago Assign sun-yibo1 as responsible person
陈思陈思Collaborators
25 days ago Assign xieshengwei1024 as responsible person
陈思陈思Collaborators
25 days ago Assign liuyang_806 as responsible person
陈思陈思Collaborators
25 days ago Removed responsible person sun-yibo1
陈思陈思Collaborators
25 days ago Removed responsible person xieshengwei1024
陈思陈思Collaborators
25 days ago Removed responsible person liuyang_806
CANN-robotCANN-robotCollaborators
22 days ago Close issue
CANN-robotCANN-robotCollaborators
22 days ago Added label: resolved