flashcodec-kv 1.0.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.
- flashcodec_kv-1.0.0/LICENSE +21 -0
- flashcodec_kv-1.0.0/PKG-INFO +430 -0
- flashcodec_kv-1.0.0/README.md +403 -0
- flashcodec_kv-1.0.0/flashcodec/__init__.py +37 -0
- flashcodec_kv-1.0.0/flashcodec/attention.py +140 -0
- flashcodec_kv-1.0.0/flashcodec/cache.py +654 -0
- flashcodec_kv-1.0.0/flashcodec/codecs/__init__.py +13 -0
- flashcodec_kv-1.0.0/flashcodec/codecs/jpeg.py +159 -0
- flashcodec_kv-1.0.0/flashcodec/codecs/png.py +107 -0
- flashcodec_kv-1.0.0/flashcodec/patch.py +136 -0
- flashcodec_kv-1.0.0/flashcodec/patch_attention.py +151 -0
- flashcodec_kv-1.0.0/flashcodec/triton_attention.py +222 -0
- flashcodec_kv-1.0.0/flashcodec_kv.egg-info/PKG-INFO +430 -0
- flashcodec_kv-1.0.0/flashcodec_kv.egg-info/SOURCES.txt +22 -0
- flashcodec_kv-1.0.0/flashcodec_kv.egg-info/dependency_links.txt +1 -0
- flashcodec_kv-1.0.0/flashcodec_kv.egg-info/requires.txt +8 -0
- flashcodec_kv-1.0.0/flashcodec_kv.egg-info/top_level.txt +1 -0
- flashcodec_kv-1.0.0/pyproject.toml +34 -0
- flashcodec_kv-1.0.0/setup.cfg +4 -0
- flashcodec_kv-1.0.0/tests/test_attention.py +129 -0
- flashcodec_kv-1.0.0/tests/test_causal.py +249 -0
- flashcodec_kv-1.0.0/tests/test_hf_integration.py +250 -0
- flashcodec_kv-1.0.0/tests/test_jpeg.py +129 -0
- flashcodec_kv-1.0.0/tests/test_png.py +96 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 feifei
|
|
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,430 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flashcodec-kv
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Compress Transformer time-series KV cache tensors as high-dimensional multimedia image signals: JPEG-style lossy frequency domain (DCT) plus PNG-style lossless time-domain deltas
|
|
5
|
+
Author-email: feifei <aifeifei798@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/aifeifei798/flashcodec-kv
|
|
8
|
+
Project-URL: Repository, https://github.com/aifeifei798/flashcodec-kv
|
|
9
|
+
Project-URL: Issues, https://github.com/aifeifei798/flashcodec-kv/issues
|
|
10
|
+
Project-URL: Changelog, https://github.com/aifeifei798/flashcodec-kv/blob/main/CHANGELOG.md
|
|
11
|
+
Keywords: kv-cache,transformer,dct,compression,llm,inference
|
|
12
|
+
Classifier: Development Status :: 4 - Beta
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
|
+
Requires-Python: >=3.10
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
License-File: LICENSE
|
|
20
|
+
Requires-Dist: torch>=2.0
|
|
21
|
+
Provides-Extra: hf
|
|
22
|
+
Requires-Dist: transformers>=5.20; extra == "hf"
|
|
23
|
+
Provides-Extra: demo
|
|
24
|
+
Requires-Dist: gradio; extra == "demo"
|
|
25
|
+
Requires-Dist: matplotlib; extra == "demo"
|
|
26
|
+
Dynamic: license-file
|
|
27
|
+
|
|
28
|
+
# flashcodec
|
|
29
|
+
|
|
30
|
+
> Compressing Transformer KV Cache tensors as **high-dimensional multimedia image signals**.
|
|
31
|
+
|
|
32
|
+
[English](README.md) | [简体中文](README.zh-CN.md)
|
|
33
|
+
|
|
34
|
+
Conventional LLM research treats hidden-space tensors as isolated scalar
|
|
35
|
+
factors. flashcodec observes that the seq axis of a `(B, H, S, D)` KV tensor
|
|
36
|
+
behaves like the rows/columns of a high-dimensional image — adjacent token
|
|
37
|
+
representations are autocorrelated, energy concentrates in low frequencies, and
|
|
38
|
+
first-order differences shrink in variance (how strongly depends on the model,
|
|
39
|
+
layer and tensor: the measured real-KV statistics and the per-layer caveat are
|
|
40
|
+
in the Benchmarks section). Two classic
|
|
41
|
+
multimedia compression theories therefore transfer directly:
|
|
42
|
+
|
|
43
|
+
| Mode | Inspired by | Transform | Nature | Storage |
|
|
44
|
+
|---|---|---|---|---|
|
|
45
|
+
| `mode="jpeg"` | JPEG | Block-wise orthonormal 1D DCT-II, keep lowest K coeffs | **Lossy, frequency domain** | `K/r` (4×–8× less memory) |
|
|
46
|
+
| `mode="png"` | PNG | Per-block anchor (I-Token) + first-order deltas, Top-M pruning | **Lossless / near-lossless, time domain** | Dense, size-neutral; sparsity left to entropy coding |
|
|
47
|
+
|
|
48
|
+
`quant_bits=8` (jpeg only) stores the coefficients as int8 with a
|
|
49
|
+
per-(block, frequency) scale — a real byte-level saving on top of `K/r`:
|
|
50
|
+
**6.6% of raw KV at k=2 (15×), 46.9 dB PSNR**, versus 25% / 47.2 dB in fp32.
|
|
51
|
+
|
|
52
|
+
Math (block size `r`, operating axis `dim=-2`):
|
|
53
|
+
|
|
54
|
+
- **JPEG**: $X_{\text{blk}} \in \mathbb{R}^{r \times D}$, $C = M_{\text{enc}} X_{\text{blk}}$ with
|
|
55
|
+
$M_{\text{enc}} \in \mathbb{R}^{K \times r}$, $M[j,n] = c_j \cos\frac{\pi}{r}(n+\tfrac12)j$,
|
|
56
|
+
orthonormal rows ($K=r$ is per-block lossless); reconstruct $\hat X = M_{\text{enc}}^\top C$.
|
|
57
|
+
- **PNG**: $\Delta_i = x_i - x_{i-1}$ ($i \ge 1$), $x_0$ kept at full precision as the
|
|
58
|
+
anchor; rebuild $\hat x_i = x_0 + \sum_{j=1}^{i} \Delta_j$ (`cumsum`, purely
|
|
59
|
+
linear and autograd-differentiable). Autocorrelation implies
|
|
60
|
+
$\mathrm{Var}(\Delta)/\mathrm{Var}(X) \approx 0.03$ (**30× variance collapse**,
|
|
61
|
+
measured); after Top-M pruning the residual is nearly zero and sparse.
|
|
62
|
+
|
|
63
|
+
```
|
|
64
|
+
┌────────────────────────────────────────────┐
|
|
65
|
+
(B,H,S,D) ──append──▶ raw_buffer (tail < r, full fidelity) │
|
|
66
|
+
│ │ r tokens accumulated, batched │
|
|
67
|
+
│ ▼ archive (single GEMM, in-place │
|
|
68
|
+
│ slice write into) │
|
|
69
|
+
│ compressed_storage (capacity-doubled) │
|
|
70
|
+
│ jpeg: coeffs (B,H,capacity,K,D) │
|
|
71
|
+
│ png: deltas (B,H,capacity,r-1,D) │
|
|
72
|
+
│ + anchors (B,H,capacity,D) │
|
|
73
|
+
└───────────────┬────────────────────────────┘
|
|
74
|
+
│ read(): one decode of active
|
|
75
|
+
│ region + raw tail
|
|
76
|
+
▼
|
|
77
|
+
reconstructed KV view (B,H,S,D)
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
Archives live in a **single geometrically-grown buffer** (std::vector-style:
|
|
81
|
+
start at 16 blocks, double on overflow). Each full block is one in-place
|
|
82
|
+
slice write — no re-allocation churn — and `read()` decodes **only the
|
|
83
|
+
newly archived blocks** (one batched GEMM / cumsum per new block, memoized
|
|
84
|
+
in an append-only reconstruction buffer) instead of re-decoding the whole
|
|
85
|
+
history on every decode step. (Fragmented per-block archive lists made
|
|
86
|
+
per-token decode ~170× slower at S=4096 due to kernel-launch overhead.)
|
|
87
|
+
|
|
88
|
+
## Causal-safety design
|
|
89
|
+
|
|
90
|
+
Only **complete blocks** are archived; the current incomplete block always
|
|
91
|
+
stays in `raw_buffer` at original precision. During incremental decoding, no
|
|
92
|
+
future perturbation can ever rewrite archived tensors (verified bitwise in
|
|
93
|
+
tests). On the lossless tier (`mode="png", keep_ratio=1.0`), step-by-step
|
|
94
|
+
decoded attention matches full-sequence attention within ≤1e-4.
|
|
95
|
+
|
|
96
|
+
## Installation
|
|
97
|
+
|
|
98
|
+
```bash
|
|
99
|
+
pip install -e ".[hf]" # from source: torch>=2.0 (+ transformers>=5.20 for the HF backend)
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
Distribution name on PyPI: `flashcodec-kv` — the import name stays
|
|
103
|
+
`flashcodec`. (The bare `flashcodec` name there belongs to an unrelated
|
|
104
|
+
placeholder package, so don't `pip install flashcodec`.)
|
|
105
|
+
|
|
106
|
+
## Quick start
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
import torch
|
|
110
|
+
from flashcodec import FlashCodecCache, apply_flash_codec
|
|
111
|
+
|
|
112
|
+
# 1) Native: a first-class transformers 5.x Cache backend (any CausalLM)
|
|
113
|
+
cache = FlashCodecCache(mode="jpeg", block_size=8, keep_freqs=2) # 4×
|
|
114
|
+
# int8 coefficients: same shape, 15× smaller archive (quant_bits=8)
|
|
115
|
+
cache = FlashCodecCache(mode="jpeg", block_size=8, keep_freqs=2,
|
|
116
|
+
quant_bits=8) # 15×
|
|
117
|
+
out = model.generate(input_ids, max_new_tokens=64,
|
|
118
|
+
past_key_values=cache, use_cache=True)
|
|
119
|
+
print(cache.get_seq_length(), cache.storage_bytes())
|
|
120
|
+
|
|
121
|
+
# 2) One-line hot patch: every generate() defaults to the codec cache
|
|
122
|
+
apply_flash_codec(model, mode="png", block_size=8, keep_ratio=1.0)
|
|
123
|
+
out = model.generate(input_ids, max_new_tokens=64)
|
|
124
|
+
# PNG lossless tier reproduces the unpatched model's tokens EXACTLY
|
|
125
|
+
from flashcodec import remove_flash_codec
|
|
126
|
+
remove_flash_codec(model) # restore default behavior
|
|
127
|
+
|
|
128
|
+
# Materialization-free decode (coefficient-domain attention, no KV view):
|
|
129
|
+
apply_flash_codec(model, mode="jpeg", block_size=8, keep_freqs=2,
|
|
130
|
+
quant_bits=8, materialization_free=True)
|
|
131
|
+
|
|
132
|
+
# 3) Standalone (no transformers needed): direct cache manager
|
|
133
|
+
from flashcodec import CodecKVCache
|
|
134
|
+
cache = CodecKVCache(mode="jpeg", block_size=8, keep_freqs=2)
|
|
135
|
+
k, v = cache.forward(k_seq, v_seq) # full-sequence fill
|
|
136
|
+
k, v = cache.decode_step(new_k, new_v) # incremental decode
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
`FlashCodecCache` subclasses `transformers.cache_utils.Cache` (5.x layer
|
|
140
|
+
API): write/read replacement is converged at the `cache.update(k, v,
|
|
141
|
+
layer_idx)` interface — zero model surgery, works with any CausalLM
|
|
142
|
+
(Llama/Qwen/Mistral-style GQA included).
|
|
143
|
+
|
|
144
|
+
## Transformers integration (real-world tested on Qwen3.5-0.8B)
|
|
145
|
+
|
|
146
|
+
Hybrid architectures (e.g. Qwen3.5: 18 linear-attention layers + 6
|
|
147
|
+
full-attention layers) are handled automatically —
|
|
148
|
+
`FlashCodecCache.from_config(model.config)` attaches codec layers **only to
|
|
149
|
+
`full_attention` layers** and keeps native dynamic layers for everything
|
|
150
|
+
else (recurrent/conv states untouched). `apply_flash_codec` uses `from_config`
|
|
151
|
+
whenever a model config is available.
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
155
|
+
from flashcodec import apply_flash_codec
|
|
156
|
+
|
|
157
|
+
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-0.8B",
|
|
158
|
+
torch_dtype="bfloat16").cuda()
|
|
159
|
+
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B")
|
|
160
|
+
ids = tok("A quiet lake at dawn", return_tensors="pt").input_ids.cuda()
|
|
161
|
+
|
|
162
|
+
# One line — every generate() now runs on the codec cache
|
|
163
|
+
apply_flash_codec(model, mode="jpeg", block_size=8, keep_freqs=2)
|
|
164
|
+
out = model.generate(ids, max_new_tokens=48, do_sample=False)
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
Measured on Qwen3.5-0.8B (bf16, CUDA, 48 greedy tokens, `benchmarks/run_qwen_demo.py`):
|
|
168
|
+
|
|
169
|
+
| Tier | Storage vs raw KV | Output | Verdict |
|
|
170
|
+
|---|---|---|---|
|
|
171
|
+
| `png keep=1.0` | ~1.8× (fp32 deltas for exactness) | **bit-identical tokens to the unpatched model** | 100% fidelity baseline |
|
|
172
|
+
| `jpeg k=8 int8` | **55.5%** (~1.8×) | fluent, **100% token match** | safe lossy tier |
|
|
173
|
+
| `jpeg k=2` | **32.3%** (~3.1×) | fluent, 38.1% token match | lossy tier |
|
|
174
|
+
| `jpeg k=4` | 54.8% (~1.8×) | fluent, 49.2% token match | lossy tier |
|
|
175
|
+
| `jpeg k=2 int8` | **21.1%** (~4.7×) | fluent, 38.1% token match | lossy + int8 |
|
|
176
|
+
|
|
177
|
+
Hybrid-layer audit: `codec_layers=6/24` (full attention only), zero warnings.
|
|
178
|
+
Reproduce with `python benchmarks/run_qwen_demo.py` (offline; honors
|
|
179
|
+
`FLASHCODEC_MODEL` / `FLASHCODEC_NEW_TOKENS`).
|
|
180
|
+
|
|
181
|
+
dtype contracts (see `AGENTS.md`): PNG deltas are always fp32 (bf16
|
|
182
|
+
differences lose the bitwise-exactness the lossless tier promises); JPEG
|
|
183
|
+
coefficients and `read()` views stay in the stream's input dtype (SDPA
|
|
184
|
+
requires q/k/v dtype parity).
|
|
185
|
+
|
|
186
|
+
## Materialization-free decode attention
|
|
187
|
+
|
|
188
|
+
Every codec here is a **linear map along the seq axis**: the r tokens of a block
|
|
189
|
+
are `K_block = A @ S`, where `S` are the stored rows (DCT coefficients, or
|
|
190
|
+
anchor + deltas) and `A` is the synthesis matrix. So a decode query can attend
|
|
191
|
+
on `S` directly:
|
|
192
|
+
|
|
193
|
+
```
|
|
194
|
+
q . K[p] = (q . S^T) @ A[p] out = (prob @ A) @ S
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
`codec_decode_attention(q, key_stream, value_stream)` implements this for one
|
|
198
|
+
new token (jpeg k<r, jpeg int8 and png; GQA grouping included). The `(S, D)`
|
|
199
|
+
reconstruction never exists, and the result matches the materialized path to
|
|
200
|
+
≤1e-6.
|
|
201
|
+
|
|
202
|
+
Measured (`benchmarks/bench_attention.py`, B=1, H_kv=8, H_q=16, D=64, fp32):
|
|
203
|
+
|
|
204
|
+
| seq | mode | materialized | coefficient-domain |
|
|
205
|
+
|---|---|---|---|
|
|
206
|
+
| 16384 | jpeg k=2 int8 | 0.20ms / 130MB | 0.30ms / **32MB** |
|
|
207
|
+
| 65536 | jpeg k=2 int8 | 0.86ms / 520MB | **0.46ms / 128MB** |
|
|
208
|
+
| 65536 | png keep=1.0 | 0.87ms / 520MB | 1.03ms / 272MB |
|
|
209
|
+
|
|
210
|
+
Scope: **decode only** (prefill still uses the view — the coefficient form would
|
|
211
|
+
explode the score tensor). png stores r rows per block, so it is numerically
|
|
212
|
+
supported but gains nothing; the win is jpeg `k < r`.
|
|
213
|
+
|
|
214
|
+
A **Triton split-K kernel** (`flashcodec/triton_attention.py`, flash-decoding
|
|
215
|
+
style: block chunks write partial max/sum/acc, a combine kernel merges them)
|
|
216
|
+
removes the intermediate tensors and the launch overhead. One layer, B=1,
|
|
217
|
+
H_kv=8, H_q=16, D=128, `jpeg k=2 int8` (`benchmarks/bench_attention.py --cuda`):
|
|
218
|
+
|
|
219
|
+
| seq | materialized | coefficient (torch) | coefficient (Triton) |
|
|
220
|
+
|---|---|---|---|
|
|
221
|
+
| 16384 | 0.20ms / 130MB | 0.29ms / 24MB | **0.09ms / 0MB** |
|
|
222
|
+
| 65536 | 0.86ms / 520MB | 0.41ms / 96MB | **0.16ms / 0MB** |
|
|
223
|
+
|
|
224
|
+
`apply_flash_codec(model, ..., materialization_free=True)` wires it into
|
|
225
|
+
transformers: the forward of standard Llama/Qwen-style attention modules is
|
|
226
|
+
wrapped (projection + RoPE unchanged, the attention output replaced), guarded to
|
|
227
|
+
single-token / batch-1 / eval / non-sliding layers, otherwise it falls back to
|
|
228
|
+
the original path. The prefill reconstruction is released on the first decode
|
|
229
|
+
step. `remove_flash_codec` restores everything.
|
|
230
|
+
|
|
231
|
+
End-to-end (`benchmarks/bench_e2e.py`, Qwen3-0.6B, 4096-token prompt + 32 new):
|
|
232
|
+
|
|
233
|
+
| tier | tok/s | peak | steady |
|
|
234
|
+
|---|---|---|---|
|
|
235
|
+
| baseline (raw KV) | 52.1 | 581MB | 0MB |
|
|
236
|
+
| codec k=2 int8, materialized | 44.2 | 1545MB | 1535MB |
|
|
237
|
+
| codec k=2 int8, materialization-free (Triton) | 40.3 | **642MB** | **122MB** |
|
|
238
|
+
|
|
239
|
+
**12.6× less steady memory** at ~9% slower decode than the materialized path
|
|
240
|
+
(per decode step at 4096 tokens: 22.8ms vs 21.7ms materialized, 17.8ms raw).
|
|
241
|
+
The remaining gap is the per-layer split/combine launch pair; a persistent
|
|
242
|
+
kernel would close it.
|
|
243
|
+
|
|
244
|
+
## Demo (side-by-side chat on Qwen3.5-0.8B)
|
|
245
|
+
|
|
246
|
+

|
|
247
|
+
|
|
248
|
+
Same prompt, two panels: native KV cache on the left, flashcodec on the right
|
|
249
|
+
(five tiers: png lossless / jpeg k=8 int8 / k=4 / k=2 / k=2 int8), with a
|
|
250
|
+
materialization-free toggle (coefficient-domain + Triton) and per-turn stats:
|
|
251
|
+
KV footprint (archive vs materialized view), peak/steady memory, per-token
|
|
252
|
+
latency and token match.
|
|
253
|
+
|
|
254
|
+
The **Full comparison** tab runs the baseline plus every tier on one prompt and
|
|
255
|
+
reports, per tier: KV footprint vs raw, archive/view split, decode tok/s,
|
|
256
|
+
peak/steady memory, teacher-forced NLL of the continuation (a perplexity-style
|
|
257
|
+
quality number) and token match — plus storage / speed / NLL bars.
|
|
258
|
+
|
|
259
|
+
```bash
|
|
260
|
+
pip install -e ".[demo]" # gradio + matplotlib
|
|
261
|
+
python demo/qwen_gradio_chat.py # Gradio side-by-side + savings chart
|
|
262
|
+
```
|
|
263
|
+
|
|
264
|
+
## Benchmarks (measured; B=2 H=8 S=2048 D=64, random-walk autocorrelated signal, fp32)
|
|
265
|
+
|
|
266
|
+
### Reconstruction quality (`benchmarks/bench_reconstruct.py`)
|
|
267
|
+
|
|
268
|
+
| Config | PSNR (dB) | cosine | Storage |
|
|
269
|
+
|---|---|---|---|
|
|
270
|
+
| jpeg r=8 k=1 (8×) | 42.95 | 0.99943 | 12.5% |
|
|
271
|
+
| jpeg r=8 k=2 (4×) | 47.22 | 0.99985 | 25.0% |
|
|
272
|
+
| jpeg r=8 k=4 (2×) | 51.66 | 1.00001 | 50.0% |
|
|
273
|
+
| jpeg r=8 k=2 int8 | 46.92 | 0.99983 | 6.6% |
|
|
274
|
+
| jpeg r=8 k=4 int8 | 50.89 | 0.99999 | 13.3% |
|
|
275
|
+
| png r=8 keep=0.5 | 50.11 | 0.99997 | 100% (sparsity → entropy coder) |
|
|
276
|
+
| png r=8 keep=1.0 | ~196 (fp rounding) | 1.00000 | 100% |
|
|
277
|
+
|
|
278
|
+
### Real KV statistics (real models, 512 tokens)
|
|
279
|
+
|
|
280
|
+
`bench_reconstruct` uses a synthetic random walk; real caches are messier.
|
|
281
|
+
Measured on real post-RoPE KV (`Var(Δ)/Var(X)` along the seq axis, and the
|
|
282
|
+
fraction of block energy kept by the lowest k=2 DCT coefficients of a block of
|
|
283
|
+
8):
|
|
284
|
+
|
|
285
|
+
| model / layer | tensor | Var(Δ)/Var(X) | energy in k=2 |
|
|
286
|
+
|---|---|---|---|
|
|
287
|
+
| Qwen3-0.6B L0 | K | 0.16 | 0.94 |
|
|
288
|
+
| Qwen3-0.6B L0 | V | 1.87 | 0.31 |
|
|
289
|
+
| Qwen3-0.6B L27 | V | 1.02 | 0.59 |
|
|
290
|
+
| Qwen3.5-0.8B L3 | K | 0.72 | 0.71 |
|
|
291
|
+
| Qwen3.5-0.8B L23 | V | 0.09 | 0.97 |
|
|
292
|
+
|
|
293
|
+
Values are frequently **anti**-correlated (ratio > 1), so "first differences
|
|
294
|
+
collapse the variance" is a property of smooth synthetic signals, not of every
|
|
295
|
+
real layer. Practical consequence: coefficient truncation (`k < r`) is a
|
|
296
|
+
per-layer decision — on a dense model (28 compressed layers) the error compounds
|
|
297
|
+
and quality collapses; use `k = r` + `quant_bits=8` there.
|
|
298
|
+
|
|
299
|
+
### Perplexity (`benchmarks/bench_ppl.py`; 16 × 512 tokens, bf16, local corpus)
|
|
300
|
+
|
|
301
|
+
| tier | Qwen3.5-0.8B (6/24 layers) | Qwen3-0.6B (28/28 layers) | storage |
|
|
302
|
+
|---|---|---|---|
|
|
303
|
+
| baseline | 22.4803 | 16.7931 | 100% |
|
|
304
|
+
| png keep=1.0 | 22.4803 (+0.000) | 16.7931 (+0.000) | 188% |
|
|
305
|
+
| jpeg k=8 (lossless basis) | 22.4877 (+0.007) | 16.8058 (+0.013) | 100% |
|
|
306
|
+
| **jpeg k=8 int8** | **22.4991 (+0.019)** | **17.1348 (+0.34)** | **~52%** |
|
|
307
|
+
| jpeg k=4 | 22.3010 (−0.18) | 470.85 (+454) | 50% |
|
|
308
|
+
| jpeg k=2 | 24.0566 (+1.58) | 6485.5 (+6469) | 25% |
|
|
309
|
+
| jpeg k=2 int8 | 24.0590 (+1.58) | 8118.5 (+8102) | 13% |
|
|
310
|
+
|
|
311
|
+
The hybrid result is good because only 6 of 24 layers carry a codec (the 18
|
|
312
|
+
linear-attention layers keep their native state); the dense model shows why
|
|
313
|
+
per-layer truncation must not be applied blindly.
|
|
314
|
+
|
|
315
|
+
### Per-layer budgets (`benchmarks/bench_adaptive.py`, Qwen3-0.6B, 6 × 512 tokens)
|
|
316
|
+
|
|
317
|
+
Calibration: compress one layer at a time to k=4, measure the perplexity damage,
|
|
318
|
+
then give the cheapest half of the layers k=4 and leave the rest at k=r (75%
|
|
319
|
+
storage):
|
|
320
|
+
|
|
321
|
+
| schedule | ppl | Δppl | storage |
|
|
322
|
+
|---|---|---|---|
|
|
323
|
+
| baseline (k=r) | 18.9786 | — | 100% |
|
|
324
|
+
| uniform k=6 | 48.0531 | +29.07 | 75% |
|
|
325
|
+
| adaptive (14 layers k=4) | 40.8309 | +21.85 | 75% |
|
|
326
|
+
| uniform k=4 | 710.22 | +691.24 | 50% |
|
|
327
|
+
| uniform k=r int8 | 19.3627 | +0.38 | 50% |
|
|
328
|
+
|
|
329
|
+
Calibration beats a uniform budget at the same storage (−15%), but the damages
|
|
330
|
+
are **super-additive** (the cheapest 14 layers add up to +22 ppl although no
|
|
331
|
+
single one costs more than +0.5), and int8 at 50% is far better. Use
|
|
332
|
+
`per_layer={idx: {"keep_freqs": ...}}` when a budget forces structural
|
|
333
|
+
truncation; otherwise prefer `quant_bits=8`.
|
|
334
|
+
|
|
335
|
+
### Cache footprint & latency (`benchmarks/bench_vram.py`, CPU)
|
|
336
|
+
|
|
337
|
+
| seq | Mode | storage | ratio | fill | read |
|
|
338
|
+
|---|---|---|---|---|---|
|
|
339
|
+
| 4096 | jpeg k=2 | 4.0MB | 0.25 | 6.7ms | 0.0ms |
|
|
340
|
+
| 16384 | jpeg k=2 | 16.0MB | 0.25 | 34.3ms | 0.5ms |
|
|
341
|
+
| 65536 | jpeg k=2 | 64.0MB | 0.25 | 141.9ms | 3.1ms |
|
|
342
|
+
| 4096 | jpeg k=2 int8 | 1.1MB | 0.07 | 9.2ms | 0.0ms |
|
|
343
|
+
| 16384 | jpeg k=2 int8 | 4.3MB | 0.07 | 46.0ms | 0.3ms |
|
|
344
|
+
| 65536 | jpeg k=2 int8 | 17.0MB | 0.07 | 197.9ms | 3.1ms |
|
|
345
|
+
| 65536 | png keep=1.0 | 256.0MB | 1.00 | 243.3ms | 3.3ms |
|
|
346
|
+
|
|
347
|
+
CPU timings vary ~±20% between runs (min of 3); the storage/ratio columns are
|
|
348
|
+
exact. `read` = materialize the view and touch it (`.sum()`, so the number is the
|
|
349
|
+
consumer cost, not just kernel launch). Decoding is memoized per archived block,
|
|
350
|
+
so `read` costs one decode per **new** block instead of re-decoding the whole
|
|
351
|
+
history on every decode step (previously 38.7ms / 71.9ms in the last two rows).
|
|
352
|
+
|
|
353
|
+
PNG does not shrink dense storage (the fate of a lossless tier) — its value is
|
|
354
|
+
the 30× variance collapse plus near-zero sparse residuals after Top-M, i.e.
|
|
355
|
+
ideal input for an entropy coder (none is bundled).
|
|
356
|
+
|
|
357
|
+
## Repository layout
|
|
358
|
+
|
|
359
|
+
```
|
|
360
|
+
flashcodec/
|
|
361
|
+
├── flashcodec/
|
|
362
|
+
│ ├── __init__.py # apply_flash_codec / CodecKVCache / JPEGCodec / PNGCodec
|
|
363
|
+
│ ├── codecs/jpeg.py # DCT-II basis, project/reconstruct (pure GEMM, no loops)
|
|
364
|
+
│ ├── codecs/png.py # anchor + first-order delta, cumsum, Top-M filter
|
|
365
|
+
│ ├── cache.py # CodecKVCache: raw tail + capacity-doubled archives
|
|
366
|
+
│ ├── attention.py # materialization-free decode attention (coefficient domain)
|
|
367
|
+
│ └── patch.py # one-line hot patch: cache-level injection (HF Cache duck-typing)
|
|
368
|
+
├── tests/ # jpeg / png / causal / attention suites (standalone scripts)
|
|
369
|
+
├── benchmarks/ # storage + PSNR/cosine + perplexity benchmarks
|
|
370
|
+
├── demo/ # Gradio side-by-side chat
|
|
371
|
+
├── images/ # README screenshots
|
|
372
|
+
├── .github/workflows/ # CI: CPU tests + transformers integration
|
|
373
|
+
├── CHANGELOG.md
|
|
374
|
+
├── pyproject.toml
|
|
375
|
+
└── README.md
|
|
376
|
+
```
|
|
377
|
+
|
|
378
|
+
## Tests
|
|
379
|
+
|
|
380
|
+
```bash
|
|
381
|
+
python tests/test_jpeg.py # orthonormality, k=r lossless, k=2 cosine>0.99, K monotone
|
|
382
|
+
python tests/test_png.py # lossless maxdiff<1e-6, variance ratio printed, differentiable
|
|
383
|
+
python tests/test_causal.py # zero future leakage, incremental vs full ≤1e-4,
|
|
384
|
+
# memoized read == full decode, exact + repeatable crop,
|
|
385
|
+
# beam-search batch gather
|
|
386
|
+
python tests/test_hf_integration.py # native transformers 5.x generate: explicit cache / hot patch /
|
|
387
|
+
# PNG lossless exact tokens / 70-token multi-block stability /
|
|
388
|
+
# beam search + crop exactness against the unpatched model /
|
|
389
|
+
# per-layer schedule
|
|
390
|
+
python tests/test_attention.py # coefficient-domain decode attention matches the materialized
|
|
391
|
+
# path (jpeg / int8 / png, GQA); never builds the KV view
|
|
392
|
+
```
|
|
393
|
+
|
|
394
|
+
## Stability (1.x)
|
|
395
|
+
|
|
396
|
+
The public API is frozen for 1.x: `apply_flash_codec` / `remove_flash_codec`,
|
|
397
|
+
`apply_codec_attention` / `remove_codec_attention`, `FlashCodecCache`,
|
|
398
|
+
`FlashCodecLayer`, `CodecKVCache`, `JPEGCodec`, `PNGCodec`,
|
|
399
|
+
`codec_decode_attention`, `make_codec`, plus the `quant_bits` / `per_layer`
|
|
400
|
+
options. New features ship as additive arguments; existing signatures keep
|
|
401
|
+
their defaults. The caveats below still apply — this is a research library with
|
|
402
|
+
a stable interface, not a claim of production readiness.
|
|
403
|
+
|
|
404
|
+
## Boundaries (honest list)
|
|
405
|
+
|
|
406
|
+
- Research prototype: `read()` materializes the full reconstructed view for
|
|
407
|
+
attention; memory gains are at the **cache storage layer** (table above);
|
|
408
|
+
end-to-end benefits require combining with materialization-free kernels
|
|
409
|
+
(flash-attention style). That view is memoized and reused across decode
|
|
410
|
+
steps, so steady-state memory = archive + one raw-KV-sized view.
|
|
411
|
+
- Full-sequence `forward(seq)` prefill compresses complete blocks directly, so
|
|
412
|
+
mid-block attention sees the block-level reconstruction (within-block
|
|
413
|
+
approximation); the incremental `decode_step` path has no such issue.
|
|
414
|
+
- Compressed tensors are dense layouts; PNG Top-M sparsity is not yet wired to
|
|
415
|
+
an entropy coder. (Packing sparse png alone was measured to be a wash: the
|
|
416
|
+
fp32 delta values dominate, and on bf16 input the packed form can be *larger*
|
|
417
|
+
than raw — delta quantization has to come first.)
|
|
418
|
+
- Coefficient truncation (`k < r`) is a per-layer decision: on a dense model
|
|
419
|
+
every layer is compressed and the error compounds (measured: Qwen3-0.6B
|
|
420
|
+
perplexity 16.8 → 6486 at k=2). Use `k = r` + `quant_bits=8` there (~52%
|
|
421
|
+
storage, +2% ppl) and reserve aggressive ratios for hybrid models where only
|
|
422
|
+
the full-attention layers carry a codec.
|
|
423
|
+
- Same lineage as the v4.2 prototypes inside
|
|
424
|
+
[Mini-Architecture](https://github.com/aifeifei798/Mini-Architecture)
|
|
425
|
+
(there: per-position view semantics for cross-layer KV reuse; here:
|
|
426
|
+
standard cache storage semantics).
|
|
427
|
+
|
|
428
|
+
## License
|
|
429
|
+
|
|
430
|
+
[MIT](LICENSE)
|