[Transform] Attention/Cache transforms

[kylesayrs]

Purpose

• Support fully-expressive attention and kv cache quantization
• Support running kv cache quantization evals with hf transformers

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_num_attn_heads, get_num_kv_heads, get_head_dim get attention config values from 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.

[brian-dellabetta]

Just trying to figure out what the lifetime of _original_impl is, once everything is set it can basically be treated as no longer necessary? Or is it something that is important the entire duration, during model loading as well as during any forward passes.

[brian-dellabetta]

This is all net-new, so it's not like we will breaking anything pre-existing. global vars make me nervous, but this seems like a legitimate enough use case to use them and accept the risk

[kylesayrs]

If it makes you feel any better, this variable is still only scoped to this file. This is the same as any module-scoped read, only this time we're writing to it, and therefore need the global keyword

[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]

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.

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/18477528039

[kylesayrs]

Last nightly worked, but e2e failed due to model storage issues
https://github.com/neuralmagic/llm-compressor-testing/actions/runs/18483826999

[brian-dellabetta]

We can resolve the global var thread, I have another new comment we might want to consider in a follow-up but marking this as approved. Cool stuff! Excited to see it in action

[brian-dellabetta]

This is all net-new, so it's not like we will breaking anything pre-existing. global vars make me nervous, but this seems like a legitimate enough use case to use them and accept the risk

[brian-dellabetta]

This seems like something that might've been missed in the scoped quant work. If multiple statuses are found, rather than just using the last one found don't we want to set format="mixed-precision" in the returned QuantizationConfig?

[kylesayrs]

Essentially, you're right. However, this value is essentially meaningless, as it is later overridden by the model compressor.