bigsmall 1.0.0__py3-none-any.whl

bigsmall/__init__.py

@@ -0,0 +1,45 @@
+"""BigSmall - Lossless neural network weight compression.
+
+Public API:
+    bigsmall.compress(src, dst, mode="balanced")          - compress safetensors -> .bs
+    bigsmall.decompress(src, dst=None) -> dict[str, ndarray]
+    bigsmall.load(src, device="cpu") -> dict[str, torch.Tensor]
+    bigsmall.info(src) -> dict
+    bigsmall.verify(src) -> bool
+    bigsmall.compress_delta(finetune, base, dst, mode="balanced")
+    bigsmall.decompress_delta(delta_src, base_src, dst=None) -> dict[str, ndarray]
+
+HuggingFace Hub round-trip (Phase 4):
+    bigsmall.compress_for_hub(source, output_dir)   - compress any HF model
+    bigsmall.upload_to_hub(output_dir, repo_id)     - push to the Hub
+    bigsmall.from_pretrained(repo_or_path)          - download + decompress -> state_dict
+    bigsmall.install_hook()                         - monkey-patch safetensors.load_file
+
+Streaming loader (Phase 4 cont.):
+    bigsmall.StreamingLoader(path, device="cuda")   - layer-by-layer decompression
+"""
+__version__ = "1.0.0"
+
+from .encoder import compress, compress_delta
+from .decoder import decompress, decompress_delta, load
+from .verify import verify
+from .container import info
+from .hub import compress_for_hub, upload_to_hub, from_pretrained
+from .integrations.huggingface import install_hook
+from .streaming import StreamingLoader
+
+__all__ = [
+    "compress",
+    "decompress",
+    "load",
+    "info",
+    "verify",
+    "compress_delta",
+    "decompress_delta",
+    "compress_for_hub",
+    "upload_to_hub",
+    "from_pretrained",
+    "install_hook",
+    "StreamingLoader",
+    "__version__",
+]

bigsmall/cli.py

