lifeboatcache 0.1.0__py3-none-any.whl

lifeboat_cache/__init__.py

@@ -0,0 +1,7 @@
+"""Lifeboat Cache — commercial LMCache plugin (L2 storage, serde, policies).
+
+See the project brief / docs/ for architecture. This package is import-safe on a
+CPU-only Mac (the "fleet of size one") and on GPU fleets alike.
+"""
+
+__version__ = "0.0.0"

lifeboat_cache/compress/__init__.py

@@ -0,0 +1,57 @@
+"""CPU reference compression layer for Lifeboat Cache.
+
+A set of KV-cache codecs behind one :class:`Codec` interface, driven by the
+``compression`` config section. Every codec is pure-numpy (torch optional) and produces a
+**self-describing blob** so ``decode`` needs no external arguments.
+
+Codecs (each approximates a published method; see the module docstrings for exactly what is
+reproduced vs. approximated):
+
+* ``turboquant`` — group-wise affine int3/int4 with float16-kept layers.
+* ``fp8``        — e4m3 / e5m2 emulation via an exact FP8 value grid.
+* ``oscar``      — 2-bit middle + full-precision sink/recent tokens + optional rotation.
+* ``kvtc``       — PCA + quantized coefficients/residual + zlib (ANS stand-in).
+
+Importing this package registers all four codec factories, so
+:func:`build_codecs` can instantiate any ladder from config.
+"""
+from __future__ import annotations
+
+from .accuracy import cosine, logit_kl, mse, tensor_relative_error
+from .base import (
+    Codec,
+    build_codecs,
+    create_codec,
+    pack_arrays,
+    pack_blob,
+    register_codec,
+    unpack_arrays,
+    unpack_blob,
+)
+from .calibration import CalibrationProfile
+from .demote import recompress
+from .fp8 import FP8Codec
+from .kvtc import KVTCCodec
+from .oscar2 import Oscar2Codec
+from .turboquant import TurboQuantCodec
+
+__all__ = [
+    "Codec",
+    "register_codec",
+    "create_codec",
+    "build_codecs",
+    "pack_blob",
+    "unpack_blob",
+    "pack_arrays",
+    "unpack_arrays",
+    "TurboQuantCodec",
+    "FP8Codec",
+    "Oscar2Codec",
+    "KVTCCodec",
+    "CalibrationProfile",
+    "recompress",
+    "logit_kl",
+    "cosine",
+    "mse",
+    "tensor_relative_error",
+]

lifeboat_cache/compress/accuracy.py

