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

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.