64 lines
3.0 KiB
Markdown
64 lines
3.0 KiB
Markdown
# ADR-0022: Sharded per-node generation cache for distributed PyTorch routes
|
|
|
|
## Status: Accepted
|
|
|
|
## 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.
|