Metadata-Version: 2.4
Name: flashcodec-kv
Version: 1.0.0
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
Author-email: feifei <aifeifei798@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/aifeifei798/flashcodec-kv
Project-URL: Repository, https://github.com/aifeifei798/flashcodec-kv
Project-URL: Issues, https://github.com/aifeifei798/flashcodec-kv/issues
Project-URL: Changelog, https://github.com/aifeifei798/flashcodec-kv/blob/main/CHANGELOG.md
Keywords: kv-cache,transformer,dct,compression,llm,inference
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Provides-Extra: hf
Requires-Dist: transformers>=5.20; extra == "hf"
Provides-Extra: demo
Requires-Dist: gradio; extra == "demo"
Requires-Dist: matplotlib; extra == "demo"
Dynamic: license-file

# flashcodec

> Compressing Transformer KV Cache tensors as **high-dimensional multimedia image signals**.

[English](README.md) | [简体中文](README.zh-CN.md)

Conventional LLM research treats hidden-space tensors as isolated scalar
factors. flashcodec observes that the seq axis of a `(B, H, S, D)` KV tensor
behaves like the rows/columns of a high-dimensional image — adjacent token
representations are autocorrelated, energy concentrates in low frequencies, and
first-order differences shrink in variance (how strongly depends on the model,
layer and tensor: the measured real-KV statistics and the per-layer caveat are
in the Benchmarks section). Two classic
multimedia compression theories therefore transfer directly:

| Mode | Inspired by | Transform | Nature | Storage |
|---|---|---|---|---|
| `mode="jpeg"` | JPEG | Block-wise orthonormal 1D DCT-II, keep lowest K coeffs | **Lossy, frequency domain** | `K/r` (4×–8× less memory) |
| `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 |

`quant_bits=8` (jpeg only) stores the coefficients as int8 with a
per-(block, frequency) scale — a real byte-level saving on top of `K/r`:
**6.6% of raw KV at k=2 (15×), 46.9 dB PSNR**, versus 25% / 47.2 dB in fp32.

Math (block size `r`, operating axis `dim=-2`):

- **JPEG**: $X_{\text{blk}} \in \mathbb{R}^{r \times D}$, $C = M_{\text{enc}} X_{\text{blk}}$ with
  $M_{\text{enc}} \in \mathbb{R}^{K \times r}$, $M[j,n] = c_j \cos\frac{\pi}{r}(n+\tfrac12)j$,
  orthonormal rows ($K=r$ is per-block lossless); reconstruct $\hat X = M_{\text{enc}}^\top C$.
- **PNG**: $\Delta_i = x_i - x_{i-1}$ ($i \ge 1$), $x_0$ kept at full precision as the
  anchor; rebuild $\hat x_i = x_0 + \sum_{j=1}^{i} \Delta_j$ (`cumsum`, purely
  linear and autograd-differentiable). Autocorrelation implies
  $\mathrm{Var}(\Delta)/\mathrm{Var}(X) \approx 0.03$ (**30× variance collapse**,
  measured); after Top-M pruning the residual is nearly zero and sparse.

```
                    ┌────────────────────────────────────────────┐
  (B,H,S,D) ──append──▶  raw_buffer (tail < r, full fidelity)     │
                    │        │ r tokens accumulated, batched      │
                    │        ▼ archive (single GEMM, in-place     │
                    │             slice write into)               │
                    │  compressed_storage (capacity-doubled)      │
                    │   jpeg: coeffs (B,H,capacity,K,D)           │
                    │   png:  deltas (B,H,capacity,r-1,D)         │
                    │         + anchors (B,H,capacity,D)          │
                    └───────────────┬────────────────────────────┘
                                    │ read(): one decode of active
                                    │         region + raw tail
                                    ▼
                        reconstructed KV view (B,H,S,D)
```

Archives live in a **single geometrically-grown buffer** (std::vector-style:
start at 16 blocks, double on overflow). Each full block is one in-place
slice write — no re-allocation churn — and `read()` decodes **only the
newly archived blocks** (one batched GEMM / cumsum per new block, memoized
in an append-only reconstruction buffer) instead of re-decoding the whole
history on every decode step. (Fragmented per-block archive lists made
per-token decode ~170× slower at S=4096 due to kernel-launch overhead.)

## Causal-safety design

Only **complete blocks** are archived; the current incomplete block always
stays in `raw_buffer` at original precision. During incremental decoding, no
future perturbation can ever rewrite archived tensors (verified bitwise in
tests). On the lossless tier (`mode="png", keep_ratio=1.0`), step-by-step
decoded attention matches full-sequence attention within ≤1e-4.

## Installation

```bash
pip install -e ".[hf]"  # from source: torch>=2.0 (+ transformers>=5.20 for the HF backend)
```

Distribution name on PyPI: `flashcodec-kv` — the import name stays
`flashcodec`. (The bare `flashcodec` name there belongs to an unrelated
placeholder package, so don't `pip install flashcodec`.)

## Quick start

```python
import torch
from flashcodec import FlashCodecCache, apply_flash_codec

