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

ADR-0022: Sharded per-node generation cache for distributed PyTorch routes

Status: Superseded — see 0022-sharded-per-node-kv-cache.md

Draft alternate header names (X-Meshnet-Cache-Mode, X-Meshnet-Seq-Len) were not implemented. The accepted wire protocol and implementation use X-Meshnet-Cache and X-Meshnet-Past-Len per the linked ADR.

Context

The distributed PyTorch chat path previously recomputed the full prompt-so-far for every generated token. The head shard embedded the entire sequence each step, forwarded full-sequence activations through every downstream shard, and every shard called its decoder layers with use_cache=False. On a two-node Qwen2.5-0.5B route this produced the expected quadratic slowdown as output length grew.

ADR-0020 and ADR-0021 fixed route construction and start_layer semantics. They did not define the per-request cache lifecycle needed for efficient decode.

Decision

Distributed PyTorch generation now uses one stable route session id for an entire chat request. The wire protocol marks each activation hop with:

• X-Meshnet-Session: stable per generation.
• X-Meshnet-Cache-Mode: prefill, decode, or stateless.
• X-Meshnet-Seq-Len: the total sequence length represented by the step.

Step 0 is prefill: the head sends the full prompt activation through the planned route. Each shard stores only the opaque cache state returned by its own executed layer range. No shard receives or stores another shard's cache.

Later cached decode steps send only the newest token activation ([1, 1, hidden]) with the full sequence length and newest position id. The backend deliberately treats layer cache state as opaque. Standard K/V tuples, HuggingFace cache objects, and hybrid linear-attention recurrent state are stored without shape assumptions.

Cache lifecycle

Each TorchModelShard owns an in-memory LRU map keyed by (session_id, effective_start_layer, shard_end). Entries expire by TTL and by a maximum session count (MESHNET_SHARD_CACHE_TTL_SECONDS, default 600; MESHNET_SHARD_CACHE_MAX_SESSIONS, default 16).

If a decode step reaches a node after restart, eviction, TTL expiry, or route mismatch, the node returns an explicit cache_miss response. The head falls back to full prefill for the current prompt-so-far using the same session id, rebuilding the shard-local caches before continuing. Alpha route repair still does not migrate cache state across nodes; a true route change is treated as cache loss and recovered by re-prefill.

Consequences

• Healthy decode sends O(1) activation payloads per token between nodes instead of O(sequence length).
• Cache internals stay behind the model backend boundary, which keeps Qwen3.6-style hybrid recurrent cache state compatible with the same route protocol.
• Restart and eviction degrade to slower stateless/full-prefill work rather than silent output corruption.
• Cross-node cache migration, batching cache state across sessions, and speculative decoding remain future work.

Verification

Unit coverage in tests/test_real_model_backend.py verifies opaque per-layer cache storage, cached one-token decode, explicit cache-miss errors, and LRU eviction. Live two-node Qwen2.5-0.5B TPS measurement still requires the physical two-machine topology used to observe the regression.