Closed
[Requirement|需求建议]:WeightQuantPreprocess + transdata Ascend960适配 #3686
zhouhang96Add on  21 days agoClosed at  20 days ago
zhouhang96Collaborators
21 days ago Create

一、背景信息 (必填)

为 aclnnWeightQuantPreprocess 大接口增加 Ascend960dt(DAV_9201)适配,支持低比特量化权重及 scale、offset、bias 的预处理,并补充依赖的 TransData ND→NZ 转换能力。

二、价值/作用 (必填)

统一量化参数预处理入口,为下游量化矩阵乘提供符合要求的数据布局,减少调用方对权重格式转换、scale 重排和编码的处理。

三、设计方案 (必填)

通过 ACLNN 两阶段接口调用:

  1. aclnnWeightQuantPreprocessGetWorkspaceSize:校验参数、构建执行任务并返回 workspace 大小。
  2. aclnnWeightQuantPreprocess:在指定 stream 上执行预处理任务。
3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)

Aclnn直调 pytorch单算子

3.2 总体设计
3.2.1 算子支持的数据类型

本次新增 Ascend960 通路主要支持:
参数 数据类型
计算类型 xDtype FLOAT16、BFLOAT16、FLOAT8_E4M3FN,按数据流选择
weight INT4、FLOAT4_E2M1、HIFLOAT4
weightScale FLOAT16、BFLOAT16、FLOAT8_E8M0、HIFLOAT4_SCALE,按数据流匹配
offset、bias 按对应数据流校验并处理
支持 A8W4、A16W4、A16MXFP4、A16HIF4 等数据流,类型组合不能任意搭配。

3.2.2 host侧设计

按 NpuArch 选择数据流注册表,通过 judge 匹配数据流,再依次执行参数检查和处理函数。
复用已有公共模板,针对 9201 完成权重 ND→NZ 转换、必要的转置,以及 scale 的重排、编码或直拷。各子算子使用同一个 executor,由框架统一计算 workspace。保留原有 3510 数据流行为。

3.2.3 kernel侧设计

大接口不新增独立 kernel,组合调用已有子算子完成处理。
TransData 复用现有 SIMT 转换 kernel,补充 Ascend960 的构建及类型配置,并在 tiling 中将 HIFLOAT4 按两个元素打包为一个字节处理。

3.3 支持硬件

本次新增 Ascend960 支持,保留已有硬件支持。

3.4 算子约束限制

本次新增 Ascend960 通路的 weight 为二维,支持的布局按数据流校验。
输入、输出的 dtype、shape、stride 和 NZ storageShape 必须满足对应通路要求。
MX 通路的 group size 为 32,HIF4 通路为 64;其他通路按注册配置校验。
C0=64 通路要求 K、N 满足相应对齐条件。
offset、bias 的支持范围及输出布局以对应数据流约束为准。

💡 备注(选填)

likedislike
陈思
陈思Collaborators
21 days ago Comments:

/assign @zhouhang96

likedislike
CANN-robotCANN-robotCollaborators
21 days ago Assign zhouhang96 as responsible person
CANN-robotCANN-robotCollaborators
20 days ago Close issue
CANN-robotCANN-robotCollaborators
20 days ago Added label: resolved