neuron-tai/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md

Status: implemented 2026-07-08 — pending live 2-node GPU verification

Implemented in packages/node/meshnet_node/model_backend.py + torch_server.py; design in ADR-0022; tests in tests/test_kv_cache_distributed.py (11 fast tests + env-gated golden test, MESHNET_REAL_MODEL_TESTS=1).

Measured (two-shard Qwen2.5-0.5B 0-11/12-23, CPU, 44-token prompt, 40 steps): stateless 7.05 tps decaying 32% (8.09 → 5.50 first-10 vs last-10); cached 18.93 tps and FLAT (17.21 → 19.28) — 2.68× overall, gap grows quadratically with length. Remaining acceptance item: re-measure on the live 2-node GPU topology (needs both machines).

Scoped 2026-07-08 from a live two-machine distributed-inference debugging session (Qwen2.5-0.5B GPU+GPU pipeline, and Qwen3.6-35B-A3B mixed GPU/CPU). The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this issue is the next performance blocker in the same code path.

25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)

What to build

The distributed generation loop (torch_server.py:515-612, _do_chat_completions distributed path) currently has no KV cache at all: model_backend.py passes use_cache: False in every layer-forward call (lines 763, 768, 770-771), and each autoregressive step re-encodes the entire prompt-so-far from scratch (backend.encode_prompt(current_text)), re-running every layer on every node in the route for every generated token.

Observed cost of this on a live 2-node Qwen2.5-0.5B GPU pipeline (layers 0-20 / 21-23): tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation — the expected quadratic-cost signature. On the Qwen3.6-35B-A3B mixed-topology case this collapses to ~0.07 tps even after the routing fix, partly for this reason.

X-Meshnet-Session already exists on the wire (torch_server.py:707, minted fresh per token, not per generation) but today only labels one activation transfer for chunk reassembly/logging — it is not used to key any cached state.

Subtask	Owner package	Deliverable
Session lifecycle	packages/node/meshnet_node/torch_server.py	Mint session ID once per chat request (not per token); reuse across all steps of that generation; add X-Meshnet-Seq-Len / position header so a node can tell prefill from decode steps
Per-node sharded cache	packages/node/meshnet_node/model_backend.py	TorchModelShard holds a session_id → cache_state map scoped to its own layer range only (naturally sharded — no node stores another node's KV); forward_bytes takes use_cache=True and returns/reuses past_key_values (or use_cache=False for the prefill token to keep failure/eviction simple)
Prefill vs. decode split	packages/node/meshnet_node/torch_server.py	Step 0 sends the full prompt activation (current behavior); steps 1+ send only the newest token's hidden state ([1, 1, hidden]) with correct position_ids, cutting per-step payload from O(seq_len) to O(1)
MoE / hybrid-attention state	packages/node/meshnet_node/model_backend.py	Cache abstraction must hold "whatever use_cache=True returns for this layer range," not assume standard K/V tensors — Qwen3.6's linear-attention/hybrid layers (see [transformers] The fast path is not available... warning already logged at startup) cache recurrent conv/delta state, not K/V pairs. MoE expert routing itself is layer-local and needs no cross-token cache, but confirm no expert-choice state leaks across the stateless-vs-cached boundary when use_cache toggles between prefill and decode
Cache lifecycle	packages/node/meshnet_node/torch_server.py	TTL + LRU eviction per node (bounded by max_loaded_shards/memory budget); explicit "cache miss" response so a restarted/evicted node causes the head to fall back to a full re-prefill instead of a hard error — keep today's fully-stateless path as the recovery mode
Correctness parity	tests/	Golden-output test: distributed multi-token output with caching enabled must match the existing stateless path token-for-token (or within sampling tolerance) for a fixed prompt/seed

Non-goals for first landing: cross-node cache migration/rebalancing on route change (evict + re-prefill is acceptable initially); speculative decoding; batching multiple concurrent sessions' KV within one node beyond what eviction already requires.

Code refs:

• packages/node/meshnet_node/torch_server.py:515-612 — distributed generation loop (current_text = current_text + token_str, full re-encode every step)
• packages/node/meshnet_node/torch_server.py:690-789 — _run_downstream_pipeline, session minting, X-Meshnet-Session/X-Meshnet-Hop-Index/X-Meshnet-Start-Layer headers
• packages/node/meshnet_node/model_backend.py:189-201, 330-351, 763-771 — use_cache: False call sites, effective_start layer-slicing logic that any cache keying must respect
• docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md — prerequisite routing fix this issue builds on
• docs/adr/0021-dynamic-statistical-routing.md — route selection this cache must stay compatible with (a route change mid-generation should trigger cache-miss fallback, not corruption)

Acceptance criteria

• A session ID is stable across all steps of one chat generation (not re-minted per token) — minted once in _do_chat_completions, asserted in test_session_is_stable_and_decode_payloads_are_single_token
• Steps after the first prefill send only the new token's activation ([1, 1, hidden] via encode_next_token) with X-Meshnet-Cache: decode + X-Meshnet-Past-Len
• Each node caches state only for its own shard's layer range (TorchModelShard.kv_sessions; sharding falls out of per-node layer execution)
• Cache abstraction is not K/V-shaped-only: DynamicCache(config=model.config) — the same construction Qwen3.6-Next's own forward uses for hybrid linear-attention conv/delta state; store treats it as opaque; TypeError fallback disables caching per-backend
• Bounded memory: TTL (600 s, MESHNET_KV_TTL_SECONDS) + LRU (8, MESHNET_KV_MAX_SESSIONS); miss → HTTP 409 {"error": "cache_miss"} → head re-prefills (tested)
• Golden-output test: cached and stateless produce identical token ids on real two-shard Qwen2.5-0.5B (test_cached_distributed_generation_matches_stateless_golden, passed)
• Measured (CPU two-shard proxy, 40 steps): stateless 7.05 tps w/ 32% decay → cached 18.93 tps flat, 2.68×. ⚠️ still to run on the live 2-node GPU topology
• tests/test_two_node_pipeline.py and tests/test_dynamic_routing.py pass (30 passed; 6 tmp-dir fixture errors are a pre-existing Windows temp-permission env issue, identical on clean tree)
• Design captured in ADR-0022 incl. cache-miss/route-change interaction with ADR-0021

Notes

MoE routing (router + expert FFN) is layer-local per token and does not itself need a cross-token cache — it was ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the ADR-0020 start_layer double-execution). The MoE angle that does matter here is architecture-awareness in the cache design: don't hardcode a K/V tensor shape assumption that breaks on Qwen3.6's hybrid attention layers.