kv-cache int8 quant

Signed-off-by: George Ohashi <[email protected]>

Author: George Ohashi

src/llmcompressor/modifiers/quantization/cache.py

@@ -151,8 +151,15 @@ def _quantize(self, tensor, kv_type, layer_idx):
             scales = self.v_scales
             zps = self.v_zps
 
-        # tensor
-        scale, zp = observer(tensor)
+        # note: key, value states are in the shape:
+        # [batch, num_key_value_heads, seq_len, head_dim]
+
+        base_name = None  # tensor-wise quantization, shape of [1]
+        if self.quantization_args.strategy == "channel":
+            # target last dim to quantize, shape of [head_dim]
+            base_name = "kv_cache"
+
+        scale, zp = observer(tensor, base_name=base_name)
         if len(scales) <= layer_idx:
             scales.append(scale)
             zps.append(zp)

src/llmcompressor/observers/base.py

@@ -128,7 +128,7 @@ def get_qparams(
                     self._zero_point[:, group_index] = zero_point.squeeze(1)
 
             elif self.quantization_args.strategy == QuantizationStrategy.CHANNEL:
-                if base_name == "output":
+                if base_name in ("output", "kv_cache"):
                     # the last dimension is the hidden dimension
                     # shape of [1,1, num_key_value_heads * head_dim]
                     scale, zero_point = self.get_qparams_along_dim(