Add SimLayerKVPress

[SimJeg]

Add SimLayerKVPress (paper, official repository) following issue #19 and PR #22.

SimLayerKV uses a layer-wise approach to compression:
- layers identified as lazy use the Streaming LLM approach (only initial and recent KV pairs are kept)
- other layers use the full KV cache

To identify lazy layers, the last attention weights are used. If the sum of attention weights of the last tokens
over the initial and recent tokens is above the lazy_threshold, the layer is considered lazy.

As for the wrapper for layer-wise compression ratio, this press only works using flash attention (more investigation to be done on why). A notable difference from other presses is that the input of SimLayerKVPress is not a compression ratio but the lazy_threshold as defined in the paper. However I implemented a compression_ratio property that is computed dynamically.

@dame-cell, after implementing it, I found it was not necessary after all to implement a new cache class. Could you please review and comment this implementation ? I will also share it on the official repository to get a feedback from the authors.

[dame-cell]

Add SimLayerKVPress (paper, official repository) following issue #19 and PR #22.

SimLayerKV uses a layer-wise approach to compression:
- layers identified as lazy use the Streaming LLM approach (only initial and recent KV pairs are kept)
- other layers use the full KV cache

To identify lazy layers, the last attention weights are used. If the sum of attention weights of the last tokens
over the initial and recent tokens is above the lazy_threshold, the layer is considered lazy.

As for the wrapper for layer-wise compression ratio, this press only works using flash attention (more investigation to be done on why). A notable difference from other presses is that the input of SimLayerKVPress is not a compression ratio but the lazy_threshold as defined in the paper. However I implemented a compression_ratio property that is computed dynamically.

@dame-cell, after implementing it, I found it was not necessary after all to implement a new cache class. Could you please review and comment this implementation ? I will also share it on the official repository to get a feedback from the authors.

Thanks for adding this i was banging my head just trying to implement this

Right now Im actually busy I'll try reviewing tommorow if that's ok 😔

[SimJeg]

Thanks for adding this i was banging my head just trying to implement this

Right now Im actually busy I'll try reviewing tommorow if that's ok 😔

@dame-cell no pb, I also ask for a review on the official repository (see sail-sg/SimLayerKV#4)

[maxjeblick]

Thanks a lot for the PR!
In general, the PR is in a good shape. I left some comments, they should be fast to fix.
Regrading the press itself, it looks good to me. I haven't studied the original work in detail, so it may make sense to also wait if the authors give some feedback.

[maxjeblick]

Code + implementation LGTM, thanks for adding the press!

I haven't checked in detail if the press is equivalent to the original one, if there's no feedback in the next day(s), I will have a look at this, as well.

[dame-cell]

@SimJeg I have added a comment please check it out and tell me if this will work or not ?
other than that i have tested it myself seems to be working pretty good 💯

[SimJeg]

@dame-cell I don't see your comment, can you provide a link to it ?

[dame-cell]

@SimJeg forgive me here is the comment
in the original implementation they included different threshold for different models

        if 'llama3' in out_path:
            threshold = 0.9
        elif 'llama2' in out_path:
            threshold = 0.65
        elif 'mistral' in out_path:
            threshold = 0.8
        elif 'qwen' in out_path:
            threshold = 0.85
       

Adding something similar could make this even more versatile and model-aware. Just a thought—curious to hear your perspective! 😊

maybe something like this

def get_lazy_threshold(model_name: str) -> float:
    if 'llama3' in model_name:
        return 0.9
    elif 'llama2' in model_name:
        return 0.65
    elif 'mistral' in model_name:
        return 0.8
    elif 'qwen' in model_name:
        return 0.85
    else:
        return 0.7  # Default threshold

# Example  
module.config.get("model_name", "")
lazy_threshold = get_lazy_threshold(model_name)

[SimJeg]

I prefered to let the user specify the lazy_threshold argument to make it model agnostic but I will update the docstring to provide to the user help to set this value.

[maxjeblick]

LGTM, thanks!
I left two small comments.

[maxjeblick]

LGTM, thanks!

[maxjeblick]

Reapproving

[jadeCurl]

Hi,

Thanks for incorporating our SimLayerKV!

We are currently working on version 2 of our project, with a major update being its integration with flash attention to enhance efficiency. For your reference, here is the source code:

attn_out, lse = flash_attn(q, k, v, causal=True, return_lse=True)
# identification
# w_last = 32, w_sink=4, w_recent=1020
q_last = q[:, -w_last:].permute(0, 2, 1, 3)
k_comb = torch.cat([k[:, 0:w_sink], k[:, -w_recent:]], dim=1).permute(0, 2, 3, 1)
log_lazy_weight = torch.matmul(q_last, k_comb).logsumexp(dim=-1) - lse[:,:,-w_recent:]