hook with CompressedAttentionImpl

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modifiers/quantization/calibration.py

@@ -282,6 +282,22 @@ def initialize_quantized_kv_cache(module: Module):
     setattr(module, "kv_cache", quantized_kv_cache)
 
 
+def initialize_attention_observers(module: Module):
+    input_args = getattr_chain(module, "quantization_scheme.input_activations", None)
+    if input_args is not None:
+        initialize_observer(module, "q", input_args)
+        initialize_observer(module, "k", input_args)
+        initialize_observer(module, "v", input_args)
+
+
+def calibrate_attention(
+    module: Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+):
+    calibrate_activations(module, value=query, base_name="q")
+    calibrate_activations(module, value=key, base_name="k")
+    calibrate_activations(module, value=value, base_name="v")
+
+
 def apply_calibration_status(module: Module):
     scheme = getattr(module, "quantization_scheme", None)
     if not scheme:

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -232,6 +232,7 @@ def _initialize_observers(self, module: torch.nn.Module):
         # kv_cache activations. Within `apply_quantization_config`, the config is
         # modified to use attention output quantization if a kv_cache_scheme exists
         if is_attention and output:
+            # initialize_attention_observers(module)  # TODO: attnq
             initialize_quantized_kv_cache(module)
 
         # output activations
@@ -240,6 +241,11 @@ def _initialize_observers(self, module: torch.nn.Module):
 
     def _initialize_hooks(self, model: torch.nn.Module) -> Set[RemovableHandle]:
         hooks = set()
+
+        # TODO: attnq
+        # attention_impl = enable_compressed_attention(model)
+        # hooks.add(self.register_hook(attention_impl, calibrate_attention, "calib"))
+
         for module in model.modules():
             if not hasattr(module, "quantization_scheme"):
                 continue