@@ -0,0 +1,124 @@
+"""BigSmall command-line interface."""
+import argparse
+import sys
+import time
+from pathlib import Path
+
+
+def _cmd_compress(args):
+    from . import encoder
+    src = Path(args.src)
+    if args.output:
+        dst = Path(args.output)
+    else:
+        dst = src.with_suffix(".bs")
+    mode = "balanced"
+    if args.storage:
+        mode = "storage"
+    elif args.inference:
+        mode = "inference"
+
+    t0 = time.perf_counter()
+    if args.base:
+        encoder.compress_delta(src, args.base, dst, mode=mode)
+    else:
+        encoder.compress(src, dst, mode=mode)
+    elapsed = time.perf_counter() - t0
+    src_size = src.stat().st_size
+    dst_size = dst.stat().st_size
+    pct = (dst_size / src_size * 100) if src_size > 0 else 0
+    print(f"compressed {src} -> {dst}", flush=True)
+    print(f"  source:     {src_size:,} bytes", flush=True)
+    print(f"  compressed: {dst_size:,} bytes ({pct:.2f}%)", flush=True)
+    print(f"  saved:      {src_size - dst_size:,} bytes", flush=True)
+    print(f"  elapsed:    {elapsed:.1f}s", flush=True)
+
+
+def _cmd_decompress(args):
+    from . import decoder
+    src = Path(args.src)
+    if args.output:
+        dst = Path(args.output)
+    else:
+        dst = src.with_suffix(".safetensors")
+
+    t0 = time.perf_counter()
+    if args.base:
+        decoder.decompress_delta(src, args.base, dst)
+    else:
+        decoder.decompress(src, dst)
+    elapsed = time.perf_counter() - t0
+    print(f"decompressed {src} -> {dst}  ({elapsed:.1f}s)", flush=True)
+
+
+def _cmd_info(args):
+    from .container import info
+    i = info(args.src)
+    for k, v in i.items():
+        print(f"  {k:24s} {v}")
+
+
+def _cmd_verify(args):
+    from .verify import verify
+    ok = verify(args.src, source_safetensors=args.source)
+    if ok:
+        print("OK", flush=True)
+        sys.exit(0)
+    print("FAIL", flush=True)
+    sys.exit(1)
+
+
+def _cmd_benchmark(args):
+    from . import encoder, decoder
+    src = Path(args.src)
+    dst = src.with_suffix(".bs")
+    print(f"Benchmarking {src.name}...")
+    t0 = time.perf_counter(); encoder.compress(src, dst); te = time.perf_counter() - t0
+    t0 = time.perf_counter(); _ = decoder.decompress(dst); td = time.perf_counter() - t0
+    src_size = src.stat().st_size
+    dst_size = dst.stat().st_size
+    pct = (dst_size / src_size * 100) if src_size > 0 else 0
+    print(f"  encode:     {te:.1f}s  ({src_size / te / 1024 / 1024:.1f} MiB/s)")
+    print(f"  decode:     {td:.1f}s  ({src_size / td / 1024 / 1024:.1f} MiB/s)")
+    print(f"  ratio:      {pct:.2f}% ({src_size:,} -> {dst_size:,})")
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="bigsmall", description="BigSmall lossless NN weight compression")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    c = sub.add_parser("compress", help="Compress a .safetensors file to .bs")
+    c.add_argument("src")
+    c.add_argument("-o", "--output", default=None)
+    c.add_argument("--base", default=None, help="Base safetensors path - enables delta mode")
+    grp = c.add_mutually_exclusive_group()
+    grp.add_argument("--storage", action="store_true", help="Maximum compression mode")
+    grp.add_argument("--balanced", action="store_true", help="Balanced ratio+speed (default)")
+    grp.add_argument("--inference", action="store_true", help="Fastest decode mode")
+    c.set_defaults(func=_cmd_compress)
+
+    d = sub.add_parser("decompress", help="Decompress a .bs file to .safetensors")
+    d.add_argument("src")
+    d.add_argument("-o", "--output", default=None)
+    d.add_argument("--base", default=None, help="Base file path for delta decompression")
+    d.set_defaults(func=_cmd_decompress)
+
+    i = sub.add_parser("info", help="Show metadata for a .bs file")
+    i.add_argument("src")
+    i.set_defaults(func=_cmd_info)
+
+    v = sub.add_parser("verify", help="Verify md5 round-trip of a .bs file")
+    v.add_argument("src")
+    v.add_argument("--source", default=None, help="Compare against original .safetensors")
+    v.set_defaults(func=_cmd_verify)
+
+    b = sub.add_parser("benchmark", help="Encode/decode benchmark for a model")
+    b.add_argument("src")
+    b.set_defaults(func=_cmd_benchmark)
+
+    args = p.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

bigsmall/codecs/__init__.py

@@ -0,0 +1,19 @@
+"""BigSmall codecs - per-format encoder/decoder implementations."""
+from . import bf16, fp32, fp16, fp8, fp4, special, generic
+
+CODEC_REGISTRY = {
+    "bf16_se_ac":  bf16,
+    "fp16_se_ac":  fp16,
+    "fp32_se_ac":  fp32,
+    "fp8_cat_ac":  fp8,
+    "fp4_cat_ac":  fp4,
+    "special":     special,
+    "zstd":        generic,
+    "blosc2_shuffle_zstd": generic,
+}
+
+
+def get_codec(name):
+    if name not in CODEC_REGISTRY:
+        raise KeyError(f"Unknown codec: {name}")
+    return CODEC_REGISTRY[name]

bigsmall/codecs/bf16.py

