【950】Cast适配int4回合-op-plugin-AtomGit

Mop_plugin/ops/opapi/CastKernelNpuOpApi.cpp+6-4

Mop_plugin/python/meta/_meta_registrations.py+6-5

	@@ -49,6 +49,7 @@ TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP = {
49	21: torch.bits8,	49	21: torch.bits8,
50	23: torch.float8_e5m2,	50	23: torch.float8_e5m2,
51	24: torch.float8_e4m3fn,	51	24: torch.float8_e4m3fn,
		52	+ 285: torch.uint8, # torch_npu.int4 use torch.uint8
52	290: torch.uint8, # torch_npu.hifloat8 use torch.uint8	53	290: torch.uint8, # torch_npu.hifloat8 use torch.uint8
53	291: torch.float8_e5m2,	54	291: torch.float8_e5m2,
54	292: torch.float8_e4m3fn,	55	292: torch.float8_e4m3fn,
	@@ -1531,17 +1532,17 @@ def npu_dtype_cast_meta(self, dtype, input_dtype=None):
1531	if dim_num != 0:	1532	if dim_num != 0:
1532	input_shape[-1] *= 2	1533	input_shape[-1] *= 2
1533	else:	1534	else:
1534	- raise RuntimeError("Scalar input cannot be float4_e2m1 or float4_e1m2" +	1535	+ raise RuntimeError("Scalar input cannot be float4_e2m1fn_x2 or float4_e1m2fn_x2" +
1535	ops_error(ErrCode.PARAM))	1536	ops_error(ErrCode.PARAM))
1536		1537
1537	- if dtype == 296 or dtype == 297:	1538	+ if dtype == 285 or dtype == 296 or dtype == 297:
1538	if dim_num == 0 or input_shape[-1] % 2:	1539	if dim_num == 0 or input_shape[-1] % 2:
1539	- raise RuntimeError("If output dtype is float4_e2m1 or float4_e1m2, " \	1540	+ raise RuntimeError("If output dtype is float4_e2m1fn_x2, float4_e1m2fn_x2 or int4, " \
1540	"the last dim of input must be divisible by 2" +	1541	"the last dim of input must be divisible by 2" +
1541	ops_error(ErrCode.PARAM))	1542	ops_error(ErrCode.PARAM))
1542	input_shape[-1] //= 2	1543	input_shape[-1] //= 2
1543	- # torch_npu.hifloat8, torch_npu.float4_e2m1, torch_npu.float4_e1m2	1544	+ # torch_npu.hifloat8, torch_npu.float4_e2m1fn_x2, torch_npu.float4_e1m2fn_x2, torch_npu.int4
1544	- if dtype in [290, 296, 297]:	1545	+ if dtype in [285, 290, 296, 297]:
1545	output = self.new_empty(input_shape, dtype=torch.uint8)	1546	output = self.new_empty(input_shape, dtype=torch.uint8)
1546	else:	1547	else:
1547	output_dst_dtype = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP.get(dtype)	1548	output_dst_dtype = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP.get(dtype)

Mop_plugin/utils/op_api_common.h+2-2

Mtest/test_base_ops/test_dtype_cast.py+8-1

	@@ -5,7 +5,7 @@ import numpy as np
5	import torch_npu	5	import torch_npu
6		6
7	from torch_npu.testing.testcase import TestCase, run_tests	7	from torch_npu.testing.testcase import TestCase, run_tests
8	-from torch_npu.testing.common_utils import create_common_tensor	8	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
9	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	9	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
10		10
11		11
	@@ -73,6 +73,13 @@ class TestDtypeCast(TestCase):
73	y[::2].copy_(x[::2])	73	y[::2].copy_(x[::2])
74	self.assertEqual(y, [1, 0, 3, 0, 5, 0])	74	self.assertEqual(y, [1, 0, 3, 0, 5, 0])
75		75
		76	+ @SupportedDevices(['Ascend910_95'])
		77	+ def test_dtype_cast_int4(self):
		78	+ input1 = torch.tensor([1, 2, 3, 4], dtype=torch.int32).npu()
		79	+ expectOutput = torch.tensor([33, 67], dtype=torch.uint8)
		80	+ output = torch_npu.npu_dtype_cast(input1, dtype=torch_npu.int4, input_dtype=torch.int32).cpu()
		81	+ self.assertEqual(expectOutput, output)
		82	+
76		83
77	if __name__ == "__main__":	84	if __name__ == "__main__":
78	run_tests()	85	run_tests()

	@@ -33,7 +33,8 @@ at::Tensor npu_dtype_cast_impl_op_api(
33	int64_t input_dtype_tocheck = input_dtype.has_value() ? input_dtype.value() : static_cast<int64_t>(self.scalar_type());	33	int64_t input_dtype_tocheck = input_dtype.has_value() ? input_dtype.value() : static_cast<int64_t>(self.scalar_type());
34	bool special_output_type = (dtype == static_cast<int64_t>(c10_npu::DType::HIFLOAT8) \|\|	34	bool special_output_type = (dtype == static_cast<int64_t>(c10_npu::DType::HIFLOAT8) \|\|
35	dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) \|\|	35	dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) \|\|
36	- dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));	36	+ dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2) \|\|
		37	+ dtype == static_cast<int64_t>(c10_npu::DType::INT4));
