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

misc

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 23:32:51 +03:00
parent d648da3344
commit 94046f1102
6 changed files with 644 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,9 +3,12 @@
from __future__ import annotations
import base64
from collections import OrderedDict
from dataclasses import dataclass
import json
import os
from pathlib import Path
import time
from typing import Any, Literal
Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
@@ -27,6 +30,10 @@ 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."""
@dataclass(frozen=True)
class TensorPayload:
body: bytes
@@ -35,6 +42,13 @@ class TensorPayload:
position_ids_header: str | None
@dataclass
class _ShardCacheEntry:
layer_states: list[Any]
seq_len: int
last_used: float
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")
@@ -163,6 +177,9 @@ 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()
def encode_prompt(self, prompt: str) -> TensorPayload:
if not self.is_head or self._embed_tokens is None:
@@ -174,12 +191,50 @@ 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_tokens(input_ids)
if self._position_embeddings is not None:
hidden_states = hidden_states + self._position_embeddings(position_ids)
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]),
)
return self._payload(hidden_states, attention_mask, position_ids)
def encode_token_cached(self, token_id: int, seq_len: int, session_id: str) -> TensorPayload:
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")
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)
hidden_states = self._run_layers(
hidden_states,
attention_mask,
position_ids,
session_id=session_id,
cache_mode="decode",
seq_len=int(seq_len),
)
return self._payload(hidden_states, attention_mask, position_ids)
def forward_bytes(
self,
body: bytes,
@@ -187,6 +242,9 @@ class TorchModelShard:
attention_mask_header: str | None,
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:
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
self.device
@@ -198,20 +256,31 @@ class TorchModelShard:
position_ids_header, self.torch, self.device
)
hidden_states = self._run_layers(
hidden_states, attention_mask, position_ids, start_layer=start_layer
hidden_states,
attention_mask,
position_ids,
start_layer=start_layer,
session_id=session_id,
cache_mode=cache_mode,
seq_len=seq_len,
)
if self.is_tail:
return self.decode_tail(hidden_states)
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._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)
def decode_tail_token_id(self, hidden_states: Any) -> int:
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 int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
def generate_text(
self,
@@ -328,6 +397,9 @@ class TorchModelShard:
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,
) -> Any:
# start_layer overrides shard_start for overlapping-shard routing
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
@@ -337,6 +409,20 @@ class TorchModelShard:
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()
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)
position_embeddings = _rotary_position_embeddings(
self.model,
hidden_states,
@@ -348,14 +434,28 @@ class TorchModelShard:
self.torch,
)
with self.torch.inference_mode():
for layer in self.layers[effective_start:self.shard_end + 1]:
hidden_states = _call_layer(
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(
layer,
hidden_states,
layer_attention_mask,
position_ids,
position_embeddings,
use_cache=use_cache,
past_key_value=past_state,
)
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:
@@ -371,6 +471,30 @@ 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,
@@ -718,19 +842,20 @@ 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, seq_len = attention_mask.shape
if seq_len <= 1:
batch_size, key_len = attention_mask.shape
query_len = int(hidden_states.shape[-2])
if key_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(
(seq_len, seq_len),
(query_len, key_len),
min_value,
dtype=hidden_states.dtype,
device=hidden_states.device,
)
causal = torch.triu(causal, diagonal=1)
causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()
causal = torch.triu(causal, diagonal=1 + key_len - query_len)
causal = causal[None, None, :, :].expand(batch_size, 1, query_len, key_len).clone()
padding = attention_mask.to(device=hidden_states.device)
if not bool(padding.all()):
@@ -754,21 +879,27 @@ def _call_layer(
attention_mask: Any,
position_ids: Any,
position_embeddings: Any | None = None,
) -> Any:
*,
use_cache: bool = False,
past_key_value: Any | None = None,
) -> tuple[Any, Any | None]:
attempts = (
{
"attention_mask": attention_mask,
"position_ids": position_ids,
"position_embeddings": position_embeddings,
"use_cache": False,
"past_key_value": past_key_value,
"use_cache": use_cache,
},
{
"attention_mask": attention_mask,
"position_ids": position_ids,
"use_cache": False,
"past_key_value": past_key_value,
"use_cache": use_cache,
},
{"attention_mask": attention_mask, "use_cache": False},
{"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},
{},
)
last_exc: Exception | None = None
@@ -776,12 +907,28 @@ def _call_layer(
filtered = {key: value for key, value in kwargs.items() if value is not None}
try:
output = layer(hidden_states, **filtered)
return output[0] if isinstance(output, tuple) else output
return _layer_hidden_and_cache(output)
except TypeError as exc:
last_exc = exc
if last_exc is not None:
raise last_exc
return layer(hidden_states)[0]
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
def _tensor_to_bytes(tensor: Any) -> bytes:

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

@@ -19,6 +19,8 @@ from .model_backend import (
InsufficientVRAMError,
MissingModelDependencyError,
Quantization,
ShardCacheMiss,
TensorPayload,
TorchModelShard,
validate_quantization,
)
@@ -31,6 +33,16 @@ 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("/"))
@@ -94,6 +106,48 @@ 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,
@@ -353,13 +407,28 @@ 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")),
)
except ShardCacheMiss 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, str):
self._send_json(200, {"text": result})
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)
return
response_body = _compress_body(result.body, encoding)
@@ -513,9 +582,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
return
# Distributed path: autoregressive generation across shards.
# We do N single-step forward passes (no cross-node KV cache), which is slow
# but correct. Each step: head encodes current sequence → forwards through route
# → tail returns the next token string → append → repeat.
# Step 0 prefills the full prompt and creates shard-local caches. Later
# cached steps send only the previous token's activation through the route.
remaining_route = self._get_remaining_route(model_name, backend=backend)
print(
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
@@ -547,6 +615,9 @@ 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
stream_emit = None
if stream:
@@ -560,11 +631,49 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
for step in range(max_tokens):
try:
payload = backend.encode_prompt(current_text)
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
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
except Exception as exc:
print(f" [node] distributed encode error: {exc}", flush=True)
break
token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
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.
@@ -573,6 +682,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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
@@ -687,7 +799,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
return []
def _run_downstream_pipeline(self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None) -> str:
def _run_downstream_pipeline(
self,
payload: object,
route: list[dict],
*,
backend: TorchModelShard | None = None,
session_id: str | None = None,
cache_mode: str = "stateless",
seq_len: int | None = None,
) -> _PipelineResult:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
if not route:
@@ -699,12 +820,14 @@ 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]
return active_backend.decode_tail(tensor)
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)
except Exception as exc:
return f"decode error: {exc}"
return "no downstream route available for non-tail shard"
return _PipelineResult(f"decode error: {exc}")
return _PipelineResult("no downstream route available for non-tail shard")
session = str(uuid.uuid4())
session = session_id 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]
@@ -733,7 +856,10 @@ 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 current_attn:
headers["X-Meshnet-Attn-Mask"] = current_attn
if current_pos:
@@ -744,11 +870,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
relay_addr, "/forward", current_body, headers, timeout=120.0,
)
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 f"pipeline error at {node_url} via relay: status {status}"
return _PipelineResult(f"pipeline error at {node_url} via relay: status {status}")
except _PipelineCacheMiss:
raise
except Exception as exc:
print(
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
@@ -767,26 +897,34 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
with urllib.request.urlopen(req, timeout=120.0) as r:
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
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return _PipelineResult(f"pipeline error at {node_url}: {exc}")
except Exception as exc:
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return f"pipeline error at {node_url}: {exc}"
return _PipelineResult(f"pipeline error at {node_url}: {exc}")
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 text
return _PipelineResult(text, int(token_id) if token_id is not None else None)
except json.JSONDecodeError:
return resp_body.decode("utf-8", errors="replace")
return _PipelineResult(resp_body.decode("utf-8", errors="replace"))
# 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 ""
return _PipelineResult("")
def _stream_openai_response(self, token_iter, model: str) -> None:
"""Stream tokens from an iterator as SSE chunks."""