!14029 add op npu_dynamic_quant_asymmetric、npu_dynamic_quant

Merge pull request !14029 from wangqk/dev_asymmetrical_dynamic_quant_master

Author: wangqk

test/allowlist_for_publicAPI.json

@@ -2822,6 +2822,8 @@
     "npu_quant_scatter",
     "npu_scatter_nd_update_",
     "npu_swiglu",
+    "npu_dynamic_quant",
+    "npu_dynamic_quant_asymmetric",
     "npu_yolo_boxes_encode",
     "npu_yolo_boxes_encode",
     "npu_weight_quant_batchmatmul",

test/onnx/test_wrapper_onnx_ops.py

@@ -1337,6 +1337,48 @@ def export_onnx(onnx_model_name):
         export_onnx(onnx_model_name)
         assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, onnx_model_name)))
 
+    @SupportedDevices(['Ascend910B'])
+    def test_wrapper_npu_dynamic_quant(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super(Model, self).__init__()
+
+            def forward(self, input_dummy, smooth_scales_dummy):
+                output, scale = torch_npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales_dummy)
+                return output, scale
+            
+        def export_onnx(onnx_model_name):
+            input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)
+            smooth_scales_dummy = torch.rand(512).uniform_(-3, 3).npu().to(torch.float16)
+            model = Model().to("npu")
+            model(input_dummy, smooth_scales_dummy)
+            self.onnx_export(model, (input_dummy, smooth_scales_dummy), onnx_model_name,
+                             ["input", "smooth_scale_dummy"], ["output", "scale"])
+        onnx_model_name = "model_npu_dynamic_quant.onnx"
+        export_onnx(onnx_model_name)
+        assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, onnx_model_name)))
+
+    @SupportedDevices(['Ascend910B'])
+    def test_wrapper_npu_dynamic_quant_asymmetric(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super(Model, self).__init__()
+
+            def forward(self, input_dummy, smooth_scales_dummy):
+                output, scale, offset = torch_npu.npu_dynamic_quant_asymmetric(input_dummy, smooth_scales=smooth_scales_dummy)
+                return output, scale, offset
+            
+        def export_onnx(onnx_model_name):
+            input_dummy = torch.rand(4, 1024, 512).uniform_(-3, 3).npu().to(torch.float16)
+            smooth_scales_dummy = torch.rand(512).uniform_(-3, 3).npu().to(torch.float16)
+            model = Model().to("npu")
+            model(input_dummy, smooth_scales_dummy)
+            self.onnx_export(model, (input_dummy, smooth_scales_dummy), onnx_model_name,
+                             ["input", "smooth_scale_dummy"], ["output", "scale", "offset"])
+        onnx_model_name = "model_npu_dynamic_quant_asymmetric.onnx"
+        export_onnx(onnx_model_name)
+        assert (os.path.isfile(os.path.join(TestOnnxOps.test_onnx_path, onnx_model_name)))
+
     @SupportedDevices(['Ascend910B'])
     def test_wrapper_npu_weight_quant_batchmatmul(self):
         class Model(torch.nn.Module):

test/test_fake_tensor.py

@@ -1595,6 +1595,48 @@ def test_npu_ffn_meta(self):
             self.assertTrue(x.shape == res.shape)
 
 
+class TestNpuDynamicQuant(TestCase):
+    def test_npu_dynamic_quant(self):
+        with FakeTensorMode():
+            input_npu = torch.randn((4, 2048, 1024)).npu().to(torch.float16)
+            smooth_scales_npu = torch.randn((1024)).npu().to(torch.float16)
+
+            output = torch.randn((4, 2048, 1024)).npu().to(torch.int8)
+            scale = torch.randn((4, 2048)).npu().to(torch.float32)
+
+            actual_output, actual_scale = torch_npu.npu_dynamic_quant(input_npu, smooth_scales=smooth_scales_npu)
+
+            self.assertEqual(actual_output.dtype, output.dtype)
+            self.assertEqual(actual_output.shape, output.shape)
+            self.assertEqual(actual_output.device, output.device)
+            self.assertEqual(actual_scale.dtype, scale.dtype)
+            self.assertEqual(actual_scale.shape, scale.shape)
+            self.assertEqual(actual_scale.device, scale.device)
+
+
+class TestDynamicQuantAsymmetric(TestCase):
+    def test_npu_dynamic_quant_asymmetric(self):
+        with FakeTensorMode():
+            input_npu = torch.randn((4, 2048, 1024)).npu().to(torch.float16)
+            smooth_scales_npu = torch.randn((1024)).npu().to(torch.float16)
+
+            output = torch.randn((4, 2048, 1024)).npu().to(torch.int8)
+            scale = torch.randn((4, 2048)).npu().to(torch.float32)
+            offset = torch.randn((4, 2048)).npu().to(torch.float32)
+
+            actual_output, actual_scale, actual_offset = torch_npu.npu_dynamic_quant_asymmetric(input_npu, smooth_scales=smooth_scales_npu)
+
+            self.assertEqual(actual_output.dtype, output.dtype)
+            self.assertEqual(actual_output.shape, output.shape)
+            self.assertEqual(actual_output.device, output.device)
+            self.assertEqual(actual_scale.dtype, scale.dtype)
+            self.assertEqual(actual_scale.shape, scale.shape)
+            self.assertEqual(actual_scale.device, scale.device)
+            self.assertEqual(actual_offset.dtype, offset.dtype)
+            self.assertEqual(actual_offset.shape, offset.shape)
+            self.assertEqual(actual_offset.device, offset.device)
+
+
 class TestGroupedMatmul(TestCase):
     def test_npu_grouped_matmul_meta_0(self):
         with FakeTensorMode():

torch_npu/meta/_meta_registrations.py

@@ -627,6 +627,27 @@ def npu_quantize_meta(self, scales, zero_points, dtype, axis=1, div_mode=True):
     return torch.empty_like(self, dtype=torch.int8)
 
 
+@impl(m, "npu_dynamic_quant")
+def npu_dynamic_quant(input_dummy, *, smooth_scales=None):
+    dim_num = input_dummy.dim()
+    scale_shape = []
+    for dim in range(dim_num - 1):
+        scale_shape.append(input_dummy.size(dim))
+    return (torch.empty_like(input_dummy, dtype=torch.int8),
+             input_dummy.new_empty(scale_shape, dtype=torch.float32))
+
+
+@impl(m, "npu_dynamic_quant_asymmetric")
+def npu_dynamic_quant_asymmetric(input_dummy, *, smooth_scales=None, group_index=None, dst_type=torch.int8):
+    dim_num = input_dummy.dim()
+    scale_offset_shape = []
+    for dim in range(dim_num - 1):
+        scale_offset_shape.append(input_dummy.size(dim))
+    return (torch.empty_like(input_dummy, dtype=torch.int8),
+             input_dummy.new_empty(scale_offset_shape, dtype=torch.float32),
+             input_dummy.new_empty(scale_offset_shape, dtype=torch.float32))
+
+
 @impl(m, "npu_moe_compute_expert_tokens")
 def npu_moe_compute_expert_tokens_meta(sorted_experts, num_experts=1):
     out = torch.zeros(num_experts, dtype=torch.int32, device='meta')

torch_npu/onnx/wrapper_onnx_ops.py

@@ -730,6 +730,37 @@ def symbolic(g, x1: torch.Tensor, x2: torch.Tensor, hcom: str,
                     dequant_scale, pertoken_scale, comm_quant_scale_1, comm_quant_scale_2, antiquant_group_size, comm_turn)
 
 
+class _NPUDynamicQuantOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input_dummy, smooth_scales):
+        return torch.ops.npu.npu_dynamic_quant(input_dummy, smooth_scales=smooth_scales)
+
+    @staticmethod
+    def symbolic(g, input_dummy: Tensor, smooth_scales: Optional[Tensor] = None):
+        if smooth_scales is None:
+            smooth_scales = g.op("Constant", value_t=torch.tensor([]).to(input_dummy.type().dtype()))
+        return g.op("npu::NPUDynamicQuant", input_dummy, smooth_scales, outputs=2)
+
+
+class _NPUDynamicQuantV2Op(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input_dummy, smooth_scales, group_index, dst_type):
+        return torch.ops.npu.npu_dynamic_quant_asymmetric(input_dummy, smooth_scales=smooth_scales,
+                                                            group_index=group_index, dst_type=dst_type)
+
+    @staticmethod
+    def symbolic(g, input_dummy: Tensor, smooth_scales: Optional[Tensor] = None,
+                 group_index: Optional[Tensor] = None, dst_type: torch.dtype = torch.int8):
+        if smooth_scales is None:
+            smooth_scales = g.op("Constant", value_t=torch.tensor([]).to(input_dummy.type().dtype()))
+        if group_index is None:
+            group_index = g.op("Constant", value_t=torch.tensor([]).to(torch.int32))
+        dst_type_i = 2 # 当前仅支持int8
+        return g.op("npu::NPUDynamicQuantV2", input_dummy, smooth_scales,
+                    group_index, dst_type_i=dst_type_i, outputs=3)
+
 
 class _NPUWeightQuantBatchMatmulOP(torch.autograd.Function):
 