@@ -0,0 +1,53 @@
+"""Accuracy metrics for the compression accuracy-guard.
+
+The guard samples a fraction of writes, decodes them, and compares against the original to
+decide whether a codec is degrading quality (config ``compression.accuracy_guard.metric``).
+The headline metric is softmax KL on logits — the distributional distance that actually
+matters for generation quality — plus cheaper tensor-space helpers.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+_EPS = 1e-12
+
+
+def _log_softmax(z: np.ndarray) -> np.ndarray:
+    z = z - z.max(axis=-1, keepdims=True)
+    return z - np.log(np.exp(z).sum(axis=-1, keepdims=True) + _EPS)
+
+
+def logit_kl(p_logits: np.ndarray, q_logits: np.ndarray) -> float:
+    """Mean KL(softmax(p) || softmax(q)) over the last axis.
+
+    Zero iff the two logit tensors induce the same distribution (identical logits, or logits
+    differing only by a per-row additive constant), strictly positive otherwise. Computed in
+    log-space for numerical stability.
+    """
+    lp = _log_softmax(np.asarray(p_logits, dtype=np.float64))
+    lq = _log_softmax(np.asarray(q_logits, dtype=np.float64))
+    p = np.exp(lp)
+    kl = (p * (lp - lq)).sum(axis=-1)
+    return float(np.maximum(kl, 0.0).mean())
+
+
+def cosine(a: np.ndarray, b: np.ndarray) -> float:
+    """Mean cosine similarity in [-1, 1] over flattened vectors (1.0 == identical)."""
+    a = np.asarray(a, dtype=np.float64).reshape(-1)
+    b = np.asarray(b, dtype=np.float64).reshape(-1)
+    denom = (np.linalg.norm(a) * np.linalg.norm(b)) + _EPS
+    return float(np.dot(a, b) / denom)
+
+
+def mse(a: np.ndarray, b: np.ndarray) -> float:
+    """Mean squared error between two arrays."""
+    a = np.asarray(a, dtype=np.float64)
+    b = np.asarray(b, dtype=np.float64)
+    return float(np.mean((a - b) ** 2))
+
+
+def tensor_relative_error(a: np.ndarray, b: np.ndarray) -> float:
+    """Relative reconstruction error ``||a - b|| / (||a|| + eps)`` (Frobenius norm)."""
+    a = np.asarray(a, dtype=np.float64)
+    b = np.asarray(b, dtype=np.float64)
+    return float(np.linalg.norm(a - b) / (np.linalg.norm(a) + _EPS))

lifeboat_cache/compress/base.py

@@ -0,0 +1,229 @@
+"""Codec base class, self-describing blob format, and the codec registry.
+
+Every codec in this package is a *CPU reference* implementation: it uses only numpy so
+the whole compression layer runs on a laptop with no CUDA. The on-GPU production codecs
+(real TurboQuant / OSCAR / KVTC kernels) would plug in behind the same :class:`Codec`
+interface — these reference versions exist to validate the format, the round-trip math,
+and the ladder/demote plumbing.
+
+Blob self-description
+---------------------
+``decode`` takes *only* the blob — no shape, dtype, or param arguments. We achieve that
+with a tiny length-prefixed JSON header followed by a raw binary payload::
+
+    u32 little-endian  json_len
+    json_len bytes     UTF-8 JSON header (scalar params + an ``_arrays`` manifest)
+    rest               concatenated raw array bytes (order = manifest order)
+
+The ``_arrays`` manifest records each stored numpy array's ``dtype`` and ``shape`` so the
+payload can be sliced back apart deterministically (dicts preserve insertion order, and
+JSON preserves it too). This keeps each codec's serialization declarative: build a dict of
+numpy arrays + a dict of scalar params, and the helpers do the rest.
+"""
+from __future__ import annotations
+
+import json
+import struct
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable
+
+import numpy as np
+
+if TYPE_CHECKING:  # avoid a hard import cycle; only needed for type hints
+    from .calibration import CalibrationProfile
+
+
+# --------------------------------------------------------------------------- #
+# blob (de)serialization
+# --------------------------------------------------------------------------- #
+def pack_blob(header: dict[str, Any], payload: bytes) -> bytes:
+    """Serialize a JSON ``header`` + raw ``payload`` into one self-describing blob."""
+    hb = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    return struct.pack("<I", len(hb)) + hb + payload
+
+
+def unpack_blob(blob: bytes) -> tuple[dict[str, Any], bytes]:
+    """Inverse of :func:`pack_blob` — return (header dict, payload bytes)."""
+    (n,) = struct.unpack_from("<I", blob, 0)
+    header = json.loads(blob[4 : 4 + n].decode("utf-8"))
+    return header, blob[4 + n :]
+
+
+def pack_arrays(arrays: dict[str, np.ndarray]) -> tuple[dict[str, Any], bytes]:
+    """Concatenate ``arrays`` to bytes and return a manifest describing each one.
+
+    The manifest (dtype + shape per key) goes in the JSON header; the bytes become the
+    payload. Order is preserved so :func:`unpack_arrays` can slice them back out.
+    """
+    meta: dict[str, Any] = {}
+    parts: list[bytes] = []
+    for key, arr in arrays.items():
+        a = np.ascontiguousarray(arr)
+        meta[key] = {"dtype": str(a.dtype), "shape": list(a.shape)}
+        parts.append(a.tobytes())
+    return meta, b"".join(parts)
+
+
+def unpack_arrays(meta: dict[str, Any], payload: bytes) -> dict[str, np.ndarray]:
+    """Inverse of :func:`pack_arrays` using the manifest produced there."""
+    out: dict[str, np.ndarray] = {}
+    off = 0
+    for key, m in meta.items():
+        dt = np.dtype(m["dtype"])
+        shape = tuple(int(s) for s in m["shape"])
+        nbytes = int(np.prod(shape)) * dt.itemsize if shape else dt.itemsize
+        buf = payload[off : off + nbytes]
+        out[key] = np.frombuffer(buf, dtype=dt).reshape(shape).copy()
+        off += nbytes
+    return out
+
+
+# --------------------------------------------------------------------------- #
+# bit packing (for sub-byte integer codes)
+# --------------------------------------------------------------------------- #
+def pack_bits(codes: np.ndarray, bits: int) -> np.ndarray:
+    """Pack a 1-D array of small unsigned ints (0..2**bits-1) into a dense bit stream.
+
+    Each code contributes ``bits`` bits (LSB-first within the code); the stream is then
+    byte-packed with :func:`numpy.packbits`. Returned as a uint8 array so it can be stored
+    directly via :func:`pack_arrays`.
+    """
+    codes = codes.astype(np.uint8).reshape(-1)
+    if bits == 8:
+        return codes.copy()
+    # (N, bits) matrix: column j is bit j of each code.
+    shifts = np.arange(bits, dtype=np.uint8)
+    bit_mat = ((codes[:, None] >> shifts) & 1).astype(np.uint8)
+    return np.packbits(bit_mat.reshape(-1))
+
+
+def unpack_bits(packed: np.ndarray, n: int, bits: int) -> np.ndarray:
+    """Inverse of :func:`pack_bits`: recover ``n`` codes each of width ``bits``."""
+    packed = np.asarray(packed, dtype=np.uint8)
+    if bits == 8:
+        return packed[:n].astype(np.uint8).copy()
+    flat = np.unpackbits(packed)[: n * bits].reshape(n, bits)
+    weights = (1 << np.arange(bits)).astype(np.uint16)
+    return (flat.astype(np.uint16) * weights).sum(axis=1).astype(np.uint8)
+
+
+# --------------------------------------------------------------------------- #
+# group-wise affine integer quantization (shared by turboquant / oscar / kvtc)
+# --------------------------------------------------------------------------- #
+def quantize_rows(x: np.ndarray, bits: int, group_size: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Per-group affine quantize a 2-D array along its last axis.
+
+    For each row and each contiguous group of ``group_size`` columns we store a scale and
+    a zero-point (the group minimum): ``q = round((x - min) / scale)`` clipped to
+    ``[0, 2**bits - 1]`` with ``scale = (max - min) / (2**bits - 1)``.
+
+    Because every value lies within its group's ``[min, max]``, no clipping is ever needed
+    and round-to-nearest bounds the reconstruction error by ``scale / 2`` per group — the
+    property the tests assert. A constant group gets ``scale = 1`` and quantizes exactly.
+
+    Returns ``(codes, scales, zeros)`` with ``codes`` shape ``(rows, D)`` uint8 and
+    ``scales``/``zeros`` shape ``(rows, n_groups)`` float32.
+    """
+    rows, D = x.shape
+    qmax = (1 << bits) - 1
+    n_groups = (D + group_size - 1) // group_size
+    codes = np.empty((rows, D), dtype=np.uint8)
+    scales = np.empty((rows, n_groups), dtype=np.float32)
+    zeros = np.empty((rows, n_groups), dtype=np.float32)
+    for g in range(n_groups):
+        s, e = g * group_size, min((g + 1) * group_size, D)
+        blk = x[:, s:e]
+        mn = blk.min(axis=1)
+        mx = blk.max(axis=1)
+        scale = (mx - mn) / qmax
+        # Constant group -> scale 0; set to 1 so (x-min)/scale == 0 and dequant is exact.
+        scale = np.where(scale > 0, scale, 1.0).astype(np.float32)
+        q = np.rint((blk - mn[:, None]) / scale[:, None])
+        codes[:, s:e] = np.clip(q, 0, qmax).astype(np.uint8)
+        scales[:, g] = scale
+        zeros[:, g] = mn.astype(np.float32)
+    return codes, scales, zeros
+
+
+def dequantize_rows(codes: np.ndarray, scales: np.ndarray, zeros: np.ndarray, group_size: int) -> np.ndarray:
+    """Inverse of :func:`quantize_rows`."""
+    rows, D = codes.shape
+    n_groups = scales.shape[1]
+    out = np.empty((rows, D), dtype=np.float32)
+    for g in range(n_groups):
+        s, e = g * group_size, min((g + 1) * group_size, D)
+        out[:, s:e] = codes[:, s:e].astype(np.float32) * scales[:, g : g + 1] + zeros[:, g : g + 1]
+    return out
+
+
+def orthogonal_matrix(dim: int, seed: int) -> np.ndarray:
+    """Deterministic random orthogonal matrix via QR (used for OSCAR rotations).
+
+    The sign of ``Q`` is pinned by the diagonal of ``R`` so the factorization is unique
+    and reproducible from ``seed`` alone — encode and decode regenerate the same matrix
+    with no need to store it.
+    """
+    rng = np.random.default_rng(seed)
+    a = rng.standard_normal((dim, dim))
+    q, r = np.linalg.qr(a)
+    q *= np.sign(np.diag(r))
+    return q.astype(np.float32)
+
+
+# --------------------------------------------------------------------------- #
+# Codec interface
+# --------------------------------------------------------------------------- #
+class Codec(ABC):
+    """A compression codec. Encode a numpy array to a self-describing blob and back.
+
+    ``decode`` must not need any external arguments — everything required to reconstruct
+    (shape, dtype, params, scales, bases, rotation seeds, ...) lives in the blob header.
+    """
+
+    #: Instance name (the key under which the codec appears in config, e.g. ``turboquant4``).
+    name: str = "codec"
+
+    @abstractmethod
+    def encode(self, arr: np.ndarray, *, layer_id: int = 0, profile: "CalibrationProfile | None" = None) -> bytes:
+        """Compress ``arr`` to a blob. ``layer_id``/``profile`` tune per-layer behavior."""
+
+    @abstractmethod
+    def decode(self, blob: bytes) -> np.ndarray:
+        """Reconstruct the array from a blob produced by :meth:`encode`."""
+
+
+# --------------------------------------------------------------------------- #
+# registry
+# --------------------------------------------------------------------------- #
+# Factories are keyed by codec *type* (``turboquant``/``oscar``/``kvtc``/``fp8``) — the
+# ``Codec.type`` field in config. A factory takes a config object (anything exposing the
+# codec params as attributes) and returns a ready Codec instance.
+_REGISTRY: dict[str, Callable[[Any], Codec]] = {}
+
+
+def register_codec(name: str, factory: Callable[[Any], Codec]) -> None:
+    """Register a codec factory under its type ``name`` (idempotent / overwriting)."""
+    _REGISTRY[name] = factory
+
+
+def create_codec(name: str, codec_cfg: Any) -> Codec:
+    """Instantiate the codec registered under type ``name`` from ``codec_cfg``."""
+    try:
+        factory = _REGISTRY[name]
+    except KeyError as exc:  # pragma: no cover - defensive
+        raise KeyError(f"no codec registered for type '{name}' (have {sorted(_REGISTRY)})") from exc
+    return factory(codec_cfg)
+
+
+def build_codecs(compression_cfg: Any) -> dict[str, Codec]:
+    """Instantiate every codec named in ``compression_cfg.codecs``.
+
+    Returns a mapping of config key (e.g. ``turboquant4``) -> Codec, with each codec's
+    ``.name`` set to that key so blobs/telemetry can report which ladder entry produced them.
+    """
+    out: dict[str, Codec] = {}
+    for key, cfg in compression_cfg.codecs.items():
+        codec = create_codec(cfg.type, cfg)
+        codec.name = key
+        out[key] = codec
+    return out

