adding kv_cache quantization

[HDCharles]

Adding support for kv_cache quantization, we are using simple symmetric quantization, though using the full precision of the k and v values of the current token.

we see tok/s reduction of 3-5 tok/s depending on context length.

and a reduction in peak memory

We expect this reduction to scale to large context lengths, in the model memory trace we can see the point where we replace the bf16 cache with the int8 cache which visually saves about half of the used memory

at longer context lengths both quantized and non-quantized kv_cache models start outputing weird stuff but otherwise accuracy of the kv_cache quant looks reasonable though e.g. for 2048 context length:

<|begin_of_text|>Hello, my name is Richard Brown and I have been a professional musician for over 25 years. I have played in a number of bands, doing a wide variety of genres (soul/funk, rock, jazz, blues, latin, world). I have played on over a hundred albums so far.
I have played with many different singers, as well as instrumentalists (guitarists, sax players, brass players, etc.). I love to play and try to learn as much as I can from others. I have become an all-round musician - playing keyboards, drums, programming, arranging; as well as writing songs myself. I have my own studio, and I can do sessions online.
I also have my own website, where you can find out more about me and my music.
I hope that you will find the music that you are looking for here.

Otherwise there are some fixes in generate.py to get things working for large context lengths without overflowing beyond the model limit.

test plan:

sh benchmarks.sh

(specifically the last 6 rows of benchmark_results.txt)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/532

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6c0c9b7 with merge base 0cc1c4d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[msaroufim]

Were the memory decreases also minimal at larger sequence lengths? I'd increase it substantially to 100K, 200K, 500K, 1M and see how results change then

[HDCharles]

Were the memory decreases also minimal at larger sequence lengths? I'd increase it substantially to 100K, 200K, 500K, 1M and see how results change then

i'll test that

[gau-nernst]

I'm guessing the bfloat16 kv cache is materialized in global memory, thus no memory saving. need to inspect it further.

[msaroufim]

discussed followup work offline