已合并
[feat]: A5 Support for groupmatmul s8s4 requirements #5672
bynshard创建于 16 天前
[feat]: A5 Support for groupmatmul s8s4 requirements #5672
已合并
共 2 个文件变更+36-1
| @@ -313,13 +313,15 @@ std::vector<at::Tensor> npu_grouped_matmul(const at::TensorList x, | |||
| 313 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2) || | 313 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2) || |
| 314 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_4) || | 314 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_4) || |
| 315 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_16); | 315 | (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_16); |
| 316 | + const bool is_a8w4 = x_wrapper.dtype == ACL_INT8 && weight_wrapper.dtype == ACL_INT32 && | ||
| 317 | + (scale_wrapper.dtype == ACL_UINT64 || scale_wrapper.dtype == ACL_INT64) && !scale_real.empty(); | ||
| 316 | if (is_weight_nz) { | 318 | if (is_weight_nz) { |
| 317 | static const bool is_weight_nz_available = check_aclnn_kernel_available("aclnnGroupedMatmulWeightNz"); | 319 | static const bool is_weight_nz_available = check_aclnn_kernel_available("aclnnGroupedMatmulWeightNz"); |
| 318 | TORCH_CHECK(is_weight_nz_available, | 320 | TORCH_CHECK(is_weight_nz_available, |
| 319 | "Format of weight in npu_grouped_matmul is FRACTAL_NZ, current CANN version " | 321 | "Format of weight in npu_grouped_matmul is FRACTAL_NZ, current CANN version " |
| 320 | "do not support with this format. Please try to update the version of CANN." | 322 | "do not support with this format. Please try to update the version of CANN." |
| 321 | + OPS_ERROR(ErrCode::PARAM)); | 323 | + OPS_ERROR(ErrCode::PARAM)); |
| 322 | - int64_t quant_per_group_size = 0; | 324 | + int64_t quant_per_group_size = is_a8w4 && offset_real.empty() ? 256 : 0; |
C | |||
| 323 | EXEC_NPU_CMD(aclnnGroupedMatmulWeightNz, x_wrapper, weight_wrapper, bias_real, scale_wrapper, offset_real, antiquant_scale_wrapper, | 325 | EXEC_NPU_CMD(aclnnGroupedMatmulWeightNz, x_wrapper, weight_wrapper, bias_real, scale_wrapper, offset_real, antiquant_scale_wrapper, |
| 324 | antiquant_offset_real, per_token_scale_wrapper, group_list_real, activation_input_real, | 326 | antiquant_offset_real, per_token_scale_wrapper, group_list_real, activation_input_real, |
| 325 | activation_quant_scale_real, activation_quant_offset_real, split_item_value, group_type_value, | 327 | activation_quant_scale_real, activation_quant_offset_real, split_item_value, group_type_value, |
| @@ -881,6 +881,39 @@ class TestGroupedMatmul(TestCase): | |||
| 881 | self.assertEqual(out_nz_dim1, golden_dim1) | 881 | self.assertEqual(out_nz_dim1, golden_dim1) |
| 882 | self.assertEqual(out_nz[0][:golden_dim0, :], out_golden.npu()) | 882 | self.assertEqual(out_nz[0][:golden_dim0, :], out_golden.npu()) |
| 883 | 883 | ||
| 884 | + | ||
| 885 | + def test_npu_grouped_matmul_a8w4_per_group_weight_nz(self): | ||
| 886 | + """A8W4 WeightNz per-group dispatches quantGroupSize=256.""" | ||
| 887 | + expert_num = 2 | ||
| 888 | + m = 8 | ||
| 889 | + k = 512 | ||
| 890 | + n = 64 | ||
| 891 | + quant_group_size = 256 | ||
| 892 | + | ||
| 893 | + x = torch.zeros((m, k), dtype=torch.int8, device="npu") | ||
| 894 | + weight = torch.zeros((expert_num, k, n), dtype=torch.int32, device="npu") | ||
| 895 | + weight_nz = torch_npu.npu_format_cast(weight, 29, customize_dtype=torch.int8) | ||
| 896 | + weight_int4_nz = torch_npu.npu_convert_weight_to_int4pack(weight_nz) | ||
| 897 | + | ||
| 898 | + scale_fp32 = np.ones((expert_num, k // quant_group_size, n), dtype=np.float32) | ||
| 899 | + scale_uint64 = np.zeros(scale_fp32.shape, dtype=np.uint64) | ||
| 900 | + scale_uint64 |= scale_fp32.view(np.uint32).astype(np.uint64) | ||
| 901 | + scale = torch.from_numpy(scale_uint64.view(np.int64)).npu() | ||
| 902 | + bias = torch.zeros((expert_num, n), dtype=torch.float32, device="npu") | ||
| 903 | + per_token_scale = torch.ones((m,), dtype=torch.float32, device="npu") | ||
| 904 | + group_list = torch.tensor([m // 2, m], dtype=torch.int64, device="npu") | ||
| 905 | + | ||
| 906 | + output = torch_npu.npu_grouped_matmul( | ||
| 907 | + [x], [weight_int4_nz], bias=[bias], scale=[scale], offset=None, | ||
| 908 | + antiquant_scale=None, antiquant_offset=None, per_token_scale=[per_token_scale], | ||
| 909 | + group_list=group_list, activation_input=None, activation_quant_scale=None, | ||
| 910 | + activation_quant_offset=None, split_item=3, group_type=0, group_list_type=1, | ||
| 911 | + act_type=0, output_dtype=torch.bfloat16) | ||
| 912 | + | ||
| 913 | + torch_npu.npu.synchronize() | ||
| 914 | + self.assertEqual(output[0].shape, (m, n)) | ||
| 915 | + self.assertEqual(output[0], torch.zeros((m, n), dtype=torch.bfloat16, device="npu")) | ||
| 916 | + | ||
| 884 | 917 | ||
| 885 | def test_npu_grouped_matmul_group_list_none(self): | 918 | def test_npu_grouped_matmul_group_list_none(self): |
| 886 | torch.manual_seed(0) | 919 | torch.manual_seed(0) |
补充下UT