Pokémon TCG AI Battle — Hierarchical KV Cache

Purpose

The trainer holds row groups from the day-partitioned Parquet corpus (see pokemon_tcg_parquet_dataset) resident across microbatches so a TBPTT lane touching the same row group across chunks does not pay the decode cost every time. The cache is called "KV cache" by naming convention only — it is a page cache over parquet row groups, not an attention-side KV cache.

Terminology (row_group, hot tier, transient tier, SSD spill, in_opt_step, microbatch, opt_step) is defined in pokemon_tcg_glossary.

Three-tier design

   +-----------------------------+
   |    HOT tier (pinned)        |    RAM, promoted on ≥ 2 hits.
   |    ≤ HOT_ZONE * capacity    |    never evicted while capacity holds.
   +-----------------------------+
                 ▲   promotion @ HOT_PROMOTION_HITS
                 |
   +-----------------------------+
   |    TRANSIENT tier (LRU)     |    RAM, LRU under pressure.
   |    remaining RAM slots      |    demoted from hot? no — evicted straight.
   +-----------------------------+
                 ▲   miss
                 |
   +-----------------------------+
   |    SSD spill (.npz)         |    disk, per training run.
   |    <checkpoint>/.cache_spill|    written on RAM eviction.
   +-----------------------------+
                 ▲   miss
                 |
   Parquet row_group (pyarrow decode)

Fixed policy constants live at the top of _ParquetRowGroupCache:

_HIGH_WATERMARK_PCT = 85.0    # psutil.virtual_memory().percent triggers eviction
_HOT_PROMOTION_HITS = 2       # entries reach hot after this many hits in transient
_HOT_ZONE_FRACTION  = 0.6     # 60% of resident capacity reserved for hot

Cache instances

The trainer creates two independent instances so the val loop cannot evict what the train loop still needs, and vice versa:

• _tbptt_row_group_cache — training path. Wrapped by in_opt_step() around forward+backward.
• _val_row_group_cache — validation path. Streamed only (never materialized to RAM). Its SSD spill dir lives at <checkpoint_dir>/.cache_spill/val/ so the train and val spill sets do not stomp each other.

Both caches carry the same set of counters (see below) and both are reported at end of run.

Pressure-driven eviction

Eviction is not scheduled — it is triggered by memory pressure. Each cache access probes psutil.virtual_memory().percent and, if above _HIGH_WATERMARK_PCT=85%, evicts the LRU entry from the transient tier before proceeding. On darwin/M-series the percent metric tracks the same signal Activity Monitor's "Memory Pressure" band uses, so a value of 85 corresponds roughly to the yellow-to-red transition.

The hot tier is never evicted by the transient LRU pressure loop — hot entries only leave when the entire cache is torn down at end of run.

LFU-style promotion

When a transient entry gets a hit:

1. Increment its hit counter in _hits_by_key.
2. If hits_by_key[key] >= _HOT_PROMOTION_HITS and the hot tier has room, move the entry from transient to hot (drops it from the transient OrderedDict; inserts into the hot OrderedDict). Reset its counter.
3. Emit a promotion counter increment for the run report.

Hot capacity is _HOT_ZONE_FRACTION * resident_capacity (currently 60% of total). The remaining slots stay in transient. In practice hot fills up to a small stable working set (2–3 row groups on the largest ablation runs), and transient recycles through the rest.

SSD spill tier

When a transient entry is evicted under pressure, its decoded numpy arrays are pickled into <checkpoint_dir>/.cache_spill/<key>.npz before being dropped from RAM. On the next miss for that key, the SSD file is read back and re-hydrated into transient — cheaper than decoding the source Parquet again.

Spill files are per-training-run and are cleaned by the driver script after training finishes:

rm -rf "$OUT_DIR/.cache_spill"

This cleanup is essential — a 10-run ablation without it accumulated ~70 GiB across nine .cache_spill/ dirs (25 GiB alone for suite_5d_1ep_ON) and filled the disk mid-epoch-1 of the 10th run. See pokemon_tcg_bc_curriculum_ablation for the incident writeup.

in_opt_step() — reentrant eviction suppression

Without protection, a memory spike inside the forward or backward pass (activation buffers materialize, MLX peak grows) can trip the 85% watermark while the current microbatch is still reading from a resident row group — evicting the very block being used and forcing a re-decode on the next slice, thrashing under load.

The cache exposes a reentrant context manager that suppresses eviction during a critical section:

@contextmanager
def in_opt_step(self):
    with self._lock:
        self._opt_step_depth += 1
    try:
        yield
    finally:
        with self._lock:
            self._opt_step_depth = max(0, self._opt_step_depth - 1)

