103 lines
5.4 KiB
Markdown
103 lines
5.4 KiB
Markdown
# 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.
|