已关闭
[Requirement|需求建议]: WeightQuantPreprocess 新增 A16W4(INT4/FP4/MXFP4)紧凑排布数据流支持 #2851
马琦钧创建于  8月25日关闭于  8月27日
马琦钧
马琦钧成员
8月25日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

一、背景信息 (必填)

WeightQuantBatchMatmulV2 在 Ascend950 上新增 A16W4(INT4/FP4/MXFP4)4-bit 紧凑排布(uint8 载体)weight 支持(ops-nn 侧:cann/ops-nn#8545),需要 WeightQuantPreprocess 算子配套支持对应的预处理数据流。当前该算子仅支持 MX A8W4 / int8 等存量流程,不支持 A16S4(INT4 per-tensor/per-channel/per-group)与 A16F4(FP4 per-group/MX)的 4-bit 紧凑 weight 预处理。

二、价值/作用 (必填)

补齐 A16W4 紧凑排布 weight 的预处理通路,使 Ascend950 上 A16S4/A16F4 场景可以经 preprocess 完成 ND→NZ_C0_16 转换或 ND 直拷,支撑 wqbmmv2 的 4-bit 紧凑 weight 输入,降低 weight 显存占用与搬运带宽。

三、设计方案 (必填)

  • 新增 A16S4 INT4(per-tensor / per-channel / per-group)与 A16F4 FP4(per-group / MX)数据流
  • ND/NZ 内部路由:per-tensor 一律 ND 直拷(不支持 NZ);per-channel / per-group 非转置 → NZ_C0_16 转换,转置 → ND 直拷,均不支持则报错
  • FORMAT_FAKE_TO_REAL 映射扩展:NZ_C0_8→NZ_C0_16(A16 4-bit),保留 NZ_C0_16→NZ_C0_32(A8)
  • 4-bit TransData dstFormat 统一为 NZ_C0_16;TransData GE proto 补充 DT_INT4
  • registry 校验加固:outWeight 空指针防护、kGroupSize 校验、weightOffset 在 A16S4 数据流透传
3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)
  • Aclnn 直调:aclnnWeightQuantPreprocess
  • PyTorch:torch_npu npu_weight_quant_preprocess(经 op-plugin,见 Ascend/op-plugin#5676,对应 issue Ascend/pytorch#4358)
3.2 总体设计
3.2.1 算子支持的数据类型
  • weight:INT4 / FLOAT4_E2M1(4-bit 紧凑排布,uint8 载体)
  • 激活侧:FLOAT16 / BF16
3.2.2 host侧设计
  • judge 按 transpose 与输出 format 拆分;ND/NZ 不体现在数据流命名上,由 preprocess 内部分支处理
  • A16S4 非转置 preprocess 输出限定 NZ_C0_16;INT4 紧凑 weight 直拷时按 uint8 view 处理
3.2.3 kernel侧设计
  • 本需求仅涉及 host 侧 preprocess 与 registry 改动,无新增 kernel
3.3 支持硬件

Ascend950

对应 PR:https://gitcode.com/cann/ops-math/pulls/4495

likedislike
Ssunchun成员
8月25日 将 sunchun 设为负责人
马琦钧马琦钧成员
8月25日 关联了pull request:docs: weight_quant_preprocess 补充 A16W4 数据流文档与 example
sunchun成员
8月25日 评论:

/assign @马琦钧

likedislike
CANN-robot
CANN-robot成员
8月25日 评论:

Notice

This issue can not be assigned to 马琦钧. Please try to assign to the repository members.

likedislike
Ssunchun成员
8月25日 将 maqijun 设为负责人
Ssunchun成员
8月25日 移除了负责人 sunchun
CANN-robotCANN-robot成员
8月27日 关闭了 issue
CANN-robotCANN-robot成员
8月27日 添加了label:resolved