remove attn head

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

Author: kylesayrs

tests/observer.py

@@ -158,9 +158,6 @@ 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}"
 
 
@@ -185,9 +182,6 @@ 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}"
 
 
@@ -209,8 +203,4 @@ 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

@@ -302,25 +302,6 @@ 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(