[Kunlunxin]support DS R1/V3 w4a8int8 per-channel quantization (#226)

Co-authored-by: lishaohao <lsh862702688@163.com>

Author: lishaoao

angelslim/compressor/quant/core/__init__.py

@@ -20,6 +20,7 @@
 from .sample_func import EMASampler, MultiStepSampler  # noqa: F401
 from .save import DeepSeekV3PTQSaveMulti  # noqa: F401
 from .save import DeepSeekV3PTQSaveSingle  # noqa: F401
+from .save import DeepSeekV3W4A8Int8Save  # noqa: F401
 from .save import PTQOnlyScaleSave  # noqa: F401
 from .save import PTQPTMSave  # noqa: F401
 from .save import PTQSaveVllmHF  # noqa: F401

angelslim/compressor/quant/core/config.py

@@ -103,6 +103,13 @@ def __init__(self, config, global_config=None):
             self.low_memory = config.quantization.low_memory
             self.quant_analyse = config.quantization.quant_analyse
             self.quant_vit = config.quantization.quant_vit
+        elif "w4a8i8" in self.quant_algo:
+            group_size = quantization_args.quant_method["group_size"]
+            self.quant_algo_info = {
+                "group_size": group_size,
+                "ignore_layers": quantization_args.ignore_layers,
+            }
+            self.low_memory = config.quantization.low_memory
         elif "int8" in self.quant_algo:
             is_dynamic = "dynamic" if "dynamic" in self.quant_algo else "static"
             assert (

angelslim/compressor/quant/core/quant_func.py

@@ -474,3 +474,17 @@ def reduce_block_padding(input: torch.Tensor, block_sizes: dict, pad_value: floa
                 padded_tensor = F.pad(padded_tensor, pad, value=pad_value)
 
         return padded_tensor
+
+
+class Int8PerChannelQuantizer:
+    """Per-channel symmetric int8 quantizer."""
+
+    @torch.no_grad()
+    def quantize(self, tensor: torch.Tensor):
+        assert tensor.dtype == torch.bfloat16
+        qmax = 127.0
+        abs_max = torch.abs(tensor).max(dim=1, keepdim=True)[0]
+        scale = abs_max / qmax
+        quantized = torch.round(tensor / scale)
+        quantized = torch.clamp(quantized, -qmax, qmax)
+        return quantized.to(torch.int8), scale.to(torch.float32)