TensorCache ⚡

Ultra-fast, high-fidelity block-wise INT8 feature & pixel cache engine for PyTorch.

tensorcache eliminates two bottlenecks:

1. Feature Cache Bloat: AMO-BQ asymmetric MSE-optimal G32 1.09B 1.83x vs BF16 0.47% rel RMSE (sym 1.06B 0.54%) — near G16 floor 0.39%.
2. JPEG/PNG CPU Decode: Zero-copy mmap + GPU stream prefetch >2,000 MB/s, ring-buffer 6.8MB VRAM batch8 128x768.

🚀 Key Features

• AMO-BQ (Asymmetric MSE-Optimal, G32): min-max + zp uint8 + 48× clipping search [0.95,1.10] — 0.47% hetero 0.54% vs sym 0.74% (-26%), per-token 1.73% (cache\dtype_comparison.json). Presets fast/balanced/accurate/max.
• Microsecond Dequant: Triton (q-zp)*scale -> BF16 0.06ms 5.4M 337GB/s (codec.py:40), FusedDequantLinear 0 intermediate fused_ops.py:119.
• Minimal VRAM: q 5.22MB + scales 0.32MB + zp 0.16MB + out 10.45MB 5.4M; Streamer fixed ~43MB double-buffer out_buffer reuse, G64 halves scales/zp.
• Cross-Platform: CUDA/ROCm Triton else PyTorch fallback, Windows mmap safe close().
• CLI + Python one-liners: tc.compress / tc.benchmark_tensor / tensorcache benchmark.

📦 Installation

pip install tensor-cache          # PyPI (import tensorcache)
pip install -e .          # dev
pip install -e ".[fast]"  # triton+zstd+blosc2 (Linux)

⚡ Quick Start

1. In-Memory (one-liners)

import torch, tensorcache as tc

x = torch.randn(16,446,768, dtype=torch.bfloat16, device="cuda")

# AMO-BQ presets: fast (16,0.95-1.05) 6.9ms 0.49%, balanced (32,0.95-1.05) 13ms 0.478% (default), accurate (48,0.95-1.10) 49ms 0.473%
q,s,zp,shape = tc.compress(x, mode="balanced")  # or "fast"/"accurate"/"max"/"sym"/"adaptive"
rec = tc.decompress(q,s,shape,zp)               # <0.1ms BF16
tc.benchmark_tensor(x)                          # rich table
tc.estimate_compression(x.shape, group_size=32) # 1.09375 B 1.83x
tc.auto_select_mode(x, target_rmse=0.5)         # -> "balanced"
tc.help()                                       # python help

# Codec object
codec = tc.BlockwiseInt8Codec(group_size=32, amo_bq=True, amo_mode="balanced")
print(codec) # G=32, amo_bq=balanced 1.0938B 1:1.83x

2. Feature Cache to Disk (mmap)

import tensorcache as tc
from torch.utils.data import DataLoader

# Write (amo_bq, G32 default balanced, G64 for minimal VRAM 1.046B 0.55%)
writer = tc.FeatureCacheWriter("./cache/dinov3", 10000,446,768, group_size=32, amo_bq=True, amo_mode="balanced")
writer.append(x) # [seq,dim] or [B,seq,dim]
writer.close()   # writes _int8.bin (uint8) _scales.bin _zp.bin _meta.json

# Load
ds = tc.FeatureCacheDataset("./cache/dinov3") # -> (q uint8, s BF16, zp uint8)
ds = tc.FeatureCacheDataset("./cache/dinov3", auto_dequant_device="cuda") # -> BF16 directly
for q,s,zp in tc.AsyncGPUPrefetcher(DataLoader(ds,batch_size=256,pin_memory=True), device="cuda"):
    batch = tc.dequantize_int8_amo_bq(q,s,zp, shape, group_size=32)

# Minimal VRAM streamer (fixed ~43MB, zero alloc)
streamer = tc.ZeroCopyTensorStreamer("./cache/dinov3", batch_size=32, device="cuda")
for batch in streamer: # BF16 [B,seq,dim] from ring buffer out_bf16_0/1
    train(batch)
streamer.close()

# Fused head (no BF16 intermediate)
from tensorcache import FusedDequantLinear
head = FusedDequantLinear(768, num_classes, group_size=32).cuda()
logits = head(q, s) # dequant+GEMM in regs

3. CLI

python -m tensorcache info
python -m tensorcache benchmark --shape 16,446,768 --device cuda
python -m tensorcache cache-info --prefix ./cache/dinov3
python -m tensorcache compress-demo --shape 4,197,768 --mode balanced
tensorcache --help

📊 Benchmark (DINOv3 ViT-Base, 5.4M randn + RSNA hetero)

📜 License

Apache 2.0 — see LICENSE.