@@ -0,0 +1,170 @@
+"""BF16 codec: per-tensor (sign,exp) joint AC + per-tensor (mantissa | exp) AC.
+
+BF16 layout: 1 sign bit | 8 exp bits | 7 mantissa bits = 16 bits total.
+
+Encoding strategy (mirrors cc10_v2.py but operates per-tensor so it works on
+ANY safetensors model - no GPT-2-specific handling here):
+
+  1. Split each weight into (sign, exp, mantissa).
+  2. Encode (sign, exp) jointly as a 9-bit alphabet per tensor with Categorical AC.
+  3. Encode mantissa per-tensor, sorted by exp, with one Cat AC bucket per exp value.
+
+This per-tensor approach generalises cleanly. The cc10_v2 script encoded the
+mantissa stream globally across all non-special tensors; that gives a tiny
+extra ratio improvement on GPT-2 but requires a fixed 'special tensor set' to
+be known across the encode/decode path. Per-tensor encoding is simpler,
+generalises, and the per-tensor overhead is negligible relative to the AC
+codeword cost.
+"""
+import struct
+import io
+import numpy as np
+import constriction as c
+
+# Bit layout for BF16
+SIGN_SHIFT = 15
+EXP_SHIFT = 7
+EXP_MASK = 0xFF
+MANT_MASK = 0x7F
+SE_ALPHABET = 512    # 2 sign * 256 exp
+MANT_ALPHABET = 128  # 7 mantissa bits
+
+
+def _encode_cat(values: np.ndarray, alphabet: int) -> tuple[bytes, np.ndarray, np.ndarray]:
+    """Encode an int array with per-symbol Categorical AC.
+
+    Returns (codeword_bytes, nonzero_indices, nonzero_freqs).
+    """
+    fp = np.bincount(values, minlength=alphabet).astype(np.int64)
+    nz_idx = np.nonzero(fp)[0]
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    enc = c.stream.queue.RangeEncoder()
+    enc.encode(values.astype(np.int32), m)
+    cw = enc.get_compressed().tobytes()
+    return cw, nz_idx, fp[nz_idx].astype(np.int64)
+
+
+def _decode_cat(cw_bytes: bytes, nz_idx: np.ndarray, freqs: np.ndarray, alphabet: int, n: int) -> np.ndarray:
+    fp = np.zeros(alphabet, dtype=np.int64)
+    fp[nz_idx] = freqs
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    cw = np.frombuffer(cw_bytes, dtype=np.uint32)
+    dec = c.stream.queue.RangeDecoder(cw)
+    return dec.decode(m, n)
+
+
+def encode(raw: bytes) -> tuple[bytes, dict]:
+    """Encode a single BF16 tensor's raw bytes.
+
+    Args:
+        raw: little-endian bytes of the tensor (length must be even)
+
+    Returns:
+        (compressed_blob, extras_dict)
+
+    The extras dict contains nothing - the codec is self-describing inside the blob.
+    """
+    if len(raw) % 2 != 0:
+        raise ValueError(f"BF16 tensor byte length must be even, got {len(raw)}")
+    u16 = np.frombuffer(raw, dtype=np.uint16)
+    n = len(u16)
+    if n == 0:
+        return b"", {}
+
+    sign = ((u16 >> SIGN_SHIFT) & 1).astype(np.uint16)
+    exp = ((u16 >> EXP_SHIFT) & EXP_MASK).astype(np.uint16)
+    mant = (u16 & MANT_MASK).astype(np.uint16)
+    se = ((sign << 8) | exp).astype(np.int32)
+
+    # SE block
+    se_cw, se_nz_idx, se_freqs = _encode_cat(se, SE_ALPHABET)
+
+    # Sort mantissa by exp -> per-exp buckets
+    order_e = np.argsort(exp, kind="stable")
+    mant_sorted = mant[order_e]
+    exp_sorted = exp[order_e]
+    counts = np.bincount(exp_sorted, minlength=256)
+    nonzero_exps = np.nonzero(counts)[0].astype(np.int32)
+    bstart = np.zeros(257, dtype=np.int64)
+    bstart[1:] = np.cumsum(counts)
+
+    # Encode mantissa per nonzero-exp bucket
+    m_buf = io.BytesIO()
+    m_buf.write(struct.pack("<I", len(nonzero_exps)))
+    m_buf.write(nonzero_exps.astype(np.uint16).tobytes())
+    for ev in nonzero_exps:
+        bs = bstart[ev]
+        be = bstart[ev + 1]
+        bucket = mant_sorted[bs:be].astype(np.int32)
+        cw, nz_idx, freqs = _encode_cat(bucket, MANT_ALPHABET)
+        m_buf.write(struct.pack("<IB", be - bs, len(nz_idx)))
+        m_buf.write(nz_idx.astype(np.uint8).tobytes())
+        m_buf.write(freqs.astype(np.uint32).tobytes())
+        m_buf.write(struct.pack("<I", len(cw)))
+        m_buf.write(cw)
+    m_blob = m_buf.getvalue()
+
+    # Container blob
+    out = io.BytesIO()
+    out.write(struct.pack("<I", n))                         # tensor weight count
+    out.write(struct.pack("<H", len(se_nz_idx)))            # SE nonzero count
+    out.write(se_nz_idx.astype(np.uint16).tobytes())        # SE nonzero indices
+    out.write(se_freqs.astype(np.uint32).tobytes())         # SE freqs
+    out.write(struct.pack("<I", len(se_cw)))                # SE codeword length
+    out.write(se_cw)
+    out.write(struct.pack("<I", len(m_blob)))               # M block length
+    out.write(m_blob)
+    return out.getvalue(), {}
+
+
+def decode(blob: bytes, extras: dict, n_weights: int) -> bytes:
+    """Decode a BF16 blob back to raw bytes."""
+    if n_weights == 0 or len(blob) == 0:
+        return b""
+    inp = io.BytesIO(blob)
+    n, = struct.unpack("<I", inp.read(4))
+    if n != n_weights:
+        raise ValueError(f"BF16 decode: weight count mismatch ({n} vs {n_weights})")
+    if n == 0:
+        return b""
+
+    se_n_nz, = struct.unpack("<H", inp.read(2))
+    se_nz_idx = np.frombuffer(inp.read(se_n_nz * 2), dtype=np.uint16).astype(np.int32)
+    se_freqs = np.frombuffer(inp.read(se_n_nz * 4), dtype=np.uint32)
+    se_cw_len, = struct.unpack("<I", inp.read(4))
+    se_cw = inp.read(se_cw_len)
+    se = _decode_cat(se_cw, se_nz_idx, se_freqs, SE_ALPHABET, n).astype(np.uint16)
+    sign = ((se >> 8) & 1).astype(np.uint16)
+    exp = (se & 0xFF).astype(np.uint16)
+
+    m_blob_len, = struct.unpack("<I", inp.read(4))
+    m_inp = io.BytesIO(inp.read(m_blob_len))
+    n_nz_exp, = struct.unpack("<I", m_inp.read(4))
+    nonzero_exps = np.frombuffer(m_inp.read(n_nz_exp * 2), dtype=np.uint16)
+
+    # Reconstruct sort order from decoded exp
+    order_e = np.argsort(exp, kind="stable")
+    counts = np.bincount(exp.astype(np.int64), minlength=256)
+    bstart = np.zeros(257, dtype=np.int64)
+    bstart[1:] = np.cumsum(counts)
+
+    mant_sorted = np.empty(n, dtype=np.uint16)
+    for ev in nonzero_exps:
+        nb, n_nz = struct.unpack("<IB", m_inp.read(5))
+        nz_idx = np.frombuffer(m_inp.read(n_nz), dtype=np.uint8).astype(np.int32)
+        freqs = np.frombuffer(m_inp.read(n_nz * 4), dtype=np.uint32)
+        cw_len, = struct.unpack("<I", m_inp.read(4))
+        cw = m_inp.read(cw_len)
+        mant_sorted[bstart[ev]:bstart[ev + 1]] = _decode_cat(
+            cw, nz_idx, freqs, MANT_ALPHABET, nb
+        ).astype(np.uint16)
+
+    mant = np.empty(n, dtype=np.uint16)
+    mant[order_e] = mant_sorted
+
+    out = ((sign << SIGN_SHIFT) | (exp << EXP_SHIFT) | mant).astype(np.uint16)
+    return out.tobytes()

