flashcodec-kv 1.0.0__py3-none-any.whl

flashcodec/codecs/jpeg.py

@@ -0,0 +1,159 @@
+"""JPEG-style lossy frequency-domain compression: blockwise orthonormal 1D
+DCT-II.
+
+Physical intuition: the hidden vectors of adjacent tokens form a strongly
+autocorrelated time-series signal whose energy concentrates in low
+frequencies. Each block is projected onto the first K orthonormal DCT-II basis
+vectors (DC mean + AC linear trend, ...), filtering out high-frequency random
+noise and yielding a geometric r/K compression ratio. Basis rows are
+orthonormal, so keep=r is lossless per block.
+
+Tensor contract: input (..., S, D); the compression/filtering axis is the
+second-to-last dimension (seq axis, dim=-2). Everything is batched matmul
+(GEMM), with no explicit Python loops.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+_BASIS_CACHE = {}
+
+
+def dct_basis(r, k, device=None, dtype=None):
+    """Precompute the orthonormal DCT-II basis M_enc (k, r):
+    M[j, n] = c_j · cos(π/r · (n+½)·j), c_0 = sqrt(1/r), c_j = sqrt(2/r).
+    Rows are orthonormal (M @ M.T = I); k == r gives the complete orthonormal
+    basis (lossless reconstruction).
+    Module-level constant cache, keyed by (r, k)."""
+    key = (r, k)
+    if key not in _BASIS_CACHE:
+        n = torch.arange(r, dtype=torch.float32)
+        j = torch.arange(k, dtype=torch.float32)
+        c = torch.where(j == 0, math.sqrt(1.0 / r), math.sqrt(2.0 / r))
+        m = c[:, None] * torch.cos(math.pi / r * (n[None, :] + 0.5) * j[:, None])
+        _BASIS_CACHE[key] = m
+    m = _BASIS_CACHE[key]
+    if device is not None or dtype is not None:
+        m = m.to(device=device, dtype=dtype or m.dtype)
+    return m
+
+
+def _pad_to_blocks(x, r):
+    """Zero-pad the seq axis (-2) to a multiple of r; returns (block tensor,
+    original length)."""
+    S = x.shape[-2]
+    n = (S + r - 1) // r
+    pad = n * r - S
+    if pad:
+        x = F.pad(x, (0, 0, 0, pad))          # (..., S, D) pads the -2 axis only
+    blocks = x.reshape(*x.shape[:-2], n, r, x.shape[-1])
+    return blocks, S
+
+
+def dct_compress(x, r, k):
+    """Blockwise DCT projection: x (..., S, D) -> coeffs (..., n_blocks, k, D)
+    (x dtype).
+    The seq axis is split into blocks of r (zero-padded at the tail); each
+    block is projected along the within-block axis onto the first k DCT-II
+    basis vectors (computed in fp32, then cast back to the input dtype — lossy
+    modes need no higher precision, and the storage ratio stays honest).
+    Compression storage ratio = k/r (same dtype, plus tail padding overhead)."""
+    blocks, _ = _pad_to_blocks(x, r)
+    basis = dct_basis(r, k, device=x.device, dtype=torch.float32)
+    hp = blocks.float().transpose(-2, -1)              # (..., n, D, r)
+    coef = torch.matmul(hp, basis.t())                 # (..., n, D, k)
+    return coef.transpose(-2, -1).to(x.dtype)          # (..., n, k, D)
+
+
+def dct_decompress(coeffs, r, out_len=None):
+    """Blockwise DCT reconstruction: coeffs (..., n_blocks, k, D) ->
+    x (..., n_blocks*r, D).
+    Low-frequency coefficients are multiplied by the basis to reconstruct the
+    within-block waveform; when out_len is given, the tail padding is trimmed."""
+    n, k = coeffs.shape[-3], coeffs.shape[-2]
+    basis = dct_basis(r, k, device=coeffs.device, dtype=torch.float32)
+    hp = torch.matmul(coeffs.transpose(-2, -1).float(), basis)   # (...,n,D,r)
+    rec = hp.transpose(-2, -1).reshape(
+        *coeffs.shape[:-3], n * r, coeffs.shape[-1])
+    if out_len is not None:
+        rec = rec[..., :out_len, :]
+    return rec
+
+
+def quantize_coeffs(coeffs, bits=8):
+    """Symmetric per-(block, frequency) quantization of DCT coefficients.
+
+    coeffs (..., n, k, D) -> (q int8, scale fp32 of shape (..., n, k, 1)).
+    One scale per block **and** frequency row: the DC row and the AC rows have
+    very different magnitudes, and a row-wise scale costs only 4 bytes per row
+    (vs D*1 byte for the int8 row itself).
+    """
+    if bits != 8:
+        raise ValueError(f"only 8-bit coefficient quantization is supported, "
+                         f"got bits={bits}")
+    amax = coeffs.abs().amax(dim=-1, keepdim=True)          # (...,n,k,1)
+    scale = amax / 127.0
+    scale = torch.where(scale > 0, scale, torch.ones_like(scale))
+    q = torch.round(coeffs.float() / scale).clamp_(-127, 127).to(torch.int8)
+    return q, scale.to(torch.float32)
+
+
+def dequantize_coeffs(q, scale, dtype=torch.float32):
+    """int8 coefficients + per-row scale -> coefficients in ``dtype``."""
+    return q.to(dtype) * scale.to(dtype)
+
+
+class JPEGCodec:
+    """JPEG-style lossy frequency-domain codec (stateless, shareable across
+    caches). ``quant_bits=8`` additionally stores the DCT coefficients as int8
+    with a per-(block, frequency) scale — a real byte-level saving on top of the
+    k/r geometric ratio (the scales cost 4 bytes per row)."""
+
+    name = "jpeg"
+
+    def __init__(self, block_size=8, keep_freqs=2, quant_bits=0):
+        assert 1 <= keep_freqs <= block_size, \
+            f"keep_freqs must be within [1, {block_size}]"
+        assert quant_bits in (0, 8), \
+            f"quant_bits must be 0 (store in input dtype) or 8, got {quant_bits}"
+        self.block_size = block_size
+        self.keep_freqs = keep_freqs
+        self.quant_bits = quant_bits
+
+    @property
+    def quantized(self):
+        return self.quant_bits > 0
+
+    def compress(self, x):
+        return dct_compress(x, self.block_size, self.keep_freqs)
+
+    def quantize(self, coeffs):
+        """Coefficients -> (int8, scale) when quantized, else (None, None)."""
+        if not self.quantized:
+            return None, None
+        return quantize_coeffs(coeffs, self.quant_bits)
+
+    def decompress(self, coeffs, out_len=None, scales=None):
+        if self.quantized:
+            assert scales is not None, "quantized codec needs the scales"
+            coeffs = dequantize_coeffs(coeffs, scales)
+        return dct_decompress(coeffs, self.block_size, out_len)
+
+    def storage_ratio(self):
+        """Archive bytes / original bytes **assuming fp32 input** (excluding the
+        raw tail). Non-quantized: k/r. Quantized: (k*D + k*4) / (r*D*4)."""
+        if not self.quantized:
+            return self.keep_freqs / self.block_size
+        # D cancels out only in the element-count sense; report per fp32 element
+        # with D=1 for the generic bound and make the byte count exact in
+        # storage_bytes(); here we use the k/r geometric ratio scaled by the
+        # int8-vs-fp32 element ratio, ignoring the (tiny) scale overhead.
+        return (self.keep_freqs / self.block_size) / 4.0
+
+    def __repr__(self):
+        q = f", quant_bits={self.quant_bits}" if self.quantized else ""
+        return (f"JPEGCodec(block_size={self.block_size}, "
+                f"keep_freqs={self.keep_freqs}{q})")
+

flashcodec/codecs/png.py

@@ -0,0 +1,107 @@
+"""PNG-style lossless/near-lossless time-domain compression: causal blockwise
+prediction delta filtering.
+
+Physical intuition: the hidden vectors of adjacent tokens are highly
+autocorrelated; after a first-order delta Δ = x_i - x_{i-1} the variance drops
+sharply and the values cluster near 0 (very low information entropy). Each
+block stores a full-precision anchor at its head (I-Token) and deltas
+elsewhere; cumsum gives causal reconstruction that is bitwise invertible at
+keep_ratio=1.0 (up to fp rounding).
+When keep_ratio<1.0, the Top-M steps by |Δ| magnitude are kept per
+(block, feature), emulating a near-lossless mode.
+
+Tensor contract: input (..., S, D); the operating axis is the second-to-last
+dimension (seq axis, dim=-2). Purely linear operators + discrete selection, so
+autograd works naturally (gradients flow only to the retained delta steps).
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+def _pad_to_blocks(x, r):
+    """Zero-pad the seq axis (-2) to a multiple of r; returns (block tensor,
+    original length)."""
+    S = x.shape[-2]
+    n = (S + r - 1) // r
+    pad = n * r - S
+    if pad:
+        x = F.pad(x, (0, 0, 0, pad))
+    blocks = x.reshape(*x.shape[:-2], n, r, x.shape[-1])
+    return blocks, S
+
+
+def png_filter(x, block_size, keep_ratio=1.0):
+    """Delta filtering: x (..., S, D) -> (deltas (..., n, r-1, float32),
+    anchor (..., n, D, input dtype)).
+    Position 0 of every block is forced to a full-precision anchor; position
+    i>=1 stores Δ_i = x_i - x_{i-1}.
+    When keep_ratio<1.0, the Top-M steps by |Δ| magnitude are kept per
+    (block, feature) (m = ceil((r-1)*keep_ratio), m>=1), the rest are set to 0.
+    dtype contract: **deltas are always fp32** — computing the deltas of bf16
+    inputs in bf16 loses ~0.4% precision and breaks the bitwise exactness of
+    lossless reconstruction; Δ is a "portable representation" whose precision
+    must not be lower than that of the original signal (bf16 deltas are exactly
+    representable in fp32, so cumsum restores them bitwise).
+    The anchor keeps the input dtype."""
+    assert 0 < keep_ratio <= 1.0, "keep_ratio must be within (0, 1]"
+    blocks, _ = _pad_to_blocks(x, block_size)
+    anchor = blocks[..., 0, :]                                # (..., n, D)
+    bf = blocks.float()
+    delta = bf[..., 1:, :] - bf[..., :-1, :]                  # (..., n, r-1, D)
+    if keep_ratio < 1.0:
+        m = max(1, math.ceil(delta.shape[-2] * keep_ratio))
+        topi = delta.abs().topk(m, dim=-2).indices
+        mask = torch.zeros_like(delta, dtype=torch.bool)
+        mask.scatter_(-2, topi, True)
+        delta = delta * mask
+    return delta, anchor
+
+
+def png_reconstruct(deltas, anchor):
+    """Causal reconstruction: anchor + cumsum (accumulated in fp32, then cast
+    back to the anchor dtype).
+    Position 0 = anchor, position i>=1 = anchor + the first i delta steps.
+    deltas (..., n, r-1, D), anchor (..., n, D) -> x (..., n*r, D).
+    With bf16 input + keep_ratio=1.0 the result is bitwise lossless (deltas are
+    exact in fp32)."""
+    base = anchor.float().unsqueeze(-2)                       # (..., n, 1, D)
+    rec = torch.cat([base, base + torch.cumsum(deltas.float(), dim=-2)], dim=-2)
+    return rec.to(anchor.dtype).reshape(*rec.shape[:-3],
+                                        rec.shape[-3] * rec.shape[-2],
+                                        rec.shape[-1])
+
+
+class PNGCodec:
+    """PNG-style prediction-filter codec (stateless, shareable across caches)."""
+
+    name = "png"
+
+    def __init__(self, block_size=8, keep_ratio=1.0):
+        assert 0 < keep_ratio <= 1.0, "keep_ratio must be within (0, 1]"
+        self.block_size = block_size
+        self.keep_ratio = keep_ratio
+
+    def filter(self, x):
+        return png_filter(x, self.block_size, self.keep_ratio)
+
+    def reconstruct(self, deltas, anchor):
+        return png_reconstruct(deltas, anchor)
+
+    def storage_ratio(self):
+        """Dense tensor storage ratio: 1.0 for fp32 input (anchor + (r-1) delta
+        steps are the same size).
+        With bf16 input, deltas are stored as fp32 (the dtype contract for
+        bitwise losslessness), so the archive is about 1.9× — bandwidth traded
+        for exactness, with the compression headroom left to an entropy coder.
+        PNG's value lies at the information-entropy level: delta variance drops
+        ~30×, and after Top-M truncation the residual is highly sparse and
+        near zero — that is the headroom for an entropy coder
+        (run-length/arithmetic coding)."""
+        return 1.0
+
+    def __repr__(self):
+        return (f"PNGCodec(block_size={self.block_size}, "
+                f"keep_ratio={self.keep_ratio})")

