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

distributd cache

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 22:53:03 +02:00
parent 436e872abe
commit daddbaa4a3
8 changed files with 966 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,8 +3,12 @@
from __future__ import annotations
import base64
from collections import OrderedDict
from dataclasses import dataclass
import json
import os
import threading
import time
from pathlib import Path
from typing import Any, Literal
@@ -27,12 +31,128 @@ class PartialModelLoadUnsupported(ModelBackendError):
"""Raised when a shard cannot be materialized from a local snapshot subset."""
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)
class TensorPayload:
body: bytes
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 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:
@@ -163,8 +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
# 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")
@@ -177,9 +303,44 @@ class TorchModelShard:
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)
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_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("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)
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, 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,
)
def forward_bytes(
self,
body: bytes,
@@ -187,7 +348,10 @@ class TorchModelShard:
attention_mask_header: str | None,
position_ids_header: str | None,
start_layer: int | None = None,
) -> TensorPayload | str:
session_id: str | None = None,
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
)
@@ -197,21 +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
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:
return self.decode_tail(hidden_states)
return self.decode_tail_token(hidden_states)
return self._payload(hidden_states, attention_mask, position_ids)
def decode_tail(self, hidden_states: Any) -> str:
return self.decode_tail_token(hidden_states).text
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)
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
return self.tokenizer.decode([token_id], skip_special_tokens=True)
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,
@@ -322,21 +511,108 @@ class TorchModelShard:
)
return dict(self.tokenizer(prompt, return_tensors="pt"))
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: str | None = None,
past_len: int | None = None,
) -> Any:
"""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":
# 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:
# 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
)
effective_start = self._effective_start(start_layer)
position_embeddings = _rotary_position_embeddings(
self.model,
hidden_states,
@@ -347,6 +623,12 @@ 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():
for layer in self.layers[effective_start:self.shard_end + 1]:
hidden_states = _call_layer(
@@ -355,6 +637,8 @@ class TorchModelShard:
layer_attention_mask,
position_ids,
position_embeddings,
cache=cache,
cache_position=cache_position,
)
return hidden_states.to(self.torch.bfloat16)
@@ -754,6 +1038,8 @@ def _call_layer(
attention_mask: Any,
position_ids: Any,
position_embeddings: Any | None = None,
cache: Any = None,
cache_position: Any = None,
) -> Any:
attempts = (
{
@@ -774,6 +1060,14 @@ def _call_layer(
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 output[0] if isinstance(output, tuple) else output