# 1) Native: a first-class transformers 5.x Cache backend (any CausalLM)
cache = FlashCodecCache(mode="jpeg", block_size=8, keep_freqs=2)  # 4×
#    int8 coefficients: same shape, 15× smaller archive (quant_bits=8)
cache = FlashCodecCache(mode="jpeg", block_size=8, keep_freqs=2,
                        quant_bits=8)                              # 15×
out = model.generate(input_ids, max_new_tokens=64,
                     past_key_values=cache, use_cache=True)
print(cache.get_seq_length(), cache.storage_bytes())

# 2) One-line hot patch: every generate() defaults to the codec cache
apply_flash_codec(model, mode="png", block_size=8, keep_ratio=1.0)
out = model.generate(input_ids, max_new_tokens=64)
# PNG lossless tier reproduces the unpatched model's tokens EXACTLY
from flashcodec import remove_flash_codec
remove_flash_codec(model)          # restore default behavior

# Materialization-free decode (coefficient-domain attention, no KV view):
apply_flash_codec(model, mode="jpeg", block_size=8, keep_freqs=2,
                  quant_bits=8, materialization_free=True)

# 3) Standalone (no transformers needed): direct cache manager
from flashcodec import CodecKVCache
cache = CodecKVCache(mode="jpeg", block_size=8, keep_freqs=2)
k, v = cache.forward(k_seq, v_seq)          # full-sequence fill
k, v = cache.decode_step(new_k, new_v)      # incremental decode
```

`FlashCodecCache` subclasses `transformers.cache_utils.Cache` (5.x layer
API): write/read replacement is converged at the `cache.update(k, v,
layer_idx)` interface — zero model surgery, works with any CausalLM
(Llama/Qwen/Mistral-style GQA included).

## Transformers integration (real-world tested on Qwen3.5-0.8B)

Hybrid architectures (e.g. Qwen3.5: 18 linear-attention layers + 6
full-attention layers) are handled automatically —
`FlashCodecCache.from_config(model.config)` attaches codec layers **only to
`full_attention` layers** and keeps native dynamic layers for everything
else (recurrent/conv states untouched). `apply_flash_codec` uses `from_config`
whenever a model config is available.

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from flashcodec import apply_flash_codec

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-0.8B",
                                             torch_dtype="bfloat16").cuda()
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B")
ids = tok("A quiet lake at dawn", return_tensors="pt").input_ids.cuda()

# One line — every generate() now runs on the codec cache
apply_flash_codec(model, mode="jpeg", block_size=8, keep_freqs=2)
out = model.generate(ids, max_new_tokens=48, do_sample=False)
```

Measured on Qwen3.5-0.8B (bf16, CUDA, 48 greedy tokens, `benchmarks/run_qwen_demo.py`):

| Tier | Storage vs raw KV | Output | Verdict |
|---|---|---|---|
| `png keep=1.0` | ~1.8× (fp32 deltas for exactness) | **bit-identical tokens to the unpatched model** | 100% fidelity baseline |
| `jpeg k=8 int8` | **55.5%** (~1.8×) | fluent, **100% token match** | safe lossy tier |
| `jpeg k=2` | **32.3%** (~3.1×) | fluent, 38.1% token match | lossy tier |
| `jpeg k=4` | 54.8% (~1.8×) | fluent, 49.2% token match | lossy tier |
| `jpeg k=2 int8` | **21.1%** (~4.7×) | fluent, 38.1% token match | lossy + int8 |

Hybrid-layer audit: `codec_layers=6/24` (full attention only), zero warnings.
Reproduce with `python benchmarks/run_qwen_demo.py` (offline; honors
`FLASHCODEC_MODEL` / `FLASHCODEC_NEW_TOKENS`).

dtype contracts (see `AGENTS.md`): PNG deltas are always fp32 (bf16
differences lose the bitwise-exactness the lossless tier promises); JPEG
coefficients and `read()` views stay in the stream's input dtype (SDPA
requires q/k/v dtype parity).

## Materialization-free decode attention