flashcodec/patch.py

@@ -0,0 +1,136 @@
+"""One-step hot patch: make any HuggingFace CausalLM use FlashCodecCache.
+
+Mechanism (zero intrusion): in standard HF models, KV writes/reads are all
+funnelled through the Cache interface — after the attention layer computes
+key/value it calls ``cache.update(k, v, layer_idx, ...)`` and gets the
+full-history KV back. Swapping the Cache object for FlashCodecCache (a native
+transformers 5.x Cache subclass) therefore completes the "replacement of the
+KV write and read logic" without model surgery for every attention
+implementation (GQA/MQA/RoPE variants).
+
+- `apply_flash_codec(model, ...)`: wraps `model.generate` and `model.forward`;
+  when `use_cache` is requested and no `past_key_values` is passed explicitly,
+  it injects the configured FlashCodecCache; it still walks model.modules() to
+  find standard KV-projection attention modules (with k_proj/v_proj) for
+  validation and reporting;
+- `remove_flash_codec(model)`: fully restores both wrappers.
+"""
+
+import functools
+import warnings
+
+from .cache import HAS_HF_CACHE, CodecKVCache, FlashCodecCache
+
+__all__ = ["apply_flash_codec", "remove_flash_codec", "undo_flash_codec",
+           "find_attention_modules"]
+
+
+def find_attention_modules(model):
+    """Find standard Self-Attention / KV projection modules (with k_proj and
+    v_proj)."""
+    return [m for m in model.modules()
+            if hasattr(m, "k_proj") and hasattr(m, "v_proj")]
+
+
+def _make_factory(mode, block_size, kwargs, sink, config=None):
+    codec_kw = dict(keep_freqs=kwargs.get("keep_freqs", 2),
+                    keep_ratio=kwargs.get("keep_ratio", 1.0),
+                    quant_bits=kwargs.get("quant_bits", 0))
+    per_layer = kwargs.get("per_layer")
+    if HAS_HF_CACHE:
+        def factory():
+            c = None
+            if config is not None:      # hybrid architecture (e.g. Qwen3.5):
+                try:                    # only full-attention layers get a codec
+                    c = FlashCodecCache.from_config(        # per layer_types
+                        config, mode=mode, block_size=block_size,
+                        per_layer=per_layer, **codec_kw)
+                except Exception:
+                    c = None
+            if c is None:
+                c = FlashCodecCache(mode=mode, block_size=block_size,
+                                    **codec_kw)
+            sink.append(c)
+            return c
+    else:                                       # fallback without transformers
+        def factory():
+            c = CodecKVCache(mode=mode, block_size=block_size, **codec_kw)
+            sink.append(c)
+            return c
+    return factory
+
+
+def apply_flash_codec(model, mode="jpeg", block_size=8,
+                      materialization_free=False, **kwargs):
+    """One-line hot patch: inject a codec KV cache into model (native
+    transformers 5.x).
+
+    Supports mode="jpeg" (kwargs: keep_freqs, default 2; quant_bits, default 0
+    or 8 for int8 coefficients) and mode="png" (kwargs: keep_ratio, default
+    1.0).
+    Afterwards every `model.generate(..., use_cache=True)` call that does not
+    pass past_key_values explicitly uses FlashCodecCache. With
+    ``materialization_free=True`` the forward of standard Llama/Qwen-style
+    attention modules is patched too, so decode steps attend on the stored rows
+    directly (no reconstructed KV view). Returns model;
+    `remove_flash_codec(model)` restores it at any time."""
+    if getattr(model, "_flashcodec", None) is not None:
+        warnings.warn("model is already patched with flashcodec, "
+                      "call remove_flash_codec first")
+        return model
+
+    attns = find_attention_modules(model)
+    caches = []                    # lets tests/diagnostics inspect injected caches
+    factory = _make_factory(mode, block_size, kwargs, caches,
+                            config=getattr(model, "config", None))
+
+    base_generate = model.generate
+    base_forward = model.forward
+
+    @functools.wraps(base_generate)
+    def wrapped_generate(*args, **kw):
+        use_cache = kw.get("use_cache")
+        if (use_cache is None or use_cache) and kw.get("past_key_values") is None:
+            kw["past_key_values"] = factory()
+        return base_generate(*args, **kw)
+
+    @functools.wraps(base_forward)
+    def wrapped_forward(*args, **kw):
+        if kw.get("use_cache") and kw.get("past_key_values") is None:
+            kw["past_key_values"] = factory()
+        return base_forward(*args, **kw)
+
+    model.generate = wrapped_generate
+    model.forward = wrapped_forward
+    n_wrapped = 0
+    if materialization_free:
+        from .patch_attention import apply_codec_attention
+        n_wrapped = apply_codec_attention(model)
+    model._flashcodec = {
+        "mode": mode, "caches": caches, "attention_modules": attns,
+        "codec_attention_modules": n_wrapped,
+        "undo": lambda: _restore(model, base_generate, base_forward)}
+    if not attns:
+        warnings.warn("no attention module with k_proj/v_proj found — if the "
+                      "model uses a non-standard KV interface, the patch may "
+                      "have no effect")
+    return model
+
+
+def _restore(model, base_generate, base_forward):
+    from .patch_attention import remove_codec_attention
+    remove_codec_attention(model)
+    model.generate = base_generate
+    model.forward = base_forward
+    model._flashcodec = None
+
+
+def remove_flash_codec(model):
+    """Undo the patch applied by apply_flash_codec (restores both generate and
+    forward)."""
+    info = getattr(model, "_flashcodec", None)
+    if info is not None:
+        info["undo"]()
+
+
+undo_flash_codec = remove_flash_codec          # legacy alias

