feat: WeightQuantPreprocess 新增 A16W4 紧凑排布数据流支持（INT4/FP4/MXFP4）-op-plugin-AtomGit

Mop_plugin/ops/opapi/WeightQuantBatchMatmulV2KernelNpuOpApi.cpp+43-9

	@@ -20,6 +20,7 @@ namespace op_api {
20	using npu_preparation = at_npu::native::OpPreparation;	20	using npu_preparation = at_npu::native::OpPreparation;
21	constexpr int MINIMUM_SHAPE_SIZE = 2;	21	constexpr int MINIMUM_SHAPE_SIZE = 2;
22	const int64_t INT4_NUMS_IN_INT32 = 8;	22	const int64_t INT4_NUMS_IN_INT32 = 8;
		23	+const int64_t B4_IN_UINT8 = 2; // uint8 载体每字节打包 2 个 4-bit（INT4/FP4）
		马琦钧8月25日临时意见
23	at::Tensor npu_weight_quant_batchmatmul(	24	at::Tensor npu_weight_quant_batchmatmul(
24	const at::Tensor& x,	25	const at::Tensor& x,
25	const at::Tensor& weight,	26	const at::Tensor& weight,
	@@ -49,9 +50,35 @@ at::Tensor npu_weight_quant_batchmatmul(
49	"x dim is not the same as weight dim",	50	"x dim is not the same as weight dim",
50	OPS_ERROR(ErrCode::PARAM));	51	OPS_ERROR(ErrCode::PARAM));
51	auto x_k_dim = x.size(x_dim_num - 1);	52	auto x_k_dim = x.size(x_dim_num - 1);
52	- auto weight_k_dim = ((weight.dtype() == at::kInt \|\| weight.dtype() == at::kFloat) && trans_weight)	53	+
53	- ? weight.size(weight_dim_num - MINIMUM_SHAPE_SIZE) * INT4_NUMS_IN_INT32	54	+ // 计算 weight 的 K 维度
54	- : weight.size(weight_dim_num - MINIMUM_SHAPE_SIZE);	55	+ // 对于 INT32/FLOAT 类型的 4-bit 打包（INT4_NUMS_IN_INT32 = 8）
		56	+ // uint8 载体 4-bit 紧凑排布（A16S4 链路）：NZ_C0_16 的 view 是逻辑 [K, N]，NZ_C0_8 的 view 是
		57	+ // 物理打包形状 [K, N/2]；ND 的 view 是物理打包形状——非转置 [K, N/2]（沿 N 打包，行主序连续，
		58	+ // stride(-1)==1 且 stride(-2)==size(-1)）、转置 [K/2, N]（沿 K 打包）。K=2 时转置视图退化为
		59	+ // [1, N] strides [1, 1]，需靠 stride(-2)==size(-1) 排除误判，K/N 需按打包方向还原
		60	+ bool is_int32_float_packed = (weight.dtype() == at::kInt \|\| weight.dtype() == at::kFloat);
		61	+
		62	+ int64_t weight_format = at_npu::native::custom_ops::get_npu_format(weight);
		63	+ aclDataType weight_acl_dtype =
		64	+ weight_dtype.has_value() ? c10_npu::GetAclDataType(weight_dtype.value()) : ACL_DT_UNDEFINED;
		65	+ bool is_4bit_acl_dtype =
		66	+ (weight_acl_dtype == ACL_INT4 \|\| weight_acl_dtype == ACL_FLOAT4_E2M1 \|\| weight_acl_dtype == ACL_FLOAT4_E1M2);
		67	+ bool is_uint8_4bit_nd = (weight.dtype() == at::kByte) && is_4bit_acl_dtype && (weight_format == ACL_FORMAT_ND);
		68	+ bool is_uint8_4bit_nz_c08 =
		69	+ (weight.dtype() == at::kByte) && is_4bit_acl_dtype && (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_8);
		70	+ bool uint8_pack_along_n = is_uint8_4bit_nd &&
		71	+ (weight.stride(weight_dim_num - 1) == 1 &&
		72	+ weight.stride(weight_dim_num - MINIMUM_SHAPE_SIZE) == weight.size(weight_dim_num - 1));
		73	+ bool uint8_pack_along_k = is_uint8_4bit_nd && !uint8_pack_along_n;
		74	+
		75	+ int64_t weight_k_dim = weight.size(weight_dim_num - MINIMUM_SHAPE_SIZE);
		76	+ if (is_int32_float_packed && trans_weight) {
		77	+ weight_k_dim *= INT4_NUMS_IN_INT32; // 8
		78	+ } else if (uint8_pack_along_k) {
		79	+ weight_k_dim *= B4_IN_UINT8; // 沿 K 打包，K 按每字节 2 个 4-bit 还原
		80	+ }
		81	+
55	TORCH_CHECK(	82	TORCH_CHECK(
56	x_k_dim == weight_k_dim,	83	x_k_dim == weight_k_dim,
57	"The k of x and weight should be equal. but x_k_dim is ",	84	"The k of x and weight should be equal. but x_k_dim is ",
	@@ -61,12 +88,17 @@ at::Tensor npu_weight_quant_batchmatmul(
61	OPS_ERROR(ErrCode::PARAM));	88	OPS_ERROR(ErrCode::PARAM));
62	auto out_dim_num = std::max(x_dim_num, weight_dim_num);	89	auto out_dim_num = std::max(x_dim_num, weight_dim_num);
63	auto output_size = op_infer::array_to_small_vector(x.sizes());	90	auto output_size = op_infer::array_to_small_vector(x.sizes());
		91	+ // weight.dim() > x.dim() 时 output_size 需预留 batch 维（batch 维由下方较长 tensor 填充），否则越界写
		92	+ output_size.resize(out_dim_num, 1);
64	output_size[out_dim_num - MINIMUM_SHAPE_SIZE] = x.size(x_dim_num - MINIMUM_SHAPE_SIZE);	93	output_size[out_dim_num - MINIMUM_SHAPE_SIZE] = x.size(x_dim_num - MINIMUM_SHAPE_SIZE);
65	auto weight_size_base = weight.size(weight_dim_num - MINIMUM_SHAPE_SIZE + 1);	94	auto weight_size_base = weight.size(weight_dim_num - MINIMUM_SHAPE_SIZE + 1);
66	- output_size[out_dim_num - MINIMUM_SHAPE_SIZE + 1] =	95	+ if (is_int32_float_packed && !trans_weight) {
67	- ((weight.dtype() == at::kInt \|\| weight.dtype() == at::kFloat) && !trans_weight)	96	+ output_size[out_dim_num - MINIMUM_SHAPE_SIZE + 1] = weight_size_base * INT4_NUMS_IN_INT32;
68	- ? weight_size_base * INT4_NUMS_IN_INT32	97	+ } else if (uint8_pack_along_n \|\| is_uint8_4bit_nz_c08) {
69	- : weight_size_base;	98	+ output_size[out_dim_num - MINIMUM_SHAPE_SIZE + 1] = weight_size_base * B4_IN_UINT8;
		99	+ } else {
		100	+ output_size[out_dim_num - MINIMUM_SHAPE_SIZE + 1] = weight_size_base;
		101	+ }
70	if (x_dim_num == weight_dim_num) {	102	if (x_dim_num == weight_dim_num) {
71	for (auto i = 0; i < out_dim_num - MINIMUM_SHAPE_SIZE; i++) {	103	for (auto i = 0; i < out_dim_num - MINIMUM_SHAPE_SIZE; i++) {
72	TORCH_CHECK(x.size(i) == weight.size(i), "batch of x is diff from batch of weight", OPS_ERROR(ErrCode::PARAM));	104	TORCH_CHECK(x.size(i) == weight.size(i), "batch of x is diff from batch of weight", OPS_ERROR(ErrCode::PARAM));
	@@ -104,10 +136,12 @@ at::Tensor npu_weight_quant_batchmatmul(
104	quant_scale.has_value() ? x.options().dtype(at::kChar) : x.options().dtype(x.scalar_type());	136	quant_scale.has_value() ? x.options().dtype(at::kChar) : x.options().dtype(x.scalar_type());
105	at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);	137	at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);
106		138
		139	+ // uint8 载体 4-bit 紧凑排布（NZ_C0_8）同属 NZ 分形，统一走 aclnnWeightQuantBatchMatmulNz
		140	+ const bool is_weight_nz = (weight_format == ACL_FORMAT_FRACTAL_NZ) \|\|
		141	+ (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2) \|\| (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_8);
		142	+
107	TensorWrapper weight_wrapper = make_wrapper(weight, weight_dtype);	143	TensorWrapper weight_wrapper = make_wrapper(weight, weight_dtype);
108		144
109	- int64_t weight_format = at_npu::native::custom_ops::get_npu_format(weight);
110	- const bool is_weight_nz = (weight_format == ACL_FORMAT_FRACTAL_NZ) \|\| (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2);
111	if (is_weight_nz && c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950) {	145	if (is_weight_nz && c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950) {
112	static const bool is_weight_quant_matmul_nz_available =	146	static const bool is_weight_quant_matmul_nz_available =
113	check_aclnn_kernel_available("aclnnWeightQuantBatchMatmulNz");	147	check_aclnn_kernel_available("aclnnWeightQuantBatchMatmulNz");

Mop_plugin/ops/opapi/WeightQuantPreprocessKernelNpuOpApi.cpp+215-0

	@@ -25,7 +25,9 @@ using npu_preparation = at_npu::native::OpPreparation;
25		25
26	namespace {	26	namespace {
27	constexpr int64_t NZ_16 = 16;	27	constexpr int64_t NZ_16 = 16;
		28	+constexpr int64_t NZ_C0_8 = 8;
28	constexpr int64_t NZ_C0_16 = 16;	29	constexpr int64_t NZ_C0_16 = 16;
		30	+constexpr int64_t DIMS_1 = 1;
29	constexpr int64_t DIMS_2 = 2;	31	constexpr int64_t DIMS_2 = 2;
30	constexpr int64_t DIMS_3 = 3;	32	constexpr int64_t DIMS_3 = 3;
31	constexpr int64_t DIMS_4 = 4;	33	constexpr int64_t DIMS_4 = 4;
	@@ -75,6 +77,16 @@ static bool is_transpose_certain_two_dims(const at::Tensor& tensor, int64_t firs
75	return tensor.stride(first_dim + 1) == tensor.stride(first_dim) * tensor.size(first_dim);	77	return tensor.stride(first_dim + 1) == tensor.stride(first_dim) * tensor.size(first_dim);
76	}	78	}
77		79
		80	+// A16W4（INT4 / FP4 E2M1，非 MX）torch 侧统一用 uint8 载体物理打包（每字节 2 个 4-bit 元素），
		81	+// 与 aclnnWeightQuantBatchMatmulV2 的 uint8 packed 输入约定一致；其他载体直接拒绝
		82	+static void check_a16w4_uint8_carrier(const QuantContext& ctx) {
		83	+ TORCH_CHECK(
		84	+ ctx.weight.scalar_type() == at::kByte,
		85	+ "A16W4 4-bit weight must be packed into a uint8 tensor (2 elements per byte), but got ",
		86	+ ctx.weight.scalar_type(),
		87	+ OPS_ERROR(ErrCode::PARAM));
		88	+}
		89	+
78	bool judge_mm_mx_a8w4(QuantContext& ctx) {	90	bool judge_mm_mx_a8w4(QuantContext& ctx) {
79	aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);	91	aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
80	aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);	92	aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
	@@ -151,6 +163,141 @@ bool judge_gmm_mx_a8w4(QuantContext& ctx) {
151	return false;	163	return false;
152	}	164	}
153		165
		166	+bool judge_mm_a16s4_per_tensor(QuantContext& ctx) {
		167	+ aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
		168	+ aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
		169	+
		170	+ bool x_dtype_match = (x_acl_dtype == ACL_FLOAT16 \|\| x_acl_dtype == ACL_BF16);
		171	+
		172	+ // A16S4 per-tensor：scale 仅含单个元素（{1}/{1,1}），不支持 NZ，一律 ND 直拷（转置/非转置均支持，
		173	+ // ND 直拷为物理透传，与转置状态无关）
		174	+ if (x_dtype_match && weight_acl_dtype == ACL_INT4 && ctx.weight.dim() == DIMS_2 &&
		175	+ ctx.weight_scale.numel() == 1) {
		176	+ check_a16w4_uint8_carrier(ctx);
		177	+ ctx.is_weight_trans = is_transpose_certain_two_dims(ctx.weight, 0);
		178	+ return true;
		179	+ }
		180	+ return false;
		181	+}
		182	+
		183	+bool judge_mm_a16s4_per_channel(QuantContext& ctx) {
		184	+ aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
		185	+ aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
		186	+
		187	+ OP_LOG_DEBUG(
		188	+ "judge_mm_a16s4_per_channel: x_acl_dtype=%d, weight_acl_dtype=%d, "
		189	+ "weight_dim=%d, weight_stride=[%d,%d], weight_size=[%d,%d]",
		190	+ static_cast<int>(x_acl_dtype),
		191	+ static_cast<int>(weight_acl_dtype),
		192	+ static_cast<int>(ctx.weight.dim()),
		193	+ ctx.weight.dim() >= 2 ? static_cast<int>(ctx.weight.stride(0)) : -1,
		194	+ ctx.weight.dim() >= 2 ? static_cast<int>(ctx.weight.stride(1)) : -1,
		195	+ ctx.weight.dim() >= 2 ? static_cast<int>(ctx.weight.size(0)) : -1,
		196	+ ctx.weight.dim() >= 2 ? static_cast<int>(ctx.weight.size(1)) : -1);
		197	+
		198	+ bool x_dtype_match = (x_acl_dtype == ACL_FLOAT16 \|\| x_acl_dtype == ACL_BF16);
		199	+
		200	+ if (x_dtype_match && weight_acl_dtype == ACL_INT4 && ctx.weight.dim() == DIMS_2) {
		201	+ int64_t scale_dim = ctx.weight_scale.dim();
		202	+ bool is_per_channel = (scale_dim == DIMS_1) \|\| (scale_dim == DIMS_2 && ctx.weight_scale.size(0) == 1);
		203	+ // 转置状态内部分流：转置 → ND 直拷，非转置 → NZ 转换（prepare_out_weight_a16s4）
		204	+ if (is_per_channel && ctx.weight_scale.numel() > 1) {
		205	+ check_a16w4_uint8_carrier(ctx);
		206	+ ctx.is_weight_trans = is_transpose_certain_two_dims(ctx.weight, 0);
		207	+ return true;
		208	+ }
		209	+ }
		210	+ return false;
		211	+}
		212	+
		213	+bool judge_mm_a16s4_per_group(QuantContext& ctx) {
		214	+ aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
		215	+ aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
		216	+ aclDataType x_scale_acl_dtype =
		217	+ ctx.x_scale_dtype.has_value() ? c10_npu::GetAclDataType(ctx.x_scale_dtype.value()) : ACL_DT_UNDEFINED;
		218	+
		219	+ bool x_dtype_match = (x_acl_dtype == ACL_FLOAT16 \|\| x_acl_dtype == ACL_BF16);
		220	+
		221	+ // A16S4 per-group：转置状态内部分流，转置 → ND 直拷，非转置 → NZ 转换（prepare_out_weight_a16s4）
		222	+ if (x_dtype_match && weight_acl_dtype == ACL_INT4 && x_scale_acl_dtype == ACL_DT_UNDEFINED &&
		223	+ ctx.weight.dim() == DIMS_2 && ctx.weight_scale.dim() == DIMS_2 && ctx.weight_scale.size(0) > 1) {
		224	+ check_a16w4_uint8_carrier(ctx);
		225	+ ctx.is_weight_trans = is_transpose_certain_two_dims(ctx.weight, 0);
		226	+ return true;
		227	+ }
		228	+ return false;
		229	+}
		230	+
		231	+bool judge_mm_a16f4_nz_pergroup(QuantContext& ctx) {
		232	+ aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
		233	+ aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
		234	+ aclDataType weight_scale_acl_dtype = c10_npu::GetAclDataType(ctx.weight_scale_dtype);
		235	+ aclDataType x_scale_acl_dtype =
		236	+ ctx.x_scale_dtype.has_value() ? c10_npu::GetAclDataType(ctx.x_scale_dtype.value()) : ACL_DT_UNDEFINED;
		237	+
		238	+ bool x_dtype_match = (x_acl_dtype == ACL_FLOAT16 \|\| x_acl_dtype == ACL_BF16);
		239	+ bool scale_dtype_match = (weight_scale_acl_dtype == ACL_FLOAT16 \|\| weight_scale_acl_dtype == ACL_BF16);
		240	+
		241	+ // A16F4 per-group NZ：FP4 weight + per-group scale [G, N]（G > 1），仅支持非转置 weight
		242	+ if (x_dtype_match && scale_dtype_match && weight_acl_dtype == ACL_FLOAT4_E2M1 &&
		243	+ x_scale_acl_dtype == ACL_DT_UNDEFINED && ctx.weight.dim() == DIMS_2 && ctx.weight_scale.dim() == DIMS_2 &&
		244	+ ctx.weight_scale.size(0) > 1 && !is_transpose_certain_two_dims(ctx.weight, 0)) {
		245	+ check_a16w4_uint8_carrier(ctx);
		246	+ ctx.is_weight_trans = false;
		247	+ return true;
		248	+ }
		249	+ return false;
		250	+}
		251	+
		252	+bool judge_mm_a16f4_mx(QuantContext& ctx) {
		253	+ aclDataType x_acl_dtype = c10_npu::GetAclDataType(ctx.x_dtype);
		254	+ aclDataType weight_acl_dtype = c10_npu::GetAclDataType(ctx.weight_dtype);
		255	+ aclDataType weight_scale_acl_dtype = c10_npu::GetAclDataType(ctx.weight_scale_dtype);
		256	+ aclDataType x_scale_acl_dtype =
		257	+ ctx.x_scale_dtype.has_value() ? c10_npu::GetAclDataType(ctx.x_scale_dtype.value()) : ACL_DT_UNDEFINED;
		258	+
		259	+ bool x_dtype_match = (x_acl_dtype == ACL_FLOAT16 \|\| x_acl_dtype == ACL_BF16);
		260	+
		261	+ // A16 MXFP4：FP4 weight + MX scale（E8M0，2D [K/32, N] 连续），仅支持非转置 weight
		262	+ if (x_dtype_match && weight_acl_dtype == ACL_FLOAT4_E2M1 && weight_scale_acl_dtype == ACL_FLOAT8_E8M0 &&
		263	+ x_scale_acl_dtype == ACL_DT_UNDEFINED && ctx.weight.dim() == DIMS_2 && ctx.weight_scale.dim() == DIMS_2 &&
		264	+ !is_transpose_certain_two_dims(ctx.weight, 0)) {
		265	+ check_a16w4_uint8_carrier(ctx);
		266	+ ctx.is_weight_trans = false;
		267	+ return true;
		268	+ }
		269	+ return false;
		270	+}
		271	+
		272	+// 校验镜像 strides 的寻址包络不超出按 numel 最小分配的 buffer；
		273	+// 连续/确定转置视图必过，padding/重叠/expand 视图直接拒绝而不是写出界
		274	+static void check_strides_envelope(c10::IntArrayRef sizes, c10::IntArrayRef strides, int64_t numel) {
		275	+ int64_t max_offset = 0;
		276	+ for (size_t i = 0; i < sizes.size(); ++i) {
		277	+ TORCH_CHECK(strides[i] > 0, "expanded/overlapped view is not supported here", OPS_ERROR(ErrCode::PARAM));
		278	+ max_offset += (sizes[i] - 1) * strides[i];
		279	+ }
		280	+ TORCH_CHECK(max_offset < numel, "mirrored strides exceed output allocation: max offset ", max_offset,
		281	+ " vs numel ", numel, OPS_ERROR(ErrCode::PARAM));
		282	+}
		283	+
		284	+static void prepare_out_weight_scale(QuantContext& ctx) {
		285	+ auto scale_view_shape = op_infer::array_to_small_vector(ctx.weight_scale.sizes());
		286	+ ctx.out_weight_scale = npu_preparation::apply_tensor_without_format(scale_view_shape, ctx.weight_scale.options());
		287	+ check_strides_envelope(ctx.weight_scale.sizes(), ctx.weight_scale.strides(), ctx.out_weight_scale.numel());
		288	+ // 保持与输入 scale 相同的 strides（ND per-group 转置场景需配转置 scale），连续输入时为空操作
		289	+ ctx.out_weight_scale =
		290	+ ctx.out_weight_scale.as_strided_(scale_view_shape, op_infer::array_to_small_vector(ctx.weight_scale.strides()));
		291	+}
		292	+
		293	+static void prepare_out_weight_nd(QuantContext& ctx) {
		294	+ auto weight_view_shape = op_infer::array_to_small_vector(ctx.weight.sizes());
		295	+ ctx.out_weight = npu_preparation::apply_tensor_without_format(weight_view_shape, ctx.weight.options());
		296	+ check_strides_envelope(ctx.weight.sizes(), ctx.weight.strides(), ctx.out_weight.numel());
		297	+ // ND 直拷是物理透传，out_weight 必须与 weight 保持相同的 sizes/strides（转置场景尤为关键）
		298	+ ctx.out_weight = ctx.out_weight.as_strided_(weight_view_shape, op_infer::array_to_small_vector(ctx.weight.strides()));
		299	+}
		300	+
154	template <bool IsGmm, int64_t NzC0, aclFormat OutWeightFormat>	301	template <bool IsGmm, int64_t NzC0, aclFormat OutWeightFormat>
155	static void prepare_out_weight_nz(QuantContext& ctx) {	302	static void prepare_out_weight_nz(QuantContext& ctx) {
156	auto weight_sizes = ctx.weight.sizes();	303	auto weight_sizes = ctx.weight.sizes();
	@@ -202,6 +349,43 @@ static void prepare_out_weight_nz(QuantContext& ctx) {
202	static_cast<aclFormat>(OutWeightFormat));	349	static_cast<aclFormat>(OutWeightFormat));
203	}	350	}
204		351
		352	+// A16W4（INT4/FP4）NZ（非转置）专用 prepare：torch 侧 uint8 [K, N/2] 非转置输入
		353	+// out_weight: torch 侧保持 uint8 物理打包视图 [K, N/2] + NZ_C0_8 标签（C0=8 字节=16 个 4-bit），
		354	+// storage 为物理分形 [N/16, K/16, 16, 8]（uint8 字节数精确）；经 ConvertType 统一 ×2 还原为
		355	+// ACL 4-bit view [K, N]、storage [N/16, K/16, 16, 16]、format NZ_C0_16，与 ops-math 校验一致。
		356	+static void prepare_out_weight_nz_a16w4(QuantContext& ctx) {
		357	+ auto weight_sizes = ctx.weight.sizes();
		358	+ int64_t weight_dim = ctx.weight.dim();
		359	+ int64_t k = weight_sizes[weight_dim - IDX_2];
		360	+ int64_t n_packed = weight_sizes[weight_dim - IDX_1]; // uint8 物理打包字节数（N/2）
		361	+ int64_t n_logical = n_packed * 2; // 每字节 2 个 INT4
		362	+
		363	+ c10::SmallVector<int64_t, DIMS_MAX> storage_shape = {ceil_div(n_logical, NZ_16), ceil_div(k, NZ_16), NZ_16, NZ_C0_8};
		364	+
		365	+ int64_t storage_size = 1;
		366	+ for (auto dim : storage_shape) {
		367	+ storage_size *= dim;
		368	+ }
		369	+
		370	+ ctx.out_weight = npu_preparation::apply_tensor_without_format({storage_size}, ctx.weight.options());
		371	+
		372	+ // 物理打包视图 [K, N/2]：4-bit 的 dtype/format 语义由 ConvertType 与 fake format 标签承担
		373	+ c10::SmallVector<int64_t, DIMS_MAX> out_view_shape = {k, n_packed};
		374	+ ctx.out_weight.unsafeGetTensorImpl()->set_sizes_contiguous(out_view_shape);
		375	+
		376	+ at_npu::native::StorageDescHelper::SetDesc(
		377	+ ctx.out_weight, ctx.out_weight.sizes(), storage_shape, ctx.out_weight.strides(), ACL_FORMAT_FRACTAL_NZ_C0_8);
		378	+}
		379	+
		380	+// A16S4 内部分流：转置 weight → ND 直拷（物理透传，与转置状态无关）；非转置 → NZ 转换
		381	+static void prepare_out_weight_a16s4(QuantContext& ctx) {
		382	+ if (ctx.is_weight_trans) {
		383	+ prepare_out_weight_nd(ctx);
		384	+ } else {
		385	+ prepare_out_weight_nz_a16w4(ctx);
		386	+ }
		387	+}
		388	+
205	template <bool IsGmm>	389	template <bool IsGmm>
206	static void prepare_out_weight_scale_mx(QuantContext& ctx) {	390	static void prepare_out_weight_scale_mx(QuantContext& ctx) {
207	constexpr int64_t expected_dim = IsGmm ? DIMS_4 : DIMS_3;	391	constexpr int64_t expected_dim = IsGmm ? DIMS_4 : DIMS_3;
	@@ -239,6 +423,12 @@ static void prepare_out_weight_offset(QuantContext& ctx) {
239	auto offset_sizes = op_infer::array_to_small_vector(ctx.weight_offset.value().sizes());	423	auto offset_sizes = op_infer::array_to_small_vector(ctx.weight_offset.value().sizes());
240	ctx.out_weight_offset =	424	ctx.out_weight_offset =
241	npu_preparation::apply_tensor_without_format(offset_sizes, ctx.weight_offset.value().options());	425	npu_preparation::apply_tensor_without_format(offset_sizes, ctx.weight_offset.value().options());
		426	+ check_strides_envelope(
		427	+ ctx.weight_offset.value().sizes(), ctx.weight_offset.value().strides(), ctx.out_weight_offset.numel());
		428	+ // 与 prepare_out_weight_scale 同理：保持与输入 offset 相同的 strides，
		429	+ // 转置场景下游 wqbmmv2 要求 antiquantOffset 与 weight 连续/转置状态一致，连续输入时为空操作
		430	+ ctx.out_weight_offset = ctx.out_weight_offset.as_strided_(
		431	+ offset_sizes, op_infer::array_to_small_vector(ctx.weight_offset.value().strides()));
242	}	432	}
243	}	433	}
244		434
	@@ -268,6 +458,31 @@ static const std::unordered_map<c10_npu::SocVersion, std::vector<DataFlowConfig>
268	prepare_out_weight_nz<true, NZ_C0_16, ACL_FORMAT_FRACTAL_NZ_C0_16>,	458	prepare_out_weight_nz<true, NZ_C0_16, ACL_FORMAT_FRACTAL_NZ_C0_16>,
269	prepare_out_weight_scale_mx<true>,	459	prepare_out_weight_scale_mx<true>,
270	prepare_out_weight_offset,	460	prepare_out_weight_offset,
		461	+ prepare_out_bias}},
		462	+ {judge_mm_a16s4_per_tensor,
		463	+ {prepare_out_weight_nd,
		464	+ prepare_out_weight_scale,
		465	+ prepare_out_weight_offset,
		466	+ prepare_out_bias}},
		467	+ {judge_mm_a16s4_per_channel,
		468	+ {prepare_out_weight_a16s4,
		469	+ prepare_out_weight_scale,
		470	+ prepare_out_weight_offset,
		471	+ prepare_out_bias}},
		472	+ {judge_mm_a16s4_per_group,
		473	+ {prepare_out_weight_a16s4,
		474	+ prepare_out_weight_scale,
		475	+ prepare_out_weight_offset,
		476	+ prepare_out_bias}},
		477	+ {judge_mm_a16f4_nz_pergroup,
		478	+ {prepare_out_weight_nz_a16w4,
		479	+ prepare_out_weight_scale,
		480	+ prepare_out_weight_offset,
		481	+ prepare_out_bias}},
		482	+ {judge_mm_a16f4_mx,
		483	+ {prepare_out_weight_nz_a16w4,
		484	+ prepare_out_weight_scale,
		485	+ prepare_out_weight_offset,
271	prepare_out_bias}}}},	486	prepare_out_bias}}}},
272	};	487	};
273		488

