[Transform] Attention/Cache transforms

[kylesayrs]

Purpose

• Support attention quantization
• Support kv cache quantization within huggingface for lm_eval tests

Prerequisites

• Must be merged at the same time as [Quantization] Attention/ KV Cache Refactor vllm-project/llm-compressor#1651

Changes

New Classes

• Add hookable attention and kvcache implementations which are registered to the attention module as submodules
  • QuantizedAttentionImpl injects itself into the model by registering a new attention implementation called ct_hooked_attention overriding model.config._attn_implementation to be the new implementation name
  • QuantizedKVCache injects itself into the model by overriding the past_key_values input kwarg to attention, and wrapping the functionality of the original cache
  • Calibration and transform hooks can be added to these modules via the hook functions
    • register_query_hook,
    • register_key_hook
    • register_value_hook

Quantization Lifecycle Changes

• Apply
  • The kv_cache_scheme field of the quantization config is now used to call initialize_hooked_kv_cache
  • Attention modules can now be targeted, and are used to call initialize_hooked_attention if attention modules are explicitly targted (see is_narrow_match)
  • Remove logic for "merging" kv cache schemes (this doesn't really make any sense, I'm not sure why it was ever included)
• Initiailize
  • Hooked kv cache and attention modules have their quantization parameters initialized by initialize_module_for_quantization
  • The presence of attention or kvcache submodules is what determines whether attention or kv cache only quantization is being applied
• Serialization
  • QuantizationConfig. from_pretrained was cleaned up with additional comments
  • The kv_cache_scheme field is added if there are any attention modules with a quantization_scheme attached

Helpers

• is_narrow_match is used to check that attention modules are being specifically targeted (rather than targeting all modules in a layer)
• get_head_dim is used to get the attention head_dim from a config

Testing

• Added tests for is_narrow_match
• Added tests for added attention and kvcache classes

[brian-dellabetta]

This looks good, though i have a number of questions and minor suggestions

[brian-dellabetta]

😬 i don't want to delay things, but we should briefly consider if there are alternative solutions

[kylesayrs]

I spent 20 minutes exploring this, it requires creating specialized _ct_hooked_attention functions and specialized QuantizedAttentionImpl, which is more complexity than value added imho

[brian-dellabetta]

can _original_impl be registered on the module level (i.e. each self_attn block) instead of setting a global var?

[kylesayrs]

Sure, but in order to register the _original_impl, it needs to be gotten from somewhere.

The first time, you "get" it from model.config. However on subsequent calls, model.config is overridden. This means that in order to "get" the original implementation, you'd have to go find the last Attention module you registered it to, or else store it in some global store.

You could register it to the model module itself or something like that, but I think that that's less reliable than just a a global store. If it's functionality you're after, we can turn it into a hash table or something, keyed by model hash.

[dsikka]

If the goal is to use this generally for kv_cache and attn quantize, can we move the initialize_hooked_attention and initialize_hooked_kv_cache to initialize.py?

I understand we haven't hooked them in yet for those workflows but I think these belong there.

[dsikka]

do a pass through on any missing docstring, otherwise lgtm.
nice work

The base branch was changed.

[brian-dellabetta]

Following for the most part. A few clarifications, but this makes sense to me

[brian-dellabetta]

Does this wrap every single attention block? If so, global _original_impl will be re-set multiple times, though if the same attention function is used throughout the entire model that's probably ok?

[kylesayrs]

We guard against multiple sets using if model.config._attn_implementation != HOOKED_ATTENTION_NAME:

[brian-dellabetta]

can _original_impl be registered on the module level (i.e. each self_attn block) instead of setting a global var?

The base branch was changed.