transformers: quantized KV-cache storage op (int8/int4)#2458
Open
czoli1976 wants to merge 3 commits into
Open
Conversation
Add QuantizedDynKeyValueCache, a drop-in for DynKeyValueCache that keeps the resident KV cache in int8 or int4 (KIVI layout: Keys quantized per channel in frozen fixed-size blocks, Values per token) and dequantizes back to the input dtype on read, so downstream RoPE/mask/attention are unchanged and the past-length symbol still resolves. The opt-in quantize_kv_storage transform walks each Sdpa's K and V inputs back through the cache-read plumbing and swaps those caches for the quantized op. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a cost() reporting values-per-token and resident quantized bytes-per-token, so the memory footprint of a quantized cache shows up in `dump` like the plain DynKeyValueCache it replaces. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e cache The first few tokens act as attention sinks whose keys and values carry outsized weight, so quantizing them costs disproportionate accuracy. Keep the first SINK_TOKENS in f32 in both the per-token Value store and the block-wise per-channel Key store — a constant, tiny cost that makes int4 essentially lossless (perplexity-matched, ~99% top-1 agreement on OpenELM, the remaining disagreement being a statistical tie in the reference distribution). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
QuantizedDynKeyValueCache, a drop-in forDynKeyValueCachethat stores the resident KV cache in int8/int4 and dequantizes on read — cutting decode KV memory (~2× int8, ~4× int4 vs an f16 cache) while leaving RoPE/mask/attention untouched. It's opt-in via thequantize_kv_storage/quantize_kv_storage_int4transforms, which walk eachSdpa's K/V inputs back through the cache-read plumbing (broadcast/reshape/cast/on-read RoPE) and swap those caches; because it stays a cache-shaped op, the past-length symbol still resolves.KIVI layout: Keys per channel (quantized in fixed-size blocks with frozen scales for streaming consistency), Values per token. The first few tokens are kept in f32 (attention sinks carry outsized weight), a constant tiny cost that matters most for int4. Distinct from the
block_quantweight-quantization path — this quantizes the activation cache, not the weights.Validated end-to-end on OpenELM-270M (f16): the transform fires on all layers and runs, and the dequant-on-read adds only ~1–4% to a decode step (the step is FFN-dominated). Perplexity and next-token agreement over an 80-token greedy continuation:
Both are perplexity-lossless; int4's single top-1 disagreement is a statistical tie (the reference distribution puts the two candidates at p=0.019 vs 0.018). Unit tests cover accumulation vs a plain concat, block-wise per-channel consistency over a growing range, and NNEF round-trip.
🍍
🤖 Generated with Claude Code