aither-kvcache 2.2.0__tar.gz → 2.4.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.
- {aither_kvcache-2.2.0/aither_kvcache.egg-info → aither_kvcache-2.4.0}/PKG-INFO +101 -3
- aither_kvcache-2.2.0/PKG-INFO → aither_kvcache-2.4.0/README.md +95 -37
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/__init__.py +2 -2
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/block_selector.py +0 -2
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/fused_kv_update.py +0 -1
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/hybrid_quantizer.py +0 -7
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/__init__.py +70 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/capture.py +648 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/fit.py +257 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/geometry.py +472 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/mapper.py +407 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/pack.py +521 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/rope.py +417 -0
- aither_kvcache-2.4.0/aither_kvcache/kvtransfer/transfer.py +432 -0
- aither_kvcache-2.4.0/aither_kvcache/persistence.py +253 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/attention.py +0 -4
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/cache.py +1 -1
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/calibration.py +4 -2
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/config.py +1 -1
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/encoder.py +0 -1
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/scorer.py +0 -2
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/backend.py +823 -115
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/engine.py +122 -19
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/eviction_plugin.py +0 -1
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/hooks.py +5 -5
- aither_kvcache-2.4.0/aither_kvcache/vllm/isolation_status_endpoint.py +128 -0
- aither_kvcache-2.4.0/aither_kvcache/vllm/kv_warmth_client.py +317 -0
- aither_kvcache-2.4.0/aither_kvcache/vllm/plugin.py +159 -0
- aither_kvcache-2.4.0/aither_kvcache/vllm/strata_shadow_plugin.py +387 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/tenant_isolation.py +55 -5
- aither_kvcache-2.2.0/README.md → aither_kvcache-2.4.0/aither_kvcache.egg-info/PKG-INFO +135 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache.egg-info/SOURCES.txt +16 -7
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache.egg-info/requires.txt +4 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/pyproject.toml +9 -3
- aither_kvcache-2.4.0/tests/test_block_selector.py +68 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/tests/test_core.py +3 -3
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/tests/test_fused.py +3 -3
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/tests/test_hybrid_hooks.py +3 -3
- aither_kvcache-2.4.0/tests/test_persistence.py +157 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/tests/test_triattention.py +0 -4
- aither_kvcache-2.4.0/tests/test_vllm_plugin.py +74 -0
- aither_kvcache-2.4.0/turboquant/gb10_fp8_rescue.py +317 -0
- aither_kvcache-2.2.0/aither_kvcache/bench.py +0 -140
- aither_kvcache-2.2.0/aither_kvcache/codebook.py +0 -166
- aither_kvcache-2.2.0/aither_kvcache/fused_attention.py +0 -748
- aither_kvcache-2.2.0/aither_kvcache/quantizer.py +0 -370
- aither_kvcache-2.2.0/aither_kvcache/vllm/plugin.py +0 -71
- aither_kvcache-2.2.0/turboquant/packing.py +0 -131
- aither_kvcache-2.2.0/turboquant/rotation.py +0 -108
- aither_kvcache-2.2.0/turboquant/triton_ops.py +0 -259
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/LICENSE +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/eviction_advisor.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/kvcache_graph.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/__init__.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/triattention/spectral.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/__init__.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/cache.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache/vllm/sitecustomize.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache.egg-info/dependency_links.txt +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache.egg-info/entry_points.txt +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/aither_kvcache.egg-info/top_level.txt +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/setup.cfg +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/tests/test_kvcache_graph.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/turboquant/__init__.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/turboquant/bench.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/turboquant/codebook.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/turboquant/fused_attention.py +0 -0
- {aither_kvcache-2.2.0/aither_kvcache → aither_kvcache-2.4.0/turboquant}/packing.py +0 -0
- {aither_kvcache-2.2.0 → aither_kvcache-2.4.0}/turboquant/quantizer.py +0 -0
- {aither_kvcache-2.2.0/aither_kvcache → aither_kvcache-2.4.0/turboquant}/rotation.py +0 -0
- {aither_kvcache-2.2.0/aither_kvcache → aither_kvcache-2.4.0/turboquant}/triton_ops.py +0 -0
|
@@ -1,13 +1,13 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: aither-kvcache
|
|
3
|
-
Version: 2.
|
|
4
|
-
Summary:
|
|
3
|
+
Version: 2.4.0
|
|
4
|
+
Summary: KV cache tooling for LLM inference — TurboQuant (vector quantization), TriAttention (spectral compression), and cross-model KV transfer (skip prefill by converting one model's cache into another's)
|
|
5
5
|
Author: Aitherium
|
|
6
6
|
License-Expression: CC-BY-4.0
|
|
7
7
|
Project-URL: Paper, https://arxiv.org/abs/2504.19874
|
|
8
8
|
Project-URL: Repository, https://github.com/Aitherium/aitherkvcache
|
|
9
9
|
Project-URL: Changelog, https://github.com/Aitherium/aitherkvcache/releases
|
|
10
|
-
Keywords: kv-cache,quantization,llm,inference,compression,turboquant,triattention,spectral,vllm,attention
|
|
10
|
+
Keywords: kv-cache,quantization,llm,inference,compression,turboquant,triattention,spectral,vllm,attention,kv-transfer,cross-model,prefill,rope,mla
|
|
11
11
|
Classifier: Development Status :: 4 - Beta
|
|
12
12
|
Classifier: Intended Audience :: Science/Research
|
|
13
13
|
Classifier: Programming Language :: Python :: 3
|
|
@@ -29,10 +29,13 @@ Provides-Extra: dev
|
|
|
29
29
|
Requires-Dist: pytest; extra == "dev"
|
|
30
30
|
Provides-Extra: vllm
|
|
31
31
|
Requires-Dist: vllm>=0.15; extra == "vllm"
|
|
32
|
+
Provides-Extra: transfer
|
|
33
|
+
Requires-Dist: transformers>=4.40; extra == "transfer"
|
|
32
34
|
Provides-Extra: all
|
|
33
35
|
Requires-Dist: triton>=2.0; extra == "all"
|
|
34
36
|
Requires-Dist: scipy; extra == "all"
|
|
35
37
|
Requires-Dist: pytest; extra == "all"
|
|
38
|
+
Requires-Dist: transformers>=4.40; extra == "all"
|
|
36
39
|
Dynamic: license-file
|
|
37
40
|
|
|
38
41
|
# aither-kvcache
|
|
@@ -46,12 +49,18 @@ Near-optimal KV cache compression for LLM inference. Two compression engines:
|
|
|
46
49
|
Retains top RoPE frequency pairs, scores via trig series without materializing full K/V.
|
|
47
50
|
~10× compression with bounded approximation error. Calibrated for Qwen3.5 family.
|
|
48
51
|
|
|
52
|
+
- **KVTransfer** *(NEW in v2.4)* — Cross-model KV transfer. Convert one model's KV cache
|
|
53
|
+
into another model's so the receiving model **skips prefill entirely**. Per-(layer, head)
|
|
54
|
+
linear maps fitted by ridge; source layers chosen per target layer by held-out R².
|
|
55
|
+
Same model family only — the two models must share a tokenizer.
|
|
56
|
+
|
|
49
57
|
## Installation
|
|
50
58
|
|
|
51
59
|
```bash
|
|
52
60
|
pip install aither-kvcache # core library
|
|
53
61
|
pip install aither-kvcache[vllm] # + vLLM plugin (v0.15+)
|
|
54
62
|
pip install aither-kvcache[triton] # + fused GPU kernels
|
|
63
|
+
pip install aither-kvcache[transfer] # + cross-model KV transfer
|
|
55
64
|
pip install aither-kvcache[all] # everything
|
|
56
65
|
```
|
|
57
66
|
|
|
@@ -126,6 +135,95 @@ instead of 256 — a ~10× reduction with bounded approximation error.
|
|
|
126
135
|
| F=16, int4 | 4 | 34 | 7.53× |
|
|
127
136
|
| F=8, int4 | 4 | 18 | **14.2×** |
|
|
128
137
|
|
|
138
|
+
## Quick Start — KVTransfer (v2.4+)
|
|
139
|
+
|
|
140
|
+
Prefill is the tax you pay before a model says anything, and a KV cache only works on the
|
|
141
|
+
model that produced it. Route a conversation to a different model and the accumulated
|
|
142
|
+
cache becomes dead weight. KVTransfer fits a map that converts one model's cache into
|
|
143
|
+
another's, so the receiving model can skip prefill.
|
|
144
|
+
|
|
145
|
+
```bash
|
|
146
|
+
# 1. capture aligned KV from both models over one corpus
|
|
147
|
+
python -m aither_kvcache.kvtransfer.capture \
|
|
148
|
+
--source Qwen/Qwen3-0.6B --target Qwen/Qwen3-4B \
|
|
149
|
+
--out ./cap --corpus ./my-corpus --seq-len 256 --train-seqs 64 --val-seqs 24
|
|
150
|
+
|
|
151
|
+
# 2. fit the mapper pack
|
|
152
|
+
python -m aither_kvcache.kvtransfer.fit --capture ./cap --out ./pack --top-k 8
|
|
153
|
+
|
|
154
|
+
# 3. measure what the RECEIVING model does with a translated cache, and record it
|
|
155
|
+
python -m aither_kvcache.kvtransfer.transfer \
|
|
156
|
+
--pack ./pack --source Qwen/Qwen3-0.6B --target Qwen/Qwen3-4B \
|
|
157
|
+
--corpus ./held-out --sequences 24 --cuts 32 64 96 128 160 192 224 --record
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
```python
|
|
161
|
+
from aither_kvcache.kvtransfer import load_pack
|
|
162
|
+
from aither_kvcache.kvtransfer.transfer import translate_cache
|
|
163
|
+
|
|
164
|
+
pack = load_pack("./pack", live_source=src_geom, live_target=tgt_geom)
|
|
165
|
+
layers = translate_cache(src_k, src_v, pack, positions) # -> the target's KV cache
|
|
166
|
+
```
|
|
167
|
+
|
|
168
|
+
### Step 3 is not optional
|
|
169
|
+
|
|
170
|
+
`load_pack` **refuses** any pack that carries no downstream acceptance evidence. That is
|
|
171
|
+
the core design decision here, and it is not defensive programming — it is the only thing
|
|
172
|
+
standing between an interesting R² and a broken deployment.
|
|
173
|
+
|
|
174
|
+
A converted cache does not fail loudly. An approximate one makes the receiving model
|
|
175
|
+
produce fluent, on-topic, **confidently wrong** text: no exception, no error status, no
|
|
176
|
+
unhealthy process, nothing in a log. Every cheap signal is green.
|
|
177
|
+
|
|
178
|
+
So a pack must carry a measurement, and the measurement has **four arms**:
|
|
179
|
+
|
|
180
|
+
| arm | what it is |
|
|
181
|
+
|---|---|
|
|
182
|
+
| `reference` | the target prefills the context itself — the ceiling |
|
|
183
|
+
| `translated` | the target is handed the converted cache — the candidate |
|
|
184
|
+
| `control` | the map is run on a **different document** — catches a mapper ignoring its input |
|
|
185
|
+
| `nocontext` | no cache at all — the floor |
|
|
186
|
+
|
|
187
|
+
The control arm is the one that matters. A mapper that has quietly learned the target's
|
|
188
|
+
*average* key/value statistics scores respectably against `reference` alone. Without a
|
|
189
|
+
floor, "68% agreement" is unreadable — it could be excellent, or exactly what a dead
|
|
190
|
+
mapper gets.
|
|
191
|
+
|
|
192
|
+
A minimum sample size is enforced too. Top-1 agreement is a proportion, so a mean over a
|
|
193
|
+
few dozen positions has a standard error near ten percentage points — and once recorded it
|
|
194
|
+
is indistinguishable from a mean over ten thousand.
|
|
195
|
+
|
|
196
|
+
### Preconditions that are definitional, not quality knobs
|
|
197
|
+
|
|
198
|
+
- **The two models must tokenize identically.** The map sends position *i* to position
|
|
199
|
+
*i*; different tokenizers mean row *i* is not the same token and the fit is regressing
|
|
200
|
+
misaligned data. Nothing downstream can see this — it reads as "this pair transfers
|
|
201
|
+
poorly" rather than "these rows do not correspond". Enforced by vocabulary digest, and
|
|
202
|
+
by comparing token ids per document at capture time.
|
|
203
|
+
- **The RoPE schedule must be exactly reproducible.** Keys carry a position rotation, so
|
|
204
|
+
the map is fitted in position-free space: stripped with the *source* schedule, re-applied
|
|
205
|
+
with the *target's* (not a no-op — the two models often disagree on `rope_theta`).
|
|
206
|
+
`default`, `linear` and `yarn` are implemented; anything else is refused rather than
|
|
207
|
+
approximated.
|
|
208
|
+
|
|
209
|
+
`check_pair()` separates definitional refusals from merely-unproven regimes — mismatched
|
|
210
|
+
KV-head counts, differing head dimension, cross-family pairs. The latter are recorded as
|
|
211
|
+
flags and decided by measurement, because a gate keyed on "heads must match" would refuse
|
|
212
|
+
the very experiment that tests whether heads need to match.
|
|
213
|
+
|
|
214
|
+
### Latent attention (MLA)
|
|
215
|
+
|
|
216
|
+
Models such as DeepSeek V2/V3/V4 cache a compressed latent plus a decoupled RoPE key
|
|
217
|
+
rather than per-head K and V. That is a different cache **layout**, not an obstacle — the
|
|
218
|
+
latent is dense, and often smaller than a conventional cache for the same context.
|
|
219
|
+
`KVGeometry.roles()` returns `("c", "kr")` for those models, and only `kr` carries a
|
|
220
|
+
rotation because the latent is position-free by construction.
|
|
221
|
+
|
|
222
|
+
The latent is captured at the compression projection rather than from `past_key_values`:
|
|
223
|
+
HuggingFace materialises per-head K/V before caching, so reading the cache would yield
|
|
224
|
+
tensors a real serving engine never stores, and a mapper fitted to them could not drive a
|
|
225
|
+
production deployment while looking entirely correct end to end.
|
|
226
|
+
|
|
129
227
|
## vLLM Integration
|
|
130
228
|
|
|
131
229
|
### Native kv-cache-dtype (recommended)
|
|
@@ -1,40 +1,3 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: aither-kvcache
|
|
3
|
-
Version: 2.2.0
|
|
4
|
-
Summary: Near-optimal KV cache compression for LLM inference — TurboQuant (vector quantization) + TriAttention (spectral compression)
|
|
5
|
-
Author: Aitherium
|
|
6
|
-
License-Expression: CC-BY-4.0
|
|
7
|
-
Project-URL: Paper, https://arxiv.org/abs/2504.19874
|
|
8
|
-
Project-URL: Repository, https://github.com/Aitherium/aitherkvcache
|
|
9
|
-
Project-URL: Changelog, https://github.com/Aitherium/aitherkvcache/releases
|
|
10
|
-
Keywords: kv-cache,quantization,llm,inference,compression,turboquant,triattention,spectral,vllm,attention
|
|
11
|
-
Classifier: Development Status :: 4 - Beta
|
|
12
|
-
Classifier: Intended Audience :: Science/Research
|
|
13
|
-
Classifier: Programming Language :: Python :: 3
|
|
14
|
-
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
-
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
19
|
-
Requires-Python: >=3.10
|
|
20
|
-
Description-Content-Type: text/markdown
|
|
21
|
-
License-File: LICENSE
|
|
22
|
-
Requires-Dist: torch>=2.0
|
|
23
|
-
Requires-Dist: numpy
|
|
24
|
-
Provides-Extra: triton
|
|
25
|
-
Requires-Dist: triton>=2.0; extra == "triton"
|
|
26
|
-
Provides-Extra: scipy
|
|
27
|
-
Requires-Dist: scipy; extra == "scipy"
|
|
28
|
-
Provides-Extra: dev
|
|
29
|
-
Requires-Dist: pytest; extra == "dev"
|
|
30
|
-
Provides-Extra: vllm
|
|
31
|
-
Requires-Dist: vllm>=0.15; extra == "vllm"
|
|
32
|
-
Provides-Extra: all
|
|
33
|
-
Requires-Dist: triton>=2.0; extra == "all"
|
|
34
|
-
Requires-Dist: scipy; extra == "all"
|
|
35
|
-
Requires-Dist: pytest; extra == "all"
|
|
36
|
-
Dynamic: license-file
|
|
37
|
-
|
|
38
1
|
# aither-kvcache
|
|
39
2
|
|
|
40
3
|
Near-optimal KV cache compression for LLM inference. Two compression engines:
|
|
@@ -46,12 +9,18 @@ Near-optimal KV cache compression for LLM inference. Two compression engines:
|
|
|
46
9
|
Retains top RoPE frequency pairs, scores via trig series without materializing full K/V.
|
|
47
10
|
~10× compression with bounded approximation error. Calibrated for Qwen3.5 family.
|
|
48
11
|
|
|
12
|
+
- **KVTransfer** *(NEW in v2.4)* — Cross-model KV transfer. Convert one model's KV cache
|
|
13
|
+
into another model's so the receiving model **skips prefill entirely**. Per-(layer, head)
|
|
14
|
+
linear maps fitted by ridge; source layers chosen per target layer by held-out R².
|
|
15
|
+
Same model family only — the two models must share a tokenizer.
|
|
16
|
+
|
|
49
17
|
## Installation
|
|
50
18
|
|
|
51
19
|
```bash
|
|
52
20
|
pip install aither-kvcache # core library
|
|
53
21
|
pip install aither-kvcache[vllm] # + vLLM plugin (v0.15+)
|
|
54
22
|
pip install aither-kvcache[triton] # + fused GPU kernels
|
|
23
|
+
pip install aither-kvcache[transfer] # + cross-model KV transfer
|
|
55
24
|
pip install aither-kvcache[all] # everything
|
|
56
25
|
```
|
|
57
26
|
|
|
@@ -126,6 +95,95 @@ instead of 256 — a ~10× reduction with bounded approximation error.
|
|
|
126
95
|
| F=16, int4 | 4 | 34 | 7.53× |
|
|
127
96
|
| F=8, int4 | 4 | 18 | **14.2×** |
|
|
128
97
|
|
|
98
|
+
## Quick Start — KVTransfer (v2.4+)
|
|
99
|
+
|
|
100
|
+
Prefill is the tax you pay before a model says anything, and a KV cache only works on the
|
|
101
|
+
model that produced it. Route a conversation to a different model and the accumulated
|
|
102
|
+
cache becomes dead weight. KVTransfer fits a map that converts one model's cache into
|
|
103
|
+
another's, so the receiving model can skip prefill.
|
|
104
|
+
|
|
105
|
+
```bash
|
|
106
|
+
# 1. capture aligned KV from both models over one corpus
|
|
107
|
+
python -m aither_kvcache.kvtransfer.capture \
|
|
108
|
+
--source Qwen/Qwen3-0.6B --target Qwen/Qwen3-4B \
|
|
109
|
+
--out ./cap --corpus ./my-corpus --seq-len 256 --train-seqs 64 --val-seqs 24
|
|
110
|
+
|
|
111
|
+
# 2. fit the mapper pack
|
|
112
|
+
python -m aither_kvcache.kvtransfer.fit --capture ./cap --out ./pack --top-k 8
|
|
113
|
+
|
|
114
|
+
# 3. measure what the RECEIVING model does with a translated cache, and record it
|
|
115
|
+
python -m aither_kvcache.kvtransfer.transfer \
|
|
116
|
+
--pack ./pack --source Qwen/Qwen3-0.6B --target Qwen/Qwen3-4B \
|
|
117
|
+
--corpus ./held-out --sequences 24 --cuts 32 64 96 128 160 192 224 --record
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
from aither_kvcache.kvtransfer import load_pack
|
|
122
|
+
from aither_kvcache.kvtransfer.transfer import translate_cache
|
|
123
|
+
|
|
124
|
+
pack = load_pack("./pack", live_source=src_geom, live_target=tgt_geom)
|
|
125
|
+
layers = translate_cache(src_k, src_v, pack, positions) # -> the target's KV cache
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
### Step 3 is not optional
|
|
129
|
+
|
|
130
|
+
`load_pack` **refuses** any pack that carries no downstream acceptance evidence. That is
|
|
131
|
+
the core design decision here, and it is not defensive programming — it is the only thing
|
|
132
|
+
standing between an interesting R² and a broken deployment.
|
|
133
|
+
|
|
134
|
+
A converted cache does not fail loudly. An approximate one makes the receiving model
|
|
135
|
+
produce fluent, on-topic, **confidently wrong** text: no exception, no error status, no
|
|
136
|
+
unhealthy process, nothing in a log. Every cheap signal is green.
|
|
137
|
+
|
|
138
|
+
So a pack must carry a measurement, and the measurement has **four arms**:
|
|
139
|
+
|
|
140
|
+
| arm | what it is |
|
|
141
|
+
|---|---|
|
|
142
|
+
| `reference` | the target prefills the context itself — the ceiling |
|
|
143
|
+
| `translated` | the target is handed the converted cache — the candidate |
|
|
144
|
+
| `control` | the map is run on a **different document** — catches a mapper ignoring its input |
|
|
145
|
+
| `nocontext` | no cache at all — the floor |
|
|
146
|
+
|
|
147
|
+
The control arm is the one that matters. A mapper that has quietly learned the target's
|
|
148
|
+
*average* key/value statistics scores respectably against `reference` alone. Without a
|
|
149
|
+
floor, "68% agreement" is unreadable — it could be excellent, or exactly what a dead
|
|
150
|
+
mapper gets.
|
|
151
|
+
|
|
152
|
+
A minimum sample size is enforced too. Top-1 agreement is a proportion, so a mean over a
|
|
153
|
+
few dozen positions has a standard error near ten percentage points — and once recorded it
|
|
154
|
+
is indistinguishable from a mean over ten thousand.
|
|
155
|
+
|
|
156
|
+
### Preconditions that are definitional, not quality knobs
|
|
157
|
+
|
|
158
|
+
- **The two models must tokenize identically.** The map sends position *i* to position
|
|
159
|
+
*i*; different tokenizers mean row *i* is not the same token and the fit is regressing
|
|
160
|
+
misaligned data. Nothing downstream can see this — it reads as "this pair transfers
|
|
161
|
+
poorly" rather than "these rows do not correspond". Enforced by vocabulary digest, and
|
|
162
|
+
by comparing token ids per document at capture time.
|
|
163
|
+
- **The RoPE schedule must be exactly reproducible.** Keys carry a position rotation, so
|
|
164
|
+
the map is fitted in position-free space: stripped with the *source* schedule, re-applied
|
|
165
|
+
with the *target's* (not a no-op — the two models often disagree on `rope_theta`).
|
|
166
|
+
`default`, `linear` and `yarn` are implemented; anything else is refused rather than
|
|
167
|
+
approximated.
|
|
168
|
+
|
|
169
|
+
`check_pair()` separates definitional refusals from merely-unproven regimes — mismatched
|
|
170
|
+
KV-head counts, differing head dimension, cross-family pairs. The latter are recorded as
|
|
171
|
+
flags and decided by measurement, because a gate keyed on "heads must match" would refuse
|
|
172
|
+
the very experiment that tests whether heads need to match.
|
|
173
|
+
|
|
174
|
+
### Latent attention (MLA)
|
|
175
|
+
|
|
176
|
+
Models such as DeepSeek V2/V3/V4 cache a compressed latent plus a decoupled RoPE key
|
|
177
|
+
rather than per-head K and V. That is a different cache **layout**, not an obstacle — the
|
|
178
|
+
latent is dense, and often smaller than a conventional cache for the same context.
|
|
179
|
+
`KVGeometry.roles()` returns `("c", "kr")` for those models, and only `kr` carries a
|
|
180
|
+
rotation because the latent is position-free by construction.
|
|
181
|
+
|
|
182
|
+
The latent is captured at the compression projection rather than from `past_key_values`:
|
|
183
|
+
HuggingFace materialises per-head K/V before caching, so reading the cache would yield
|
|
184
|
+
tensors a real serving engine never stores, and a mapper fitted to them could not drive a
|
|
185
|
+
production deployment while looking entirely correct end to end.
|
|
186
|
+
|
|
129
187
|
## vLLM Integration
|
|
130
188
|
|
|
131
189
|
### Native kv-cache-dtype (recommended)
|
|
@@ -40,7 +40,7 @@ Usage:
|
|
|
40
40
|
candidates = advisor.get_eviction_candidates(n=16)
|
|
41
41
|
"""
|
|
42
42
|
|
|
43
|
-
from
|
|
43
|
+
from turboquant import TurboQuant, TurboQuantConfig
|
|
44
44
|
from .hybrid_quantizer import HybridTurboQuant, HybridLayout, GroupLayout
|
|
45
45
|
from .kvcache_graph import KVCacheGraph, KVBlockNode, KVEdge, EdgeType, get_kvcache_graph
|
|
46
46
|
from .eviction_advisor import GraphEvictionAdvisor, reorder_by_ranking
|
|
@@ -56,4 +56,4 @@ __all__ = [
|
|
|
56
56
|
# TriAttention — lazy-loaded via subpackage
|
|
57
57
|
# from aither_kvcache.triattention import TriAttention, TriAttentionConfig
|
|
58
58
|
]
|
|
59
|
-
__version__ = "2.0
|
|
59
|
+
__version__ = "2.1.0"
|
|
@@ -20,7 +20,6 @@ Integration:
|
|
|
20
20
|
"""
|
|
21
21
|
|
|
22
22
|
import math
|
|
23
|
-
import os
|
|
24
23
|
import torch
|
|
25
24
|
from typing import Optional, Tuple
|
|
26
25
|
|
|
@@ -107,7 +106,6 @@ class BlockRepresentativeCache:
|
|
|
107
106
|
If block_indices is None, finalizes all blocks with count > 0.
|
|
108
107
|
"""
|
|
109
108
|
if block_indices is not None:
|
|
110
|
-
mask = self._counts[block_indices] > 0
|
|
111
109
|
safe_count = self._counts[block_indices].clamp(min=1).unsqueeze(-1).float()
|
|
112
110
|
self.rep_even[block_indices] = self._acc_even[block_indices] / safe_count
|
|
113
111
|
self.rep_odd[block_indices] = self._acc_odd[block_indices] / safe_count
|
|
@@ -114,7 +114,6 @@ if HAS_TRITON:
|
|
|
114
114
|
# --- Step 3: compute L2 norm and normalize ---
|
|
115
115
|
norm_sq = tl.sum(x_vec * x_vec)
|
|
116
116
|
norm = tl.sqrt(norm_sq + 1e-20)
|
|
117
|
-
x_unit = x_vec / norm
|
|
118
117
|
|
|
119
118
|
# --- Step 4: tiled rotation y = rotation @ x_unit ---
|
|
120
119
|
# rotation is [HEAD_DIM, HEAD_DIM] stored row-major.
|
|
@@ -36,7 +36,6 @@ from __future__ import annotations
|
|
|
36
36
|
|
|
37
37
|
import functools
|
|
38
38
|
import math
|
|
39
|
-
import time
|
|
40
39
|
from dataclasses import dataclass
|
|
41
40
|
from typing import Dict, Optional, Tuple
|
|
42
41
|
|
|
@@ -405,7 +404,6 @@ def _pack_sign_bits(x: torch.Tensor) -> torch.Tensor:
|
|
|
405
404
|
Returns:
|
|
406
405
|
[..., ceil(dim / 8)] uint8 tensor.
|
|
407
406
|
"""
|
|
408
|
-
dim = x.shape[-1]
|
|
409
407
|
signs = (x >= 0).to(torch.uint8) # 1 for positive, 0 for negative
|
|
410
408
|
return pack_multibit(signs, bits=1)
|
|
411
409
|
|
|
@@ -679,8 +677,6 @@ class HybridTurboQuant:
|
|
|
679
677
|
Returns:
|
|
680
678
|
[N, packed_dim] uint8
|
|
681
679
|
"""
|
|
682
|
-
N = x.shape[0]
|
|
683
|
-
device = x.device
|
|
684
680
|
|
|
685
681
|
# Gather group dimensions
|
|
686
682
|
x_g0 = x[:, g0_idx] # [N, g0_dim]
|
|
@@ -733,9 +729,6 @@ class HybridTurboQuant:
|
|
|
733
729
|
Returns:
|
|
734
730
|
[N, layout.total_bytes] uint8 packed tensor.
|
|
735
731
|
"""
|
|
736
|
-
N = x_group.shape[0]
|
|
737
|
-
G = x_group.shape[1]
|
|
738
|
-
device = x_group.device
|
|
739
732
|
|
|
740
733
|
# 1. Compute L2 norm
|
|
741
734
|
vec_norm = x_group.norm(dim=-1) # [N]
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""Cross-model KV cache transfer.
|
|
2
|
+
|
|
3
|
+
Convert one model's KV cache into another model's, so the receiving model can
|
|
4
|
+
skip prefill entirely. Per-(target layer, head) affine maps fitted by ridge
|
|
5
|
+
regression, with source layers chosen per target layer by held-out R².
|
|
6
|
+
|
|
7
|
+
capture -> fit -> measure -> translate
|
|
8
|
+
|
|
9
|
+
Each stage is a CLI as well as an API::
|
|
10
|
+
|
|
11
|
+
python -m aither_kvcache.kvtransfer.capture --source A --target B --out CAP
|
|
12
|
+
python -m aither_kvcache.kvtransfer.fit --capture CAP --out PACK
|
|
13
|
+
python -m aither_kvcache.kvtransfer.transfer --pack PACK --source A --target B --record
|
|
14
|
+
|
|
15
|
+
**The measure step is not optional.** ``load_pack`` refuses any pack that does
|
|
16
|
+
not carry downstream acceptance evidence over a large enough sample, because a
|
|
17
|
+
converted cache fails silently: an approximate one does not raise, it makes the
|
|
18
|
+
receiving model produce fluent, on-topic, confidently wrong text with a green
|
|
19
|
+
healthcheck and nothing in a log.
|
|
20
|
+
|
|
21
|
+
Two preconditions are definitional rather than quality knobs. The two models
|
|
22
|
+
must **tokenize identically** — the map sends position *i* to position *i*, so
|
|
23
|
+
different tokenizers mean the fit is regressing misaligned rows — and the RoPE
|
|
24
|
+
schedule must be exactly reproducible. Both are enforced.
|
|
25
|
+
|
|
26
|
+
Only ``numpy`` is needed for the mapper, the pack format and the rotation math.
|
|
27
|
+
``capture`` and ``transfer``'s evaluation additionally need ``torch`` and
|
|
28
|
+
``transformers``; they are imported lazily so the rest of the package works
|
|
29
|
+
without them.
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
from __future__ import annotations
|
|
33
|
+
|
|
34
|
+
from .geometry import KVGeometry, PairVerdict, check_pair, source_roles_for
|
|
35
|
+
from .mapper import GramStats, LayerMap, apply_layer_map, fit_layer_map, r2_from_gram
|
|
36
|
+
from .pack import (
|
|
37
|
+
AcceptanceMetrics,
|
|
38
|
+
LoadedPack,
|
|
39
|
+
PackManifest,
|
|
40
|
+
PackRefusedError,
|
|
41
|
+
load_pack,
|
|
42
|
+
record_acceptance,
|
|
43
|
+
write_pack,
|
|
44
|
+
)
|
|
45
|
+
from .rope import RopeSpec, apply_rope, rotate_to_target, strip_rope
|
|
46
|
+
|
|
47
|
+
__all__ = [
|
|
48
|
+
"AcceptanceMetrics",
|
|
49
|
+
"GramStats",
|
|
50
|
+
"KVGeometry",
|
|
51
|
+
"LayerMap",
|
|
52
|
+
"LoadedPack",
|
|
53
|
+
"PackManifest",
|
|
54
|
+
"PackRefusedError",
|
|
55
|
+
"PairVerdict",
|
|
56
|
+
"RopeSpec",
|
|
57
|
+
"apply_layer_map",
|
|
58
|
+
"apply_rope",
|
|
59
|
+
"check_pair",
|
|
60
|
+
"fit_layer_map",
|
|
61
|
+
"load_pack",
|
|
62
|
+
"r2_from_gram",
|
|
63
|
+
"record_acceptance",
|
|
64
|
+
"rotate_to_target",
|
|
65
|
+
"source_roles_for",
|
|
66
|
+
"strip_rope",
|
|
67
|
+
"write_pack",
|
|
68
|
+
]
|
|
69
|
+
|
|
70
|
+
__version__ = "2.4.0"
|