Every codec here is a **linear map along the seq axis**: the r tokens of a block
are `K_block = A @ S`, where `S` are the stored rows (DCT coefficients, or
anchor + deltas) and `A` is the synthesis matrix. So a decode query can attend
on `S` directly:

```
q . K[p] = (q . S^T) @ A[p]        out = (prob @ A) @ S
```

`codec_decode_attention(q, key_stream, value_stream)` implements this for one
new token (jpeg k<r, jpeg int8 and png; GQA grouping included). The `(S, D)`
reconstruction never exists, and the result matches the materialized path to
≤1e-6.

Measured (`benchmarks/bench_attention.py`, B=1, H_kv=8, H_q=16, D=64, fp32):

| seq | mode | materialized | coefficient-domain |
|---|---|---|---|
| 16384 | jpeg k=2 int8 | 0.20ms / 130MB | 0.30ms / **32MB** |
| 65536 | jpeg k=2 int8 | 0.86ms / 520MB | **0.46ms / 128MB** |
| 65536 | png keep=1.0 | 0.87ms / 520MB | 1.03ms / 272MB |

Scope: **decode only** (prefill still uses the view — the coefficient form would
explode the score tensor). png stores r rows per block, so it is numerically
supported but gains nothing; the win is jpeg `k < r`.

A **Triton split-K kernel** (`flashcodec/triton_attention.py`, flash-decoding
style: block chunks write partial max/sum/acc, a combine kernel merges them)
removes the intermediate tensors and the launch overhead. One layer, B=1,
H_kv=8, H_q=16, D=128, `jpeg k=2 int8` (`benchmarks/bench_attention.py --cuda`):

| seq | materialized | coefficient (torch) | coefficient (Triton) |
|---|---|---|---|
| 16384 | 0.20ms / 130MB | 0.29ms / 24MB | **0.09ms / 0MB** |
| 65536 | 0.86ms / 520MB | 0.41ms / 96MB | **0.16ms / 0MB** |

`apply_flash_codec(model, ..., materialization_free=True)` wires it into
transformers: the forward of standard Llama/Qwen-style attention modules is
wrapped (projection + RoPE unchanged, the attention output replaced), guarded to
single-token / batch-1 / eval / non-sliding layers, otherwise it falls back to
the original path. The prefill reconstruction is released on the first decode
step. `remove_flash_codec` restores everything.

End-to-end (`benchmarks/bench_e2e.py`, Qwen3-0.6B, 4096-token prompt + 32 new):

| tier | tok/s | peak | steady |
|---|---|---|---|
| baseline (raw KV) | 52.1 | 581MB | 0MB |
| codec k=2 int8, materialized | 44.2 | 1545MB | 1535MB |
| codec k=2 int8, materialization-free (Triton) | 40.3 | **642MB** | **122MB** |

**12.6× less steady memory** at ~9% slower decode than the materialized path
(per decode step at 4096 tokens: 22.8ms vs 21.7ms materialized, 17.8ms raw).
The remaining gap is the per-layer split/combine launch pair; a persistent
kernel would close it.

## Demo (side-by-side chat on Qwen3.5-0.8B)

![flashcodec side-by-side Gradio demo](images/demo.png)

Same prompt, two panels: native KV cache on the left, flashcodec on the right
(five tiers: png lossless / jpeg k=8 int8 / k=4 / k=2 / k=2 int8), with a
materialization-free toggle (coefficient-domain + Triton) and per-turn stats:
KV footprint (archive vs materialized view), peak/steady memory, per-token
latency and token match.

The **Full comparison** tab runs the baseline plus every tier on one prompt and
reports, per tier: KV footprint vs raw, archive/view split, decode tok/s,
peak/steady memory, teacher-forced NLL of the continuation (a perplexity-style
quality number) and token match — plus storage / speed / NLL bars.

```bash
pip install -e ".[demo]"   # gradio + matplotlib
python demo/qwen_gradio_chat.py    # Gradio side-by-side + savings chart
```

## Benchmarks (measured; B=2 H=8 S=2048 D=64, random-walk autocorrelated signal, fp32)

### Reconstruction quality (`benchmarks/bench_reconstruct.py`)

