---
type: reference
title: "Pokémon TCG AI Battle — Hierarchical KV Cache"
description: "Design and mechanics of the three-tier row-group cache (hot pinned, transient LRU, SSD spill) with pressure-driven eviction, LFU promotion, opt-step protection and separate train/val instances."
tags: [pokemon-tcg, kv-cache, hierarchical, ssd-spill, memory-pressure, parquet, row-group]
timestamp: "2026-08-07T11:00:00-03:00"
---

# 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

```text
   +-----------------------------+
   |    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`:

```python
_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:

```bash
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:

```python
@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:

```text
[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):

```python
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