Mop_plugin/python/meta/_meta_registrations.py+8-0

	@@ -4167,6 +4167,14 @@ def npu_weight_quant_batchmatmul_meta(x, weight, antiquant_scale, antiquant_offs
4167	dim_m = x.size(0)	4167	dim_m = x.size(0)
4168	if (weight.dtype == torch.int32 or weight.dtype == torch.float32) and weight.is_contiguous():	4168	if (weight.dtype == torch.int32 or weight.dtype == torch.float32) and weight.is_contiguous():
4169	dim_n = weight.size(1) * 8	4169	dim_n = weight.size(1) * 8
		4170	+ elif weight.dtype == torch.uint8 and weight_dtype in (torch_npu.int4, torch_npu.float4_e2m1fn_x2, torch_npu.float4_e1m2fn_x2):
		4171	+ # 4-bit 紧凑打包载体（仅 uint8，与 C++ 侧 at::kByte 判定一致）：与 C++ 输出 size 推齐，
		4172	+ # 以 stride(-1)==1 判定沿 N 打包（ND 非转置视图 / NZ_C0_8 视图末维为 N/2，逻辑 N 需 ×2）；
		4173	+ # 转置视图（stride(-2)==1）末维已是逻辑 N
		4174	+ if weight.stride(-1) == 1:
		4175	+ dim_n = weight.size(1) * 2
		4176	+ else:
		4177	+ dim_n = weight.size(1)
4170	else:	4178	else:
4171	dim_n = weight.size(1)	4179	dim_n = weight.size(1)
4172	if quant_scale is not None:	4180	if quant_scale is not None:

Mop_plugin/utils/op_api_common.cpp+1-0

	@@ -30,6 +30,7 @@ constexpr int OWNER_ROOT_UID = 0;
30	typedef void (AddTensorAddrToCachedList)(void addr);	30	typedef void (AddTensorAddrToCachedList)(void addr);
31		31
32	static std::unordered_map<aclFormat, aclFormat> FORMAT_FAKE_TO_REAL{	32	static std::unordered_map<aclFormat, aclFormat> FORMAT_FAKE_TO_REAL{
		33	+ {ACL_FORMAT_FRACTAL_NZ_C0_8, ACL_FORMAT_FRACTAL_NZ_C0_16},
33	{ACL_FORMAT_FRACTAL_NZ_C0_16, ACL_FORMAT_FRACTAL_NZ_C0_32},	34	{ACL_FORMAT_FRACTAL_NZ_C0_16, ACL_FORMAT_FRACTAL_NZ_C0_32},
34	{ACL_FORMAT_FRACTAL_NZ, ACL_FORMAT_FRACTAL_NZ}};	35	{ACL_FORMAT_FRACTAL_NZ, ACL_FORMAT_FRACTAL_NZ}};
35		36

Mtest/test_custom_ops/test_npu_weight_quant_batchmatmul.py+65-0

	@@ -19,6 +19,23 @@ class TestNPUWeightQuantBatchMatmul(TestCase):
19	def custom_op_exec(self, x, weight, antiquant_scale, antiquant_offset, weight_dtype=None):	19	def custom_op_exec(self, x, weight, antiquant_scale, antiquant_offset, weight_dtype=None):
20	return torch_npu.npu_weight_quant_batchmatmul(x, weight, antiquant_scale, antiquant_offset, weight_dtype=weight_dtype)	20	return torch_npu.npu_weight_quant_batchmatmul(x, weight, antiquant_scale, antiquant_offset, weight_dtype=weight_dtype)
21		21
		22	+ @staticmethod
		23	+ def pack_int4_along_n(weight_int4):
		24	+ # [K, N] int4 值（int8 存储，-8~7）-> uint8 载体 [K, N/2]，低 4 位为偶数列
		25	+ w = (weight_int4.numpy() & 0xF).astype(np.uint8)
		26	+ return torch.from_numpy((w[:, 1::2] << 4) \| w[:, 0::2])
		27	+
		28	+ @staticmethod
		29	+ def pack_int4_along_k(weight_int4):
		30	+ # [K, N] int4 值 -> uint8 载体物理 [N, K/2]（沿 K 打包，低 4 位为偶数行）
		31	+ w = (weight_int4.numpy() & 0xF).astype(np.uint8)
		32	+ return torch.from_numpy(np.ascontiguousarray((w[1::2, :].T << 4) \| w[0::2, :].T))
		33	+
		34	+ @staticmethod
		35	+ def gen_fp16_exact_values(shape, pool):
		36	+ # 从 fp16 精确可表示的值池中取样，保证与 CPU 反量化参考比对时基本无舍入差
		37	+ return pool[torch.randint(0, len(pool), shape)]
		38	+
22	@SupportedDevices(['Ascend310P'])	39	@SupportedDevices(['Ascend310P'])
23	def test_npu_weight_quant_batchmatmul2(self, device="npu"):	40	def test_npu_weight_quant_batchmatmul2(self, device="npu"):
24	torch.manual_seed(0)	41	torch.manual_seed(0)
	@@ -79,6 +96,54 @@ class TestNPUWeightQuantBatchMatmul(TestCase):
79		96
80	self.assertRtolEqual(supported_output, custom_output, 0.001)	97	self.assertRtolEqual(supported_output, custom_output, 0.001)
81		98
		99	+ @SupportedDevices(['Ascend950'])
		100	+ def test_npu_weight_quant_batchmatmul2_with_A16W4_nd_uint8_perchannel(self, device="npu"):
		101	+ # uint8 载体紧凑打包 int4（沿 N 打包，ND 非转置视图 [K, N/2]，stride(-1)==1），
		102	+ # weight_dtype=torch_npu.int4 时 N 维按每字节 2 个 4-bit 还原
		103	+ torch.manual_seed(0)
		104	+ m, k, n = 16, 256, 128
		105	+ x_pool = torch.tensor([0.5, -0.5, 1.0, -1.0], dtype=torch.float16)
		106	+ scale_pool = torch.tensor([0.5, 1.0, 1.5, 2.0], dtype=torch.float16)
		107	+ cpu_x = self.gen_fp16_exact_values((m, k), x_pool)
		108	+ cpu_weight = torch.randint(low=-8, high=8, size=(k, n), dtype=torch.int8)
		109	+ cpu_antiquant_scale = self.gen_fp16_exact_values((n,), scale_pool)
		110	+
		111	+ npu_x = cpu_x.clone().npu()
		112	+ npu_weight = self.pack_int4_along_n(cpu_weight).npu()
		113	+ npu_antiquant_scale = cpu_antiquant_scale.clone().npu()
		114	+
		115	+ supported_output = torch.matmul(cpu_x.float(), cpu_weight.float() * cpu_antiquant_scale.float())
		116	+ custom_output = self.custom_op_exec(
		117	+ npu_x, npu_weight, npu_antiquant_scale, None, torch_npu.int4)
		118	+
		119	+ self.assertEqual(custom_output.shape, supported_output.shape)
		120	+ self.assertEqual(custom_output.dtype, torch.float16)
		121	+ self.assertRtolEqual(supported_output, custom_output.cpu().float(), 0.001)
		122	+
		123	+ @SupportedDevices(['Ascend950'])
		124	+ def test_npu_weight_quant_batchmatmul2_with_A16W4_nd_uint8_transposed(self, device="npu"):
		125	+ # uint8 载体紧凑打包 int4（沿 K 打包，转置视图 [K/2, N]，stride(-2)==1），
		126	+ # weight_dtype=torch_npu.int4 时 K 维按每字节 2 个 4-bit 还原
		127	+ torch.manual_seed(0)
		128	+ m, k, n = 16, 256, 128
		129	+ x_pool = torch.tensor([0.5, -0.5, 1.0, -1.0], dtype=torch.float16)
		130	+ scale_pool = torch.tensor([0.5, 1.0, 1.5, 2.0], dtype=torch.float16)
		131	+ cpu_x = self.gen_fp16_exact_values((m, k), x_pool)
		132	+ cpu_weight = torch.randint(low=-8, high=8, size=(k, n), dtype=torch.int8)
		133	+ cpu_antiquant_scale = self.gen_fp16_exact_values((n,), scale_pool)
		134	+
		135	+ npu_x = cpu_x.clone().npu()
		136	+ npu_weight = self.pack_int4_along_k(cpu_weight).npu().transpose(0, 1)
		137	+ npu_antiquant_scale = cpu_antiquant_scale.clone().npu()
		138	+
		139	+ supported_output = torch.matmul(cpu_x.float(), cpu_weight.float() * cpu_antiquant_scale.float())
		140	+ custom_output = self.custom_op_exec(
		141	+ npu_x, npu_weight, npu_antiquant_scale, None, torch_npu.int4)
		142	+
		143	+ self.assertEqual(custom_output.shape, supported_output.shape)
		144	+ self.assertEqual(custom_output.dtype, torch.float16)
		145	+ self.assertRtolEqual(supported_output, custom_output.cpu().float(), 0.001)
		146	+
82		147
83	if __name__ == "__main__":	148	if __name__ == "__main__":
84	run_tests()	149	run_tests()

Mtest/test_custom_ops/test_npu_weight_quant_preprocess.py+15-0

	@@ -34,6 +34,21 @@ class TestNpuWeightQuantPreprocess(TestCase):
34	if torch_npu._C._npu_getOption("ALLOW_INTERNAL_FORMAT") == b"enable":	34	if torch_npu._C._npu_getOption("ALLOW_INTERNAL_FORMAT") == b"enable":
35	self.assertEqual(torch_npu.get_npu_format(out_weight), FRACTAL_NZ_C0_16)	35	self.assertEqual(torch_npu.get_npu_format(out_weight), FRACTAL_NZ_C0_16)
36		36
		37	+ @SupportedDevices(['Ascend950'])
		38	+ def test_npu_weight_quant_preprocess_a16w4_int8_carrier_rejected(self):
		39	+ # A16W4 的 4-bit weight 只接受 uint8 载体（每字节 2 个 4-bit），int8 载体应被拒绝
		40	+ k, n = 256, 128
		41	+ weight = torch.zeros((k, n // 2), dtype=torch.int8).npu()
		42	+ weight_scale = torch.ones((n,), dtype=torch.float16).npu()
		43	+
		44	+ with self.assertRaisesRegex(RuntimeError, "must be packed into a uint8 tensor"):
		45	+ torch_npu.npu_weight_quant_preprocess(
		46	+ weight,
		47	+ weight_scale,
		48	+ x_dtype=torch.float16,
		49	+ weight_dtype=torch_npu.int4,
		50	+ weight_scale_dtype=torch.float16)
		51	+
37		52
38	if __name__ == "__main__":	53	if __name__ == "__main__":
39	run_tests()	54	run_tests()