bigsmall/codecs/fp16.py

@@ -0,0 +1,135 @@
+"""FP16 codec: per-tensor (sign,exp) joint AC + per-tensor (mantissa | exp) AC.
+
+FP16 layout: 1 sign bit | 5 exp bits | 10 mantissa bits = 16 bits total.
+
+Same structure as bf16.py with different alphabet sizes.
+"""
+import struct
+import io
+import numpy as np
+import constriction as c
+
+SIGN_SHIFT = 15
+EXP_SHIFT = 10
+EXP_MASK = 0x1F        # 5 bits
+MANT_MASK = 0x3FF      # 10 bits
+SE_ALPHABET = 64       # 2 sign * 32 exp
+MANT_ALPHABET = 1024   # 10 bits
+
+
+def _encode_cat(values: np.ndarray, alphabet: int):
+    fp = np.bincount(values, minlength=alphabet).astype(np.int64)
+    nz_idx = np.nonzero(fp)[0]
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    enc = c.stream.queue.RangeEncoder()
+    enc.encode(values.astype(np.int32), m)
+    cw = enc.get_compressed().tobytes()
+    return cw, nz_idx, fp[nz_idx].astype(np.int64)
+
+
+def _decode_cat(cw_bytes, nz_idx, freqs, alphabet, n):
+    fp = np.zeros(alphabet, dtype=np.int64)
+    fp[nz_idx] = freqs
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    cw = np.frombuffer(cw_bytes, dtype=np.uint32)
+    dec = c.stream.queue.RangeDecoder(cw)
+    return dec.decode(m, n)
+
+
+def encode(raw: bytes) -> tuple[bytes, dict]:
+    if len(raw) % 2 != 0:
+        raise ValueError(f"FP16 tensor byte length must be even, got {len(raw)}")
+    u16 = np.frombuffer(raw, dtype=np.uint16)
+    n = len(u16)
+    if n == 0:
+        return b"", {}
+
+    sign = ((u16 >> SIGN_SHIFT) & 1).astype(np.uint16)
+    exp = ((u16 >> EXP_SHIFT) & EXP_MASK).astype(np.uint16)
+    mant = (u16 & MANT_MASK).astype(np.uint16)
+    se = ((sign << 5) | exp).astype(np.int32)
+
+    se_cw, se_nz_idx, se_freqs = _encode_cat(se, SE_ALPHABET)
+
+    order_e = np.argsort(exp, kind="stable")
+    mant_sorted = mant[order_e]
+    exp_sorted = exp[order_e]
+    counts = np.bincount(exp_sorted, minlength=32)
+    nonzero_exps = np.nonzero(counts)[0].astype(np.int32)
+    bstart = np.zeros(33, dtype=np.int64)
+    bstart[1:] = np.cumsum(counts)
+
+    m_buf = io.BytesIO()
+    m_buf.write(struct.pack("<I", len(nonzero_exps)))
+    m_buf.write(nonzero_exps.astype(np.uint8).tobytes())
+    for ev in nonzero_exps:
+        bs = bstart[ev]; be = bstart[ev + 1]
+        bucket = mant_sorted[bs:be].astype(np.int32)
+        cw, nz_idx, freqs = _encode_cat(bucket, MANT_ALPHABET)
+        m_buf.write(struct.pack("<IH", be - bs, len(nz_idx)))
+        m_buf.write(nz_idx.astype(np.uint16).tobytes())
+        m_buf.write(freqs.astype(np.uint32).tobytes())
+        m_buf.write(struct.pack("<I", len(cw)))
+        m_buf.write(cw)
+    m_blob = m_buf.getvalue()
+
+    out = io.BytesIO()
+    out.write(struct.pack("<I", n))
+    out.write(struct.pack("<B", len(se_nz_idx)))
+    out.write(se_nz_idx.astype(np.uint8).tobytes())
+    out.write(se_freqs.astype(np.uint32).tobytes())
+    out.write(struct.pack("<I", len(se_cw)))
+    out.write(se_cw)
+    out.write(struct.pack("<I", len(m_blob)))
+    out.write(m_blob)
+    return out.getvalue(), {}
+
+
+def decode(blob: bytes, extras: dict, n_weights: int) -> bytes:
+    if n_weights == 0 or len(blob) == 0:
+        return b""
+    inp = io.BytesIO(blob)
+    n, = struct.unpack("<I", inp.read(4))
+    if n != n_weights:
+        raise ValueError(f"FP16 decode: weight count mismatch ({n} vs {n_weights})")
+    if n == 0:
+        return b""
+
+    se_n_nz, = struct.unpack("<B", inp.read(1))
+    se_nz_idx = np.frombuffer(inp.read(se_n_nz), dtype=np.uint8).astype(np.int32)
+    se_freqs = np.frombuffer(inp.read(se_n_nz * 4), dtype=np.uint32)
+    se_cw_len, = struct.unpack("<I", inp.read(4))
+    se_cw = inp.read(se_cw_len)
+    se = _decode_cat(se_cw, se_nz_idx, se_freqs, SE_ALPHABET, n).astype(np.uint16)
+    sign = ((se >> 5) & 1).astype(np.uint16)
+    exp = (se & 0x1F).astype(np.uint16)
+
+    m_blob_len, = struct.unpack("<I", inp.read(4))
+    m_inp = io.BytesIO(inp.read(m_blob_len))
+    n_nz_exp, = struct.unpack("<I", m_inp.read(4))
+    nonzero_exps = np.frombuffer(m_inp.read(n_nz_exp), dtype=np.uint8)
+
+    order_e = np.argsort(exp, kind="stable")
+    counts = np.bincount(exp.astype(np.int64), minlength=32)
+    bstart = np.zeros(33, dtype=np.int64)
+    bstart[1:] = np.cumsum(counts)
+
+    mant_sorted = np.empty(n, dtype=np.uint16)
+    for ev in nonzero_exps:
+        nb, n_nz = struct.unpack("<IH", m_inp.read(6))
+        nz_idx = np.frombuffer(m_inp.read(n_nz * 2), dtype=np.uint16).astype(np.int32)
+        freqs = np.frombuffer(m_inp.read(n_nz * 4), dtype=np.uint32)
+        cw_len, = struct.unpack("<I", m_inp.read(4))
+        cw = m_inp.read(cw_len)
+        mant_sorted[bstart[ev]:bstart[ev + 1]] = _decode_cat(
+            cw, nz_idx, freqs, MANT_ALPHABET, nb
+        ).astype(np.uint16)
+
+    mant = np.empty(n, dtype=np.uint16)
+    mant[order_e] = mant_sorted
+    out = ((sign << SIGN_SHIFT) | (exp << EXP_SHIFT) | mant).astype(np.uint16)
+    return out.tobytes()

