popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-09 00:06:01 +03:00
parent 599aa44d97 5feb5b96f8
commit 08826f6ace
8 changed files with 952 additions and 323 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.scratch/alpha-hardening/README.md
View File

@@ -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

30
.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md
View File

@@ -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

8
.scratch/alpha-hardening/prd.json
View File

@@ -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"
}
}
}

1
QUICKSTART.md
View File

@@ -82,6 +82,7 @@ python -m venv .venv
.\.venv\Scripts\meshnet-node.exe --help
```
</details>
<details>

102
docs/adr/0022-sharded-per-node-kv-cache.md Normal file
View File

@@ -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=<cache>, 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.

471
packages/node/meshnet_node/model_backend.py
View File

@@ -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:

277
packages/node/meshnet_node/torch_server.py
View File

@@ -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."""

384
tests/test_kv_cache_distributed.py Normal file
View File

@@ -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