Python Module as a drop-in replacement for transformers using Attention Sinks

[tomaarsen]

Hello!

This work is looking extremely promising! Great job. I wanted to notify you that I've created a project for a drop-in replacement of transformers using the attention sinks approach. For example, loading Llama 2 with attention sinks is now as simple as:

from attention_sinks import AutoModel

model = AutoModel.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

If you're interested, you can check out the repository here: https://github.com/tomaarsen/attention_sinks

I ran some experiments over there using the attention_sinks Python module, and was able to get some extremely promising results, much like your paper:

1. transformers: Linear VRAM usage as it doesn't do any windowing. Performance fails after ~4096
2. window_attention: Constant VRAM usage due to the windowing at 1024 tokens. Fails after ~1024 tokens.
3. attention_sinks: Constant VRAM usage due to windowing at 4 attention sink tokens + 1020 most recent tokens. Never fails despite constant VRAM usage.

• Tom Aarsen

[Guangxuan-Xiao]

Great implementation! I'm looking forward to integrating our work into Huggingface transformers!