misc

packages/node/meshnet_node/model_backend.py

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