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 数据流支持。
weight_quant_preprocess
torch_npu.npu_weight_quant_preprocess
Aclnn 直调(aclnnWeightQuantPreprocess);PyTorch 单算子(torch_npu.npu_weight_quant_preprocess,op-plugin 侧配套适配,另仓交付)。图模式不在本期范围。
aclnnWeightQuantPreprocess
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。
DataFlowEntry{judge, checks, processes}
NpuArch::DAV_3510
数据流枚举新增 1 个:MM_A8F4_TCG(15)。
MM_A8F4_TCG(15)
{K, N}
{ceildiv(K, kGroupSize), N}
{ceildiv(N, 32), ceildiv(K, 16), 16, 32}
IsMMA8F4TcgTransDataFlow
IsMMA8F4TcgNonTransDataFlow
{ceildiv(K,32),N}
INPUT_4BIT_BASE_CHECKS
CheckOutWeightNz<NZ_C0_32, FORMAT_FRACTAL_NZ_C0_32, NON_TRANS_WEIGHT>
ProcessWeightDirectCopy
ProcessWeightNd2Nz<IDX_0, NON_TRANS_WEIGHT>
无新增 kernel:转置复用 ViewCopy ND 物理字节级直拷;非转置复用既有 FLOAT4_E2M1 ND→FRACTAL_NZ_C0_32 TransData 注册(非转置喂入使 c0 轴落在 N),仅输出重标签 format。
FLOAT4_E2M1 ND→FRACTAL_NZ_C0_32
Ascend 950PR / 950DT(NpuArch::DAV_3510)
💡 备注(选填)
配套交付: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
docs/aclnnWeightQuantPreprocess.md
examples/arch35/test_aclnn_weight_quant_preprocess_f8f4_tcg.cpp
tests/ut/op_host/test_aclnn_weight_quant_preprocess_l2.cpp
一、背景信息 (必填)
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 数据流支持。二、价值/作用 (必填)
torch_npu.npu_weight_quant_preprocess)两条通路,支撑 F8F4 T-CG 量化模型在 Ascend 950 系列部署三、设计方案 (必填)
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 算子支持的数据类型
{K, N}){ceildiv(K, kGroupSize), N}且分组数大于 1;weightOffset:可选透传;bias:可选透传{ceildiv(N, 32), ceildiv(K, 16), 16, 32}(N 块在前,区别于 MX A8W4 转置 NZ 的 K 块在前)3.2.2 host侧设计
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 互斥INPUT_4BIT_BASE_CHECKS(含打包维 N 偶数校验)+ 输出 NZ 校验模板CheckOutWeightNz<NZ_C0_32, FORMAT_FRACTAL_NZ_C0_32, NON_TRANS_WEIGHT>+ 输出一致性尾部ProcessWeightDirectCopy;非转置ProcessWeightNd2Nz<IDX_0, NON_TRANS_WEIGHT>(既有模板实例化);scale/offset/bias 均 DirectCopy 透传3.2.3 kernel侧设计
无新增 kernel:转置复用 ViewCopy ND 物理字节级直拷;非转置复用既有
FLOAT4_E2M1 ND→FRACTAL_NZ_C0_32TransData 注册(非转置喂入使 c0 轴落在 N),仅输出重标签 format。3.3 支持硬件
Ascend 950PR / 950DT(NpuArch::DAV_3510)
3.4 算子约束限制
💡 备注(选填)
配套交付:
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