bigsmall/codecs/fp32.py

@@ -0,0 +1,119 @@
+"""FP32 codec: per-tensor (sign,exp) joint AC + per-tensor (mantissa | exp) AC.
+
+FP32 layout: 1 sign | 8 exp | 23 mantissa = 32 bits total.
+
+The 23-bit mantissa cannot be efficiently AC-encoded as a single 8M-symbol
+alphabet, so we encode mantissa per-(exp) bucket using 3 byte-streams
+(low, mid, high) compressed with zstd. This mirrors cc10_v2 byte-transpose
+intuition while staying simple.
+
+Strategy per tensor:
+  1. Split into sign(1), exp(8), mant_lo8 + mant_mid8 + mant_hi7.
+  2. Encode (sign, exp) as 9-bit alphabet (512) per tensor with Cat AC.
+  3. Compress mant_lo / mant_mid / mant_hi as three byte streams with zstd L9
+     (the high 7 bits of mantissa have lower entropy than the low 16).
+"""
+import struct
+import io
+import numpy as np
+import constriction as c
+import zstandard as zstd
+
+SE_ALPHABET = 512  # 2 sign * 256 exp
+
+
+def _encode_cat(values: np.ndarray, alphabet: int):
+    fp = np.bincount(values, minlength=alphabet).astype(np.int64)
+    nz_idx = np.nonzero(fp)[0]
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    enc = c.stream.queue.RangeEncoder()
+    enc.encode(values.astype(np.int32), m)
+    cw = enc.get_compressed().tobytes()
+    return cw, nz_idx, fp[nz_idx].astype(np.int64)
+
+
+def _decode_cat(cw_bytes, nz_idx, freqs, alphabet, n):
+    fp = np.zeros(alphabet, dtype=np.int64)
+    fp[nz_idx] = freqs
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    cw = np.frombuffer(cw_bytes, dtype=np.uint32)
+    dec = c.stream.queue.RangeDecoder(cw)
+    return dec.decode(m, n)
+
+
+def encode(raw: bytes) -> tuple[bytes, dict]:
+    if len(raw) % 4 != 0:
+        raise ValueError(f"FP32 tensor byte length must be % 4 == 0, got {len(raw)}")
+    u32 = np.frombuffer(raw, dtype=np.uint32)
+    n = len(u32)
+    if n == 0:
+        return b"", {}
+
+    sign = ((u32 >> 31) & 1).astype(np.uint16)
+    exp = ((u32 >> 23) & 0xFF).astype(np.uint16)
+    mant = (u32 & 0x7FFFFF).astype(np.uint32)
+    se = ((sign << 8) | exp).astype(np.int32)
+
+    se_cw, se_nz_idx, se_freqs = _encode_cat(se, SE_ALPHABET)
+
+    # Split mantissa into low 8, mid 8, high 7 bits and zstd-compress each
+    # (these 3 byte-streams approximate per-byte entropy of the mantissa)
+    mant_lo = (mant & 0xFF).astype(np.uint8)
+    mant_mid = ((mant >> 8) & 0xFF).astype(np.uint8)
+    mant_hi = ((mant >> 16) & 0x7F).astype(np.uint8)
+
+    cctx = zstd.ZstdCompressor(level=9)
+    blob_lo = cctx.compress(mant_lo.tobytes())
+    blob_mid = cctx.compress(mant_mid.tobytes())
+    blob_hi = cctx.compress(mant_hi.tobytes())
+
+    out = io.BytesIO()
+    out.write(struct.pack("<I", n))
+    out.write(struct.pack("<H", len(se_nz_idx)))
+    out.write(se_nz_idx.astype(np.uint16).tobytes())
+    out.write(se_freqs.astype(np.uint32).tobytes())
+    out.write(struct.pack("<I", len(se_cw)))
+    out.write(se_cw)
+    out.write(struct.pack("<I", len(blob_lo)))
+    out.write(blob_lo)
+    out.write(struct.pack("<I", len(blob_mid)))
+    out.write(blob_mid)
+    out.write(struct.pack("<I", len(blob_hi)))
+    out.write(blob_hi)
+    return out.getvalue(), {}
+
+
+def decode(blob: bytes, extras: dict, n_weights: int) -> bytes:
+    if n_weights == 0 or len(blob) == 0:
+        return b""
+    inp = io.BytesIO(blob)
+    n, = struct.unpack("<I", inp.read(4))
+    if n != n_weights:
+        raise ValueError(f"FP32 decode: weight count mismatch ({n} vs {n_weights})")
+    if n == 0:
+        return b""
+
+    se_n_nz, = struct.unpack("<H", inp.read(2))
+    se_nz_idx = np.frombuffer(inp.read(se_n_nz * 2), dtype=np.uint16).astype(np.int32)
+    se_freqs = np.frombuffer(inp.read(se_n_nz * 4), dtype=np.uint32)
+    se_cw_len, = struct.unpack("<I", inp.read(4))
+    se_cw = inp.read(se_cw_len)
+    se = _decode_cat(se_cw, se_nz_idx, se_freqs, SE_ALPHABET, n).astype(np.uint32)
+    sign = ((se >> 8) & 1).astype(np.uint32)
+    exp = (se & 0xFF).astype(np.uint32)
+
+    blob_lo_len, = struct.unpack("<I", inp.read(4)); blob_lo = inp.read(blob_lo_len)
+    blob_mid_len, = struct.unpack("<I", inp.read(4)); blob_mid = inp.read(blob_mid_len)
+    blob_hi_len, = struct.unpack("<I", inp.read(4)); blob_hi = inp.read(blob_hi_len)
+    dctx = zstd.ZstdDecompressor()
+    mant_lo = np.frombuffer(dctx.decompress(blob_lo), dtype=np.uint8).astype(np.uint32)
+    mant_mid = np.frombuffer(dctx.decompress(blob_mid), dtype=np.uint8).astype(np.uint32)
+    mant_hi = np.frombuffer(dctx.decompress(blob_hi), dtype=np.uint8).astype(np.uint32)
+    mant = mant_lo | (mant_mid << 8) | (mant_hi << 16)
+
+    out = ((sign << 31) | (exp << 23) | mant).astype(np.uint32)
+    return out.tobytes()