| Config | PSNR (dB) | cosine | Storage |
|---|---|---|---|
| jpeg r=8 k=1 (8×) | 42.95 | 0.99943 | 12.5% |
| jpeg r=8 k=2 (4×) | 47.22 | 0.99985 | 25.0% |
| jpeg r=8 k=4 (2×) | 51.66 | 1.00001 | 50.0% |
| jpeg r=8 k=2 int8 | 46.92 | 0.99983 | 6.6% |
| jpeg r=8 k=4 int8 | 50.89 | 0.99999 | 13.3% |
| png r=8 keep=0.5 | 50.11 | 0.99997 | 100% (sparsity → entropy coder) |
| png r=8 keep=1.0 | ~196 (fp rounding) | 1.00000 | 100% |

### Real KV statistics (real models, 512 tokens)

`bench_reconstruct` uses a synthetic random walk; real caches are messier.
Measured on real post-RoPE KV (`Var(Δ)/Var(X)` along the seq axis, and the
fraction of block energy kept by the lowest k=2 DCT coefficients of a block of
8):

| model / layer | tensor | Var(Δ)/Var(X) | energy in k=2 |
|---|---|---|---|
| Qwen3-0.6B L0 | K | 0.16 | 0.94 |
| Qwen3-0.6B L0 | V | 1.87 | 0.31 |
| Qwen3-0.6B L27 | V | 1.02 | 0.59 |
| Qwen3.5-0.8B L3 | K | 0.72 | 0.71 |
| Qwen3.5-0.8B L23 | V | 0.09 | 0.97 |

Values are frequently **anti**-correlated (ratio > 1), so "first differences
collapse the variance" is a property of smooth synthetic signals, not of every
real layer. Practical consequence: coefficient truncation (`k < r`) is a
per-layer decision — on a dense model (28 compressed layers) the error compounds
and quality collapses; use `k = r` + `quant_bits=8` there.

### Perplexity (`benchmarks/bench_ppl.py`; 16 × 512 tokens, bf16, local corpus)

| tier | Qwen3.5-0.8B (6/24 layers) | Qwen3-0.6B (28/28 layers) | storage |
|---|---|---|---|
| baseline | 22.4803 | 16.7931 | 100% |
| png keep=1.0 | 22.4803 (+0.000) | 16.7931 (+0.000) | 188% |
| jpeg k=8 (lossless basis) | 22.4877 (+0.007) | 16.8058 (+0.013) | 100% |
| **jpeg k=8 int8** | **22.4991 (+0.019)** | **17.1348 (+0.34)** | **~52%** |
| jpeg k=4 | 22.3010 (−0.18) | 470.85 (+454) | 50% |
| jpeg k=2 | 24.0566 (+1.58) | 6485.5 (+6469) | 25% |
| jpeg k=2 int8 | 24.0590 (+1.58) | 8118.5 (+8102) | 13% |

The hybrid result is good because only 6 of 24 layers carry a codec (the 18
linear-attention layers keep their native state); the dense model shows why
per-layer truncation must not be applied blindly.

### Per-layer budgets (`benchmarks/bench_adaptive.py`, Qwen3-0.6B, 6 × 512 tokens)

Calibration: compress one layer at a time to k=4, measure the perplexity damage,
then give the cheapest half of the layers k=4 and leave the rest at k=r (75%
storage):

| schedule | ppl | Δppl | storage |
|---|---|---|---|
| baseline (k=r) | 18.9786 | — | 100% |
| uniform k=6 | 48.0531 | +29.07 | 75% |
| adaptive (14 layers k=4) | 40.8309 | +21.85 | 75% |
| uniform k=4 | 710.22 | +691.24 | 50% |
| uniform k=r int8 | 19.3627 | +0.38 | 50% |

Calibration beats a uniform budget at the same storage (−15%), but the damages
are **super-additive** (the cheapest 14 layers add up to +22 ppl although no
single one costs more than +0.5), and int8 at 50% is far better. Use
`per_layer={idx: {"keep_freqs": ...}}` when a budget forces structural
truncation; otherwise prefer `quant_bits=8`.

### Cache footprint & latency (`benchmarks/bench_vram.py`, CPU)

| seq | Mode | storage | ratio | fill | read |
|---|---|---|---|---|---|
| 4096 | jpeg k=2 | 4.0MB | 0.25 | 6.7ms | 0.0ms |
| 16384 | jpeg k=2 | 16.0MB | 0.25 | 34.3ms | 0.5ms |
| 65536 | jpeg k=2 | 64.0MB | 0.25 | 141.9ms | 3.1ms |
| 4096 | jpeg k=2 int8 | 1.1MB | 0.07 | 9.2ms | 0.0ms |
| 16384 | jpeg k=2 int8 | 4.3MB | 0.07 | 46.0ms | 0.3ms |
| 65536 | jpeg k=2 int8 | 17.0MB | 0.07 | 197.9ms | 3.1ms |
| 65536 | png keep=1.0 | 256.0MB | 1.00 | 243.3ms | 3.3ms |