The trainer wraps every optimizer_step(_accum_grads, _accum_examples) call in _tbptt_row_group_cache.in_opt_step() (or nullcontext() when TBPTT is disabled). This is safe because the working set during an optimizer step is bounded: capacity was already sized to fit it. Evictions still happen outside the critical section on the next fetch.

Counters (tensorboard-visible)

Every access updates a small counter set. The trainer emits the whole set per optimizer step to tensorboard under cache_train/* and cache_val/*, and prints the final tally per cache at end of run:

• hits / misses / hit_rate_pct — self-explanatory
• promotions — transient → hot moves
• evictions — transient → SSD spill events
• ssd_hits / ssd_spills / ssd_resident — the disk tier's own hits, writes, current resident count
• resident_hot / resident_transient — current entry counts in each RAM tier
• bytes_loaded — total decoded bytes drawn from Parquet (excludes SSD reads)

Sample from the end of a 5d_10ep_OFF run:

[cache-train] hits=85705 misses=5 hit_rate=100.0%
  resident=5 rg (hot=3 / transient=2) promotions=916
  evictions=1 ssd_hits=1 ssd_spills=1 ssd_resident=1 decoded=7151.2 MiB
[cache-val]   hits=9585 misses=5 hit_rate=99.9%
  resident=5 rg (hot=3 / transient=2) promotions=413
  evictions=0 ssd_hits=0 ssd_spills=0 ssd_resident=0 decoded=7151.4 MiB

Interpretation: the training loop touched five distinct row groups the first time (misses=5) and served every subsequent access from cache. Only one eviction event happened over the whole run — good sign that the hot tier absorbs the working set. Val's 0 evictions is the expected shape (val is a fixed episode set, small and stable).

Threading

The cache is thread-safe: _lock guards every mutation (touch, evict, promote, spill, re-hydrate). Decoded arrays are the expensive part, so the decode path releases the lock before calling pyarrow to avoid stalling other threads during I/O. The MLX training loop currently uses a single ThreadPoolExecutor prefetch worker; the cache is designed to handle N workers correctly.

What lives in _ParquetRowGroupCache

For each cached row group, the value stored in the tier is a dict of column-name → numpy array, projected to only the columns the caller requested when the cache was created (train projection ≠ val projection, hence two caches with different column sets):

class _ParquetRowGroupCache:
    def __init__(self, file_paths, columns, shapes, int_keys, *, ssd_spill_dir=None):
        self._transient: OrderedDict[Key, dict] = OrderedDict()
        self._hot: OrderedDict[Key, dict] = OrderedDict()
        self._hits_by_key: dict[Key, int] = {}
        self._ssd_keys: set[Key] = set()
        self._opt_step_depth = 0
        ...

Key = (file_idx, row_group_idx). columns is the projection list. shapes and int_keys are metadata used to reconstruct fixed_size_list columns as N-dim numpy arrays.

Rules of thumb

• On the M3 Pro 24 GB, steady-state resident capacity for a 30k rows/day × 5 days run is 5 row groups (3 hot + 2 transient). Do not try to grow beyond ~10 without measuring host memory pressure.
• If hits/misses after warmup is < 90%, either the projection is fetching too much or the working set doesn't fit in RAM.
• If evictions climbs high while ssd_hits stays 0, the working set is churning and SSD spill is dead weight — either raise resident_capacity or reduce column projection.
• Always delete .cache_spill/ between training runs. The suite script does this automatically; ad-hoc runs must do it manually.

Provenance / anchor commits

• e772fe0 (2026-08-06) — train loader: KV-style parquet cache, async prefetch, top-elo/per-day filters, tensorboard — introduced the first-generation flat single-tier cache with FIFO eviction, plus a single-worker ThreadPoolExecutor prefetch.
• fa38caa (2026-08-06) — train loader: hierarchical KV cache (hot/transient + SSD spill), opt-step protection, val streaming — replaced the flat cache with the current three-tier design, added the reentrant in_opt_step() suppression, and moved val to stream through its own cache instance (eliminating the ~22 GiB val materialization).
• e8d6a4a (2026-08-07) — experiments: purge KV cache SSD spill after each training run — established the rule that .cache_spill/ is training-scoped storage and must be deleted between runs; landed the driver-script cleanup after the disk-fill incident.

Full timeline in pokemon_tcg_repository_timeline.

Cross-references

• pokemon_tcg_glossary — cache-specific vocabulary
• pokemon_tcg_parquet_dataset — the row groups the cache holds
• pokemon_tcg_training_pipeline — where in_opt_step() is called
• pokemon_tcg_bc_curriculum_ablation — SSD-spill disk-fill incident and fix