KV cache modes from Python

TinyLoop has five KV cache storage modes (default FP16, INT8, store-h FP16, store-h INT8, store-h INT4) selected via environment variables. They trade decode latency against VRAM: INT8 KV saves ~49 % KV memory at a small (~8 %) latency cost, the three h modes save 43–78 % at a much larger (~3×) decode cost under the scalar reconstruction path — or become faster than the FP16 KV baseline with TINYLOOP_EXPERIMENTAL_FP16_BODY=1.

This page covers the Python-side surface. For storage layouts, full memory math, composability, and decision guidance, see the framework-level KV Cache Modes reference.

Quick mode table

Mode	Enable	VRAM saving (at L=32, seq=1024)	K/V noise	Decode vs baseline (+fp16_body)
FP16 KV	(default)	—	0	— (−22 % with fp16_body)
INT8 KV	TINYLOOP_KV_INT8=1	49 %	~0 for ~40 tok, drift after	+13 %
FP16 h	TINYLOOP_KV_H_MODE=1	43 %	0 (bit-exact reconstruction)	−38 %
INT8 h	TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1	66 %	~0.6 % L2-rel	−37 %
INT4 h	TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1	78 %	~10 % L2-rel	−37 %

Setting the mode

Set the env vars before importing tinyloop_py:

import os
os.environ["TINYLOOP_KV_H_MODE"] = "1"      # enable store-h mode
os.environ["TINYLOOP_KV_INT4"]   = "1"      # ... pack h as signed INT4 nibbles
import tinyloop_py as tl
model = tl.Model("model.tinyloop", max_seq_len=2048)

The env vars are probed at first cache-allocation call and cached in a module-level static. Changing them after the first generate / build_prefix_cache / build_resume_handle has no effect — that's why swapping modes in-process is not supported.

Precedence

If both TINYLOOP_KV_INT8 and TINYLOOP_KV_INT4 are set under TINYLOOP_KV_H_MODE=1, INT4 wins — the more aggressive compression takes the h storage. TINYLOOP_KV_INT4=1 without TINYLOOP_KV_H_MODE=1 is silently ignored (there is no INT4 K/V path in TinyLoop; INT4 is only wired for the stored h).

TINYLOOP_KV_INT8=1 without TINYLOOP_KV_H_MODE=1 selects the INT8 K/V path (first row above), completely independent of h mode.

Recommended combos

• Quality-first, small models (≤ 1B), VRAM abundant: leave env vars unset, optionally add TINYLOOP_EXPERIMENTAL_FP16_BODY=1 for the 22 % decode speedup.
• Latency + quality, weight VRAM headroom: TINYLOOP_KV_H_MODE=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 — FP16-h + tensor-core reconstruction is −38 % decode at 43 % KV saving with zero K/V noise.
• Long context, willing to tolerate small noise: TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 — 66 % KV saving at 0.6 % L2-rel noise, same −37 % decode.
• Long context, extremely VRAM-constrained: TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 — 78 % KV saving (37 % smaller than INT8-h) at 10 % L2-rel noise.

Python-visible behaviour

From the Python caller's perspective, the storage mode is transparent: every API on these pages (score_*, generate, prefix cache, warm-start) works under any mode. What changes is:

• VRAM usage — model.vram_usage_mb() reports lower numbers under INT8 / h modes as cache grows.
• Decode latency — visible in generate_stream's decode_steady_ms field.
• Numerical tolerance — lossy modes produce K/V that match the FP16 reference to within their L2-rel bound. Bit-exact decode parity tests only pass under FP16 modes; lossy modes have their own parity tests with mode-specific tolerances.
• Cache construction speed — INT8 / INT4 packing adds a small per-token overhead on the prefill append path (one reduction + one quant per head), negligible next to the GEMMs.

Composition with other APIs

API	h mode	INT8 h	INT4 h	Notes
score, score_last	✓	✓	✓	Output affected only by the reconstruction noise envelope.
score_logit_lens	✓	✓	✓	Each iteration's logit lens uses the cache-mode-appropriate K/V.
score_trajectory	✓	✓	✓	Trajectory captures the FP32 residual regardless of cache mode; the hidden array is cache-independent.
score_with_uncertainty / score_with_consistency	✓	✓	✓	KL / argmax comparison is robust to small K/V noise.
generate	✓	✓	✓	Sampling is on top of the cache's logits — noise can occasionally flip close-call argmaxes under INT4-h.
generate_stream	✓	✓	✓	Same as generate. Timing fields include the reconstruction cost.
generate_speculative	✓	✓	✓	Both draft and verifier run under the same cache mode.
build_prefix_cache / generate_from_prefix_cache	✓	✓	✓	Prefix cache inherits the mode. h modes make the prefix cache ~2× smaller.
build_resume_handle / resume_generate	✓	✓	✓	Resume handle inherits the mode. Bit-exact parity guaranteed under FP16 modes; lossy modes' parity is up to the mode's noise bound.

