[Attention] Attention head quantization strategy

[kylesayrs]

Purpose

• Support attention head quantization
  • This will be used by [Transform] Attention/Cache transforms #436

Prerequisites

• [Tests] Mock Observers, Static Lifecycle Tests #482

Changes

• Add attention head quantization strategy

Testing

• Added attention head quantization test and validated that generated scales and zero points make sense

[brian-dellabetta]

how are we expecting users to use QuantizationStrategy.ATTN_HEAD in a recipe? If i'm understanding correctly, it would look something like this?

quant_stage:
  quant_modifiers:
    QuantizationModifier:
      config_groups:
        group0:
          targets: ["re:.*self_attn$"]
          weights:
            strategy: attn_head
            ...
        group1:
          targets: ["re:.*(q|k|v)_proj$"]
          weights:
            strategy: group
            ...

[kylesayrs]

@brian-dellabetta I’ve decided that giving per-attention strategy its own strategy (rather than reusing group) makes more sense.

quant_stage:
  quant_modifiers:
    QuantizationModifier:
      config_groups:
        group0:
          targets: ["re:.*self_attn$"]
          input_activations:
            strategy: attn_head
            ...