Closed
[Requirement|需求建议]: weight_quant_preprocess 算子新增 A8W4(F8F4) T-CG 数据流支持 #3839
马琦钧Add on  19 days agoClosed at  19 days ago
马琦钧
马琦钧Collaborators
19 days ago Create

一、背景信息 (必填)

weight_quant_preprocess 是伪量化 Matmul 的前置参数预处理算子,负责 weight 格式处理与 weightScale/weightOffset/bias 的参数同步,当前已支持 A8W4(MX)、A16S4、A16F4、A16MXF4、A16INT8、A16FP8HIF8 数据流。为支撑 QuantBatchMatmulV4 在 Ascend 950 上的 T-CG 伪量化场景(x 为 FLOAT8_E4M3FN、weight 为 FLOAT4_E2M1 紧凑 4-bit、per-group 量化),需为该算子新增 A8W4(F8F4) T-CG 数据流支持。

二、价值/作用 (必填)

  1. 覆盖 MM_A8F4_TCG 2 个子场景:weight 转置(ND 直拷透传)+ weight 非转置(ND→FRACTAL_NZ_C0_32 分形转换)
  2. 打通 aclnn 直调与 PyTorch 单算子(torch_npu.npu_weight_quant_preprocess)两条通路,支撑 F8F4 T-CG 量化模型在 Ascend 950 系列部署
  3. Host 侧改动为注册表纯增量插入,存量数据流(A8W4/A16S4/A16F4/A16MXF4/A16INT8/A16FP8HIF8)行为零改动

三、设计方案 (必填)

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

Aclnn 直调(aclnnWeightQuantPreprocess);PyTorch 单算子(torch_npu.npu_weight_quant_preprocess,op-plugin 侧配套适配,另仓交付)。图模式不在本期范围。

3.2 总体设计

Host 侧沿用"数据流注册表"声明式架构(每个数据流 = 一个 DataFlowEntry{judge, checks, processes},按序匹配),在 NpuArch::DAV_3510 注册表追加 2 个条目:转置条目为 ViewCopy ND 物理直拷 + scale/offset/bias 透传;非转置条目复用 MX 的 TransData FLOAT4_E2M1 ND→FRACTAL_NZ_C0_32 注册(纯字节级 permutation),无新增 kernel。

数据流枚举新增 1 个:MM_A8F4_TCG(15)。

3.2.1 算子支持的数据类型
  • x:FLOAT8_E4M3FN
  • weight:FLOAT4_E2M1(紧凑 4-bit,每字节打包 2 个值,2-D {K, N})
  • weightScale:FP16 / BF16,per-group 2-D {ceildiv(K, kGroupSize), N} 且分组数大于 1;weightOffset:可选透传;bias:可选透传
  • outWeight:转置场景为 ND(直拷透传,与输入 weight 同 view shape/strides);非转置场景为 FRACTAL_NZ_C0_32,storage shape {ceildiv(N, 32), ceildiv(K, 16), 16, 32}(N 块在前,区别于 MX A8W4 转置 NZ 的 K 块在前)
  • outWeightScale/outWeightOffset/outBias:与对应输入直拷透传(格式/形状/数据类型一致)
3.2.2 host侧设计
  • judge 按转置拆 2 个(IsMMA8F4TcgTransDataFlow / IsMMA8F4TcgNonTransDataFlow,与 A16MXF4 同机制):FP4_E2M1 weight + FP16/BF16 per-group scale({ceildiv(K,32),N},G>1)+ x 为 FP8_E4M3FN + xScaleDtype == ACL_DT_UNDEFINED,与 MX(scale=E8M0)、A16(x=F16/BF16)按 dtype 互斥
  • checks 全部复用既有组件,TCG 零新增:转置条目为「输入侧非空/ND/4-bit 打包视图校验 + 输出四项 SameAsInput 一致性」;非转置条目复用 INPUT_4BIT_BASE_CHECKS(含打包维 N 偶数校验)+ 输出 NZ 校验模板 CheckOutWeightNz<NZ_C0_32, FORMAT_FRACTAL_NZ_C0_32, NON_TRANS_WEIGHT> + 输出一致性尾部
  • processes:转置 ProcessWeightDirectCopy;非转置 ProcessWeightNd2Nz<IDX_0, NON_TRANS_WEIGHT>(既有模板实例化);scale/offset/bias 均 DirectCopy 透传
3.2.3 kernel侧设计

无新增 kernel:转置复用 ViewCopy ND 物理字节级直拷;非转置复用既有 FLOAT4_E2M1 ND→FRACTAL_NZ_C0_32 TransData 注册(非转置喂入使 c0 轴落在 N),仅输出重标签 format。

3.3 支持硬件

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

3.4 算子约束限制

  • kGroupSize 须为 32:取值及与 weightScale 量化粒度的匹配约束由下游 QuantBatchMatmulV4 校验,本算子不拦截
  • 紧凑 4-bit 打包维须为偶数:非转置时 N 为偶数(本算子拦截);转置时 K 为偶数(由下游 QuantBatchMatmulV4 校验)
  • offset 仅透传,语义约束由下游 QuantBatchMatmulV4 校验(下游 T-CG 场景当前不支持非空 offset,建议传 nullptr)
  • xScaleDtype 必须为 ACL_DT_UNDEFINED(T-CG 无独立 xScale 入参)
  • judge 未命中(如 dtype 组合不支持)报 161002(dataflow is UNKNOWN),不允许静默落入其它数据流;非转置场景 outWeight format/storage 不满足 C0_32 约束由 checks 拦截报错
  • 存量数据流(A8W4/A16S4/A16F4/A16MXF4/A16INT8/A16FP8HIF8)行为零改动

💡 备注(选填)

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

likedislike
马琦钧马琦钧Collaborators
19 days ago Added label: requirement
CANN-robotCANN-robotCollaborators
19 days ago Close issue
CANN-robotCANN-robotCollaborators
19 days ago Added label: resolved