@@ -1083,6 +1114,14 @@ def _wrapper_npu_stride_add(self, other, offset1, offset2, c1_len):
     return _NPUStrideAddOP.apply(self, other, offset1, offset2, c1_len)
 
 
+def _wrapper_npu_dynamic_quant(input_dummy, smooth_scales=None):
+    return _NPUDynamicQuantOp.apply(input_dummy, smooth_scales)
+
+
+def _wrapper_npu_dynamic_quant_asymmetric(input_dummy, smooth_scales=None, group_index=None, dst_type=torch.int8):
+    return _NPUDynamicQuantV2Op.apply(input_dummy, smooth_scales, group_index, dst_type)
+
+
 def _wrapper_npu_gru(inputs, hx, weight_input, weight_hidden, bias_input, bias_hidden,
                     seq_length, has_biases, num_layers, dropout, train, bidirectional, batch_first):
     return _NPUGruOP.apply(inputs, hx, weight_input, weight_hidden, bias_input, bias_hidden,
@@ -1189,6 +1228,8 @@ def _add_onnx_ops():
     torch_npu.npu_scatter = _wrapper_npu_scatter
     torch_npu.npu_scatter_nd_update = _wrapper_npu_scatter_nd_update
     torch_npu.npu_lstm = _wrapper_npu_lstm
+    torch_npu.npu_dynamic_quant = _wrapper_npu_dynamic_quant
+    torch_npu.npu_dynamic_quant_asymmetric = _wrapper_npu_dynamic_quant_asymmetric
     torch_npu.npu_rms_norm = _wrapper_npu_rms_norm
     torch_npu.npu_add_rms_norm = _wrapper_npu_add_rms_norm
     torch_npu.npu_lstm_cell = _wrapper_npu_lstm_cell