MoR-TurboQuant: Recursion-Aware Compressed KV Cache

A PyTorch module bridging Mixture of Recursions (MoR) adaptive compute with TurboQuant-style KV cache quantization. Tokens that exit early skip KV allocation entirely; the KV entries that do exist are stored as bit-packed 3-bit codes (Walsh-Hadamard rotation + per-group asymmetric quantization).

The two savings are orthogonal, so they multiply:

early exit (fewer KV entries)  x  3-bit packing (smaller entries)  =  combined
        ~2.7x                  x          4.57x                    =   ~12x

Why "real storage" matters

PyTorch has no 3-bit dtype. A quantize→dequantize pipeline that stores results in fp16 tensors reports compression analytically while allocating the same (or more) VRAM — an early version of this repo had exactly that bug, plus dense cache allocation with masking for exited tokens. Hardware measurement showed ~1.0x. Both are fixed in mor_tq/cache/:

• Bit packing — 8 three-bit codes are packed into 3 uint8 bytes (pack3/unpack3, exact roundtrip). With one fp16 scale + zero per group of 64, effective storage is 3.5 bits/value = 4.57x vs fp16.
• Sparse allocation — each recursion depth stores only its active tokens plus an int32 position map. Exited tokens consume zero bytes (skipped, not masked).
• Honest measurement — cache_bytes() sums element_size() * nelement() over every tensor the cache owns. No formulas.

Results

Measured cache bytes (tensor storage, device-agnostic)

Config: 8 heads, head_dim 64 (d_model=512), 8 recursions, expert-choice routing (capacity 0.5), 3-bit + WHT, group=64. Reproduce with python benchmarks/bench_cache_bytes.py.

Compression depends on the router's active fraction: a strict halving schedule keeps 24.9% of token-depth slots (18.2x); the router trained on WikiText-103 kept 37.3% (≈12.5x). Match a measured router with --active-fraction 0.373.

Hardware validation (torch.cuda.memory_allocated() deltas isolating cache growth) is in benchmarks/bench_decode_memory.py — run it on a GPU box and report both numbers.

Language modeling (WikiText-103, 2 epochs, batch 64, seq 256)

Caveats (read before citing): the MoR model has ~22% more parameters than the baseline (router + shared-block bookkeeping), so part of the PPL gap is parameter count, not the method — iso-parameter baselines are on the roadmap. Results are small-scale (35M params, 2 epochs) and demonstrate the mechanism, not SOTA quality. Per-entry 3-bit compression is comparable to published KV quantization work (KIVI, KVQuant); the contribution here is the recursion-aware combination — quantizing a cache whose size already depends on learned exit depths — and the measured interaction between the two.

Architecture

Token Embeddings
      │
      ▼
┌─────────────────────────┐
│   Recursive Block (Φ)   │◄── Same weights, applied N times
│   Attention + FFN       │
└─────────┬───────────────┘
          ▼
┌─────────────────────────┐
│   Adaptive Router       │── g_t = σ(θᵀ · h_t)
└─────┬─────────┬─────────┘
   continue    exit ──► freeze h_t, skip KV (zero bytes)
      ▼
┌──────────────────────────────────────────┐
│ PackedKVCache (per-depth buckets)        │
│  • bit-packed 3-bit codes (uint8)        │
│  • fp16 scale+zero per group of 64       │
│  • int32 position map, active tokens only│
└──────────────────────────────────────────┘

Quick start

import torch
from mor_tq.cache import PackedKVCache

cache = PackedKVCache(n_recursions=8, n_heads=8, head_dim=64,
                      group_size=64, use_wht=True)

# Inside the model, per recursion step d:
#   active_idx = router.active_indices(d)            # tokens still alive
#   cache.append(d, k[active_idx], v[active_idx], positions=active_idx)

# At attention time (returned tensors are TEMPORARY — never store them back):
K, V, pos = cache.gather(depth=0)

print(cache.cache_bytes())   # actual allocated bytes
print(cache.stats())         # tokens per depth, total entries

No retraining is required to adopt the packed cache — it is inference-side storage for an already-trained checkpoint.

Repository layout

mor_tq/
├── recursive_block.py        # weight-shared transformer block
├── router.py                 # adaptive depth router (token/expert choice)
├── model.py                  # full MoR model
├── config.py                 # configuration dataclass
└── cache/
    ├── quantizer.py          # pack3/unpack3, WHT, GroupQuantizer3bit
    └── packed_kv_cache.py    # PackedKVCache, StandardKVCache
benchmarks/
├── bench_cache_bytes.py      # real allocated bytes (CPU or GPU)
└── bench_decode_memory.py    # CUDA memory_allocated() validation
tests/
└── test_packed_cache.py      # packing roundtrip, fidelity, sparse bytes
train.py                      # WikiText-103 training

Install & test

pip install -e ".[dev]"
pytest tests/ -v
python benchmarks/bench_cache_bytes.py
python benchmarks/bench_decode_memory.py   # GPU box

Roadmap

• Iso-parameter and iso-FLOP baselines for the PPL comparison
• KIVI / KVQuant baselines at matched compression
• Fused Triton dequant-attention kernel (throughput; memory claim already holds)
• Quality-vs-compression curves at 2/3/4 bits across exit depths

References

• Mixture of Recursions (MoR) — Google DeepMind
• TurboQuant paper — KV cache quantization
• KIVI — 2-bit asymmetric KV quantization
• KVQuant — long-context KV quantization
• HIGGS — Hadamard-based weight quantization

License

MIT