Inspecting the active mode

Model.config() does not directly report the active cache mode, but you can observe it indirectly:

import os
import time
import numpy as np

prompt = np.arange(64, dtype=np.int32)
t0 = time.perf_counter()
_ = model.generate(prompt, max_tokens=32, loops=8, temperature=0.0)
dt = time.perf_counter() - t0

print(f"{dt * 1000:.1f} ms for 32 tokens (config: "
      f"h_mode={os.environ.get('TINYLOOP_KV_H_MODE', '0')}, "
      f"int8={os.environ.get('TINYLOOP_KV_INT8', '0')}, "
      f"int4={os.environ.get('TINYLOOP_KV_INT4', '0')}, "
      f"fp16_body={os.environ.get('TINYLOOP_EXPERIMENTAL_FP16_BODY', '0')})")

If you enabled TINYLOOP_KV_H_MODE=1 without TINYLOOP_EXPERIMENTAL_FP16_BODY=1 and the wall time is close to the FP16 KV baseline, the env vars did not take effect (likely because they were set after the first Model construction). Restart the process.

Benchmarking modes side-by-side

In-process mode toggling is unsafe (env vars are cached). The correct pattern is subprocess-per-mode:

import os, subprocess, sys

def bench(mode_env, model_path):
    code = f'''
import os
for k, v in {mode_env!r}.items(): os.environ[k] = v
import numpy as np, time, tinyloop_py as tl
model = tl.Model({model_path!r}, max_seq_len=2048)
prompt = np.arange(512, dtype=np.int32)
_ = model.generate(prompt, max_tokens=8, loops=8, temperature=0.0)  # warmup
t0 = time.perf_counter()
r = model.generate_stream(prompt, lambda t: True,
    max_tokens=64, loops=8, temperature=0.0)
print(f"decode_ms={{r['decode_steady_ms']:.2f}}")
'''
    return subprocess.check_output([sys.executable, "-c", code]).decode()

for label, env in [
    ("FP16 KV", {}),
    ("FP16 h + fp16_body", {"TINYLOOP_KV_H_MODE": "1",
                              "TINYLOOP_EXPERIMENTAL_FP16_BODY": "1"}),
    ("INT4 h + fp16_body", {"TINYLOOP_KV_H_MODE": "1",
                              "TINYLOOP_KV_INT4": "1",
                              "TINYLOOP_EXPERIMENTAL_FP16_BODY": "1"}),
]:
    print(f"{label}: {bench(env, '/workspace/model_1b_effective.tinyloop').strip()}")

The ready-made version is tests/bench_kv_h_mode.py, which runs nine configs (three cache modes × two reconstruction paths, plus controls) and prints a summary table.

Invariants and correctness

Under greedy decoding (temperature=0.0) all five modes produce the same token stream for the same prompt on the same model, up to the quality drift noted in the table above:

• FP16 KV and FP16 h are bit-exact: every token id is identical. If you see a divergence, file a bug.
• INT8 KV matches FP16 KV's argmax for roughly the first 40 tokens and may diverge after; both continuations are valid samples.
• INT8 h matches FP16 KV's argmax on the prompt position and the first several decode steps; the 0.6 % L2-relative K/V reconstruction error may eventually flip an argmax on a close-call logit.
• INT4 h matches FP16 KV's argmax on the prompt position in most cases; the 10 % L2-relative K/V reconstruction error flips close-call argmaxes occasionally. Not yet semantic-tested at scale — PPL sweep on 1B WikiText-103 val is the next validation item.

If you see divergence on the first decoded token between FP16 KV and any other mode, that is a bug — please file an issue.

Gotchas

• Order of imports. Set env vars before import tinyloop_py. Setting them after has no effect.
• Mode is process-wide. You cannot have two Model instances in the same process running different cache modes. Use subprocesses.
• head_dim must be even for INT4-h. head_dim = dim / n_heads. The alloc path logs a warning and falls back to FP16-h if this is violated. Current reference 407M / 1B models all satisfy head_dim = 128 (even).
• Scratch is shared across layers. The h modes allocate one set of reconstruction scratch buffers per RuntimeKVCache, not per layer. Memory-accounting expectations should reflect this — at low n_layers the scratch dominates.