已关闭
[Feature]: npu_weight_quant_preprocess 支持 A16W4(INT4/FP4/MXFP4)紧凑排布数据流 #444
马琦钧创建于  23 天前关闭于  22 天前
马琦钧
马琦钧成员
23 天前 创建

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

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

需求背景:WeightQuantBatchMatmulV2 在 Ascend950 上新增 A16W4(INT4/FP4/MXFP4)4-bit 紧凑排布(uint8 载体)weight 支持(ops-nn 侧见 cann/ops-nn#8545),需要 PTA 侧 npu_weight_quant_preprocess 配套支持对应数据流。

当前现状npu_weight_quant_preprocess 仅支持 MX A8W4 / int8 等存量流程,不支持 A16S4(INT4 per-tensor/per-channel/per-group)与 A16F4(FP4 per-group/MX)的 4-bit 紧凑 weight 预处理。

期望实现的功能

  • A16S4 INT4:per-tensor / per-channel / per-group 全场景支持
  • A16F4 FP4:per-group / MX 支持
  • torch 接口保持原 9 参形式,无新增参数

设计方案

  • 入口按 A16S4 + scale 形状分流:per-tensor(scale 单元素)/ per-channel({N}/{1,N})/ per-group({G,N},G>1)
  • per-tensor:不支持 NZ,转置/非转置一律 ND 直拷
  • 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)
  • npu_weight_quant_batchmatmul 支持 uint8 载体 4-bit ND weight 按打包方向还原逻辑 K/N;is_weight_nz_4bit_compact 识别 NZ_C0_8/NZ_C0_16
  • meta 推导:uint8 载体 + 4-bit weight_dtype 时按 scale 末维推导逻辑 N

测试方案

  • eager:per-tensor/per-channel/per-group、转置自动 ND、负向用例,与 CPU golden bit-exact 比对
  • 图模式(torchair 静态+动态)全数据流覆盖
  • 存量 MX A8W4 / int8 流程回归不受影响

替代方案

补充说明

对应 PR:https://gitcode.com/Ascend/op-plugin/pulls/5676

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

likedislike
ascend-robotascend-robot成员
23 天前 添加了label:feature
马琦钧
马琦钧成员
23 天前 评论:

按 op-plugin 仓库约定(README:问题请提交至 Ascend/pytorch Issues),本 issue 已迁移至 https://gitcode.com/Ascend/pytorch/issues/4358 ,此处关闭。

likedislike
ascend-robotascend-robot成员
22 天前 关闭了 issue
ascend-robotascend-robot成员
22 天前 添加了label:resolved