メインコンテンツまでスキップ

KV Cache Modes

TinyLoop supports five KV cache storage modes. They differ in what they store, how large the cache is in VRAM, and how much extra compute each decoded token pays. This page is the exhaustive reference — if you are choosing a cache mode for a deployment, everything you need to decide should be here.

This page assumes you already have a TinyLoop runtime working (Model.generate, Model.score_last, etc.). For the surrounding API, see Python API.

The five modes

NameEnv var to enableWhat is stored per token per layerQuality cost
FP16 KV (default)Two FP16 vectors: K ([D]), V ([D])None (reference)
INT8 KVTINYLOOP_KV_INT8=1Two INT8 vectors + one FP16 scale per head~0% PPL drift up to ~40 tokens, then slow divergence
FP16 hTINYLOOP_KV_H_MODE=1One FP16 vector: h = LN(residual) ([D])None (bit-exact K/V reconstruction)
INT8 hTINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1One INT8 vector of h + FP16 per-head scales~0.6 % L2-relative K/V error, no measurable attention-argmax drift
INT4 hTINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1One packed-nibble vector of h + FP16 per-head scales~10 % L2-relative K/V error (15 symmetric levels vs INT8's 255)

The three h modes are architecturally exclusive to weight-shared looped transformers. They work because every loop iteration applies the same W_k / W_v, so the K and V for each cached position can be reconstructed on demand as W_k · h / W_v · h. Standard deep transformers cannot do this because each layer has its own W_k.

TINYLOOP_KV_INT4 takes precedence over TINYLOOP_KV_INT8 for h storage — if both are set, INT4-h wins. INT4 without TINYLOOP_KV_H_MODE=1 is silently ignored (there is no INT4 K/V path; INT4 is only supported for the stored h).

Memory math

Let D be the hidden dimension, n_heads the number of attention heads, max_seq_len the cache capacity in tokens, and n_layers = n_pre_blocks + L the total cached-attention layer count. Per token per layer, with D = 2048 and n_heads = 16:

ModeBytes per token per layer formulaBytes at D=2048, n_heads=16
FP16 KV2 × D × 28192
INT8 KV2 × D × 1 + 2 × n_heads × 24160
FP16 hD × 24096
INT8 hD × 1 + n_heads × 22080
INT4 hD / 2 + n_heads × 21056

Multiply by n_layers × max_seq_len for the per-cache total. At max_seq_len=1024 and L=32 (so n_layers=34) the cache footprints are (measured 2026-04-18 on H100):

ModeVRAM (seq=1024, L=32)Savings vs FP16 KV
FP16 KV272 MB
INT8 KV138 MB49 %
FP16 h136 MB + 20 MB scratch = 156 MB43 %
INT8 h68 MB + 24 MB scratch = 92 MB66 %
INT4 h34 MB + 24 MB scratch = 59 MB78 %

The scratch buffers in the three h modes are a constant ~20–24 MB allocation shared across all layers (used to reconstruct K/V at attention time). They are independent of L, so they matter less and less as L grows — at L=64 n_layers=66 the INT4-h saving reaches 82 %, approaching the 87 % theoretical floor where per-layer cost dominates the scratch.

Latency characteristics

Measured on H100 / 407M GPTQ-INT4 / L=8 / prompt=32 / decode=64 (2026-04-18, after the skip-Q + partial-fp16-attn_qkv optimisations):

Store modescalar recon+ partial-attn_qkv+ full fp16_body
FP16 KV (baseline)9.48
FP16 h16.48 (+74 %)9.02 (−5 %)
INT8 h16.48 (+74 %)9.06 (−4 %)
INT4 h16.54 (+74 %)9.09 (−4 %)2.92 (−69 %)
Weight VRAM delta on 407M0+75 MB+406 MB

The recommended h-mode production config is TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 TINYLOOP_H_MODE_FP16_ATTN_QKV=1: matches baseline latency within measurement noise, pays only ~75 MB extra weight VRAM on the 407M (vs 406 MB for full fp16_body), preserves the full 78 % KV cache saving of INT4-h. Net VRAM is positive for any seq_len ≥ ~512 because the KV cache savings exceed the weight overhead. The same flag works for FP16-h (43 % cache saving, bit-exact K/V) and INT8-h (66 % cache saving, ~0.6 % K/V error).

How it works. When h-mode is active, every attention call reconstructs K and V from the cached h via W_k · h / W_v · h — a GEMM against the attn_qkv weight. On a scalar INT4 path this is compute-bound and dominates decode time (+74 % regression). With TINYLOOP_H_MODE_FP16_ATTN_QKV=1, only attn_qkv is pre-dequantised to FP16 at model load, and the reconstruction GEMM runs on cuBLAS tensor cores — the same fast path as full fp16_body, but at a fraction of the VRAM cost (attn_qkv is ~1/5 of a block's weight footprint and there are typically only 3 unique blocks on a weight-shared loop architecture — n_pre_blocks=2 + 1 loop — so on 407M this is 3 × 25 MB = 75 MB).

The earlier +108 % regression was cut to +74 % by a separate 2026-04-18 change: the reconstruction GEMM now skips the Q third of attn_qkv (we only need K, V to rebuild the cache) and a single kernel handles the K/V split, replacing two cudaMemcpy2D calls. That optimisation is memory-cost-free and always on; partial-fp16-attn_qkv stacks on top.

Ephemeral attn_qkv dequant

TINYLOOP_H_MODE_FP16_ATTN_QKV_EPHEMERAL=1 allocates a single shared FP16 scratch (sized 3D × D × 2 bytes = 25 MB on 407M) in the buffer pool and re-dequantises attn_qkv on entry to every reconstruction call. Compared to the persistent variant:

Modems/tokVRAM cost1B projected7B projected
persistent (…FP16_ATTN_QKV=1)9.09 (−4 %)+75 MB+300 MB+1.2 GB
ephemeral (…FP16_ATTN_QKV_EPHEMERAL=1)10.39 (+9 %)+25 MB+100 MB+400 MB

Use ephemeral on larger models where the per-block persistent cache footprint gets expensive (the trade flips at scale: ephemeral saves 200 MB on 1B, 800 MB on 7B, at the same ~10 % decode-latency cost). On 407M the persistent path is typically preferable because +75 MB is negligible and saves 1.3 ms/tok.

The scratch is shared across blocks — dequantised into in place at the start of each block's reconstruction, overwritten by the next. Parity with the INT4 scalar path is preserved because the dequant kernel is the same one that populates fp16_cache when TINYLOOP_EXPERIMENTAL_FP16_BODY=1 is set globally.

Hopper mma.sync INT4 tensor-core kernel (§4 P1) would close the gap with zero memory overhead — matching full-fp16_body latency on INT4 weights directly. Multi-day kernel work; tracked as a parallel track.

Historical measurement on 1B-effective (before skip-Q)

These were captured before the 2026-04-18 skip-Q optimization and are preserved here for the absolute numbers. On the updated code the h-mode rows would be ~17 % faster:

ModeDecode wall timeThroughputvs baseline
FP16 KV (baseline)674 ms93 tok/s
FP16 KV + fp16_body527 ms120 tok/s−22 %
INT8 KV760 ms83 tok/s+13 %
FP16 h (scalar INT4 reconstruction)2117 ms30 tok/s+214 %
INT8 h (scalar INT4 reconstruction)2107 ms30 tok/s+213 %
INT4 h (scalar INT4 reconstruction)2108 ms30 tok/s+213 %
FP16 h + fp16_body415 ms152 tok/s−38 %
INT8 h + fp16_body423 ms149 tok/s−37 %
INT4 h + fp16_body425 ms148 tok/s−37 %

Two distinct regimes for h mode:

  • Without fp16_body, the reconstruction quant_gemm runs on the scalar INT4 kernel and costs ~3× baseline decode. Use this only when VRAM is so tight you cannot afford the FP16 body cache (which triples weight VRAM).
  • With fp16_body, the reconstruction runs through cuBLAS WMMA/WGMMA tensor cores and h mode flips from "+215 %" to "−38 %" — faster than the FP16 KV baseline. This is the recommended production configuration when weight VRAM headroom exists.

Combined with the memory tables above, the production tradeoffs are:

Deployment scenarioRecommended mode
VRAM-tight edge / many-tenant / small GPUFP16 KV baseline, or INT8 KV
Latency-critical server with weight VRAM headroomFP16 h + fp16_body (quality-identical, 43 % KV saving, −38 % decode)
Long-context (>2 k) inference where cache dominates memoryINT4 h + fp16_body (78 % KV saving, matches INT8-h latency, ~10 % K/V L2-rel noise)
Quality-sensitive long-contextINT8 h + fp16_body (66 % saving, ~0.6 % K/V L2-rel noise)
Research / ablation runs exploring h dynamicsAny h mode; score_trajectory gives the hidden state trajectory directly

The fused INT8-h kernel (TINYLOOP_KV_H_FUSED=1) is provided for experimentation but is currently slower than the unfused path on Hopper. It is retained because the idea is correct (skip the scratch write/read, skip the Q third of the projection) and a future vectorised-load rewrite should flip the sign. On Hopper it is superseded by the fp16_body composition above.

When to pick which mode

Default (FP16 KV). Pick this unless you have a concrete reason not to. Lowest latency, zero quality cost, fewest knobs. Every feature in TinyLoop is tested here first.

INT8 KV (TINYLOOP_KV_INT8=1). Pick this when you need to fit roughly twice the context length in the same VRAM budget and can tolerate small quality drift past 40 generated tokens. Latency regression is small (+8 %).

FP16 h (TINYLOOP_KV_H_MODE=1). Pick this when you need ~50 % less KV VRAM and zero quality drift. Decode takes ~3× longer, which typically matters only in interactive UIs. Good fit for batch evaluation on small GPUs.

INT8 h (TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1). Pick this when VRAM is genuinely scarce (edge device, many-tenant serving) and a ~0.6 % L2-relative K/V reconstruction error is acceptable. Gives ~66 % KV saving at the same ~3× decode cost as FP16 h.

INT4 h (TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1). Pick this when VRAM is extremely scarce (long context on consumer GPUs, >50 concurrent tenants) and you're running evaluation / ranking / perplexity workloads, not token-level generation. The 2026-04-18 six-axis stress test (see paper/store_h_stress/PAPER.md) shows:

  • Eval-safe. On 407M GPTQ-INT4 at seq_len up to 4096, INT4-h adds ≤ +0.08 PPL and < 1 % argmax flip rate vs FP16 baseline on WikiText-103 val. On a 10-item multiple-choice benchmark, INT4-h agrees with the FP16 baseline's preferred option 10/10.
  • Generate-divergent. On greedy decoding, INT4-h selects a different coherent attractor than baseline within 1-3 tokens on 75-100 % of tested prompts. The outputs are well-formed but differ; use INT8-h or FP16 KV for production token-level generation.
  • One adversarial pathology. On highly-repetitive inputs ([tok_a, tok_b] * N) INT4-h can collapse into a degenerate 2-token loop. INT8-h does not have this failure mode.

Usage rule of thumb: INT4-h for eval / RAG retrieval / embedding-sum ranking; INT8-h (or FP16 KV) for chat / text generation / code completion.

Fused INT8 h (+TINYLOOP_KV_H_FUSED=1). Do not pick this for production today. It is currently slower than the unfused INT8-h path; kept as an opt-in for kernel development.

Python usage

Enabling a mode

Cache modes are controlled through environment variables read once per process at first cache allocation. Set them before importing tinyloop_py:

import os
os.environ["TINYLOOP_KV_INT8"] = "1" # 49 % KV saving, small latency hit
# or
os.environ["TINYLOOP_KV_H_MODE"] = "1" # 43 % KV saving, zero quality cost, 3× latency
# or INT8-h:
os.environ["TINYLOOP_KV_H_MODE"] = "1"
os.environ["TINYLOOP_KV_INT8"] = "1" # 66 % KV saving, ~0.6 % L2 K/V error
# or INT4-h — aggressive VRAM savings:
os.environ["TINYLOOP_KV_H_MODE"] = "1"
os.environ["TINYLOOP_KV_INT4"] = "1" # 78 % KV saving, ~10 % L2 K/V error

import tinyloop_py as tl
model = tl.Model("model.tinyloop", max_seq_len=2048)

Setting a flag after the first cache has been allocated has no effect for the rest of the process — the env-var reads are cached.

Measuring the active mode

Model.config() does not directly report the active cache mode, but you can observe it via the decode wall clock and the reported memory-controller duty cycle on nvidia-smi. A quick sanity check:

import numpy as np, time
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')})")

If you enabled TINYLOOP_KV_H_MODE=1 and the wall time is close to the FP16 KV baseline, something is misconfigured — h_mode always adds measurable compute.

Invariant: all modes produce the same argmax stream

Under greedy decoding (temperature=0.0) all four 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.
  • 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 few decode steps; the 0.6 % L2-relative K/V reconstruction error may eventually flip an argmax on a close-call logit.

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

Composing with other modes

CombinationSupported?Notes
TINYLOOP_KV_INT8 onlyYes§2 path. Default K/V storage becomes INT8 + per-head FP16 scale.
TINYLOOP_KV_H_MODE onlyYes§16.14.b Phase 1C path. K/V not stored; reconstructed per attention call.
TINYLOOP_KV_H_MODE + TINYLOOP_KV_INT8Yes§16.14.b Phase 1D-INT8. h stored as per-head INT8 + FP16 scales.
TINYLOOP_KV_H_MODE + TINYLOOP_KV_INT4Yes§16.14.b Phase 1G-INT4. h stored as packed signed nibbles + FP16 scales. Takes precedence over TINYLOOP_KV_INT8 when both are set.
Any of the above + TINYLOOP_EXPERIMENTAL_FP16_BODY=1YesRuns body GEMMs through cuBLAS FP16 tensor cores. Costs ~3× weight VRAM, halves or better decode latency.
Any of the above + TINYLOOP_CROSS_ITER_ROUND=1 (§16.9 stoch rounding)Yes for FP16 KV / INT8 KV; h modes fall through to the unfused reconstruction path because the fused kernel only handles plain INT4 weights.
TINYLOOP_KV_H_FUSED=1 (opt-in fused kernel)Yes but slower todayTakes effect only when h_mode + INT8 h are active and the attn_qkv weight is on the plain INT4 path.

End-to-end example: picking the best mode for a batch eval

import os, time, numpy as np

def bench(mode_env):
for k, v in mode_env.items():
os.environ[k] = v
import importlib, sys
if "tinyloop_py" in sys.modules:
# Env vars are cached at first import; a fresh run is needed to compare.
raise RuntimeError("Restart the process to change cache mode")
import tinyloop_py as tl
model = tl.Model("model.tinyloop", max_seq_len=2048)
prompt = np.arange(512, dtype=np.int32)
t0 = time.perf_counter()
r = model.generate_stream(prompt, lambda t: True,
max_tokens=64, loops=8, temperature=0.0)
return time.perf_counter() - t0, r["decode_steady_ms"]

# Call `bench({...})` in separate Python invocations for each mode and
# pick the one whose (wall time, VRAM) pair fits your constraint.

The env-var caching inside TinyLoop means you cannot switch modes in-process. The helper above is therefore structured as a single-shot per-process benchmark; run it once per mode in subprocesses if you want a side-by-side.

Hidden-state trajectory API

When h_mode is active, the per-position hidden state is a first-class cached quantity. If you want to inspect it directly (for example to research iteration-axis compression variants), use Model.score_trajectory, which runs a single forward pass and returns the entire stage-by-stage last-position hidden trajectory without needing a cache.

When not to use any of these

The four modes above are the runtime-level KV compression options. For the orthogonal weight compression (INT2 / INT4 / FP16-body cache / Marlin tensor-core GEMM), see the runtime modes page and the TINYLOOP_EXPERIMENTAL_FP16_BODY section of the CLI reference. Combining weight and KV compression is supported — they are independent knobs.

Known issues and upcoming work

  1. Fused INT8-h kernel (TINYLOOP_KV_H_FUSED=1) is slower than the unfused path (~4 s vs ~2.1 s on the 1B-effective bench). The fused kernel correctness is verified by tests/test_kv_h_mode, and it is kept opt-in for future vectorised-load + shared-memory-scale optimisation. Do not enable it for latency-sensitive workloads.
  2. Tensor-core reconstruction for h_mode is the natural next step for making decode latency competitive. Path: enable TINYLOOP_EXPERIMENTAL_FP16_BODY=1 so the extra attn_qkv projection flows through cuBLAS FP16 matmul; or, once SM_80 hardware is available, wire gptq-marlin into the h_mode reconstruction. See the CHECKLIST §16.14.b Phase 1F entry.
  3. INT8-h sliding-window parity test exists for the FP16-h variant (tests/test_kv_h_mode test 7) but has not been run end-to-end for INT8-h. The wrap path is wired but untested; expect it to work, but do not rely on it for production until a test lands.
  4. INT4-h variant would be another 2× saving on top of INT8-h (memory down to ~13 % of baseline). Tracked as Phase 1G; not yet implemented. Shipped 2026-04-18 — enable with TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1. See the INT4 h row in the modes table above.

See also