flashcodec/patch_attention.py

@@ -0,0 +1,151 @@
+"""Route decode steps through the coefficient-domain attention kernel.
+
+`apply_flash_codec` swaps the cache, but transformers still computes attention
+on the reconstructed KV view. This patch replaces the forward of *standard*
+Llama/Qwen-style attention modules (q_proj / k_proj / v_proj / o_proj, optional
+q_norm / k_norm) so that, for a single-token decode step on a codec layer, the
+projection + RoPE part runs as usual and the attention itself is computed by
+`codec_decode_attention` — the (S, D) reconstructed view is never built.
+
+Scope / guards (anything else falls back to the original forward):
+- single-token query, batch size 1, eval mode
+- the cache layer must be a ``FlashCodecLayer`` (i.e. a full-attention layer)
+- no sliding-window attention
+- ``position_embeddings`` must be provided (5.x standard signature)
+
+The fallback is silent by design (mixed models have linear-attention layers that
+never hit this path); `codec_attention_stats` reports how many layers are wrapped.
+"""
+
+import importlib
+
+import torch
+
+from .attention import codec_decode_attention
+from .cache import FlashCodecLayer
+
+_ROPE_CACHE = {}
+
+
+def _rope_fn(module):
+    """The `apply_rotary_pos_emb` of the model file that defines this module."""
+    mod_name = type(module).__module__
+    if mod_name not in _ROPE_CACHE:
+        mod = importlib.import_module(mod_name)
+        _ROPE_CACHE[mod_name] = getattr(mod, "apply_rotary_pos_emb", False)
+    return _ROPE_CACHE[mod_name] or None
+
+
+def _attention_forward(module, orig, hidden_states,
+                       position_embeddings=None, attention_mask=None,
+                       past_key_values=None, **kwargs):
+    layer = None
+    if past_key_values is not None and hasattr(past_key_values, "layers"):
+        try:
+            layer = past_key_values.layers[module.layer_idx]
+        except (IndexError, TypeError):
+            layer = None
+
+    usable = (
+        isinstance(layer, FlashCodecLayer)
+        and not module.training
+        and hidden_states.shape[0] == 1          # batch 1: no padding mask
+        and hidden_states.shape[-2] == 1         # decode step only
+        and position_embeddings is not None
+        and getattr(module, "sliding_window", None) is None
+        and _rope_fn(module) is not None
+    )
+    if not usable:
+        return orig(hidden_states, position_embeddings=position_embeddings,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values, **kwargs)
+
+    input_shape = hidden_states.shape[:-1]
+    hidden_shape = (*input_shape, -1, module.head_dim)
+    o_in = module.o_proj.in_features
+    gated = module.q_proj.out_features == 2 * o_in      # Qwen3.5-style Q gate
+    q_raw = module.q_proj(hidden_states)
+    gate = None
+    if gated:
+        q, gate = torch.chunk(
+            q_raw.view(*input_shape, -1, module.head_dim * 2), 2, dim=-1)
+        gate = gate.reshape(*input_shape, -1)
+    else:
+        q = q_raw
+    q = q.view(hidden_shape)
+    k = module.k_proj(hidden_states).view(hidden_shape)
+    v = module.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    # Qwen3-style QK-norm is applied on (heads, head_dim) before the transpose
+    if hasattr(module, "q_norm"):
+        q = module.q_norm(q)
+    if hasattr(module, "k_norm"):
+        k = module.k_norm(k)
+    q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+    cos, sin = position_embeddings
+    q, k = _rope_fn(module)(q, k, cos, sin)
+
+    if not layer.is_initialized:
+        layer.lazy_initialization(k, v)
+    # prefill left a full reconstructed view behind; decode does not need it
+    for st in (layer._key_stream, layer._value_stream):
+        if st._rec is not None:
+            st.discard_view()
+    layer.keys = layer.values = None            # drop the prefill view reference
+    layer._key_stream.append(k)                  # append only: no read()
+    layer._value_stream.append(v)
+    out = codec_decode_attention(q, layer._key_stream, layer._value_stream,
+                                 scale=module.scaling)
+    out = out.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+    if gated:
+        out = out * torch.sigmoid(gate)
+    return module.o_proj(out), None
+
+
+def apply_codec_attention(model):
+    """Patch every supported attention module so decode steps use the
+    coefficient-domain kernel. Idempotent; returns the number of wrapped
+    modules."""
+    wrapped = 0
+    for module in model.modules():
+        if not (hasattr(module, "q_proj") and hasattr(module, "k_proj")
+                and hasattr(module, "v_proj") and hasattr(module, "o_proj")
+                and hasattr(module, "layer_idx") and hasattr(module, "head_dim")):
+            continue
+        # only the standard layout: head_dim divides the output width, and the
+        # query projection is either plain (Llama/Qwen3) or gated (Qwen3.5)
+        o_in = module.o_proj.in_features
+        if module.head_dim <= 0 or o_in % module.head_dim:
+            continue
+        if module.q_proj.out_features not in (o_in, 2 * o_in):
+            continue
+        if getattr(module, "_flashcodec_orig_forward", None) is not None:
+            wrapped += 1
+            continue
+        orig = module.forward
+
+        def forward(*args, _module=module, _orig=orig, **kw):
+            return _attention_forward(_module, _orig, *args, **kw)
+
+        module._flashcodec_orig_forward = orig
+        module.forward = forward
+        wrapped += 1
+    return wrapped
+
+
+def remove_codec_attention(model):
+    """Restore the original forwards. Returns the number restored."""
+    restored = 0
+    for module in model.modules():
+        orig = getattr(module, "_flashcodec_orig_forward", None)
+        if orig is not None:
+            module.forward = orig
+            del module._flashcodec_orig_forward
+            restored += 1
+    return restored
+
+
+def codec_attention_stats(model):
+    """Number of attention modules currently wrapped."""
+    return sum(1 for m in model.modules()
+               if getattr(m, "_flashcodec_orig_forward", None) is not None)