bigsmall/codecs/fp4.py

@@ -0,0 +1,76 @@
+"""FP4 codec: per-tensor Categorical AC on 4-bit indices (alphabet 16).
+
+The codec operates on UNPACKED FP4 (one byte per 4-bit value). The encoder
+takes raw bytes representing already-unpacked 4-bit indices in the low nibble.
+"""
+import struct
+import io
+import numpy as np
+import constriction as c
+
+
+def encode(raw: bytes) -> tuple[bytes, dict]:
+    """Encode FP4 unpacked stream (low nibble of each byte)."""
+    if len(raw) == 0:
+        return b"", {}
+    u8 = np.frombuffer(raw, dtype=np.uint8)
+    # Defensive: ensure values are in [0, 15]
+    if u8.max() > 15:
+        # Caller passed packed FP4 - unpack
+        n_total = len(u8) * 2
+        unp = np.empty(n_total, dtype=np.uint8)
+        unp[0::2] = u8 & 0x0F
+        unp[1::2] = (u8 >> 4) & 0x0F
+        u8 = unp
+        was_packed = True
+    else:
+        was_packed = False
+    n = len(u8)
+    vals = u8.astype(np.int32)
+    fp = np.bincount(vals, minlength=16).astype(np.int64)
+    nz_idx = np.nonzero(fp)[0]
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    enc = c.stream.queue.RangeEncoder()
+    enc.encode(vals, m)
+    cw = enc.get_compressed().tobytes()
+
+    out = io.BytesIO()
+    out.write(struct.pack("<IBB", n, len(nz_idx), 1 if was_packed else 0))
+    out.write(nz_idx.astype(np.uint8).tobytes())
+    out.write(fp[nz_idx].astype(np.uint32).tobytes())
+    out.write(struct.pack("<I", len(cw)))
+    out.write(cw)
+    return out.getvalue(), {"was_packed": was_packed}
+
+
+def decode(blob: bytes, extras: dict, n_weights: int) -> bytes:
+    """Decode FP4 to unpacked stream (low nibble of each byte).
+
+    Note: n_weights is interpreted as the unpacked count.
+    If extras['was_packed'], the original was packed; we re-pack on output.
+    """
+    if n_weights == 0 or len(blob) == 0:
+        return b""
+    inp = io.BytesIO(blob)
+    n, n_nz, was_packed_flag = struct.unpack("<IBB", inp.read(6))
+    nz_idx = np.frombuffer(inp.read(n_nz), dtype=np.uint8)
+    freqs = np.frombuffer(inp.read(n_nz * 4), dtype=np.uint32)
+    cw_len, = struct.unpack("<I", inp.read(4))
+    cw_bytes = inp.read(cw_len)
+    fp = np.zeros(16, dtype=np.int64)
+    fp[nz_idx] = freqs
+    probs = fp.astype(np.float64) + 0.01
+    probs /= probs.sum()
+    m = c.stream.model.Categorical(probs, perfect=True)
+    cw = np.frombuffer(cw_bytes, dtype=np.uint32)
+    dec = c.stream.queue.RangeDecoder(cw)
+    out = dec.decode(m, n).astype(np.uint8)
+
+    if was_packed_flag:
+        if n % 2 != 0:
+            raise ValueError("FP4 packed decode requires even unpacked count")
+        packed = (out[0::2] | (out[1::2] << 4)).astype(np.uint8)
+        return packed.tobytes()
+    return out.tobytes()