37	at::SmallVector<int64_t, op_infer::SIZE> input_shape;	38	at::SmallVector<int64_t, op_infer::SIZE> input_shape;
38	at::SmallVector<int64_t, op_infer::SIZE> output_shape;	39	at::SmallVector<int64_t, op_infer::SIZE> output_shape;
39	int32_t input_dim = self.dim();	40	int32_t input_dim = self.dim();
	@@ -50,12 +51,13 @@ at::Tensor npu_dtype_cast_impl_op_api(
50	input_shape.push_back(self.size(index));	51	input_shape.push_back(self.size(index));
51	}	52	}
52		53
53	- // float4 shape check	54	+ // float4/int4 shape check
54	if (dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2) \|\|	55	if (dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2) \|\|
55	- dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1)) {	56	+ dtype == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) \|\|
		57	+ dtype == static_cast<int64_t>(c10_npu::DType::INT4)) {
56	TORCH_CHECK(input_shape[index] % FP4_IN_UINT8_NUM == 0,	58	TORCH_CHECK(input_shape[index] % FP4_IN_UINT8_NUM == 0,
57	"The last dim input shape must be divisible by 2 if "	59	"The last dim input shape must be divisible by 2 if "
58	- "output dtype is torch_npu.float4_e2m1 or torch_npu.float4_e1m2" + OPS_ERROR(ErrCode::PARAM));	60	+ "output dtype is torch_npu.float4_e2m1fn_x2, torch_npu.float4_e1m2fn_x2 or torch_npu.int4" + OPS_ERROR(ErrCode::PARAM));
		C C chuboning2025年12月3日实现在torch2.8版本b=a.to(torch.float4_e2m1fn_x2)的调用方法
59	output_shape.push_back(input_shape[index] / FP4_IN_UINT8_NUM);	61	output_shape.push_back(input_shape[index] / FP4_IN_UINT8_NUM);
60	} else {	62	} else {
61	output_shape.push_back(input_shape[index]);	63	output_shape.push_back(input_shape[index]);

	@@ -488,7 +488,7 @@ inline aclTensor *ConvertType(const TensorWrapper &tensor_r)
488	if (acl_data_type != ACL_STRING) {	488	if (acl_data_type != ACL_STRING) {
489	TORCH_CHECK(at_tensor.itemsize() > 0, "the itemsize of tensor must be greater than 0.",	489	TORCH_CHECK(at_tensor.itemsize() > 0, "the itemsize of tensor must be greater than 0.",
490	OPS_ERROR(ErrCode::VALUE));	490	OPS_ERROR(ErrCode::VALUE));
491	- if (acl_data_type == ACL_FLOAT4_E2M1 \|\| acl_data_type == ACL_FLOAT4_E1M2) {	491	+ if (acl_data_type == ACL_FLOAT4_E2M1 \|\| acl_data_type == ACL_FLOAT4_E1M2 \|\| acl_data_type == ACL_INT4) {
492	storageDims.push_back(at_tensor.storage().nbytes() / at_tensor.itemsize() * FP4_IN_INT8);	492	storageDims.push_back(at_tensor.storage().nbytes() / at_tensor.itemsize() * FP4_IN_INT8);
493	if (at_tensor.sizes().size() == 1) {	493	if (at_tensor.sizes().size() == 1) {
494	wrapperShape[0] = wrapperShape[0] * FP4_IN_INT8;	494	wrapperShape[0] = wrapperShape[0] * FP4_IN_INT8;
	@@ -611,7 +611,7 @@ inline aclTensor *ConvertTypeV2(TensorStructPtr at_tensor)
611	if (acl_data_type != ACL_STRING) {	611	if (acl_data_type != ACL_STRING) {
612	TORCH_CHECK((*at_tensor).itemsize > 0, "the itemsize of tensor must be greater than 0.",	612	TORCH_CHECK((*at_tensor).itemsize > 0, "the itemsize of tensor must be greater than 0.",
613	OPS_ERROR(ErrCode::VALUE));	613	OPS_ERROR(ErrCode::VALUE));
614	- if (acl_data_type == ACL_FLOAT4_E2M1 \|\| acl_data_type == ACL_FLOAT4_E1M2) {	614	+ if (acl_data_type == ACL_FLOAT4_E2M1 \|\| acl_data_type == ACL_FLOAT4_E1M2 \|\| acl_data_type == ACL_INT4) {
615	storageDims.push_back((at_tensor).nbytes / (at_tensor).itemsize * FP4_IN_INT8);	615	storageDims.push_back((at_tensor).nbytes / (at_tensor).itemsize * FP4_IN_INT8);
616	if ((*at_tensor).sizes.size() == 1) {	616	if ((*at_tensor).sizes.size() == 1) {
617	wrapperShape[0] = wrapperShape[0] * FP4_IN_INT8;	617	wrapperShape[0] = wrapperShape[0] * FP4_IN_INT8;