KV Cache Design — cells, format, sessions

Scope. The contract for the KV cache: how a run resolves its cells, what the cell store guarantees, what MINFER_CACHE_TYPE is allowed to do, and what a session file contains. The graph-side invariants are in COMPUTE-GRAPH-DESIGN.md; the per-ticket history (C1–C8, #99, #130, #153, #306, #310) is in ARCHITECTURE-EXECUTION-PLAN.md §5.

AGENTS.md states each rule below as a one-line invariant and links here for the elaboration.

1. Positions and cells

  1. KV positions are data, not structure — topology never depends on n_past (precondition for decode reuse). So is the allowed window: attn_span (E1) carries each query's [lo, hi) cell range, resolved by KvCache from its per-sequence span list (position base, first cell, length) — not start + position. A read path refuses a multi-span sequence loudly; a set-valued window (shared prefix + private run) goes through the kv_map input instead (C8b S2/S4, gathered by all three backends — Metal's since #362). A backend expresses its window in one of three modes, selected by the size of the window input: causal positions, one [lo, hi) pair per query, or KV_MAP_MAX_SPANS (cell, len) runs. Op::Attn { explicit_span } marks a node positions cannot bound (more than one sequence, or a window not starting at cell 0); only a backend with supports_attn_span() (CPU, CUDA and — since #44 part (a) landed on a Mac 2026-10-06 — Metal, via kernel_gqa_attn_window_f32/_f16 in src/metal/kernels/attn_window.metal) may take it. Metal reads both explicit layouts — the one-range attn_span since #44 part (a) and the set-valued kv_map since #362, through the sibling kernel_gqa_attn_map_f32/_f16 — and a packed q8_0 KV cache is read on Metal since #310 (mechanism A's native packed decode plus mechanism B's f32 stage, crate::metal::packed_attn_route). The write/move side landed in #44 part (b) on a Mac 2026-10-06: copy_cells (C3 compaction), the copy_kv_to_cpu Metal arm (C2 shift / C5 sessions, f32-only) and the per-engine kv_format. graph/batch.rs composes such batches; forward_cached is its one-sequence case. A store may never write into a prefix a sequence reads in place: production has one fill entry point, GraphAllocator::fill_batch_inputs, and it runs the copy-on-write (GraphAllocator::kv_private_row_for, C8b S3) for every group before the first cell is resolved — the share shrinks to the first position that forward writes and the run's own rows shift up inside it. The kv_cells_for_seq refusal of a position still inside the share is a belt-and-braces invariant check on the store resolver, not a second entry point: production cannot reach it with a shared position, because fill_batch_inputs copied first (E1's test-only fill_attn_inputs and its C1 remnant kv_note_used/own_prefix were deleted in #228). GraphAllocator::kv_cell_of is that same read side, but test-only (#236): production reads a sharing sequence's rows as windows (the kv_map input above) and snapshots the whole arena (kv_save*, C5), and the one consumer is the C8b S3 gate's kv_rows_of in server::batch::tests. The windowed Metal path now has fast prefill families (kernel_flash_attn_window_blk_* for the one-range attn_span, #359; kernel_flash_attn_window_map_* for the set-valued kv_map, #369 — both in src/metal/kernels/fa_window.metal) selected for nt > 1 at hd ∈ {64,128}, alongside the #44/#362 correctness families (which stay byte-untouched and serve the other shapes), so a batched multi-sequence prefill is no longer the slow path — #315 measured the correctness families, #359 and #369 the fast ones.

2. The cell store: ownership, removal, compaction, sharing

  1. Each layer owns two persistent KV regions (K/V) via kv_pair(layer); they survive rebuilds (the allocator lives in GraphCache). kvcache.rs tracks the owner of every cell (C1), can drop a row position in place (C2), can compact the arena (C3), lets a sequence read another's prefix in place (C8b S2) and copies a row out of that prefix the moment a store would land on it (C8b S3). C6 splits position from cell: positions is a token's index within its sequence (what RoPE rotates by); the allocator resolves cells — the row to write — from the run's span list. They coincide only while a run starts at cell 0, which is why the single-sequence path is bitwise unchanged. The fused decode QKV family (Op::FusedQKV, Op::QkvBiasRopeStore) consumes cells too — positions ropes, cells stores ((cuda_on || !explicit_span); Metal keeps the pre-C6 gate, G5) — and whoever feeds a batch must pass sequence-relative positions, never run.start + pos. A physical kv_rm/kv_shift (C2) changes positions, so it still re-ropes; a compaction (C3) changes only cells, so rows move verbatim (Backend::copy_cells: CPU copy_within; CUDA's kv_move_rows — one row at a time with a barrier, ascending when the run slides down and descending when it slides up (C7b), no staging buffer, because overlapping device-to-device cudaMemcpy is undefined; Metal's arm walks one row at a time via MTLBlitCommandEncoder in the same order, since #44 part (b) 2026-10-06) and kv_defrag(need) takes no rope. Whoever copies must use kvcache::order_moves — upward moves top-down, downward bottom-up, upward first — because a destination may never land on a row that has not been copied yet. Whoever caches a run's start (E2's server keeps one per slot) must apply the moves kv_defrag / kv_reserve_seq_with_defrag return.

3. The storage format is a gate, never a guess

  1. The KV storage format is a gate, never a guess (C4), and it is per engine (#99, CUDA device half #153). MINFER_CACHE_TYPE is parsed strictly into KvFormat { f32, f16, q8_0 } (graph/kvformat.rs, the single authority) and resolved once per load against the device in models::load_model_configured; the resolver folds in the GPU's own auto policy (kvformat::auto_device_format: f16 for the 7B class, f32 for small models, never on the CPU), so one answer drives both the region width and the kernel. The answer is stored on the loaded engine (ModelDef::kv_format/set_kv_format) and reaches the graph through CParams::kv_format (part of the reuse identity: the builder stamps each KV node's KvcacheMeta::row_elems from it) and the CPU, CUDA and Metal kernels through GraphAllocator::set_kv_format. There is deliberately no process-wide kvformat global (the decision, and the measured failure that forced it, are in ADR-0006), and no process-wide CUDA layout tag either: cuda::layout_of/format_of are the only KvFormat ↔ FFI-tag binding, and the tag is part of the captured-graph identity — cuda_backend.rs::graph_replay_step refuses an exec whose recorded kv_layout moved, exactly like a pool_gen change. An unknown value fails the load on every device; q8_0 fails on a backend whose attention kernel has no packed read — the registry's BackendCaps::reads_packed_kv, true for CPU, CUDA and Metal (#310 enabled it on Metal via mechanisms A and B). A Q8_0 cell is a whole number of Q8_0 blocks rounded up to whole f32 words, so elems / n_ctx is one cell's width and every copy_cells move stays verbatim; KvcacheMeta::row_elems carries that width while the node's shape stays logical, and ensure_kv refuses a packed width on a backend without the capability. The store quantizes with one quantizer on both backends (store_kv_q8_0: amax/127, f16 scale, round-ties-even), and attention reads the packed blocks directly — on the CPU the K score is dot_q8_0_q8_0 against the Q8_0-quantized query row and V accumulates out of the cell (kvformat::accumulate_q8_0_row; MINFER_NO_FUSED_Q8_KV restores S1's dequantize-into-scratch A/B) — while on CUDA the layout-tagged kv4<KV_LAYOUT_Q8_0> dequantizes each 4-element group. A physical kv_rm/kv_shift works on a packed region: survivors move verbatim and each is mapped through kvformat::map_q8_0_cells (dequantize → re-rope → requantize). Honest scope: a physical shift on an f16 region refuses loudly, and that refusal is pinned by a gate (graph::alloc::tests::kv_shift::an_f16_region_refuses_the_physical_shift_and_the_other_formats_take_it, with f32 and packed q8_0 controls on the same fixture) (#306 — the host round trip has no hunk→f32→hunk map; CUDA is exposed too), so the CLI --cnv overflow falls back to re-rendering the retained window (measured dgxspark (aarch64, GB10 sm_121), 2026-10-07, 7B Q4_K_M MINFER_CACHE_TYPE=f16 ./target/release/minfer --cnv --n-ctx 512 -n 8: 499–505 tokens re-prefilled per overflow in 0.30–0.35 s, ≈ +0.09 s over the f32 shift's 22-token / 0.26 s delta, and cheaper than the f32 fallback's own 0.72–0.76 s), and a speculative session refuses a packed cache (spec::SpecEngine::new) because the batched split kernel's bitwise identity with sequential decode is what its greedy contract rests on.

4. A session is a file with a header, never a memory dump

  1. A KV session is a file with a header, never a memory dump (C5). graph/kvsession.rs writes a versioned, checksummed container: shape, backend, KV element type (the header's flags word: FLAG_PACKED for Q8_0, FLAG_F16 for f16, 0 for f32; the two element-type bits are mutually exclusive and an unknown bit is refused loudly, so a pre-#130 flags == 0 file still loads as f32 and an older build refuses a newer f16 file instead of decoding it as f32) — one K/V blob per layer as pool words, then the owner table + run table + span lists + written extents, then — version 2, C5 S2 — an opaque, length-prefixed host-state blob inside the checksum: the KV rows belong to a host state, so the container carries both or neither and a version-1 file is refused loudly. GraphAllocator::kv_save/kv_load stream it through the backends' host I/O and enable the pool a file names if the graph has not been built yet. kv_load verifies the whole file before it applies anything, so a truncated, corrupted or foreign file is a no-op, and the header must describe the run that loads it (backend, n_ctx, n_embd, element type). KvCache::restore_session validates the bookkeeping — arena capacity, owner-table length, every reservation and span inside the arena, every live sequence carrying a span list — before applying. The real-model gate asserts the resumed session continues bitwise (max |Δlogit| = 0). CLI (--session FILE under --cnv): keeps the JSON history and writes FILE.kv beside it on exit; a matching companion is resumed with no prefill (the host state — messages, stream_tokens, current_pos, turn_pos, prev_tokens, need_insert_eot — rides in the container as a versioned JSON blob), and everything else (another --n-ctx/model/MINFER_CACHE_TYPE, an edited history, an unknown snapshot version, an engine that cannot hand its KV to the host) prints the reason and re-renders. --slots-file <PATH> is the server twin (S2b): it snapshots the batched engine's slot table after every completed request and resumes it at startup, so a request whose prompt matches a restored slot prefills only its delta; the in-flight request is deliberately not in the snapshot. A mixed offload plan has KV on two backends, so kv_load refuses a session (one arena per file).

Decisions governing this document

This page is the current contract; the decisions behind it are frozen in the ADR corpus:

  • ADR-0002 — Topology is a function of GraphParams alone, so positions cannot be structure
  • ADR-0006 — The KV storage format is a per-engine gate, not a process-wide global
  • ADR-0014 — A KV session is a versioned, checksummed file — never a memory dump
  • ADR-0005 — Metal becomes a first-class backend
  • ADR-0011 — Backend ids are a file-format contract: appended, never renumbered