CPU timings vary ~±20% between runs (min of 3); the storage/ratio columns are
exact. `read` = materialize the view and touch it (`.sum()`, so the number is the
consumer cost, not just kernel launch). Decoding is memoized per archived block,
so `read` costs one decode per **new** block instead of re-decoding the whole
history on every decode step (previously 38.7ms / 71.9ms in the last two rows).

PNG does not shrink dense storage (the fate of a lossless tier) — its value is
the 30× variance collapse plus near-zero sparse residuals after Top-M, i.e.
ideal input for an entropy coder (none is bundled).

## Repository layout

```
flashcodec/
├── flashcodec/
│   ├── __init__.py        # apply_flash_codec / CodecKVCache / JPEGCodec / PNGCodec
│   ├── codecs/jpeg.py     # DCT-II basis, project/reconstruct (pure GEMM, no loops)
│   ├── codecs/png.py      # anchor + first-order delta, cumsum, Top-M filter
│   ├── cache.py           # CodecKVCache: raw tail + capacity-doubled archives
│   ├── attention.py       # materialization-free decode attention (coefficient domain)
│   └── patch.py           # one-line hot patch: cache-level injection (HF Cache duck-typing)
├── tests/                 # jpeg / png / causal / attention suites (standalone scripts)
├── benchmarks/            # storage + PSNR/cosine + perplexity benchmarks
├── demo/                  # Gradio side-by-side chat
├── images/                # README screenshots
├── .github/workflows/     # CI: CPU tests + transformers integration
├── CHANGELOG.md
├── pyproject.toml
└── README.md
```

## Tests

```bash
python tests/test_jpeg.py           # orthonormality, k=r lossless, k=2 cosine>0.99, K monotone
python tests/test_png.py            # lossless maxdiff<1e-6, variance ratio printed, differentiable
python tests/test_causal.py         # zero future leakage, incremental vs full ≤1e-4,
                                   # memoized read == full decode, exact + repeatable crop,
                                   # beam-search batch gather
python tests/test_hf_integration.py # native transformers 5.x generate: explicit cache / hot patch /
                                    # PNG lossless exact tokens / 70-token multi-block stability /
                                    # beam search + crop exactness against the unpatched model /
                                    # per-layer schedule
python tests/test_attention.py      # coefficient-domain decode attention matches the materialized
                                    # path (jpeg / int8 / png, GQA); never builds the KV view
```

## Stability (1.x)

The public API is frozen for 1.x: `apply_flash_codec` / `remove_flash_codec`,
`apply_codec_attention` / `remove_codec_attention`, `FlashCodecCache`,
`FlashCodecLayer`, `CodecKVCache`, `JPEGCodec`, `PNGCodec`,
`codec_decode_attention`, `make_codec`, plus the `quant_bits` / `per_layer`
options. New features ship as additive arguments; existing signatures keep
their defaults. The caveats below still apply — this is a research library with
a stable interface, not a claim of production readiness.

## Boundaries (honest list)

- Research prototype: `read()` materializes the full reconstructed view for
  attention; memory gains are at the **cache storage layer** (table above);
  end-to-end benefits require combining with materialization-free kernels
  (flash-attention style). That view is memoized and reused across decode
  steps, so steady-state memory = archive + one raw-KV-sized view.
- Full-sequence `forward(seq)` prefill compresses complete blocks directly, so
  mid-block attention sees the block-level reconstruction (within-block
  approximation); the incremental `decode_step` path has no such issue.
- Compressed tensors are dense layouts; PNG Top-M sparsity is not yet wired to
  an entropy coder. (Packing sparse png alone was measured to be a wash: the
  fp32 delta values dominate, and on bf16 input the packed form can be *larger*
  than raw — delta quantization has to come first.)
- Coefficient truncation (`k < r`) is a per-layer decision: on a dense model
  every layer is compressed and the error compounds (measured: Qwen3-0.6B
  perplexity 16.8 → 6486 at k=2). Use `k = r` + `quant_bits=8` there (~52%
  storage, +2% ppl) and reserve aggressive ratios for hybrid models where only
  the full-attention layers carry a codec.
- Same lineage as the v4.2 prototypes inside
  [Mini-Architecture](https://github.com/aifeifei798/Mini-Architecture)
  (there: per-position view semantics for cross-layer KV reuse; here:
  standard cache storage semantics).

## License

[MIT](LICENSE)
