kvquant-hf 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kvquant_hf-0.2.0/LICENSE +21 -0
- kvquant_hf-0.2.0/PKG-INFO +154 -0
- kvquant_hf-0.2.0/README.md +136 -0
- kvquant_hf-0.2.0/kvquant/__init__.py +18 -0
- kvquant_hf-0.2.0/kvquant/bench.py +91 -0
- kvquant_hf-0.2.0/kvquant/core.py +55 -0
- kvquant_hf-0.2.0/kvquant/doesitbite.py +72 -0
- kvquant_hf-0.2.0/kvquant/live.py +103 -0
- kvquant_hf-0.2.0/kvquant/rankconf.py +58 -0
- kvquant_hf-0.2.0/kvquant/storage.py +74 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/PKG-INFO +154 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/SOURCES.txt +17 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/dependency_links.txt +1 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/entry_points.txt +2 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/requires.txt +9 -0
- kvquant_hf-0.2.0/kvquant_hf.egg-info/top_level.txt +1 -0
- kvquant_hf-0.2.0/pyproject.toml +23 -0
- kvquant_hf-0.2.0/setup.cfg +4 -0
- kvquant_hf-0.2.0/tests/test_all.py +76 -0
kvquant_hf-0.2.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Sardor Razikov
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: kvquant-hf
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Tested KV-cache quantization for Hugging Face transformers >= 5 (with a silent-no-op stress test).
|
|
5
|
+
License: MIT
|
|
6
|
+
Keywords: llm,kv-cache,quantization,transformers,inference,efficiency
|
|
7
|
+
Requires-Python: >=3.9
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: torch>=2.1
|
|
11
|
+
Provides-Extra: bench
|
|
12
|
+
Requires-Dist: transformers>=5.0; extra == "bench"
|
|
13
|
+
Requires-Dist: numpy; extra == "bench"
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest; extra == "dev"
|
|
16
|
+
Requires-Dist: numpy; extra == "dev"
|
|
17
|
+
Dynamic: license-file
|
|
18
|
+
|
|
19
|
+
# kvquant
|
|
20
|
+
|
|
21
|
+
**Tested KV-cache quantization for Hugging Face `transformers >= 5` — with a stress test so a silent no-op can't fool you.**
|
|
22
|
+
|
|
23
|
+
Most "quantize the KV cache" snippets you'll find mutate `cache.key_cache[i]`. On
|
|
24
|
+
`transformers 5.x` that is a **non-persisting compatibility accessor**: your edit is
|
|
25
|
+
thrown away and the model keeps using the original fp tensors. Everything *looks* like
|
|
26
|
+
it worked — same code, no error — but perplexity doesn't move because nothing changed.
|
|
27
|
+
The real tensors live at `cache.layers[i].keys` / `.values`. kvquant mutates the right
|
|
28
|
+
ones and ships a guard that **proves** the quantization actually bit.
|
|
29
|
+
|
|
30
|
+
## Install
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
pip install kvquant-hf # core
|
|
34
|
+
pip install "kvquant-hf[bench]" # + transformers/numpy for the perplexity benchmark
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Quickstart
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
from kvquant import quantize_cache # in-place int-N quant of a transformers>=5 Cache
|
|
41
|
+
change = quantize_cache(past_key_values, bits=8) # returns mean |change|; 0.0 == a no-op bug
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
Measure the quality cost yourself:
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
48
|
+
from kvquant import kv_ppl_sweep
|
|
49
|
+
m = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m-it")
|
|
50
|
+
t = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")
|
|
51
|
+
print(kv_ppl_sweep(m, t)) # {None/8/4/2: {ppl, delta_pct, change}}
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
or from the shell:
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
kvquant-bench --model google/gemma-3-270m-it
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
## What int8 KV actually costs (measured)
|
|
61
|
+
|
|
62
|
+
`google/gemma-3-270m-it`, 300 tokens, per-channel symmetric int quant of K and V:
|
|
63
|
+
|
|
64
|
+
| KV cache | perplexity | Δ vs fp | note |
|
|
65
|
+
|---|---|---|---|
|
|
66
|
+
| fp (baseline) | 4.5479 | — | |
|
|
67
|
+
| **int8** | 4.5437 | **−0.09%** | within noise — effectively lossless |
|
|
68
|
+
| int4 | 4.7528 | +4.5% | visible degradation |
|
|
69
|
+
| int2 | 982.76 | +21509% | breaks (this is the stress check) |
|
|
70
|
+
|
|
71
|
+
int8 K/V carries **half** the bytes of fp16 (a quarter of fp32), so an int8 cache is
|
|
72
|
+
the obvious lever for longer context / larger batch at near-zero quality cost. Run the
|
|
73
|
+
sweep on *your* model — numbers vary by architecture.
|
|
74
|
+
|
|
75
|
+
### Real memory, measured (not just quality)
|
|
76
|
+
|
|
77
|
+
`kvquant.compress_cache` stores a populated cache as int8 + a per-token fp16 scale and
|
|
78
|
+
reports the actual byte reduction:
|
|
79
|
+
|
|
80
|
+
```python
|
|
81
|
+
from kvquant import compress_cache
|
|
82
|
+
stores, report = compress_cache(past_key_values)
|
|
83
|
+
print(report) # {'fp_bytes': 5_824_512, 'int8_bytes': 1_467_504, 'ratio': 3.97}
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
On the gemma cache above (158 tokens, 18 layers): **5,824,512 → 1,467,504 bytes = 3.97×
|
|
87
|
+
smaller vs fp32 (~2× vs fp16)**, ~2% per-tensor reconstruction error. Use it to compress
|
|
88
|
+
a cache for offload / serialization / holding an idle long context.
|
|
89
|
+
|
|
90
|
+
### Live int8 cache (holds int8 *during* generation)
|
|
91
|
+
|
|
92
|
+
`make_live_int8` converts a cache so every layer **stores** int8 + a per-token fp16 scale
|
|
93
|
+
and dequantizes on read inside attention — the memory held is really int8, no extra deps:
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
from kvquant import make_live_int8, held_bytes
|
|
97
|
+
|
|
98
|
+
past = None
|
|
99
|
+
for step in decode_loop:
|
|
100
|
+
out = model(next_token, past_key_values=past, use_cache=True)
|
|
101
|
+
past = make_live_int8(out.past_key_values) # cache now holds int8
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
Measured on `gemma-3-270m-it` (300 tokens, autoregressive): perplexity **+0.74%**, cache
|
|
105
|
+
held as `torch.int8` and **3.97× smaller than fp32 (~2× vs fp16)**. It subclasses each
|
|
106
|
+
layer's own class, so sliding-window and other behaviours are preserved.
|
|
107
|
+
|
|
108
|
+
## The silent-no-op guard
|
|
109
|
+
|
|
110
|
+
`kv_ppl_sweep` asserts that int2 both changes the tensors (`change > 0`) **and** wrecks
|
|
111
|
+
perplexity. If it doesn't, you're mutating the wrong thing and the whole measurement is
|
|
112
|
+
meaningless — so it raises instead of printing a comfortable lie. The same idea is
|
|
113
|
+
packaged as a standalone decorator:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
from kvquant.doesitbite import bites
|
|
117
|
+
|
|
118
|
+
@bites() # raises AssertionError if the transform doesn't change its input
|
|
119
|
+
def my_quantizer(x): ...
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Also included
|
|
123
|
+
|
|
124
|
+
- **`kvquant.doesitbite`** — `@bites` / `assert_bites` to catch transforms that silently
|
|
125
|
+
do nothing (mis-targeted in-place edits, dead paths, clamps that never fire).
|
|
126
|
+
- **`kvquant.rankconf`** — confidence utilities for rank metrics (AURC/risk-coverage):
|
|
127
|
+
`aurc`, `is_rank_noop` (any strictly-monotonic rescale is a no-op for a rank metric),
|
|
128
|
+
and `agreement_confidence` (a stronger signal than a model's self-reported confidence).
|
|
129
|
+
|
|
130
|
+
## Honest limitations (v0)
|
|
131
|
+
|
|
132
|
+
- Three tools, clear jobs: `kv_ppl_sweep` / `quantize_cache` **fake-quantize** (quant→dequant)
|
|
133
|
+
to measure the *quality* cost with the correct non-no-op hook; `compress_cache` /
|
|
134
|
+
`Int8KVStore` do **real int8 storage** for a finished cache; `make_live_int8` holds int8
|
|
135
|
+
**during** generation. All measured above.
|
|
136
|
+
- Symmetric quantization, no outlier handling or post-RoPE split yet (those shrink the
|
|
137
|
+
int8 quality gap further).
|
|
138
|
+
- Benchmarked on small models on a laptop (gemma-3-270m). Try it on your model and
|
|
139
|
+
architecture and report back — especially GQA/MQA and sliding-window variants.
|
|
140
|
+
|
|
141
|
+
## Support
|
|
142
|
+
|
|
143
|
+
kvquant is free and MIT-licensed. If it saved you GPU memory or an afternoon of
|
|
144
|
+
debugging, a tip keeps the work going — no obligation.
|
|
145
|
+
|
|
146
|
+
- **GitHub Sponsors:** the ❤ Sponsor button above (once enabled)
|
|
147
|
+
- **BTC:** `bc1q9hm4upksmvkzhl8z3zvxgz94da358qfe3tred0`
|
|
148
|
+
- **ETH / Polygon / USDT (ERC-20):** `0xa9e3b63799f5E5cc6258042D45869a5a1d39F021`
|
|
149
|
+
- **USDT (TRC-20) / TRON:** `TUyNkMXme1CYAc1wZ8SZ64X9upBc9nwTZ5`
|
|
150
|
+
- **Solana:** `j8zZCsWwM8WmUZwMSDHikNuz6E49Vn8nwo9FMQnHSoW`
|
|
151
|
+
|
|
152
|
+
## License
|
|
153
|
+
|
|
154
|
+
MIT.
|
|
@@ -0,0 +1,136 @@
|
|
|
1
|
+
# kvquant
|
|
2
|
+
|
|
3
|
+
**Tested KV-cache quantization for Hugging Face `transformers >= 5` — with a stress test so a silent no-op can't fool you.**
|
|
4
|
+
|
|
5
|
+
Most "quantize the KV cache" snippets you'll find mutate `cache.key_cache[i]`. On
|
|
6
|
+
`transformers 5.x` that is a **non-persisting compatibility accessor**: your edit is
|
|
7
|
+
thrown away and the model keeps using the original fp tensors. Everything *looks* like
|
|
8
|
+
it worked — same code, no error — but perplexity doesn't move because nothing changed.
|
|
9
|
+
The real tensors live at `cache.layers[i].keys` / `.values`. kvquant mutates the right
|
|
10
|
+
ones and ships a guard that **proves** the quantization actually bit.
|
|
11
|
+
|
|
12
|
+
## Install
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
pip install kvquant-hf # core
|
|
16
|
+
pip install "kvquant-hf[bench]" # + transformers/numpy for the perplexity benchmark
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
## Quickstart
|
|
20
|
+
|
|
21
|
+
```python
|
|
22
|
+
from kvquant import quantize_cache # in-place int-N quant of a transformers>=5 Cache
|
|
23
|
+
change = quantize_cache(past_key_values, bits=8) # returns mean |change|; 0.0 == a no-op bug
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
Measure the quality cost yourself:
|
|
27
|
+
|
|
28
|
+
```python
|
|
29
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
30
|
+
from kvquant import kv_ppl_sweep
|
|
31
|
+
m = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m-it")
|
|
32
|
+
t = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")
|
|
33
|
+
print(kv_ppl_sweep(m, t)) # {None/8/4/2: {ppl, delta_pct, change}}
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
or from the shell:
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
kvquant-bench --model google/gemma-3-270m-it
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## What int8 KV actually costs (measured)
|
|
43
|
+
|
|
44
|
+
`google/gemma-3-270m-it`, 300 tokens, per-channel symmetric int quant of K and V:
|
|
45
|
+
|
|
46
|
+
| KV cache | perplexity | Δ vs fp | note |
|
|
47
|
+
|---|---|---|---|
|
|
48
|
+
| fp (baseline) | 4.5479 | — | |
|
|
49
|
+
| **int8** | 4.5437 | **−0.09%** | within noise — effectively lossless |
|
|
50
|
+
| int4 | 4.7528 | +4.5% | visible degradation |
|
|
51
|
+
| int2 | 982.76 | +21509% | breaks (this is the stress check) |
|
|
52
|
+
|
|
53
|
+
int8 K/V carries **half** the bytes of fp16 (a quarter of fp32), so an int8 cache is
|
|
54
|
+
the obvious lever for longer context / larger batch at near-zero quality cost. Run the
|
|
55
|
+
sweep on *your* model — numbers vary by architecture.
|
|
56
|
+
|
|
57
|
+
### Real memory, measured (not just quality)
|
|
58
|
+
|
|
59
|
+
`kvquant.compress_cache` stores a populated cache as int8 + a per-token fp16 scale and
|
|
60
|
+
reports the actual byte reduction:
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
from kvquant import compress_cache
|
|
64
|
+
stores, report = compress_cache(past_key_values)
|
|
65
|
+
print(report) # {'fp_bytes': 5_824_512, 'int8_bytes': 1_467_504, 'ratio': 3.97}
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
On the gemma cache above (158 tokens, 18 layers): **5,824,512 → 1,467,504 bytes = 3.97×
|
|
69
|
+
smaller vs fp32 (~2× vs fp16)**, ~2% per-tensor reconstruction error. Use it to compress
|
|
70
|
+
a cache for offload / serialization / holding an idle long context.
|
|
71
|
+
|
|
72
|
+
### Live int8 cache (holds int8 *during* generation)
|
|
73
|
+
|
|
74
|
+
`make_live_int8` converts a cache so every layer **stores** int8 + a per-token fp16 scale
|
|
75
|
+
and dequantizes on read inside attention — the memory held is really int8, no extra deps:
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
from kvquant import make_live_int8, held_bytes
|
|
79
|
+
|
|
80
|
+
past = None
|
|
81
|
+
for step in decode_loop:
|
|
82
|
+
out = model(next_token, past_key_values=past, use_cache=True)
|
|
83
|
+
past = make_live_int8(out.past_key_values) # cache now holds int8
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
Measured on `gemma-3-270m-it` (300 tokens, autoregressive): perplexity **+0.74%**, cache
|
|
87
|
+
held as `torch.int8` and **3.97× smaller than fp32 (~2× vs fp16)**. It subclasses each
|
|
88
|
+
layer's own class, so sliding-window and other behaviours are preserved.
|
|
89
|
+
|
|
90
|
+
## The silent-no-op guard
|
|
91
|
+
|
|
92
|
+
`kv_ppl_sweep` asserts that int2 both changes the tensors (`change > 0`) **and** wrecks
|
|
93
|
+
perplexity. If it doesn't, you're mutating the wrong thing and the whole measurement is
|
|
94
|
+
meaningless — so it raises instead of printing a comfortable lie. The same idea is
|
|
95
|
+
packaged as a standalone decorator:
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
from kvquant.doesitbite import bites
|
|
99
|
+
|
|
100
|
+
@bites() # raises AssertionError if the transform doesn't change its input
|
|
101
|
+
def my_quantizer(x): ...
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
## Also included
|
|
105
|
+
|
|
106
|
+
- **`kvquant.doesitbite`** — `@bites` / `assert_bites` to catch transforms that silently
|
|
107
|
+
do nothing (mis-targeted in-place edits, dead paths, clamps that never fire).
|
|
108
|
+
- **`kvquant.rankconf`** — confidence utilities for rank metrics (AURC/risk-coverage):
|
|
109
|
+
`aurc`, `is_rank_noop` (any strictly-monotonic rescale is a no-op for a rank metric),
|
|
110
|
+
and `agreement_confidence` (a stronger signal than a model's self-reported confidence).
|
|
111
|
+
|
|
112
|
+
## Honest limitations (v0)
|
|
113
|
+
|
|
114
|
+
- Three tools, clear jobs: `kv_ppl_sweep` / `quantize_cache` **fake-quantize** (quant→dequant)
|
|
115
|
+
to measure the *quality* cost with the correct non-no-op hook; `compress_cache` /
|
|
116
|
+
`Int8KVStore` do **real int8 storage** for a finished cache; `make_live_int8` holds int8
|
|
117
|
+
**during** generation. All measured above.
|
|
118
|
+
- Symmetric quantization, no outlier handling or post-RoPE split yet (those shrink the
|
|
119
|
+
int8 quality gap further).
|
|
120
|
+
- Benchmarked on small models on a laptop (gemma-3-270m). Try it on your model and
|
|
121
|
+
architecture and report back — especially GQA/MQA and sliding-window variants.
|
|
122
|
+
|
|
123
|
+
## Support
|
|
124
|
+
|
|
125
|
+
kvquant is free and MIT-licensed. If it saved you GPU memory or an afternoon of
|
|
126
|
+
debugging, a tip keeps the work going — no obligation.
|
|
127
|
+
|
|
128
|
+
- **GitHub Sponsors:** the ❤ Sponsor button above (once enabled)
|
|
129
|
+
- **BTC:** `bc1q9hm4upksmvkzhl8z3zvxgz94da358qfe3tred0`
|
|
130
|
+
- **ETH / Polygon / USDT (ERC-20):** `0xa9e3b63799f5E5cc6258042D45869a5a1d39F021`
|
|
131
|
+
- **USDT (TRC-20) / TRON:** `TUyNkMXme1CYAc1wZ8SZ64X9upBc9nwTZ5`
|
|
132
|
+
- **Solana:** `j8zZCsWwM8WmUZwMSDHikNuz6E49Vn8nwo9FMQnHSoW`
|
|
133
|
+
|
|
134
|
+
## License
|
|
135
|
+
|
|
136
|
+
MIT.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
"""kvquant — tested KV-cache quantization for Hugging Face transformers >= 5.
|
|
2
|
+
|
|
3
|
+
Most KV-quant snippets online silently no-op on transformers 5.x because they mutate
|
|
4
|
+
``cache.key_cache[i]`` (a non-persisting compat accessor) instead of the real
|
|
5
|
+
``cache.layers[i].keys`` / ``.values``. kvquant mutates the right tensors and ships a
|
|
6
|
+
stress test so you can never be fooled.
|
|
7
|
+
"""
|
|
8
|
+
from .core import quantize_cache, quantize_dequantize
|
|
9
|
+
from .bench import kv_ppl_sweep
|
|
10
|
+
from .storage import Int8KVStore, cache_bytes, compress_cache
|
|
11
|
+
from .live import make_live_int8, held_bytes
|
|
12
|
+
|
|
13
|
+
__all__ = [
|
|
14
|
+
"quantize_cache", "quantize_dequantize", "kv_ppl_sweep",
|
|
15
|
+
"Int8KVStore", "cache_bytes", "compress_cache",
|
|
16
|
+
"make_live_int8", "held_bytes",
|
|
17
|
+
]
|
|
18
|
+
__version__ = "0.2.0"
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
"""Perplexity benchmark for KV-cache quantization, with a built-in silent-no-op guard.
|
|
2
|
+
|
|
3
|
+
Scores tokens autoregressively (so the *cache* is actually exercised) and quantizes
|
|
4
|
+
the KV cache after every step. Reports perplexity at fp / int8 / int4 / int2. The int2
|
|
5
|
+
row MUST blow up and show a non-zero quant-change — if it does not, the quantization
|
|
6
|
+
is a silent no-op and the numbers are meaningless, so :func:`kv_ppl_sweep` asserts it.
|
|
7
|
+
"""
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
import math
|
|
10
|
+
import torch
|
|
11
|
+
from .core import quantize_cache
|
|
12
|
+
|
|
13
|
+
DEFAULT_TEXT = (
|
|
14
|
+
"Perplexity measures how well a probability model predicts a sample; lower is "
|
|
15
|
+
"better. Quantizing the key-value cache to eight-bit integers halves its memory "
|
|
16
|
+
"footprint, which lets a server hold a longer context or a larger batch on the "
|
|
17
|
+
"same hardware. We test whether that compression degrades predictions. "
|
|
18
|
+
) * 4
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@torch.no_grad()
|
|
22
|
+
def _score(model, input_ids, bits):
|
|
23
|
+
past = None
|
|
24
|
+
nll = 0.0
|
|
25
|
+
n = 0
|
|
26
|
+
change = 0.0
|
|
27
|
+
T = input_ids.shape[1]
|
|
28
|
+
for t in range(T - 1):
|
|
29
|
+
out = model(input_ids[:, t : t + 1], past_key_values=past, use_cache=True)
|
|
30
|
+
past = out.past_key_values
|
|
31
|
+
if bits:
|
|
32
|
+
change += quantize_cache(past, bits)
|
|
33
|
+
logp = torch.log_softmax(out.logits[:, -1, :].float(), dim=-1)
|
|
34
|
+
nll += -logp[0, input_ids[0, t + 1]].item()
|
|
35
|
+
n += 1
|
|
36
|
+
return math.exp(nll / n), change
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def kv_ppl_sweep(model, tokenizer, text: str = DEFAULT_TEXT, bits_list=(None, 8, 4, 2),
|
|
40
|
+
max_tokens: int = 400, device: str | None = None):
|
|
41
|
+
"""Return ``{bits: {"ppl", "delta_pct", "change"}}`` and assert it isn't a no-op."""
|
|
42
|
+
if device is None:
|
|
43
|
+
device = "mps" if torch.backends.mps.is_available() else (
|
|
44
|
+
"cuda" if torch.cuda.is_available() else "cpu")
|
|
45
|
+
model = model.to(device).eval()
|
|
46
|
+
ids = tokenizer(text, return_tensors="pt").input_ids.to(device)[:, :max_tokens]
|
|
47
|
+
out, base = {}, None
|
|
48
|
+
for bits in bits_list:
|
|
49
|
+
ppl, change = _score(model, ids, bits)
|
|
50
|
+
if bits is None:
|
|
51
|
+
base = ppl
|
|
52
|
+
out[bits] = {"ppl": ppl, "delta_pct": (ppl / base - 1) * 100, "change": change}
|
|
53
|
+
# silent-no-op guard: 2-bit must both bite (change>0) and wreck ppl
|
|
54
|
+
two = out.get(2)
|
|
55
|
+
if two is not None and (two["change"] == 0.0 or two["delta_pct"] < 50):
|
|
56
|
+
raise AssertionError(
|
|
57
|
+
"int2 KV did not degrade perplexity (change=%.3f, delta=%.1f%%) — the "
|
|
58
|
+
"quantization is a SILENT NO-OP; you are probably mutating the wrong cache "
|
|
59
|
+
"attribute." % (two["change"], two["delta_pct"]))
|
|
60
|
+
return out
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def _main():
|
|
64
|
+
import argparse
|
|
65
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
66
|
+
ap = argparse.ArgumentParser(description="KV-cache quantization perplexity sweep")
|
|
67
|
+
ap.add_argument("--model", default="google/gemma-3-270m-it")
|
|
68
|
+
ap.add_argument("--tokens", type=int, default=400)
|
|
69
|
+
args = ap.parse_args()
|
|
70
|
+
tok = AutoTokenizer.from_pretrained(args.model)
|
|
71
|
+
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.float32)
|
|
72
|
+
res = kv_ppl_sweep(model, tok, max_tokens=args.tokens)
|
|
73
|
+
print(f"model: {args.model}")
|
|
74
|
+
for bits, r in res.items():
|
|
75
|
+
tag = "fp (baseline)" if bits is None else f"int{bits}"
|
|
76
|
+
print(f" {tag:14s} ppl={r['ppl']:.4f} dppl={r['delta_pct']:+.3f}% "
|
|
77
|
+
f"change={r['change']:.3f}")
|
|
78
|
+
# real, measured memory saving from int8 storage on a populated cache
|
|
79
|
+
import torch as _t
|
|
80
|
+
from .storage import compress_cache
|
|
81
|
+
with _t.no_grad():
|
|
82
|
+
ids = tok(DEFAULT_TEXT, return_tensors="pt").input_ids[:, : args.tokens]
|
|
83
|
+
dev = next(model.parameters()).device
|
|
84
|
+
cache = model(ids.to(dev), use_cache=True).past_key_values
|
|
85
|
+
_, mem = compress_cache(cache)
|
|
86
|
+
print(f" int8 storage: {mem['fp_bytes']:,} -> {mem['int8_bytes']:,} bytes "
|
|
87
|
+
f"= {mem['ratio']:.2f}x smaller vs fp32 (~{mem['ratio']/2:.2f}x vs fp16).")
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
if __name__ == "__main__":
|
|
91
|
+
_main()
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
"""Core KV-cache quantization for Hugging Face transformers >= 5.
|
|
2
|
+
|
|
3
|
+
The one thing people get wrong on transformers 5.x: the real K/V tensors live at
|
|
4
|
+
``cache.layers[i].keys`` / ``.values``. The old ``cache.key_cache[i]`` is a
|
|
5
|
+
non-persisting compatibility accessor — mutating it is a **silent no-op** (the model
|
|
6
|
+
keeps using the original fp tensors). :func:`quantize_cache` mutates the real tensors
|
|
7
|
+
and returns the mean absolute change so you can *prove* it actually bit.
|
|
8
|
+
"""
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
import torch
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def quantize_dequantize(x: torch.Tensor, bits: int = 8, axis: int = -2) -> torch.Tensor:
|
|
14
|
+
"""Symmetric per-channel fake-quantization of ``x`` to ``bits`` integers.
|
|
15
|
+
|
|
16
|
+
Scale is computed per slice along ``axis`` (default: the sequence axis of a
|
|
17
|
+
``[batch, heads, seq, head_dim]`` KV tensor). Returns a dequantized float tensor
|
|
18
|
+
with the same shape and dtype — this measures the *quality* cost of an int cache.
|
|
19
|
+
(For real memory savings the cache must *store* int8; see the README roadmap.)
|
|
20
|
+
"""
|
|
21
|
+
if bits < 2 or bits > 8:
|
|
22
|
+
raise ValueError("bits must be in [2, 8]")
|
|
23
|
+
lv = 2 ** (bits - 1) - 1
|
|
24
|
+
s = x.abs().amax(dim=axis, keepdim=True) / lv
|
|
25
|
+
s = torch.where(s == 0, torch.ones_like(s), s)
|
|
26
|
+
return torch.round(x / s).clamp(-lv, lv) * s
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def quantize_cache(cache, bits: int = 8) -> float:
|
|
30
|
+
"""In-place int-``bits`` quantization of a transformers>=5 ``Cache``.
|
|
31
|
+
|
|
32
|
+
Mutates ``cache.layers[i].keys`` / ``.values`` (the tensors attention actually
|
|
33
|
+
reads). Returns the **mean absolute change** across layers. A return value of
|
|
34
|
+
``0.0`` means the quantization did not bite — treat that as a bug, not a result.
|
|
35
|
+
"""
|
|
36
|
+
layers = getattr(cache, "layers", None)
|
|
37
|
+
if layers is None:
|
|
38
|
+
raise TypeError(
|
|
39
|
+
"expected a transformers>=5 Cache exposing `.layers`; got "
|
|
40
|
+
f"{type(cache).__name__}. On transformers 5.x do NOT mutate `.key_cache` "
|
|
41
|
+
"(a non-persisting compat accessor) — that is a silent no-op."
|
|
42
|
+
)
|
|
43
|
+
total, n = 0.0, 0
|
|
44
|
+
for layer in layers:
|
|
45
|
+
keys = getattr(layer, "keys", None)
|
|
46
|
+
if keys is None:
|
|
47
|
+
continue
|
|
48
|
+
before = keys.clone()
|
|
49
|
+
layer.keys = quantize_dequantize(layer.keys, bits)
|
|
50
|
+
layer.values = quantize_dequantize(layer.values, bits)
|
|
51
|
+
total += (layer.keys - before).abs().mean().item()
|
|
52
|
+
n += 1
|
|
53
|
+
if n == 0:
|
|
54
|
+
raise RuntimeError("cache had no populated layers to quantize")
|
|
55
|
+
return total / n
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""doesitbite — catch silent no-ops in ML pipeline transforms.
|
|
2
|
+
|
|
3
|
+
A transform that doesn't actually change its input (a mis-targeted in-place edit, a
|
|
4
|
+
dead code path, a clamp that never fires) silently degrades your results while every
|
|
5
|
+
test stays green. ``@bites`` asserts a transform MEASURABLY changes its input; use it
|
|
6
|
+
in your tests or wrap a suspicious step.
|
|
7
|
+
|
|
8
|
+
Born 2026-10-28/29: a KV-cache quantizer that mutated the wrong attribute and a
|
|
9
|
+
packing optimizer stuck in a local minimum both passed unnoticed until stress-tested.
|
|
10
|
+
"""
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
import functools
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _meanabs(x):
|
|
16
|
+
try:
|
|
17
|
+
import torch
|
|
18
|
+
if torch.is_tensor(x):
|
|
19
|
+
return (x.float()).abs().mean().item()
|
|
20
|
+
except Exception:
|
|
21
|
+
pass
|
|
22
|
+
try:
|
|
23
|
+
import numpy as np
|
|
24
|
+
a = np.asarray(x, dtype="float64")
|
|
25
|
+
return float(np.abs(a).mean())
|
|
26
|
+
except Exception:
|
|
27
|
+
return None
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def changed(before, after) -> float | None:
|
|
31
|
+
"""Mean absolute difference between two array/tensor-like objects, or None."""
|
|
32
|
+
try:
|
|
33
|
+
import torch
|
|
34
|
+
if torch.is_tensor(before) and torch.is_tensor(after):
|
|
35
|
+
return (after.float() - before.float()).abs().mean().item()
|
|
36
|
+
except Exception:
|
|
37
|
+
pass
|
|
38
|
+
try:
|
|
39
|
+
import numpy as np
|
|
40
|
+
return float(np.abs(np.asarray(after, "float64") - np.asarray(before, "float64")).mean())
|
|
41
|
+
except Exception:
|
|
42
|
+
return None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def bites(min_change: float = 1e-9):
|
|
46
|
+
"""Decorator: assert ``fn(x, ...)`` returns something that differs from ``x``.
|
|
47
|
+
|
|
48
|
+
Raises ``AssertionError`` (a silent no-op) when the output equals the input within
|
|
49
|
+
``min_change``. The first positional arg is treated as the input.
|
|
50
|
+
"""
|
|
51
|
+
def deco(fn):
|
|
52
|
+
@functools.wraps(fn)
|
|
53
|
+
def wrapper(x, *a, **k):
|
|
54
|
+
try:
|
|
55
|
+
import torch
|
|
56
|
+
before = x.clone() if hasattr(x, "clone") else x
|
|
57
|
+
except Exception:
|
|
58
|
+
before = x
|
|
59
|
+
out = fn(x, *a, **k)
|
|
60
|
+
d = changed(before, out)
|
|
61
|
+
if d is not None and d <= min_change:
|
|
62
|
+
raise AssertionError(
|
|
63
|
+
f"{fn.__name__} is a SILENT NO-OP: output differs from input by "
|
|
64
|
+
f"{d:.3g} (<= {min_change:g}). It did not bite.")
|
|
65
|
+
return out
|
|
66
|
+
return wrapper
|
|
67
|
+
return deco
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def assert_bites(fn, x, *args, min_change: float = 1e-9, **kwargs):
|
|
71
|
+
"""Imperative form: run ``fn(x, ...)`` once and assert it changed ``x``."""
|
|
72
|
+
return bites(min_change)(fn)(x, *args, **kwargs)
|
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
"""Live int8 KV cache — actually HOLDS int8 in memory during generation.
|
|
2
|
+
|
|
3
|
+
Unlike :func:`kvquant.quantize_cache` (fake-quant, measures quality) and
|
|
4
|
+
:mod:`kvquant.storage` (compress a finished cache), this converts a transformers>=5
|
|
5
|
+
cache so each layer *stores* its keys/values as int8 + a per-token fp16 scale and
|
|
6
|
+
dequantizes them on read inside attention. Memory held is ~1/4 of fp32 (~1/2 of fp16);
|
|
7
|
+
the quality cost is the int8 cost (sub-1% perplexity on the models tested).
|
|
8
|
+
|
|
9
|
+
Usage: call :func:`make_live_int8` on the cache after each forward in your decode loop
|
|
10
|
+
(or once on a prefilled cache). It is idempotent and keeps the layer's own class
|
|
11
|
+
behaviour (e.g. sliding-window) by subclassing it, so attention still works.
|
|
12
|
+
"""
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
import torch
|
|
15
|
+
|
|
16
|
+
_variants: dict = {}
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def _q(fp):
|
|
20
|
+
s = fp.abs().amax(dim=-1, keepdim=True) / 127.0
|
|
21
|
+
s = torch.where(s == 0, torch.ones_like(s), s)
|
|
22
|
+
return torch.round(fp / s).clamp(-127, 127).to(torch.int8), s.to(torch.float16)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _dq(q, s, dtype):
|
|
26
|
+
return (q.float() * s.float()).to(dtype) if q is not None else None
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _variant(base):
|
|
30
|
+
"""A subclass of ``base`` whose keys/values are int8-backed properties."""
|
|
31
|
+
if base in _variants:
|
|
32
|
+
return _variants[base]
|
|
33
|
+
|
|
34
|
+
class Int8Layer(base):
|
|
35
|
+
_kvquant_int8 = True
|
|
36
|
+
|
|
37
|
+
@property
|
|
38
|
+
def keys(self):
|
|
39
|
+
return _dq(getattr(self, "_qk", None), getattr(self, "_sk", None),
|
|
40
|
+
getattr(self, "_dt", torch.float32))
|
|
41
|
+
|
|
42
|
+
@keys.setter
|
|
43
|
+
def keys(self, v):
|
|
44
|
+
if v is None:
|
|
45
|
+
self._qk = self._sk = None
|
|
46
|
+
else:
|
|
47
|
+
self._dt = v.dtype
|
|
48
|
+
self._qk, self._sk = _q(v)
|
|
49
|
+
|
|
50
|
+
@property
|
|
51
|
+
def values(self):
|
|
52
|
+
return _dq(getattr(self, "_qv", None), getattr(self, "_sv", None),
|
|
53
|
+
getattr(self, "_dt", torch.float32))
|
|
54
|
+
|
|
55
|
+
@values.setter
|
|
56
|
+
def values(self, v):
|
|
57
|
+
if v is None:
|
|
58
|
+
self._qv = self._sv = None
|
|
59
|
+
else:
|
|
60
|
+
self._dt = v.dtype
|
|
61
|
+
self._qv, self._sv = _q(v)
|
|
62
|
+
|
|
63
|
+
def int8_bytes(self) -> int:
|
|
64
|
+
b = 0
|
|
65
|
+
for a in ("_qk", "_sk", "_qv", "_sv"):
|
|
66
|
+
t = getattr(self, a, None)
|
|
67
|
+
if t is not None:
|
|
68
|
+
b += t.element_size() * t.nelement()
|
|
69
|
+
return b
|
|
70
|
+
|
|
71
|
+
_variants[base] = Int8Layer
|
|
72
|
+
return Int8Layer
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def make_live_int8(cache):
|
|
76
|
+
"""Convert a transformers>=5 cache in place so its layers hold int8. Idempotent."""
|
|
77
|
+
layers = getattr(cache, "layers", None)
|
|
78
|
+
if layers is None:
|
|
79
|
+
raise TypeError("expected a transformers>=5 Cache exposing `.layers`")
|
|
80
|
+
for layer in layers:
|
|
81
|
+
if getattr(type(layer), "_kvquant_int8", False):
|
|
82
|
+
continue # already converted
|
|
83
|
+
k = getattr(layer, "keys", None)
|
|
84
|
+
v = getattr(layer, "values", None)
|
|
85
|
+
layer.__class__ = _variant(type(layer))
|
|
86
|
+
layer._dt = k.dtype if k is not None else torch.float32
|
|
87
|
+
layer.keys = k
|
|
88
|
+
layer.values = v
|
|
89
|
+
return cache
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def held_bytes(cache) -> int:
|
|
93
|
+
"""Live byte footprint of the cache's K/V storage (int8 if converted, else fp)."""
|
|
94
|
+
total = 0
|
|
95
|
+
for layer in getattr(cache, "layers", []):
|
|
96
|
+
if hasattr(layer, "int8_bytes"):
|
|
97
|
+
total += layer.int8_bytes()
|
|
98
|
+
else:
|
|
99
|
+
for name in ("keys", "values"):
|
|
100
|
+
t = getattr(layer, name, None)
|
|
101
|
+
if t is not None:
|
|
102
|
+
total += t.element_size() * t.nelement()
|
|
103
|
+
return total
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
"""rankconf — confidence utilities for rank-based selective-prediction metrics.
|
|
2
|
+
|
|
3
|
+
Two facts people miss on metrics that depend only on the *order* of a confidence
|
|
4
|
+
score (AURC, risk-coverage, AUROC):
|
|
5
|
+
|
|
6
|
+
1. Any **strictly monotonic** rescale of confidence (isotonic in the no-tie case,
|
|
7
|
+
Platt, temperature, sigmoid) leaves the metric unchanged — it is a provable no-op.
|
|
8
|
+
So calibrating to improve AURC is wasted effort; spend it on a *better signal*.
|
|
9
|
+
2. **Agreement-based** confidence (how much an independent second pass agrees) is a
|
|
10
|
+
much stronger ranking signal than a model's self-reported confidence.
|
|
11
|
+
|
|
12
|
+
Verified 2026-10-29: on synthetic-but-faithful data, strictly-monotonic rescales gave
|
|
13
|
+
identical AURC (delta 0), and agreement confidence cut AURC by ~34% vs self-reported.
|
|
14
|
+
"""
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
import numpy as np
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def aurc(confidence, error) -> float:
|
|
20
|
+
"""Area under the risk-coverage curve (lower is better).
|
|
21
|
+
|
|
22
|
+
Sorts items by ``confidence`` descending and averages the running mean ``error``
|
|
23
|
+
over all coverage levels. Depends only on the *order* of ``confidence``.
|
|
24
|
+
"""
|
|
25
|
+
confidence = np.asarray(confidence, dtype="float64")
|
|
26
|
+
error = np.asarray(error, dtype="float64")
|
|
27
|
+
order = np.argsort(-confidence)
|
|
28
|
+
e = error[order]
|
|
29
|
+
risk = np.cumsum(e) / np.arange(1, len(e) + 1)
|
|
30
|
+
return float(risk.mean())
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def is_rank_noop(confidence, transform, error, tol: float = 1e-9) -> bool:
|
|
34
|
+
"""True if applying ``transform`` to ``confidence`` leaves AURC unchanged."""
|
|
35
|
+
a = aurc(confidence, error)
|
|
36
|
+
b = aurc(transform(np.asarray(confidence, dtype="float64")), error)
|
|
37
|
+
return abs(a - b) <= tol
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def agreement_confidence(primary, secondary, continuous: bool = True):
|
|
41
|
+
"""Confidence from agreement between two independent predictions.
|
|
42
|
+
|
|
43
|
+
``continuous=True`` returns negative normalized distance (works for numbers or
|
|
44
|
+
equal-length strings via character distance); otherwise exact-match (1.0/0.0).
|
|
45
|
+
Higher = more agreement = more confident.
|
|
46
|
+
"""
|
|
47
|
+
out = []
|
|
48
|
+
for p, s in zip(primary, secondary):
|
|
49
|
+
if not continuous:
|
|
50
|
+
out.append(1.0 if p == s else 0.0)
|
|
51
|
+
elif isinstance(p, (int, float)) and isinstance(s, (int, float)):
|
|
52
|
+
out.append(-abs(float(p) - float(s)))
|
|
53
|
+
else:
|
|
54
|
+
p, s = str(p), str(s)
|
|
55
|
+
m = max(len(p), len(s)) or 1
|
|
56
|
+
dist = sum(a != b for a, b in zip(p, s)) + abs(len(p) - len(s))
|
|
57
|
+
out.append(-dist / m)
|
|
58
|
+
return np.asarray(out, dtype="float64")
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
"""Real int8 storage for a KV cache — actual bytes saved, not just a quality probe.
|
|
2
|
+
|
|
3
|
+
``Int8KVStore`` holds a ``[batch, heads, seq, head_dim]`` K or V tensor as int8 plus a
|
|
4
|
+
per-token fp16 scale, so its footprint is ~1/4 of fp32 (~1/2 of fp16) with <1% scale
|
|
5
|
+
overhead. Per-token scaling (over the ``head_dim`` axis) is append-friendly: every token
|
|
6
|
+
is quantized independently, so a growing cache never needs re-quantizing.
|
|
7
|
+
|
|
8
|
+
Use it to compress a populated cache for offload / serialization / holding an idle
|
|
9
|
+
long context. For int8 *during* attention (read-from-int8 every step) pair this scheme
|
|
10
|
+
with transformers' ``QuantizedCache``; kvquant's value-add is the tested scheme, the
|
|
11
|
+
stress guard, and the measured numbers.
|
|
12
|
+
"""
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
import torch
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _bytes(t: torch.Tensor) -> int:
|
|
18
|
+
return t.element_size() * t.nelement()
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class Int8KVStore:
|
|
22
|
+
"""int8 + per-token fp16 scale storage for one K or V tensor."""
|
|
23
|
+
|
|
24
|
+
def __init__(self, fp: torch.Tensor):
|
|
25
|
+
if fp.ndim < 2:
|
|
26
|
+
raise ValueError("expected a [..., seq, head_dim] tensor")
|
|
27
|
+
scale = fp.abs().amax(dim=-1, keepdim=True) / 127.0
|
|
28
|
+
scale = torch.where(scale == 0, torch.ones_like(scale), scale)
|
|
29
|
+
self.q = torch.round(fp / scale).clamp(-127, 127).to(torch.int8)
|
|
30
|
+
self.scale = scale.to(torch.float16)
|
|
31
|
+
self.dtype = fp.dtype
|
|
32
|
+
self.shape = tuple(fp.shape)
|
|
33
|
+
|
|
34
|
+
def dequantize(self) -> torch.Tensor:
|
|
35
|
+
return (self.q.float() * self.scale.float()).to(self.dtype)
|
|
36
|
+
|
|
37
|
+
@property
|
|
38
|
+
def nbytes(self) -> int:
|
|
39
|
+
return _bytes(self.q) + _bytes(self.scale)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def cache_bytes(cache) -> int:
|
|
43
|
+
"""Live byte footprint of a transformers>=5 cache's K/V tensors."""
|
|
44
|
+
total = 0
|
|
45
|
+
for layer in getattr(cache, "layers", []):
|
|
46
|
+
for name in ("keys", "values"):
|
|
47
|
+
t = getattr(layer, name, None)
|
|
48
|
+
if t is not None:
|
|
49
|
+
total += _bytes(t)
|
|
50
|
+
return total
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def compress_cache(cache):
|
|
54
|
+
"""Return ``[(Int8KVStore, Int8KVStore), ...]`` per layer and the measured savings.
|
|
55
|
+
|
|
56
|
+
Returns ``(stores, report)`` where ``report`` has ``fp_bytes``, ``int8_bytes`` and
|
|
57
|
+
``ratio`` (fp / int8). The original cache is left untouched.
|
|
58
|
+
"""
|
|
59
|
+
stores = []
|
|
60
|
+
fp_bytes = int8_bytes = 0
|
|
61
|
+
for layer in getattr(cache, "layers", []):
|
|
62
|
+
k, v = getattr(layer, "keys", None), getattr(layer, "values", None)
|
|
63
|
+
if k is None:
|
|
64
|
+
continue
|
|
65
|
+
sk, sv = Int8KVStore(k), Int8KVStore(v)
|
|
66
|
+
stores.append((sk, sv))
|
|
67
|
+
fp_bytes += _bytes(k) + _bytes(v)
|
|
68
|
+
int8_bytes += sk.nbytes + sv.nbytes
|
|
69
|
+
report = {
|
|
70
|
+
"fp_bytes": fp_bytes,
|
|
71
|
+
"int8_bytes": int8_bytes,
|
|
72
|
+
"ratio": (fp_bytes / int8_bytes) if int8_bytes else 0.0,
|
|
73
|
+
}
|
|
74
|
+
return stores, report
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: kvquant-hf
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Tested KV-cache quantization for Hugging Face transformers >= 5 (with a silent-no-op stress test).
|
|
5
|
+
License: MIT
|
|
6
|
+
Keywords: llm,kv-cache,quantization,transformers,inference,efficiency
|
|
7
|
+
Requires-Python: >=3.9
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: torch>=2.1
|
|
11
|
+
Provides-Extra: bench
|
|
12
|
+
Requires-Dist: transformers>=5.0; extra == "bench"
|
|
13
|
+
Requires-Dist: numpy; extra == "bench"
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest; extra == "dev"
|
|
16
|
+
Requires-Dist: numpy; extra == "dev"
|
|
17
|
+
Dynamic: license-file
|
|
18
|
+
|
|
19
|
+
# kvquant
|
|
20
|
+
|
|
21
|
+
**Tested KV-cache quantization for Hugging Face `transformers >= 5` — with a stress test so a silent no-op can't fool you.**
|
|
22
|
+
|
|
23
|
+
Most "quantize the KV cache" snippets you'll find mutate `cache.key_cache[i]`. On
|
|
24
|
+
`transformers 5.x` that is a **non-persisting compatibility accessor**: your edit is
|
|
25
|
+
thrown away and the model keeps using the original fp tensors. Everything *looks* like
|
|
26
|
+
it worked — same code, no error — but perplexity doesn't move because nothing changed.
|
|
27
|
+
The real tensors live at `cache.layers[i].keys` / `.values`. kvquant mutates the right
|
|
28
|
+
ones and ships a guard that **proves** the quantization actually bit.
|
|
29
|
+
|
|
30
|
+
## Install
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
pip install kvquant-hf # core
|
|
34
|
+
pip install "kvquant-hf[bench]" # + transformers/numpy for the perplexity benchmark
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Quickstart
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
from kvquant import quantize_cache # in-place int-N quant of a transformers>=5 Cache
|
|
41
|
+
change = quantize_cache(past_key_values, bits=8) # returns mean |change|; 0.0 == a no-op bug
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
Measure the quality cost yourself:
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
48
|
+
from kvquant import kv_ppl_sweep
|
|
49
|
+
m = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m-it")
|
|
50
|
+
t = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")
|
|
51
|
+
print(kv_ppl_sweep(m, t)) # {None/8/4/2: {ppl, delta_pct, change}}
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
or from the shell:
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
kvquant-bench --model google/gemma-3-270m-it
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
## What int8 KV actually costs (measured)
|
|
61
|
+
|
|
62
|
+
`google/gemma-3-270m-it`, 300 tokens, per-channel symmetric int quant of K and V:
|
|
63
|
+
|
|
64
|
+
| KV cache | perplexity | Δ vs fp | note |
|
|
65
|
+
|---|---|---|---|
|
|
66
|
+
| fp (baseline) | 4.5479 | — | |
|
|
67
|
+
| **int8** | 4.5437 | **−0.09%** | within noise — effectively lossless |
|
|
68
|
+
| int4 | 4.7528 | +4.5% | visible degradation |
|
|
69
|
+
| int2 | 982.76 | +21509% | breaks (this is the stress check) |
|
|
70
|
+
|
|
71
|
+
int8 K/V carries **half** the bytes of fp16 (a quarter of fp32), so an int8 cache is
|
|
72
|
+
the obvious lever for longer context / larger batch at near-zero quality cost. Run the
|
|
73
|
+
sweep on *your* model — numbers vary by architecture.
|
|
74
|
+
|
|
75
|
+
### Real memory, measured (not just quality)
|
|
76
|
+
|
|
77
|
+
`kvquant.compress_cache` stores a populated cache as int8 + a per-token fp16 scale and
|
|
78
|
+
reports the actual byte reduction:
|
|
79
|
+
|
|
80
|
+
```python
|
|
81
|
+
from kvquant import compress_cache
|
|
82
|
+
stores, report = compress_cache(past_key_values)
|
|
83
|
+
print(report) # {'fp_bytes': 5_824_512, 'int8_bytes': 1_467_504, 'ratio': 3.97}
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
On the gemma cache above (158 tokens, 18 layers): **5,824,512 → 1,467,504 bytes = 3.97×
|
|
87
|
+
smaller vs fp32 (~2× vs fp16)**, ~2% per-tensor reconstruction error. Use it to compress
|
|
88
|
+
a cache for offload / serialization / holding an idle long context.
|
|
89
|
+
|
|
90
|
+
### Live int8 cache (holds int8 *during* generation)
|
|
91
|
+
|
|
92
|
+
`make_live_int8` converts a cache so every layer **stores** int8 + a per-token fp16 scale
|
|
93
|
+
and dequantizes on read inside attention — the memory held is really int8, no extra deps:
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
from kvquant import make_live_int8, held_bytes
|
|
97
|
+
|
|
98
|
+
past = None
|
|
99
|
+
for step in decode_loop:
|
|
100
|
+
out = model(next_token, past_key_values=past, use_cache=True)
|
|
101
|
+
past = make_live_int8(out.past_key_values) # cache now holds int8
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
Measured on `gemma-3-270m-it` (300 tokens, autoregressive): perplexity **+0.74%**, cache
|
|
105
|
+
held as `torch.int8` and **3.97× smaller than fp32 (~2× vs fp16)**. It subclasses each
|
|
106
|
+
layer's own class, so sliding-window and other behaviours are preserved.
|
|
107
|
+
|
|
108
|
+
## The silent-no-op guard
|
|
109
|
+
|
|
110
|
+
`kv_ppl_sweep` asserts that int2 both changes the tensors (`change > 0`) **and** wrecks
|
|
111
|
+
perplexity. If it doesn't, you're mutating the wrong thing and the whole measurement is
|
|
112
|
+
meaningless — so it raises instead of printing a comfortable lie. The same idea is
|
|
113
|
+
packaged as a standalone decorator:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
from kvquant.doesitbite import bites
|
|
117
|
+
|
|
118
|
+
@bites() # raises AssertionError if the transform doesn't change its input
|
|
119
|
+
def my_quantizer(x): ...
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Also included
|
|
123
|
+
|
|
124
|
+
- **`kvquant.doesitbite`** — `@bites` / `assert_bites` to catch transforms that silently
|
|
125
|
+
do nothing (mis-targeted in-place edits, dead paths, clamps that never fire).
|
|
126
|
+
- **`kvquant.rankconf`** — confidence utilities for rank metrics (AURC/risk-coverage):
|
|
127
|
+
`aurc`, `is_rank_noop` (any strictly-monotonic rescale is a no-op for a rank metric),
|
|
128
|
+
and `agreement_confidence` (a stronger signal than a model's self-reported confidence).
|
|
129
|
+
|
|
130
|
+
## Honest limitations (v0)
|
|
131
|
+
|
|
132
|
+
- Three tools, clear jobs: `kv_ppl_sweep` / `quantize_cache` **fake-quantize** (quant→dequant)
|
|
133
|
+
to measure the *quality* cost with the correct non-no-op hook; `compress_cache` /
|
|
134
|
+
`Int8KVStore` do **real int8 storage** for a finished cache; `make_live_int8` holds int8
|
|
135
|
+
**during** generation. All measured above.
|
|
136
|
+
- Symmetric quantization, no outlier handling or post-RoPE split yet (those shrink the
|
|
137
|
+
int8 quality gap further).
|
|
138
|
+
- Benchmarked on small models on a laptop (gemma-3-270m). Try it on your model and
|
|
139
|
+
architecture and report back — especially GQA/MQA and sliding-window variants.
|
|
140
|
+
|
|
141
|
+
## Support
|
|
142
|
+
|
|
143
|
+
kvquant is free and MIT-licensed. If it saved you GPU memory or an afternoon of
|
|
144
|
+
debugging, a tip keeps the work going — no obligation.
|
|
145
|
+
|
|
146
|
+
- **GitHub Sponsors:** the ❤ Sponsor button above (once enabled)
|
|
147
|
+
- **BTC:** `bc1q9hm4upksmvkzhl8z3zvxgz94da358qfe3tred0`
|
|
148
|
+
- **ETH / Polygon / USDT (ERC-20):** `0xa9e3b63799f5E5cc6258042D45869a5a1d39F021`
|
|
149
|
+
- **USDT (TRC-20) / TRON:** `TUyNkMXme1CYAc1wZ8SZ64X9upBc9nwTZ5`
|
|
150
|
+
- **Solana:** `j8zZCsWwM8WmUZwMSDHikNuz6E49Vn8nwo9FMQnHSoW`
|
|
151
|
+
|
|
152
|
+
## License
|
|
153
|
+
|
|
154
|
+
MIT.
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
kvquant/__init__.py
|
|
5
|
+
kvquant/bench.py
|
|
6
|
+
kvquant/core.py
|
|
7
|
+
kvquant/doesitbite.py
|
|
8
|
+
kvquant/live.py
|
|
9
|
+
kvquant/rankconf.py
|
|
10
|
+
kvquant/storage.py
|
|
11
|
+
kvquant_hf.egg-info/PKG-INFO
|
|
12
|
+
kvquant_hf.egg-info/SOURCES.txt
|
|
13
|
+
kvquant_hf.egg-info/dependency_links.txt
|
|
14
|
+
kvquant_hf.egg-info/entry_points.txt
|
|
15
|
+
kvquant_hf.egg-info/requires.txt
|
|
16
|
+
kvquant_hf.egg-info/top_level.txt
|
|
17
|
+
tests/test_all.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
kvquant
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "kvquant-hf"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Tested KV-cache quantization for Hugging Face transformers >= 5 (with a silent-no-op stress test)."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
keywords = ["llm", "kv-cache", "quantization", "transformers", "inference", "efficiency"]
|
|
13
|
+
dependencies = ["torch>=2.1"]
|
|
14
|
+
|
|
15
|
+
[project.optional-dependencies]
|
|
16
|
+
bench = ["transformers>=5.0", "numpy"]
|
|
17
|
+
dev = ["pytest", "numpy"]
|
|
18
|
+
|
|
19
|
+
[project.scripts]
|
|
20
|
+
kvquant-bench = "kvquant.bench:_main"
|
|
21
|
+
|
|
22
|
+
[tool.setuptools.packages.find]
|
|
23
|
+
include = ["kvquant*"]
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import pytest
|
|
3
|
+
import torch
|
|
4
|
+
|
|
5
|
+
from kvquant.core import quantize_dequantize, quantize_cache
|
|
6
|
+
from kvquant import doesitbite as dib
|
|
7
|
+
from kvquant import rankconf as rc
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def test_quantize_dequantize_bits():
|
|
11
|
+
x = torch.randn(1, 2, 64, 16)
|
|
12
|
+
q8 = quantize_dequantize(x, 8)
|
|
13
|
+
q2 = quantize_dequantize(x, 2)
|
|
14
|
+
assert q8.shape == x.shape
|
|
15
|
+
# coarser quantization must distort more
|
|
16
|
+
assert (q2 - x).abs().mean() > (q8 - x).abs().mean() > 0
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def test_quantize_cache_rejects_wrong_object():
|
|
20
|
+
with pytest.raises(TypeError):
|
|
21
|
+
quantize_cache(object()) # no .layers -> the transformers-5.x gotcha message
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def test_doesitbite_catches_noop():
|
|
25
|
+
@dib.bites()
|
|
26
|
+
def identity(x):
|
|
27
|
+
return x.clone()
|
|
28
|
+
with pytest.raises(AssertionError):
|
|
29
|
+
identity(torch.ones(10))
|
|
30
|
+
|
|
31
|
+
@dib.bites()
|
|
32
|
+
def real(x):
|
|
33
|
+
return x + 1
|
|
34
|
+
assert real(torch.zeros(5)).sum().item() == 5
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_rankconf_monotonic_noop():
|
|
38
|
+
rng = np.random.default_rng(0)
|
|
39
|
+
conf = rng.random(2000)
|
|
40
|
+
err = rng.random(2000) ** 2
|
|
41
|
+
for f in (lambda x: 1 / (1 + np.exp(-x)), lambda x: x ** 3, np.exp):
|
|
42
|
+
assert rc.is_rank_noop(conf, f, err)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def test_rankconf_agreement_beats_self():
|
|
46
|
+
rng = np.random.default_rng(1)
|
|
47
|
+
n = 3000
|
|
48
|
+
err = rng.random(n) ** 2
|
|
49
|
+
self_conf = -0.3 * err + rng.normal(0, 0.3, n)
|
|
50
|
+
agree = -err + rng.normal(0, 0.25, n)
|
|
51
|
+
assert rc.aurc(agree, err) < rc.aurc(self_conf, err)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_int8_store_saves_bytes_and_reconstructs():
|
|
55
|
+
from kvquant.storage import Int8KVStore
|
|
56
|
+
x = torch.randn(1, 4, 128, 64) # [B, H, seq, head_dim]
|
|
57
|
+
s = Int8KVStore(x)
|
|
58
|
+
# ~4x smaller than fp32 (int8 + tiny fp16 per-token scale)
|
|
59
|
+
assert s.nbytes < x.element_size() * x.nelement() / 3.5
|
|
60
|
+
rel = (s.dequantize() - x).abs().mean() / x.abs().mean()
|
|
61
|
+
assert rel < 0.05 # int8 per-token reconstruction stays within a few percent
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_live_int8_cache_holds_int8():
|
|
65
|
+
import torch
|
|
66
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
67
|
+
from kvquant.live import make_live_int8, held_bytes
|
|
68
|
+
tok = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")
|
|
69
|
+
m = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m-it", dtype=torch.float32).eval()
|
|
70
|
+
ids = tok("live int8 cache test " * 8, return_tensors="pt").input_ids[:, :40]
|
|
71
|
+
with torch.no_grad():
|
|
72
|
+
cache = m(ids, use_cache=True).past_key_values
|
|
73
|
+
fp = held_bytes(cache)
|
|
74
|
+
make_live_int8(cache)
|
|
75
|
+
assert cache.layers[0]._qk.dtype == torch.int8
|
|
76
|
+
assert held_bytes(cache) < fp / 3 # real int8 storage, ~4x vs fp32
|