Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

docs/adr/0022-sharded-per-node-kv-cache.md

@@ -0,0 +1,102 @@
+# 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.