lifeboat_cache/compress/calibration.py

@@ -0,0 +1,81 @@
+"""Per-model calibration profile: cached state shared across encode calls.
+
+A :class:`CalibrationProfile` is fit once per model (from a handful of representative KV
+tensors) and then handed to ``Codec.encode`` to tune per-layer behavior. The reference
+codecs are mostly self-calibrating per blob, so this profile stays deliberately small; it
+carries:
+
+* ``num_layers`` — so TurboQuant can resolve negative ``keep_bf16_layers`` indices.
+* ``state`` — a dict of numpy arrays (e.g. a cached KVTC PCA basis + mean, per-codec scales).
+* ``meta`` — JSON-serializable scalars (e.g. a global scale).
+
+Persistence uses ``.npz`` (arrays) with the JSON metadata embedded as a byte array in the
+same file — no pickle, so a profile file is safe to load from untrusted locations.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Sequence
+
+import numpy as np
+
+_META_KEY = "__meta__"
+_STATE_PREFIX = "state::"
+
+
+@dataclass
+class CalibrationProfile:
+    num_layers: int = 0
+    model_hash: str = ""
+    state: dict[str, np.ndarray] = field(default_factory=dict)
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    def fit(self, samples: "np.ndarray | Sequence[np.ndarray]") -> "CalibrationProfile":
+        """Fit the profile from one or more representative tensors.
+
+        Computes a global magnitude scale and a cached PCA basis (mean + right singular
+        vectors over the last-axis feature vectors) that a KVTC codec could reuse instead
+        of recomputing SVD per blob. ``num_layers`` is inferred from the first sample's
+        leading axis when not already set.
+        """
+        arrs = list(samples) if isinstance(samples, (list, tuple)) else [samples]
+        arrs = [np.asarray(a, dtype=np.float32) for a in arrs]
+        if not arrs:
+            return self
+        if not self.num_layers:
+            self.num_layers = int(arrs[0].shape[0])
+
+        D = arrs[0].shape[-1]
+        X = np.concatenate([a.reshape(-1, D) for a in arrs], axis=0)
+        self.meta["global_scale"] = float(np.abs(X).max())
+        mean = X.mean(axis=0)
+        Xc = X - mean
+        _, _, vt = np.linalg.svd(Xc, full_matrices=False)
+        self.state["kvtc_mean"] = mean.astype(np.float32)
+        self.state["kvtc_basis"] = vt.astype(np.float32)
+        return self
+
+    # -- persistence --
+    def save(self, path: str) -> None:
+        meta = {"num_layers": self.num_layers, "model_hash": self.model_hash, "meta": self.meta}
+        meta_bytes = np.frombuffer(json.dumps(meta).encode("utf-8"), dtype=np.uint8)
+        payload = {_META_KEY: meta_bytes}
+        payload.update({f"{_STATE_PREFIX}{k}": v for k, v in self.state.items()})
+        np.savez(path, **payload)
+
+    @classmethod
+    def load(cls, path: str) -> "CalibrationProfile":
+        with np.load(path, allow_pickle=False) as data:
+            meta = json.loads(bytes(data[_META_KEY]).decode("utf-8"))
+            state = {
+                k[len(_STATE_PREFIX) :]: np.array(data[k])
+                for k in data.files
+                if k.startswith(_STATE_PREFIX)
+            }
+        return cls(
+            num_layers=meta["num_layers"],
+            model_hash=meta["model_hash"],
+            state=state,
+            meta=meta["meta"],
+        )

lifeboat_cache/compress/demote.py

@@ -0,0 +1,16 @@
+"""Recompress-on-demote: move a chunk from one codec to another.
+
+When a chunk is demoted down the storage ladder (e.g. GPU ``turboquant4`` -> archive
+``kvtc``), it is re-encoded into the destination tier's codec. The reference path is simply
+decode-then-encode through numpy; a production implementation might transcode directly in
+the quantized domain, but decode/encode is always correct and is what the ladder relies on.
+"""
+from __future__ import annotations
+
+from .base import Codec
+
+
+def recompress(blob: bytes, from_codec: Codec, to_codec: Codec, *, layer_id: int = 0) -> bytes:
+    """Decode ``blob`` with ``from_codec`` and re-encode it with ``to_codec``."""
+    arr = from_codec.decode(blob)
+    return to_codec.encode(arr, layer_id=layer_id)

lifeboat_cache/compress/fp8.py

@@ -0,0 +1,160 @@
+"""FP8 emulation (e4m3 / e5m2) on CPU via an exact 256-entry value table.
+
+Approximation note
+------------------
+Real hardware stores one byte per element in the OCP FP8 formats. We emulate that exactly
+by enumerating all 256 bit patterns of the chosen format, decoding each to its float value
+(this *is* the FP8 numeric grid, including subnormals), and encoding an input by snapping
+each element to the nearest representable value (``searchsorted`` on the sorted grid).
+Specials (e4m3 ``S.1111.111`` NaN; e5m2 ``exp==all-ones`` inf/NaN) are excluded from the
+candidate set, so finite inputs never map to them. This reproduces true round-to-nearest
+FP8 behavior with no hand-rolled mantissa bit-twiddling.
+
+A per-tensor ``scale`` maps the tensor's peak magnitude to ~half the format's max before
+quantization (and divides back out on decode). FP8's error is relative (the grid is roughly
+log-uniform), so scaling chiefly prevents overflow and lifts tiny tensors out of the
+subnormal region; it does not change the ~relative tolerance.
+
+Tolerance: e4m3 has 3 mantissa bits -> max relative step ~2**-4 = 6.25% per element; e5m2
+has 2 mantissa bits -> ~12.5%. Tensor-level relative error is well under those.
+"""
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from .base import Codec, pack_arrays, pack_blob, register_codec, unpack_arrays, unpack_blob
+
+if TYPE_CHECKING:
+    from .calibration import CalibrationProfile
+
+_FORMATS = {"e4m3": (4, 3, 7), "e5m2": (5, 2, 15)}
+
+
+@functools.lru_cache(maxsize=None)
+def _fp8_grid(fmt: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return (all_values[256], finite_mask[256], sorted_finite_values) for ``fmt``."""
+    ebits, mbits, bias = _FORMATS[fmt]
+    vals = np.empty(256, dtype=np.float32)
+    finite = np.ones(256, dtype=bool)
+    exp_all_ones = (1 << ebits) - 1
+    for b in range(256):
+        s = (b >> 7) & 1
+        exp = (b >> mbits) & ((1 << ebits) - 1)
+        man = b & ((1 << mbits) - 1)
+        is_special = (fmt == "e4m3" and exp == exp_all_ones and man == (1 << mbits) - 1) or (
+            fmt == "e5m2" and exp == exp_all_ones
+        )
+        if is_special:
+            finite[b] = False
+            vals[b] = np.nan
+            continue
+        if exp == 0:  # subnormal (or zero)
+            v = (man / (1 << mbits)) * 2.0 ** (1 - bias)
+        else:  # normal
+            v = (1.0 + man / (1 << mbits)) * 2.0 ** (exp - bias)
+        vals[b] = (-1.0) ** s * v
+    return vals, finite, np.sort(vals[finite]).astype(np.float32)
+
+
+@functools.lru_cache(maxsize=None)
+def _torch_fp8(fmt: str):
+    """(torch, fp8 dtype) when torch has native float8, else None (numpy grid path)."""
+    try:
+        import torch
+    except ImportError:
+        return None
+    dt = {"e4m3": "float8_e4m3fn", "e5m2": "float8_e5m2"}[fmt]
+    return (torch, getattr(torch, dt)) if hasattr(torch, dt) else None
+
+
+class FP8Codec(Codec):
+    def __init__(self, fmt: str = "e4m3", name: str = "fp8") -> None:
+        if fmt not in _FORMATS:
+            raise ValueError(f"fp8 format must be one of {sorted(_FORMATS)}, got {fmt}")
+        self.fmt = fmt
+        self.name = name
+
+    @classmethod
+    def from_cfg(cls, cfg: Any) -> "FP8Codec":
+        # Accept either `format` or `fmt`; default e4m3 (the common KV-cache choice).
+        return cls(fmt=getattr(cfg, "format", getattr(cfg, "fmt", "e4m3")))
+
+    def encode(self, arr: np.ndarray, *, layer_id: int = 0, profile: "CalibrationProfile | None" = None) -> bytes:
+        arr = np.asarray(arr, dtype=np.float32)
+        fast = _torch_fp8(self.fmt)
+        if fast is not None:
+            return self._encode_torch(arr, *fast)
+        vals, finite, sorted_vals = _fp8_grid(self.fmt)
+        max_val = float(np.nanmax(np.abs(vals)))
+
+        amax = float(np.abs(arr).max()) if arr.size else 0.0
+        target = max_val * 0.5
+        scale = amax / target if amax > 0 else 1.0
+
+        x = (arr.reshape(-1) / scale).astype(np.float32)
+        # Nearest representable value: searchsorted gives the right neighbor; compare both.
+        idx = np.searchsorted(sorted_vals, x)
+        idx = np.clip(idx, 1, len(sorted_vals) - 1)
+        left = sorted_vals[idx - 1]
+        right = sorted_vals[idx]
+        choose_left = np.abs(x - left) <= np.abs(x - right)
+
+        # Map chosen values back to their byte codes (value -> byte lookup).
+        finite_bytes = np.nonzero(finite)[0].astype(np.uint8)
+        order = np.argsort(vals[finite])
+        sorted_bytes = finite_bytes[order]
+        chosen_sorted_idx = np.where(choose_left, idx - 1, idx)
+        codes = sorted_bytes[chosen_sorted_idx].astype(np.uint8)
+
+        arr_meta, payload = pack_arrays({"codes": codes})
+        header = {
+            "codec": self.name,
+            "type": "fp8",
+            "fmt": self.fmt,
+            "scale": scale,
+            "shape": list(arr.shape),
+            "dtype": str(arr.dtype),
+            "_arrays": arr_meta,
+        }
+        return pack_blob(header, payload)
+
+    def _encode_torch(self, arr: np.ndarray, torch, fp8_dtype) -> bytes:
+        """Same blob as the grid path (the stored byte *is* the e4m3/e5m2 bit pattern), via
+        torch's native round-to-nearest-even cast — orders of magnitude faster on big chunks.
+        The scale keeps |x| <= max/2, so the cast never saturates."""
+        max_val = float(torch.finfo(fp8_dtype).max)
+        amax = float(np.abs(arr).max()) if arr.size else 0.0
+        scale = amax / (max_val * 0.5) if amax > 0 else 1.0
+        x = torch.from_numpy(arr.reshape(-1)) / scale
+        codes = x.to(fp8_dtype).view(torch.uint8).numpy()
+        arr_meta, payload = pack_arrays({"codes": codes})
+        header = {
+            "codec": self.name,
+            "type": "fp8",
+            "fmt": self.fmt,
+            "scale": scale,
+            "shape": list(arr.shape),
+            "dtype": str(arr.dtype),
+            "_arrays": arr_meta,
+        }
+        return pack_blob(header, payload)
+
+    def decode(self, blob: bytes) -> np.ndarray:
+        header, payload = unpack_blob(blob)
+        arrays = unpack_arrays(header["_arrays"], payload)
+        fast = _torch_fp8(header["fmt"])
+        if fast is not None:
+            torch, fp8_dtype = fast
+            codes = torch.from_numpy(arrays["codes"])
+            recon = (codes.view(fp8_dtype).to(torch.float32) * float(header["scale"])).numpy()
+            return recon.reshape(tuple(header["shape"])).astype(np.dtype(header["dtype"]))
+        vals, _, _ = _fp8_grid(header["fmt"])
+        codes = arrays["codes"]
+        recon = vals[codes.astype(np.intp)] * np.float32(header["scale"])
+        return recon.reshape(tuple(header["shape"])).astype(np.dtype(header["dtype"]))
+
+
+register_codec("fp8", FP8Codec.from_cfg)