add tests

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

Author: kylesayrs

tests/mock_observer.py

@@ -110,6 +110,9 @@ def flatten_weight_for_quantization(value: torch.Tensor, args: QuantizationArgs)
             .unsqueeze(0)
         )
 
+    if args.strategy == QuantizationStrategy.ATTN_HEAD:
+        raise ValueError("attention head quantization cannot be applied to weights")
+
     assert False, f"Unknown strategy {args.strategy}"
 
 
@@ -134,6 +137,9 @@ def flatten_activation_for_quantization(value: torch.Tensor, args: QuantizationA
     if args.strategy == QuantizationStrategy.BLOCK:
         raise ValueError("Block quantization cannot be applied to activations")
 
+    if args.strategy == QuantizationStrategy.ATTN_HEAD:
+        raise ValueError("attention head quantization cannot be applied to linear acts")
+
     assert False, f"Unknown strategy {args.strategy}"
 
 
@@ -155,4 +161,8 @@ def flatten_attention_for_quantization(value: torch.Tensor, args: QuantizationAr
     if args.strategy == QuantizationStrategy.BLOCK:
         raise ValueError("Block quantization cannot be applied to attention")
 
+    if args.strategy == QuantizationStrategy.ATTN_HEAD:
+        # (batch_size * seq_len, num_heads, 1, head_dim)
+        return value.flatten(0, 1).unsqueeze(-2)
+
     assert False, f"Unknown strategy {args.strategy}"

tests/test_quantization/lifecycle/test_static_lifecycle.py

@@ -303,6 +303,25 @@ class MockAttention(torch.nn.Module):
         # group is not supported
         # tensor group is not supported
         # block is not supported
+        (
+            QuantizationArgs(
+                num_bits=4,
+                type="int",
+                symmetric=True,
+                strategy="attn_head",
+            ),
+            torch.tensor([[0], [3]]),
+            torch.tensor([[8], [11]]),
+            torch.tensor(
+                [
+                    [
+                        [[0.0000, 1.0703, 2.1406], [2.9375, 4.4062, 4.4062]],
+                        [[6.4375, 7.5000, 7.5000], [8.8125, 10.2500, 10.2500]],
+                    ]
+                ]
+            ),
+            0.16,
+        ),
     ],
 )
 def test_static_attention_quantization(