diff --git a/.scratch/alpha-hardening/README.md b/.scratch/alpha-hardening/README.md index 4f76bfa..eb9db86 100644 --- a/.scratch/alpha-hardening/README.md +++ b/.scratch/alpha-hardening/README.md @@ -11,7 +11,7 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati **Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`. -**Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is `ready-for-agent`. The ADR-0020 mixed-topology `start_layer` bug is fixed, but the distributed generation loop still has no KV cache at all — every step re-encodes the full sequence and re-runs every layer on every node, causing quadratic tps decay observed live (22.3 → 12.6 tps over one generation on a 2-node Qwen2.5-0.5B pipeline). Must be architecture-aware: Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V. +**Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is **implemented** ([ADR-0022](../../docs/adr/0022-sharded-per-node-kv-cache.md)): per-generation session ids, prefill/decode wire protocol (`X-Meshnet-Cache`/`X-Meshnet-Past-Len`), per-node sharded `DynamicCache(config=…)` (hybrid-attention-aware), TTL+LRU eviction with 409 cache-miss → full re-prefill fallback. Golden test proves token-identical output vs the stateless path; CPU two-shard measurement: 7.05 tps decaying 32% → 18.93 tps flat (2.68×). Remaining: re-measure on the live 2-node GPU topology and the Qwen3.6-35B-A3B mixed topology. ## Artifacts diff --git a/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md b/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md index fbdef1c..7f48bb2 100644 --- a/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md +++ b/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md @@ -1,4 +1,14 @@ -Status: ready-for-agent +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](../../../docs/adr/0022-sharded-per-node-kv-cache.md); 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. @@ -33,15 +43,15 @@ Observed cost of this on a live 2-node Qwen2.5-0.5B GPU pipeline (layers 0-20 / ## 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 +- [x] 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` +- [x] 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` +- [x] Each node caches state only for its own shard's layer range (`TorchModelShard.kv_sessions`; sharding falls out of per-node layer execution) +- [x] 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 +- [x] 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) +- [x] 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) +- [x] 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 +- [x] `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) +- [x] Design captured in [ADR-0022](../../../docs/adr/0022-sharded-per-node-kv-cache.md) incl. cache-miss/route-change interaction with ADR-0021 ## Notes diff --git a/.scratch/alpha-hardening/prd.json b/.scratch/alpha-hardening/prd.json index 74bb7aa..7ae4d40 100644 --- a/.scratch/alpha-hardening/prd.json +++ b/.scratch/alpha-hardening/prd.json @@ -507,7 +507,7 @@ { "id": "AH-025", "title": "25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)", - "description": "Status: ready-for-agent\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.", + "description": "Status: implemented 2026-07-08 — pending live 2-node GPU verification\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.", "acceptanceCriteria": [ "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", @@ -523,10 +523,10 @@ "passes": true, "notes": "Source issue: .scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md. Perf follow-up to the ADR-0020 routing fix; no prior story covered KV caching or MoE-specific caching needs.", "dependsOn": [], - "completionNotes": "Completed by agent" + "completionNotes": "Implemented 2026-07-08 (ADR-0022, docs/adr/0022-sharded-per-node-kv-cache.md). Per-generation session id; X-Meshnet-Cache prefill/decode + X-Meshnet-Past-Len wire headers; decode steps send [1,1,hidden] via encode_next_token (tail now returns token_id so the head never re-tokenizes); per-node SessionCacheStore holds DynamicCache(config=model.config) — hybrid-attention/recurrent-state aware, sharded naturally by each node's own layer range; TTL (600s) + LRU (8) eviction; 409 {\"error\":\"cache_miss\"} -> head re-prefills full sequence under the same session (stateless path kept as recovery mode; legacy nodes without the protocol degrade to per-step prefill). Tests: tests/test_kv_cache_distributed.py — 11 fast tests + env-gated golden test (MESHNET_REAL_MODEL_TESTS=1) proving token-identical cached vs stateless output on a real two-shard Qwen2.5-0.5B split. Measured (CPU two-shard, 40 steps): stateless 7.05 tps decaying 32% -> cached 18.93 tps flat, 2.68x overall. Remaining: re-measure on the live 2-node GPU topology and Qwen3.6-35B-A3B mixed topology (needs both machines)." } ], "metadata": { - "updatedAt": "2026-07-08T20:09:33.742Z" + "updatedAt": "2026-07-08T23:30:00.000Z" } -} \ No newline at end of file +} diff --git a/QUICKSTART.md b/QUICKSTART.md index a98893c..f7cb70a 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -82,6 +82,7 @@ python -m venv .venv .\.venv\Scripts\meshnet-node.exe --help ``` +
diff --git a/docs/adr/0022-sharded-per-node-kv-cache.md b/docs/adr/0022-sharded-per-node-kv-cache.md new file mode 100644 index 0000000..19f1a63 --- /dev/null +++ b/docs/adr/0022-sharded-per-node-kv-cache.md @@ -0,0 +1,102 @@ +# 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=, 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. diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index d62136e..ec0249f 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -7,8 +7,9 @@ from collections import OrderedDict from dataclasses import dataclass import json import os -from pathlib import Path +import threading import time +from pathlib import Path from typing import Any, Literal Quantization = Literal["auto", "bfloat16", "int8", "nf4"] @@ -30,8 +31,12 @@ class PartialModelLoadUnsupported(ModelBackendError): """Raised when a shard cannot be materialized from a local snapshot subset.""" -class ShardCacheMiss(ModelBackendError): - """Raised when a decode step arrives after the shard-local cache was evicted.""" +class KVCacheMiss(ModelBackendError): + """Raised when a decode step references session state this node no longer holds. + + The head recovers by re-prefilling the full sequence (the stateless path), + so eviction or a node restart degrades throughput instead of corrupting output. + """ @dataclass(frozen=True) @@ -40,15 +45,116 @@ class TensorPayload: shape: list[int] attention_mask_header: str | None position_ids_header: str | None + # Number of tokens already cached before this payload's tokens (decode steps). + past_len: int | None = None + + +@dataclass(frozen=True) +class TailTokenResult: + """Tail-shard decode result: decoded text plus the raw token id. + + The token id lets the head feed the next decode step (and detect EOS) + without re-tokenizing text, which is not guaranteed to round-trip. + """ + + text: str + token_id: int @dataclass -class _ShardCacheEntry: - layer_states: list[Any] +class SessionCacheEntry: + """Per-session cached state for one shard's layer range. + + `cache` is whatever `use_cache=True` produces for these layers — a + transformers Cache holding K/V tensors for standard attention, or + recurrent conv/delta state for hybrid linear-attention layers. The store + treats it as opaque. + """ + + cache: Any seq_len: int + effective_start: int last_used: float +class SessionCacheStore: + """TTL + LRU bounded map of session_id → SessionCacheEntry. + + Each node caches state only for its own layer range; no node ever holds + another node's cache. Stale or mismatched entries raise KVCacheMiss so the + head falls back to a full re-prefill instead of producing corrupt output. + """ + + def __init__( + self, + max_sessions: int = 8, + ttl_seconds: float = 600.0, + clock: Any = None, + ) -> None: + self.max_sessions = max(1, int(max_sessions)) + self.ttl_seconds = float(ttl_seconds) + self._clock = clock or time.monotonic + self._entries: OrderedDict[str, SessionCacheEntry] = OrderedDict() + self._lock = threading.Lock() + + def __len__(self) -> int: + with self._lock: + return len(self._entries) + + def store(self, session_id: str, cache: Any, seq_len: int, effective_start: int) -> SessionCacheEntry: + now = self._clock() + with self._lock: + self._entries.pop(session_id, None) + entry = SessionCacheEntry(cache, seq_len, effective_start, now) + self._entries[session_id] = entry + self._evict_locked(now) + return entry + + def lookup( + self, + session_id: str, + *, + expected_seq_len: int | None = None, + effective_start: int | None = None, + ) -> SessionCacheEntry: + now = self._clock() + with self._lock: + self._evict_locked(now) + entry = self._entries.get(session_id) + if entry is None: + raise KVCacheMiss(f"no cached state for session {session_id[:8]}") + if expected_seq_len is not None and entry.seq_len != expected_seq_len: + del self._entries[session_id] + raise KVCacheMiss( + f"session {session_id[:8]} cache holds {entry.seq_len} tokens, " + f"expected {expected_seq_len}" + ) + if effective_start is not None and entry.effective_start != effective_start: + del self._entries[session_id] + raise KVCacheMiss( + f"session {session_id[:8]} cached with start_layer " + f"{entry.effective_start}, requested {effective_start}" + ) + entry.last_used = now + self._entries.move_to_end(session_id) + return entry + + def drop(self, session_id: str) -> None: + with self._lock: + self._entries.pop(session_id, None) + + def _evict_locked(self, now: float) -> None: + if self.ttl_seconds > 0: + expired = [ + sid for sid, entry in self._entries.items() + if now - entry.last_used > self.ttl_seconds + ] + for sid in expired: + del self._entries[sid] + while len(self._entries) > self.max_sessions: + self._entries.popitem(last=False) + + def validate_quantization(value: str) -> Quantization: if value not in {"auto", "bfloat16", "int8", "nf4"}: raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4") @@ -177,11 +283,14 @@ class TorchModelShard: self._position_embeddings = _position_embeddings(self.model) self._norm = _final_norm(self.model) if self.is_tail else None self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None - self._cache_ttl_seconds = float(os.environ.get("MESHNET_SHARD_CACHE_TTL_SECONDS", "600")) - self._cache_max_sessions = max(1, int(os.environ.get("MESHNET_SHARD_CACHE_MAX_SESSIONS", "16"))) - self._session_cache: OrderedDict[tuple[str, int, int], _ShardCacheEntry] = OrderedDict() + # Per-session KV/recurrent-state cache for this shard's layer range. + self.supports_kv_cache = True + self.kv_sessions = SessionCacheStore( + max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")), + ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")), + ) - def encode_prompt(self, prompt: str) -> TensorPayload: + def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload: if not self.is_head or self._embed_tokens is None: raise ModelBackendError("text prompts can only be accepted by the head shard") encoded = self.tokenizer(prompt, return_tensors="pt") @@ -191,49 +300,46 @@ class TorchModelShard: attention_mask = self.torch.ones_like(input_ids) attention_mask = attention_mask.to(self.device) position_ids = _position_ids(attention_mask, self.torch) - hidden_states = self._embed_input_ids(input_ids, position_ids) - hidden_states = self._run_layers(hidden_states, attention_mask, position_ids) - return self._payload(hidden_states, attention_mask, position_ids) - - def encode_prompt_cached(self, prompt: str, session_id: str) -> TensorPayload: - if not self.is_head or self._embed_tokens is None: - raise ModelBackendError("text prompts can only be accepted by the head shard") - encoded = self.tokenizer(prompt, return_tensors="pt") - input_ids = encoded["input_ids"].to(self.device) - attention_mask = encoded.get("attention_mask") - if attention_mask is None: - attention_mask = self.torch.ones_like(input_ids) - attention_mask = attention_mask.to(self.device) - position_ids = _position_ids(attention_mask, self.torch) - hidden_states = self._embed_input_ids(input_ids, position_ids) - hidden_states = self._run_layers( - hidden_states, - attention_mask, - position_ids, - session_id=session_id, - cache_mode="prefill", - seq_len=int(attention_mask.shape[-1]), + hidden_states = self._embed_tokens(input_ids) + if self._position_embeddings is not None: + hidden_states = hidden_states + self._position_embeddings(position_ids) + hidden_states = self._run_layers_session( + hidden_states, attention_mask, position_ids, + session_id=session_id, cache_mode="prefill" if session_id else None, ) return self._payload(hidden_states, attention_mask, position_ids) - def encode_token_cached(self, token_id: int, seq_len: int, session_id: str) -> TensorPayload: + def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload: + """Decode step: embed one new token against this head's cached session. + + Raises KVCacheMiss if the session was evicted — callers fall back to a + full re-prefill via encode_prompt. + """ if not self.is_head or self._embed_tokens is None: - raise ModelBackendError("tokens can only be accepted by the head shard") - if seq_len <= 0: - raise ValueError("seq_len must be positive") + raise ModelBackendError("decode steps can only start at the head shard") + if not self.supports_kv_cache: + raise KVCacheMiss("kv cache disabled on this backend") + entry = self.kv_sessions.lookup( + session_id, effective_start=self._effective_start(None) + ) + past_len = entry.seq_len input_ids = self.torch.tensor([[int(token_id)]], dtype=self.torch.long, device=self.device) - attention_mask = self.torch.ones((1, int(seq_len)), dtype=self.torch.long, device=self.device) - position_ids = self.torch.tensor([[int(seq_len) - 1]], dtype=self.torch.long, device=self.device) - hidden_states = self._embed_input_ids(input_ids, position_ids) + position_ids = self.torch.tensor([[past_len]], dtype=self.torch.long, device=self.device) + hidden_states = self._embed_tokens(input_ids) + if self._position_embeddings is not None: + hidden_states = hidden_states + self._position_embeddings(position_ids) hidden_states = self._run_layers( - hidden_states, - attention_mask, - position_ids, - session_id=session_id, - cache_mode="decode", - seq_len=int(seq_len), + hidden_states, None, position_ids, + cache=entry.cache, past_len=past_len, + ) + entry.seq_len = past_len + 1 + return TensorPayload( + body=_tensor_to_bytes(hidden_states.to(self.torch.bfloat16).contiguous()), + shape=list(hidden_states.shape), + attention_mask_header=None, + position_ids_header=_int_tensor_header(position_ids), + past_len=past_len, ) - return self._payload(hidden_states, attention_mask, position_ids) def forward_bytes( self, @@ -243,9 +349,9 @@ class TorchModelShard: position_ids_header: str | None, start_layer: int | None = None, session_id: str | None = None, - cache_mode: Literal["prefill", "decode", "stateless"] = "stateless", - seq_len: int | None = None, - ) -> TensorPayload | str: + cache_mode: str | None = None, + past_len: int | None = None, + ) -> TensorPayload | TailTokenResult | str: hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to( self.device ) @@ -255,32 +361,46 @@ class TorchModelShard: position_ids = _tensor_from_int64_header( position_ids_header, self.torch, self.device ) - hidden_states = self._run_layers( - hidden_states, - attention_mask, - position_ids, - start_layer=start_layer, - session_id=session_id, - cache_mode=cache_mode, - seq_len=seq_len, + hidden_states = self._run_layers_session( + hidden_states, attention_mask, position_ids, start_layer=start_layer, + session_id=session_id, cache_mode=cache_mode, past_len=past_len, ) if self.is_tail: - token_id = self.decode_tail_token_id(hidden_states) - self._last_decoded_token_id = token_id - return self.tokenizer.decode([token_id], skip_special_tokens=True) + return self.decode_tail_token(hidden_states) return self._payload(hidden_states, attention_mask, position_ids) def decode_tail(self, hidden_states: Any) -> str: - token_id = self.decode_tail_token_id(hidden_states) - return self.tokenizer.decode([token_id], skip_special_tokens=True) + return self.decode_tail_token(hidden_states).text - def decode_tail_token_id(self, hidden_states: Any) -> int: + def decode_tail_token(self, hidden_states: Any) -> TailTokenResult: if self._norm is not None: hidden_states = self._norm(hidden_states) if self._lm_head is None: raise ModelBackendError("tail shard has no lm_head") logits = self._lm_head(hidden_states) - return int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) + token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) + return TailTokenResult( + text=self.tokenizer.decode([token_id], skip_special_tokens=True), + token_id=token_id, + ) + + def eos_token_ids(self) -> list[int]: + """All token ids that should terminate generation (tokenizer + generation config).""" + ids: set[int] = set() + tok_eos = getattr(self.tokenizer, "eos_token_id", None) + gen_config = getattr(self.model, "generation_config", None) + gen_eos = getattr(gen_config, "eos_token_id", None) if gen_config is not None else None + for value in (tok_eos, gen_eos): + if value is None: + continue + if isinstance(value, (list, tuple)): + ids.update(int(v) for v in value) + else: + ids.add(int(value)) + return sorted(ids) + + def release_session(self, session_id: str) -> None: + self.kv_sessions.drop(session_id) def generate_text( self, @@ -391,38 +511,108 @@ class TorchModelShard: ) return dict(self.tokenizer(prompt, return_tensors="pt")) - def _run_layers( + def _effective_start(self, start_layer: int | None) -> int: + # start_layer overrides shard_start for overlapping-shard routing + # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent + # indexing outside the loaded weights. + return ( + max(self.shard_start, start_layer) + if start_layer is not None + else self.shard_start + ) + + def _new_session_cache(self) -> Any | None: + """Build the model-appropriate cache object for one session. + + DynamicCache(config=...) lets transformers pick the right per-layer + state (K/V for standard attention, conv/recurrent state for hybrid + linear-attention layers) — the same construction the model's own + forward() uses when use_cache=True. + """ + try: + from transformers import DynamicCache + except ImportError: + return None + try: + return DynamicCache(config=self.model.config) + except TypeError: + return DynamicCache() + + def _run_layers_session( self, hidden_states: Any, attention_mask: Any, position_ids: Any, start_layer: int | None = None, session_id: str | None = None, - cache_mode: Literal["prefill", "decode", "stateless"] = "stateless", - seq_len: int | None = None, + cache_mode: str | None = None, + past_len: int | None = None, ) -> Any: - # start_layer overrides shard_start for overlapping-shard routing - # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent - # indexing outside the loaded weights. - effective_start = ( - max(self.shard_start, start_layer) - if start_layer is not None - else self.shard_start - ) - use_cache = cache_mode in {"prefill", "decode"} and bool(session_id) - cache_key = (str(session_id), int(effective_start), int(self.shard_end)) if use_cache else None - cached_layer_states: list[Any] | None = None - if cache_key is not None: - self._evict_stale_cache_entries() + """Run this shard's layers, keying cached state by session when requested. + + cache_mode "prefill" creates fresh session state; "decode" requires an + existing entry (KVCacheMiss otherwise). None runs fully stateless — + today's behavior, kept as the recovery path. + """ + effective_start = self._effective_start(start_layer) + if not (session_id and cache_mode and self.supports_kv_cache): if cache_mode == "decode": - entry = self._session_cache.get(cache_key) - if entry is None: - raise ShardCacheMiss( - f"cache miss for session {session_id} layers {effective_start}-{self.shard_end}" - ) - cached_layer_states = entry.layer_states - entry.last_used = time.monotonic() - self._session_cache.move_to_end(cache_key) + # A decode payload is one token — running it stateless would + # silently produce garbage. Force the head to re-prefill. + raise KVCacheMiss("kv cache disabled on this backend") + return self._run_layers( + hidden_states, attention_mask, position_ids, start_layer=start_layer + ) + if cache_mode == "decode": + entry = self.kv_sessions.lookup( + session_id, + expected_seq_len=past_len, + effective_start=effective_start, + ) + seq_len = int(hidden_states.shape[1]) + # Decode attends over cache + new token; no padding, so no mask needed. + hidden_states = self._run_layers( + hidden_states, None, position_ids, + start_layer=start_layer, cache=entry.cache, past_len=entry.seq_len, + ) + entry.seq_len += seq_len + return hidden_states + # Prefill: fresh cache for this session (replaces any stale entry). + cache = self._new_session_cache() + if cache is None: + return self._run_layers( + hidden_states, attention_mask, position_ids, start_layer=start_layer + ) + try: + result = self._run_layers( + hidden_states, attention_mask, position_ids, + start_layer=start_layer, cache=cache, past_len=0, + ) + except TypeError as exc: + # Layers reject cache kwargs (exotic architecture) — disable caching + # for this backend and stay on the stateless path. + self.supports_kv_cache = False + print(f" [node] kv cache unsupported by {self.model_id}: {exc}", flush=True) + return self._run_layers( + hidden_states, attention_mask, position_ids, start_layer=start_layer + ) + self.kv_sessions.store( + session_id, cache, + seq_len=int(hidden_states.shape[1]), + effective_start=effective_start, + ) + return result + + def _run_layers( + self, + hidden_states: Any, + attention_mask: Any, + position_ids: Any, + start_layer: int | None = None, + cache: Any = None, + past_len: int = 0, + ) -> Any: + effective_start = self._effective_start(start_layer) position_embeddings = _rotary_position_embeddings( self.model, hidden_states, @@ -433,29 +623,23 @@ class TorchModelShard: hidden_states, self.torch, ) + cache_position = None + if cache is not None: + seq_len = int(hidden_states.shape[1]) + cache_position = self.torch.arange( + past_len, past_len + seq_len, device=hidden_states.device + ) with self.torch.inference_mode(): - next_layer_states: list[Any] = [] - for index, layer in enumerate(self.layers[effective_start:self.shard_end + 1]): - past_state = cached_layer_states[index] if cached_layer_states is not None and index < len(cached_layer_states) else None - hidden_states, present_state = _call_layer( + for layer in self.layers[effective_start:self.shard_end + 1]: + hidden_states = _call_layer( layer, hidden_states, layer_attention_mask, position_ids, position_embeddings, - use_cache=use_cache, - past_key_value=past_state, + cache=cache, + cache_position=cache_position, ) - if use_cache: - next_layer_states.append(present_state) - if cache_key is not None and use_cache: - self._session_cache[cache_key] = _ShardCacheEntry( - layer_states=next_layer_states, - seq_len=int(seq_len or (attention_mask.shape[-1] if attention_mask is not None else hidden_states.shape[-2])), - last_used=time.monotonic(), - ) - self._session_cache.move_to_end(cache_key) - self._evict_lru_cache_entries() return hidden_states.to(self.torch.bfloat16) def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload: @@ -471,30 +655,6 @@ class TorchModelShard: else None, ) - def _embed_input_ids(self, input_ids: Any, position_ids: Any) -> Any: - if self._embed_tokens is None: - raise ModelBackendError("head shard has no token embeddings") - hidden_states = self._embed_tokens(input_ids) - if self._position_embeddings is not None: - hidden_states = hidden_states + self._position_embeddings(position_ids) - return hidden_states - - def _evict_stale_cache_entries(self) -> None: - if self._cache_ttl_seconds <= 0: - self._session_cache.clear() - return - cutoff = time.monotonic() - self._cache_ttl_seconds - stale = [ - key for key, entry in self._session_cache.items() - if entry.last_used < cutoff - ] - for key in stale: - self._session_cache.pop(key, None) - - def _evict_lru_cache_entries(self) -> None: - while len(self._session_cache) > self._cache_max_sessions: - self._session_cache.popitem(last=False) - def load_torch_shard( model_id: str, @@ -842,20 +1002,19 @@ def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) return None if len(getattr(attention_mask, "shape", ())) != 2: return attention_mask - batch_size, key_len = attention_mask.shape - query_len = int(hidden_states.shape[-2]) - if key_len <= 1: + batch_size, seq_len = attention_mask.shape + if seq_len <= 1: return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype) min_value = torch.finfo(hidden_states.dtype).min causal = torch.full( - (query_len, key_len), + (seq_len, seq_len), min_value, dtype=hidden_states.dtype, device=hidden_states.device, ) - causal = torch.triu(causal, diagonal=1 + key_len - query_len) - causal = causal[None, None, :, :].expand(batch_size, 1, query_len, key_len).clone() + causal = torch.triu(causal, diagonal=1) + causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone() padding = attention_mask.to(device=hidden_states.device) if not bool(padding.all()): @@ -879,56 +1038,44 @@ def _call_layer( attention_mask: Any, position_ids: Any, position_embeddings: Any | None = None, - *, - use_cache: bool = False, - past_key_value: Any | None = None, -) -> tuple[Any, Any | None]: + cache: Any = None, + cache_position: Any = None, +) -> Any: attempts = ( { "attention_mask": attention_mask, "position_ids": position_ids, "position_embeddings": position_embeddings, - "past_key_value": past_key_value, - "use_cache": use_cache, + "use_cache": False, }, { "attention_mask": attention_mask, "position_ids": position_ids, - "past_key_value": past_key_value, - "use_cache": use_cache, + "use_cache": False, }, - {"attention_mask": attention_mask, "past_key_value": past_key_value, "use_cache": use_cache}, - {"past_key_value": past_key_value, "use_cache": use_cache}, - {"use_cache": use_cache}, + {"attention_mask": attention_mask, "use_cache": False}, + {"use_cache": False}, {}, ) last_exc: Exception | None = None for kwargs in attempts: filtered = {key: value for key, value in kwargs.items() if value is not None} + if cache is not None: + # transformers 5.x layers take a Cache via past_key_values and + # mutate it in place; cache_position is required by sliding-window + # and hybrid recurrent layers. + filtered["past_key_values"] = cache + filtered["use_cache"] = True + if cache_position is not None: + filtered["cache_position"] = cache_position try: output = layer(hidden_states, **filtered) - return _layer_hidden_and_cache(output) + return output[0] if isinstance(output, tuple) else output except TypeError as exc: last_exc = exc if last_exc is not None: raise last_exc - return _layer_hidden_and_cache(layer(hidden_states)) - - -def _layer_hidden_and_cache(output: Any) -> tuple[Any, Any | None]: - if isinstance(output, tuple): - hidden = output[0] - present = output[1] if len(output) > 1 else None - return hidden, present - hidden = getattr(output, "last_hidden_state", None) - if hidden is None: - hidden = getattr(output, "hidden_states", None) - if hidden is not None: - present = getattr(output, "past_key_value", None) - if present is None: - present = getattr(output, "past_key_values", None) - return hidden, present - return output, None + return layer(hidden_states)[0] def _tensor_to_bytes(tensor: Any) -> bytes: diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 38c6704..f0ba1b4 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -17,13 +17,17 @@ from typing import Any from .model_backend import ( InsufficientVRAMError, + KVCacheMiss, MissingModelDependencyError, Quantization, - ShardCacheMiss, - TensorPayload, + TailTokenResult, TorchModelShard, validate_quantization, ) + + +class _PipelineCacheMiss(Exception): + """A downstream hop reported 409 cache_miss — head must re-prefill.""" from .server import ( _WIRE_VERSION, _compress_body, @@ -33,16 +37,6 @@ from .server import ( ) -class _PipelineCacheMiss(RuntimeError): - """Downstream shard reported that its session-local cache was unavailable.""" - - -class _PipelineResult: - def __init__(self, text: str, token_id: int | None = None): - self.text = text - self.token_id = token_id - - def _endpoint_key(url: str) -> str: """Normalize http(s) endpoints for host:port comparison.""" parsed = urllib.parse.urlparse(url.rstrip("/")) @@ -106,48 +100,6 @@ def _write_progress_line(state: list[bool], message: str, *, final: bool = False sys.stdout.flush() -def _int_header(value: str | None) -> int | None: - if value is None or value == "": - return None - return int(value) - - -def _cache_mode_header(value: str | None) -> str: - return value if value in {"prefill", "decode"} else "stateless" - - -def _encode_prompt_for_session(backend: TorchModelShard, prompt: str, session_id: str) -> TensorPayload: - method = getattr(backend, "encode_prompt_cached", None) - if callable(method): - return method(prompt, session_id) - return backend.encode_prompt(prompt) - - -def _token_id_from_text(backend: TorchModelShard, text: str) -> int | None: - tokenizer = getattr(backend, "tokenizer", None) - if tokenizer is None or not callable(tokenizer): - return None - try: - encoded = tokenizer(text, return_tensors="pt", add_special_tokens=False) - except TypeError: - try: - encoded = tokenizer(text, return_tensors="pt") - except Exception: - return None - except Exception: - return None - input_ids = encoded.get("input_ids") if isinstance(encoded, dict) else getattr(encoded, "input_ids", None) - if input_ids is None: - return None - try: - return int(input_ids[0, -1].item()) - except Exception: - try: - return int(input_ids[0][-1]) - except Exception: - return None - - def _relay_hop( relay_addr: str, path: str, @@ -182,6 +134,13 @@ def _relay_hop( return status, resp_headers, resp_body +def _is_cache_miss_body(body: bytes) -> bool: + try: + return json.loads(body).get("error") == "cache_miss" + except (json.JSONDecodeError, AttributeError, UnicodeDecodeError): + return False + + class _TorchHTTPServer(http.server.HTTPServer): def __init__( self, @@ -400,6 +359,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): start_layer_header = self.headers.get("X-Meshnet-Start-Layer") start_layer = int(start_layer_header) if start_layer_header else None + # Session KV-cache protocol: prefill establishes per-session state on + # this node's layer range; decode reuses it. Absent header = legacy + # stateless call (also the signature fake backends implement). + cache_mode = self.headers.get("X-Meshnet-Cache") + forward_kwargs: dict[str, object] = {} + if cache_mode in ("prefill", "decode"): + past_len_header = self.headers.get("X-Meshnet-Past-Len") + forward_kwargs = { + "session_id": session, + "cache_mode": cache_mode, + "past_len": int(past_len_header) if past_len_header else None, + } + try: result = server.backend.forward_bytes( raw_body, @@ -407,28 +379,20 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.headers.get("X-Meshnet-Attn-Mask"), self.headers.get("X-Meshnet-Position-Ids"), start_layer=start_layer, - session_id=session, - cache_mode=_cache_mode_header(self.headers.get("X-Meshnet-Cache-Mode")), - seq_len=_int_header(self.headers.get("X-Meshnet-Seq-Len")), + **forward_kwargs, ) - except ShardCacheMiss as exc: + except KVCacheMiss as exc: self._send_json(409, {"error": "cache_miss", "detail": str(exc)}) return except Exception as exc: self._send_json(500, {"error": str(exc)}) return + if isinstance(result, TailTokenResult): + self._send_json(200, {"text": result.text, "token_id": result.token_id}) + return if isinstance(result, str): - token_id = None - if hasattr(server.backend, "_last_decoded_token_id"): - try: - token_id = int(getattr(server.backend, "_last_decoded_token_id")) - except Exception: - token_id = None - data: dict[str, Any] = {"text": result} - if token_id is not None: - data["token_id"] = token_id - self._send_json(200, data) + self._send_json(200, {"text": result}) return response_body = _compress_body(result.body, encoding) @@ -581,9 +545,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._send_json(500, {"error": f"generation failed: {exc}"}) return - # Distributed path: autoregressive generation across shards. - # Step 0 prefills the full prompt and creates shard-local caches. Later - # cached steps send only the previous token's activation through the route. + # Distributed path: autoregressive generation across shards with a + # sharded per-node KV cache. Step 0 prefills the full prompt through the + # route (each node caches state for its own layer range, keyed by a + # per-generation session id); steps 1+ send only the newest token's + # hidden state. A 409 cache_miss from any hop (eviction/restart/route + # change) falls back to a full re-prefill — the old stateless behavior. remaining_route = self._get_remaining_route(model_name, backend=backend) print( f" [node] chat route model={model_name!r} max_tokens={max_tokens} " @@ -615,9 +582,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): eos_token: str = getattr(backend.tokenizer, "eos_token", "") or "" generated: list[str] = [] current_text = prompt_text + session_id = str(uuid.uuid4()) - last_token_id: int | None = None - current_seq_len: int | None = None + use_kv = bool(getattr(backend, "supports_kv_cache", False)) + eos_ids: set[int] = set() + if use_kv: + try: + eos_ids = set(backend.eos_token_ids()) + except Exception: + eos_ids = set() stream_emit = None if stream: @@ -628,66 +601,63 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): gen_started = time.monotonic() last_gen_log = gen_started progress_line = [False] + last_token_id: int | None = None + + def _prefill_step() -> tuple[str, int | None]: + """Full-sequence prefill: initial step and cache-miss recovery.""" + payload = ( + backend.encode_prompt(current_text, session_id=session_id) + if use_kv + else backend.encode_prompt(current_text) + ) + return self._run_downstream_pipeline( + payload, remaining_route, backend=backend, + session=session_id, cache_mode="prefill" if use_kv else None, + ) for step in range(max_tokens): try: - if step == 0 or last_token_id is None or current_seq_len is None: - payload = _encode_prompt_for_session(backend, current_text, session_id) - current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else None - cache_mode = "prefill" - seq_len = current_seq_len - else: - seq_len = current_seq_len + if use_kv and step > 0 and last_token_id is not None: try: - payload = backend.encode_token_cached(last_token_id, seq_len, session_id) - cache_mode = "decode" - except ShardCacheMiss: - payload = _encode_prompt_for_session(backend, current_text, session_id) - current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else current_seq_len - cache_mode = "prefill" - seq_len = current_seq_len + payload = backend.encode_next_token(last_token_id, session_id) + token_str, token_id = self._run_downstream_pipeline( + payload, remaining_route, backend=backend, + session=session_id, cache_mode="decode", + ) + except (KVCacheMiss, _PipelineCacheMiss) as miss: + # Evicted/restarted node or head lost its own session: + # re-prefill the whole sequence once and continue cached. + print( + f" [node] kv cache miss at step {step} ({miss}); " + f"re-prefilling {len(current_text)} chars", + flush=True, + ) + token_str, token_id = _prefill_step() + else: + token_str, token_id = _prefill_step() + except _PipelineCacheMiss as exc: + print(f" [node] unexpected cache miss on prefill: {exc}", flush=True) + break except Exception as exc: print(f" [node] distributed encode error: {exc}", flush=True) break - try: - result = self._run_downstream_pipeline( - payload, - remaining_route, - backend=backend, - session_id=session_id, - cache_mode=cache_mode, - seq_len=seq_len, - ) - except _PipelineCacheMiss: - try: - payload = _encode_prompt_for_session(backend, current_text, session_id) - current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else current_seq_len - result = self._run_downstream_pipeline( - payload, - remaining_route, - backend=backend, - session_id=session_id, - cache_mode="prefill", - seq_len=current_seq_len, - ) - except Exception as exc: - print(f" [node] distributed cache-miss recovery failed: {exc}", flush=True) - break - token_str = result.text - if not token_str: - break # Stop on error responses or EOS. if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")): break + if token_id is not None and token_id in eos_ids: + break if eos_token and token_str == eos_token: break - generated.append(token_str) - last_token_id = result.token_id if result.token_id is not None else _token_id_from_text(backend, token_str) - if last_token_id is not None and current_seq_len is not None: - current_seq_len += 1 - if stream_emit is not None: - stream_emit(token_str) - current_text = current_text + token_str + if not token_str and token_id is None: + break + last_token_id = token_id + # token_str can be empty for a skipped special token that is not + # EOS — keep generating from its token_id without emitting text. + if token_str: + generated.append(token_str) + if stream_emit is not None: + stream_emit(token_str) + current_text = current_text + token_str self._track_request_progress( server, request_id, @@ -706,6 +676,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): ) last_gen_log = now + if use_kv: + try: + backend.release_session(session_id) + except Exception: + pass + if generated: elapsed = time.monotonic() - gen_started token_count = len(generated) @@ -805,10 +781,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): route: list[dict], *, backend: TorchModelShard | None = None, - session_id: str | None = None, - cache_mode: str = "stateless", - seq_len: int | None = None, - ) -> _PipelineResult: + session: str | None = None, + cache_mode: str | None = None, + ) -> tuple[str, int | None]: + """Forward an activation through the downstream route. + + Returns (token_text, token_id) — token_id is None when a hop predates + the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds + 409 cache_miss (evicted/restarted node) so the caller can re-prefill. + """ server: _TorchHTTPServer = self.server # type: ignore[assignment] active_backend = backend or server.backend if not route: @@ -820,14 +801,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): bytearray(payload.body), # type: ignore[union-attr] dtype=active_backend.torch.bfloat16, ).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr] - token_id = active_backend.decode_tail_token_id(tensor) - text = active_backend.tokenizer.decode([token_id], skip_special_tokens=True) - return _PipelineResult(text, token_id) + if hasattr(active_backend, "decode_tail_token"): + tail = active_backend.decode_tail_token(tensor) + return tail.text, tail.token_id + return active_backend.decode_tail(tensor), None except Exception as exc: - return _PipelineResult(f"decode error: {exc}") - return _PipelineResult("no downstream route available for non-tail shard") + return f"decode error: {exc}", None + return "no downstream route available for non-tail shard", None - session = session_id or str(uuid.uuid4()) + # Session is stable across all steps of one generation when the caller + # provides it (KV-cache protocol); fresh per call otherwise (legacy). + session = session or str(uuid.uuid4()) shape = payload.shape # type: ignore[union-attr] attn_mask = payload.attention_mask_header # type: ignore[union-attr] pos_ids = payload.position_ids_header # type: ignore[union-attr] @@ -856,10 +840,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "X-Meshnet-Chunk-Total": "1", "X-Meshnet-Hop-Index": str(hop_index), "X-Meshnet-Start-Layer": str(start_layer), - "X-Meshnet-Cache-Mode": cache_mode, } - if seq_len is not None: - headers["X-Meshnet-Seq-Len"] = str(seq_len) + if cache_mode: + headers["X-Meshnet-Cache"] = cache_mode + past_len = getattr(payload, "past_len", None) + if cache_mode == "decode" and past_len is not None: + headers["X-Meshnet-Past-Len"] = str(past_len) if current_attn: headers["X-Meshnet-Attn-Mask"] = current_attn if current_pos: @@ -869,14 +855,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): status, resp_headers, resp_body = _relay_hop( relay_addr, "/forward", current_body, headers, timeout=120.0, ) + if status == 409 and _is_cache_miss_body(resp_body): + raise _PipelineCacheMiss(node_url) if status >= 400: - if status == 409: - raise _PipelineCacheMiss(f"cache miss at {node_url}") print( f" [node] relay hop {hop_index} returned {status} from {relay_addr}", flush=True, ) - return _PipelineResult(f"pipeline error at {node_url} via relay: status {status}") + return f"pipeline error at {node_url} via relay: status {status}", None except _PipelineCacheMiss: raise except Exception as exc: @@ -898,33 +884,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): resp_body = r.read() resp_headers = {k.lower(): v for k, v in r.headers.items()} except urllib.error.HTTPError as exc: - if exc.code == 409: - raise _PipelineCacheMiss(f"cache miss at {node_url}") from exc + body = exc.read() + if exc.code == 409 and _is_cache_miss_body(body): + raise _PipelineCacheMiss(node_url) from exc print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) - return _PipelineResult(f"pipeline error at {node_url}: {exc}") + return f"pipeline error at {node_url}: {exc}", None except Exception as exc: print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) - return _PipelineResult(f"pipeline error at {node_url}: {exc}") + return f"pipeline error at {node_url}: {exc}", None content_type = resp_headers.get("content-type", "") if "application/json" in content_type: try: data = json.loads(resp_body) - if data.get("error") == "cache_miss": - raise _PipelineCacheMiss(f"cache miss at {node_url}") text = str(data.get("text", "")) token_id = data.get("token_id") if server.debug: print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True) - return _PipelineResult(text, int(token_id) if token_id is not None else None) + return text, int(token_id) if token_id is not None else None except json.JSONDecodeError: - return _PipelineResult(resp_body.decode("utf-8", errors="replace")) + return resp_body.decode("utf-8", errors="replace"), None # Binary activation — update and forward to next node shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape)) current_shape = _parse_shape(shape_header) current_body = resp_body current_attn = resp_headers.get("x-meshnet-attn-mask") current_pos = resp_headers.get("x-meshnet-position-ids") - return _PipelineResult("") + return "", None def _stream_openai_response(self, token_iter, model: str) -> None: """Stream tokens from an iterator as SSE chunks.""" diff --git a/tests/test_kv_cache_distributed.py b/tests/test_kv_cache_distributed.py new file mode 100644 index 0000000..d08fa80 --- /dev/null +++ b/tests/test_kv_cache_distributed.py @@ -0,0 +1,384 @@ +"""AH-25: sharded per-node KV cache for distributed generation. + +Covers the SessionCacheStore (TTL + LRU + mismatch handling), the HTTP +session protocol (stable session id, O(1) decode payloads, 409 cache-miss +fallback, legacy stateless compatibility), and an env-gated golden test that +proves cached and stateless distributed generation produce identical tokens +on a real two-shard Qwen2.5-0.5B split. +""" + +import json +import os +import urllib.request + +import pytest + +from meshnet_node.model_backend import ( + KVCacheMiss, + SessionCacheStore, + TailTokenResult, + TensorPayload, +) +from meshnet_node.torch_server import TorchNodeServer + + +# --------------------------------------------------------------------------- +# SessionCacheStore units +# --------------------------------------------------------------------------- + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +def test_store_lookup_roundtrip_advances_lru(): + store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0, clock=_Clock()) + store.store("s1", cache=object(), seq_len=6, effective_start=12) + entry = store.lookup("s1", expected_seq_len=6, effective_start=12) + assert entry.seq_len == 6 + entry.seq_len += 1 + assert store.lookup("s1", expected_seq_len=7).seq_len == 7 + + +def test_lookup_unknown_session_raises_cache_miss(): + store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0) + with pytest.raises(KVCacheMiss): + store.lookup("nope") + + +def test_seq_len_mismatch_drops_entry_and_raises(): + store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0) + store.store("s1", cache=object(), seq_len=6, effective_start=0) + with pytest.raises(KVCacheMiss): + store.lookup("s1", expected_seq_len=9) + # Entry must be gone — a poisoned cache is never reused. + with pytest.raises(KVCacheMiss): + store.lookup("s1") + + +def test_effective_start_mismatch_raises(): + store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0) + store.store("s1", cache=object(), seq_len=6, effective_start=12) + with pytest.raises(KVCacheMiss): + store.lookup("s1", effective_start=21) + + +def test_ttl_expiry_evicts_stale_sessions(): + clock = _Clock() + store = SessionCacheStore(max_sessions=4, ttl_seconds=60.0, clock=clock) + store.store("s1", cache=object(), seq_len=6, effective_start=0) + clock.now = 61.0 + with pytest.raises(KVCacheMiss): + store.lookup("s1") + assert len(store) == 0 + + +def test_lru_eviction_bounds_session_count(): + clock = _Clock() + store = SessionCacheStore(max_sessions=2, ttl_seconds=1000.0, clock=clock) + store.store("s1", cache=object(), seq_len=1, effective_start=0) + store.store("s2", cache=object(), seq_len=1, effective_start=0) + store.lookup("s1") # s1 becomes most recent → s2 is LRU + store.store("s3", cache=object(), seq_len=1, effective_start=0) + assert len(store) == 2 + with pytest.raises(KVCacheMiss): + store.lookup("s2") + store.lookup("s1") + store.lookup("s3") + + +def test_drop_removes_session(): + store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0) + store.store("s1", cache=object(), seq_len=1, effective_start=0) + store.drop("s1") + with pytest.raises(KVCacheMiss): + store.lookup("s1") + + +# --------------------------------------------------------------------------- +# HTTP session protocol with fake cached backends +# --------------------------------------------------------------------------- + +class _ChatTokenizer: + eos_token = "" + + def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False): + return "debug prompt" + + +class _CachedHeadBackend: + model_id = "fake-model" + total_layers = 12 + is_head = True + is_tail = False + supports_kv_cache = True + tokenizer = _ChatTokenizer() + + def __init__(self) -> None: + self.prefills: list[str | None] = [] + self.decode_calls: list[tuple[int, str]] = [] + self.released: list[str] = [] + self._seq: dict[str, int] = {} + + def eos_token_ids(self) -> list[int]: + return [99] + + def release_session(self, session_id: str) -> None: + self.released.append(session_id) + + def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload: + self.prefills.append(session_id) + if session_id: + self._seq[session_id] = 6 + return TensorPayload( + body=b"\x00" * (1 * 6 * 8 * 2), + shape=[1, 6, 8], + attention_mask_header=None, + position_ids_header=None, + ) + + def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload: + self.decode_calls.append((token_id, session_id)) + past = self._seq[session_id] + self._seq[session_id] = past + 1 + return TensorPayload( + body=b"\x00" * (1 * 1 * 8 * 2), + shape=[1, 1, 8], + attention_mask_header=None, + position_ids_header=None, + past_len=past, + ) + + +class _CachedTailBackend: + model_id = "fake-model" + total_layers = 12 + is_head = False + is_tail = True + supports_kv_cache = True + + def __init__(self, tokens, miss_on_call: int | None = None) -> None: + self._tokens = list(tokens) + self.miss_on_call = miss_on_call + self.calls: list[dict] = [] + + def forward_bytes( + self, + body, + shape, + attention_mask_header, + position_ids_header, + start_layer=None, + session_id=None, + cache_mode=None, + past_len=None, + ): + call_index = len(self.calls) + self.calls.append({ + "session": session_id, + "mode": cache_mode, + "past_len": past_len, + "shape": list(shape), + }) + if self.miss_on_call is not None and call_index == self.miss_on_call: + raise KVCacheMiss("session evicted (test)") + text, token_id = self._tokens.pop(0) + return TailTokenResult(text=text, token_id=token_id) + + +def _chat_once(head_port: int, tail_port: int, max_tokens: int) -> str: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": max_tokens, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 6}, + ]), + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=10) as resp: + body = json.loads(resp.read()) + return body["choices"][0]["message"]["content"] + + +def test_session_is_stable_and_decode_payloads_are_single_token(): + head_backend = _CachedHeadBackend() + tail_backend = _CachedTailBackend([(" a", 1), (" b", 2), (" c", 3)]) + head = TorchNodeServer(backend=head_backend, tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + content = _chat_once(head_port, tail_port, max_tokens=3) + finally: + head.stop() + tail.stop() + + assert content == " a b c" + assert len(tail_backend.calls) == 3 + # Step 0 is a full-prompt prefill; steps 1+ carry only the new token. + assert tail_backend.calls[0]["mode"] == "prefill" + assert tail_backend.calls[0]["shape"] == [1, 6, 8] + for step, call in enumerate(tail_backend.calls[1:], start=1): + assert call["mode"] == "decode" + assert call["shape"] == [1, 1, 8] + assert call["past_len"] == 6 + (step - 1) + # One session id across every step of the generation. + sessions = {call["session"] for call in tail_backend.calls} + assert len(sessions) == 1 + session_id = sessions.pop() + assert head_backend.prefills == [session_id] + assert head_backend.decode_calls == [(1, session_id), (2, session_id)] + # Head releases its own session state when the generation ends. + assert head_backend.released == [session_id] + + +def test_eos_token_id_stops_generation(): + head_backend = _CachedHeadBackend() + tail_backend = _CachedTailBackend([(" a", 1), ("", 99)]) + head = TorchNodeServer(backend=head_backend, tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + content = _chat_once(head_port, tail_port, max_tokens=8) + finally: + head.stop() + tail.stop() + + assert content == " a" + assert len(tail_backend.calls) == 2 + + +def test_downstream_cache_miss_falls_back_to_full_reprefill(): + head_backend = _CachedHeadBackend() + # Call 1 (the first decode) raises KVCacheMiss → node answers 409 → + # head re-prefills the full sequence and keeps generating. + tail_backend = _CachedTailBackend( + [(" a", 1), (" b", 2), (" c", 3)], miss_on_call=1, + ) + head = TorchNodeServer(backend=head_backend, tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + content = _chat_once(head_port, tail_port, max_tokens=3) + finally: + head.stop() + tail.stop() + + assert content == " a b c" + modes = [call["mode"] for call in tail_backend.calls] + assert modes == ["prefill", "decode", "prefill", "decode"] + # Head re-prefilled once, with the same stable session id. + assert len(head_backend.prefills) == 2 + assert len(set(head_backend.prefills)) == 1 + + +def test_kv_head_with_legacy_tail_reprefills_every_step(): + """Mixed fleet: tail predates the protocol and returns no token_id.""" + + class _LegacyTailBackend: + model_id = "fake-model" + total_layers = 12 + is_head = False + is_tail = True + + def __init__(self) -> None: + self.calls = 0 + + def forward_bytes(self, body, shape, attention_mask_header, + position_ids_header, start_layer=None, **kwargs): + self.calls += 1 + return " x" if self.calls < 3 else "" + + head_backend = _CachedHeadBackend() + tail_backend = _LegacyTailBackend() + head = TorchNodeServer(backend=head_backend, tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + content = _chat_once(head_port, tail_port, max_tokens=5) + finally: + head.stop() + tail.stop() + + assert content == " x x" + # No token_id from the tail → every step is a full prefill (legacy cost), + # never a decode against a cache the tail doesn't keep. + assert head_backend.decode_calls == [] + assert len(head_backend.prefills) == 3 + + +# --------------------------------------------------------------------------- +# Golden test on a real two-shard split (env-gated: loads Qwen2.5-0.5B twice) +# --------------------------------------------------------------------------- + +_GOLDEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" + +requires_real_model = pytest.mark.skipif( + os.environ.get("MESHNET_REAL_MODEL_TESTS") != "1", + reason="set MESHNET_REAL_MODEL_TESTS=1 to run the real-model golden test", +) + + +@requires_real_model +def test_cached_distributed_generation_matches_stateless_golden(): + pytest.importorskip("torch") + from meshnet_node.model_backend import TorchModelShard + + head = TorchModelShard(_GOLDEN_MODEL, 0, 11) + tail = TorchModelShard(_GOLDEN_MODEL, 12, 23) + steps = 12 + prompt = head.tokenizer.apply_chat_template( + [{"role": "user", "content": "Count from 1 to 5."}], + add_generation_prompt=True, + tokenize=False, + ) + + # Reference: today's stateless path — re-encode the full sequence each step. + stateless_ids: list[int] = [] + text = prompt + for _ in range(steps): + payload = head.encode_prompt(text) + result = tail.forward_bytes( + payload.body, payload.shape, + payload.attention_mask_header, payload.position_ids_header, + start_layer=12, + ) + stateless_ids.append(result.token_id) + text += result.text + + # Cached path: one prefill, then single-token decode steps. + session = "golden-session" + cached_ids: list[int] = [] + payload = head.encode_prompt(prompt, session_id=session) + result = tail.forward_bytes( + payload.body, payload.shape, + payload.attention_mask_header, payload.position_ids_header, + start_layer=12, session_id=session, cache_mode="prefill", + ) + cached_ids.append(result.token_id) + for _ in range(steps - 1): + payload = head.encode_next_token(cached_ids[-1], session) + assert payload.shape[1] == 1, "decode payload must be a single token" + result = tail.forward_bytes( + payload.body, payload.shape, + None, payload.position_ids_header, + start_layer=12, session_id=session, cache_mode="decode", + past_len=payload.past_len, + ) + cached_ids.append(result.token_id) + + assert cached_ids == stateless_ids