# ADR-0022: Sharded per-node KV cache for distributed generation

## Status: Accepted, implemented (alpha-hardening issue 25)

## Context

The distributed generation loop (`torch_server.py`, `_do_chat_completions` distributed
path) had **no KV cache**: every layer-forward call passed `use_cache: False`, and each
autoregressive step re-encoded the entire prompt-so-far from scratch, re-running every
layer on every node in the route for every generated token. Measured on a live 2-node
Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 to 12.6 within a single generation —
the quadratic-cost signature. On Qwen3.6-35B-A3B mixed GPU/CPU topology this collapsed
to ~0.07 tps even after the ADR-0020 routing fix.

`X-Meshnet-Session` existed on the wire but was minted fresh **per token** and keyed no
state.

## Decision

### Session lifecycle

The head mints one session id per chat generation (not per token) and reuses it across
every step. Two new request headers extend the `/forward` wire protocol:

- `X-Meshnet-Cache: prefill | decode` — absent means legacy stateless (unchanged
  behavior, and what old nodes send/understand).
- `X-Meshnet-Past-Len: N` — decode only: the number of tokens the node's session cache
  must already hold. A mismatch is a cache miss, never silent corruption.

Step 0 (`prefill`) sends the full prompt activation as before; each node creates fresh
session state for its own layer range. Steps 1+ (`decode`) send only the newest token's
hidden state — `[1, 1, hidden]`, cutting per-step compute and wire payload from
O(seq_len) to O(1). The head embeds the next token directly from the `token_id` the tail
now returns alongside text (`{"text": …, "token_id": …}`), avoiding text
re-tokenization drift; EOS is detected by id against tokenizer + generation-config eos
sets.

### Per-node sharded cache

`TorchModelShard.kv_sessions` is a `SessionCacheStore`: `session_id → SessionCacheEntry`
holding cache state **only for that shard's layer range** — sharding falls out naturally
because each node only executes (and therefore only caches) its own layers. No node ever
holds another node's state.

### MoE / hybrid-attention awareness

The cached object is whatever `use_cache=True` produces: a transformers
`DynamicCache(config=model.config)` — the same construction the model's own `forward()`
uses. With the config, transformers picks the right per-layer state: K/V tensors for
standard attention, conv/recurrent delta state for Qwen3.6-style hybrid linear-attention
layers, sliding-window variants, etc. The store treats it as opaque; nothing assumes a
K/V tensor shape. Cache slots are indexed by absolute `layer_idx`, so a shard updating
only layers 12–23 leaves 0–11 empty (verified: sparse `DynamicCache.update` works).
MoE expert routing is layer-local per token and needs no cross-token state.

Layers are invoked with `past_key_values=<cache>, use_cache=True, cache_position=…`
(transformers 5.x layer API; the cache is mutated in place). If a model's layers reject
those kwargs, the backend logs once, sets `supports_kv_cache = False`, and stays on the
stateless path permanently — exotic architectures degrade to today's behavior instead of
failing.

### Cache miss and route-change interaction (ADR-0021)

Any decode-mode request that cannot be served — unknown session (evicted, node
restarted), `past_len` mismatch, `start_layer` mismatch (the route or shard overlap
changed mid-generation), or caching disabled — raises `KVCacheMiss`, answered as
**HTTP 409 `{"error": "cache_miss"}`**. The head catches it and falls back to one full
re-prefill of the accumulated sequence under the same session id, which atomically
replaces every node's session state, then continues cached. The fully-stateless path is
therefore still the recovery mode: eviction and restarts cost one prefill, never
corruption or a failed generation. A decode request against a node whose caching is
disabled is also a 409 — running a single-token payload statelessly would silently
produce garbage.

Mixed fleets degrade the same way: if the tail predates the protocol and returns no
`token_id`, the head simply prefills every step (exactly the old cost).

### Bounded memory

`SessionCacheStore` enforces TTL (default 600 s, `MESHNET_KV_TTL_SECONDS`) plus LRU cap
(default 8 sessions, `MESHNET_KV_MAX_SESSIONS`), evaluated on every access. The head
additionally drops its own session explicitly when a generation completes; downstream
nodes rely on TTL/LRU (an explicit cross-node release RPC was judged not worth the
failure modes — misses are cheap).

### Non-goals (first landing)

Cross-node cache migration on route change (evict + re-prefill is acceptable),
speculative decoding, cross-session batching.

## Consequences

- Per-token cost drops from O(seq_len) layer re-execution + O(seq_len) wire transfer per
  hop to O(1) of both; tps stays flat across generation length instead of decaying.
- Golden test (`tests/test_kv_cache_distributed.py`, env-gated by
  `MESHNET_REAL_MODEL_TESTS=1`) proves cached and stateless distributed generation emit
  identical token ids on a real two-shard Qwen2.5-0.5B split.
- Nodes now hold per-session GPU/CPU memory between requests (bounded above); operators
  sizing `max_loaded_shards` should account for ~`sessions × seq_len × kv_bytes_per_token`
  per resident model.
- The wire protocol is backward- and forward-compatible: headers are additive, absent
  headers mean stateless, and 409 is only sent in reply to explicit decode-mode requests.