A test of the 'Attention is Off by One' hypothesis. Softermax implementation of llama based on HuggingFace Transformers. We name the subclassed huggingface model SofterLlama (mainly so IDE autocomplete is more convenient than "LlamaSoftermax", but it also sounds nicer).

We track development using this Google Doc

Coincidentally, our timing on adapting llama falls under a massive PR refactoring kv cache utils, partially motivated by attention sinks which is tangentially related to softermax (main difference being our inital "token" that serves as an "attention sink" doesn't have a corresponding V matrix entry).

Since this is the case, we must use transformers==4.36.1 onwards as the latest version to have access to the Cache object.