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

Status: ready-for-agent

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)
• Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes
• Each node caches past_key_values/recurrent state only for its own shard's layer range; no node ever holds another node's cache
• Cache works correctly for both standard-attention shards and Qwen3.6-style hybrid linear-attention/recurrent shards (cache abstraction is not K/V-shaped-only)
• Bounded memory: TTL + LRU eviction; eviction/restart triggers a documented cache-miss response, not silent corruption or an unhandled exception
• Golden-output regression test proves cached and uncached distributed generation produce equivalent output for a fixed prompt
• Measured tps improvement recorded on the same 2-node Qwen2.5-0.5B topology used to observe the regression (target: flat tps across generation length, not decaying)
• tests/test_two_node_pipeline.py and tests/test_dynamic_routing.py still pass
• Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction

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.