flashcodec-kv 1.0.0__py3-none-any.whl

flashcodec/triton_attention.py

@@ -0,0 +1,222 @@
+"""Triton kernel for coefficient-domain decode attention (split-K).
+
+One program per (batch, kv head, query group, block-chunk) walks its slice of
+the archive and performs the whole decode attention in registers:
+
+    rows  = dequant(coeffs)                  (S, D)   per block
+    s     = rows @ q                         (S,)
+    logits= A @ s  * scale                   (R,)
+    online-softmax update
+    w     = p @ A                            (S,)
+    acc  += w @ rows                         (D,)
+
+then a tiny combine kernel merges the per-chunk (max, sum, acc) partials
+(flash-decoding style) so the GPU is not limited to one program per head. That
+removes the five intermediate tensors of the PyTorch path and launches two
+kernels per layer instead of ~6.
+
+Falls back to the PyTorch implementation whenever the shape/dtype/device is not
+supported (see ``codec_decode_attention_triton``).
+"""
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+    HAS_TRITON = True
+except Exception:                                   # no triton -> torch path
+    HAS_TRITON = False
+
+
+if HAS_TRITON:
+
+    @triton.jit(do_not_specialize=["nb", "nr", "chunk", "ns"])
+    def _attn_split_kernel(
+        Q, K, V, KS, VS, KR, VR, A,
+        PM, PL, PA,
+        nb, nr, chunk, ns,
+        q_sb, q_sh, q_sq, q_sd,
+        k_sb, k_sh, k_sn, k_ss, k_sd,
+        v_sb, v_sh, v_sn, v_ss, v_sd,
+        ks_sb, ks_sh, ks_sn, ks_ss,
+        kr_sb, kr_sh, kr_st, kr_sd,
+        pa_s0, pa_s1,
+        SROWS: tl.constexpr, RTOK: tl.constexpr, DIM: tl.constexpr,
+        BDIM: tl.constexpr, HK: tl.constexpr, GFACT: tl.constexpr,
+        K_INT8: tl.constexpr, V_INT8: tl.constexpr, HAS_RAW: tl.constexpr,
+        SCALE: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        sid = tl.program_id(1)
+        g = pid % GFACT
+        h = (pid // GFACT) % HK
+        b = pid // (GFACT * HK)
+
+        d = tl.arange(0, BDIM)
+        dm = d < DIM
+        s = tl.arange(0, SROWS)
+        r = tl.arange(0, RTOK)
+
+        q = tl.load(Q + b * q_sb + (h * GFACT + g) * q_sh + d * q_sd,
+                    mask=dm, other=0.0).to(tl.float32)
+        a = tl.load(A + r[:, None] * SROWS + s[None, :])      # (R, S) fp32
+
+        m_i = float("-inf")
+        l_i = 0.0
+        acc = tl.zeros([BDIM], tl.float32)
+
+        start = sid * chunk
+        end = tl.minimum(start + chunk, nb)
+        for n in range(start, end):
+            kbase = b * k_sb + h * k_sh + n * k_sn
+            rows = tl.load(K + kbase + s[:, None] * k_ss + d[None, :] * k_sd,
+                           mask=dm[None, :], other=0.0).to(tl.float32)
+            if K_INT8:
+                ksc = tl.load(KS + b * ks_sb + h * ks_sh + n * ks_sn
+                              + s * ks_ss)
+                rows = rows * ksc[:, None]
+            sj = tl.sum(rows * q[None, :], axis=1)            # (S,)
+            logits = tl.sum(a * sj[None, :], axis=1) * SCALE  # (R,)
+
+            m_new = tl.maximum(m_i, tl.max(logits))
+            p = tl.exp(logits - m_new)
+            alpha = tl.exp(m_i - m_new)
+            l_i = l_i * alpha + tl.sum(p)
+
+            w = tl.sum(p[:, None] * a, axis=0)                # (S,)
+            vbase = b * v_sb + h * v_sh + n * v_sn
+            vrows = tl.load(V + vbase + s[:, None] * v_ss + d[None, :] * v_sd,
+                            mask=dm[None, :], other=0.0).to(tl.float32)
+            if V_INT8:
+                vsc = tl.load(VS + b * ks_sb + h * ks_sh + n * ks_sn
+                              + s * ks_ss)
+                vrows = vrows * vsc[:, None]
+            acc = acc * alpha + tl.sum(w[:, None] * vrows, axis=0)
+            m_i = m_new
+
+        # the raw tail lives in the last chunk only
+        if HAS_RAW and sid == ns - 1:
+            for t in range(0, nr):
+                kraw = tl.load(KR + b * kr_sb + h * kr_sh + t * kr_st
+                               + d * kr_sd, mask=dm, other=0.0).to(tl.float32)
+                logit = tl.sum(q * kraw) * SCALE
+                m_new = tl.maximum(m_i, logit)
+                p = tl.exp(logit - m_new)
+                alpha = tl.exp(m_i - m_new)
+                l_i = l_i * alpha + p
+                vraw = tl.load(VR + b * kr_sb + h * kr_sh + t * kr_st
+                               + d * kr_sd, mask=dm, other=0.0).to(tl.float32)
+                acc = acc * alpha + p * vraw
+                m_i = m_new
+
+        slot = pid * ns + sid
+        tl.store(PM + slot, m_i)
+        tl.store(PL + slot, l_i)
+        tl.store(PA + slot * pa_s0 + d * pa_s1, acc, mask=dm)
+
+    @triton.jit(do_not_specialize=["ns"])
+    def _attn_combine_kernel(
+        PM, PL, PA, OUT,
+        ns, out_sb, out_sh, out_sq, out_sd,
+        DIM: tl.constexpr, BDIM: tl.constexpr, HQT: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        d = tl.arange(0, BDIM)
+        dm = d < DIM
+        m = float("-inf")
+        l = 0.0
+        acc = tl.zeros([BDIM], tl.float32)
+        for i in range(0, ns):
+            mi = tl.load(PM + pid * ns + i)
+            li = tl.load(PL + pid * ns + i)
+            ai = tl.load(PA + (pid * ns + i) * BDIM + d, mask=dm, other=0.0)
+            m_new = tl.maximum(m, mi)
+            alpha = tl.exp(m - m_new)
+            beta = tl.exp(mi - m_new)
+            l = l * alpha + li * beta
+            acc = acc * alpha + ai * beta
+            m = m_new
+        out = acc / l
+        b = pid // HQT
+        hq = pid % HQT
+        tl.store(OUT + b * out_sb + hq * out_sh + d * out_sd,
+                 out.to(OUT.dtype.element_ty), mask=dm)
+
+
+def _strides(t):
+    return [t.stride(i) for i in range(t.dim())]
+
+
+def codec_decode_attention_triton(query, key_stream, value_stream,
+                                  scale=None, basis=None):
+    """Triton implementation; returns None if the case is unsupported."""
+    if not HAS_TRITON or not query.is_cuda:
+        return None
+    if query.shape[2] != 1:
+        return None
+    codec = key_stream.codec
+    if codec.name != value_stream.codec.name or codec.name != "jpeg":
+        return None                                    # png stores r rows: no win
+    n = key_stream.n_blocks
+    if n == 0:
+        return None
+    D = query.shape[-1]
+    BD = triton.next_power_of_2(D)
+    if D > 256:
+        return None
+    S = codec.keep_freqs
+    R = key_stream.block_size
+    B, Hq, _, _ = query.shape
+    Hk = key_stream._coeffs.shape[1]
+    if Hq % Hk != 0:
+        return None
+    G = Hq // Hk
+    scale = D ** -0.5 if scale is None else scale
+
+    krows, vrows = key_stream._coeffs, value_stream._coeffs
+    int8 = codec.quantized
+    ks_t, vs_t = key_stream._scales, value_stream._scales
+    if basis is None:
+        from .attention import block_basis
+        basis = block_basis(codec, R, device=query.device, dtype=torch.float32)
+    basis = basis.contiguous()
+
+    raw_k, raw_v = key_stream.raw, value_stream.raw
+    has_raw = raw_k is not None
+    nr = raw_k.shape[-2] if has_raw else 0
+    dummy = krows
+
+    n_prog = B * Hk * G
+    ns = max(1, min(32, (n + 63) // 64))
+    chunk = (n + ns - 1) // ns
+
+    pm = torch.empty((n_prog, ns), dtype=torch.float32, device=query.device)
+    pl = torch.empty((n_prog, ns), dtype=torch.float32, device=query.device)
+    pa = torch.empty((n_prog, ns, BD), dtype=torch.float32, device=query.device)
+    out = torch.empty_like(query)
+
+    _attn_split_kernel[(n_prog, ns)](
+        query, krows, vrows, ks_t if int8 else dummy,
+        vs_t if int8 else dummy,
+        raw_k if has_raw else dummy, raw_v if has_raw else dummy,
+        basis, pm, pl, pa,
+        n, nr, chunk, ns,
+        *_strides(query),
+        *_strides(krows),
+        *_strides(vrows),
+        *_strides(ks_t)[:4] if int8 else (0, 0, 0, 0),
+        *_strides(raw_k) if has_raw else (0, 0, 0, 0),
+        _strides(pa)[1], _strides(pa)[2],
+        SROWS=S, RTOK=R, DIM=D, BDIM=BD, HK=Hk, GFACT=G,
+        K_INT8=int8, V_INT8=int8, HAS_RAW=has_raw,
+        SCALE=float(scale),
+        num_warps=1,
+    )
+    _attn_combine_kernel[(n_prog,)](
+        pm, pl, pa, out,
+        ns, *_strides(out),
+        DIM=D, BDIM=BD, HQT=Hk * G,
+        num_warps=1,
+    )
+    return out

flashcodec_kv-1.0.0.dist-info/METADATA

@@ -0,0 +1,430 @@
+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)

flashcodec_kv-1.0.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+flashcodec/__init__.py,sha256=tTs9FLDPpFuxEfLU102DSOLut22UG8--jV2z1C20zmI,1555
+flashcodec/attention.py,sha256=IkN2AoXIMUyIqqH0QmH7UVXDjVCq9eEQrFT2XY8hSDA,6143
+flashcodec/cache.py,sha256=SXzZXN0y-5VCAkCPnQMHFsJ-cE3m7xTTET3Zi54Ovbo,28844
+flashcodec/patch.py,sha256=rMTyJa496MyJegHwJchUDSre-ajI3zLybWlDkfdtkSc,5652
+flashcodec/patch_attention.py,sha256=9bZjl-Po3aSK1UYgJmv7hOFh53xltBffy6345DbpwhM,6035
+flashcodec/triton_attention.py,sha256=bDeiF2ZGfgSGsAibf_u-dQFITvBd6SYyFbGZhwBoO0s,8085
+flashcodec/codecs/__init__.py,sha256=jboZNO_fdm_ngnqIEiyL-evI5r6N5imI_p8jGExpYzA,459
+flashcodec/codecs/jpeg.py,sha256=xNIUNbxIlmaGNifX7-hflqq3utw-xWvNSUmOrUlbIJY,6722
+flashcodec/codecs/png.py,sha256=Zxao_gtX9wkIFRt3aY2nct63yh8nEpgtulfiskHw1Yk,4708
+flashcodec_kv-1.0.0.dist-info/licenses/LICENSE,sha256=7slruha1foQbK2Y9M1SIvSs1CDfTmubSl82PZUecjCI,1063
+flashcodec_kv-1.0.0.dist-info/METADATA,sha256=VAl5BlL1dHi7dbZiMV6Fei6mWBoiGxI3Yp-oNCeE8V8,21489
+flashcodec_kv-1.0.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+flashcodec_kv-1.0.0.dist-info/top_level.txt,sha256=l4d-yumaVz4u_beYbHgdLvH8rQ5bjOcrehyWUD5j4KA,11
+flashcodec_kv-1.0.0.dist-info/RECORD,,

flashcodec_kv-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+