distributd cache

packages/node/meshnet_node/model_backend.py

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

packages/node/meshnet_node/torch_server.py

@@ -17,11 +17,17 @@ from typing import Any
 
 from .model_backend import (
     InsufficientVRAMError,
+    KVCacheMiss,
     MissingModelDependencyError,
     Quantization,
+    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,
@@ -128,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,
@@ -346,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,
@@ -353,11 +379,18 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 self.headers.get("X-Meshnet-Attn-Mask"),
                 self.headers.get("X-Meshnet-Position-Ids"),
                 start_layer=start_layer,
+                **forward_kwargs,
             )
+        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):
             self._send_json(200, {"text": result})
             return
@@ -512,10 +545,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 self._send_json(500, {"error": f"generation failed: {exc}"})
             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.
+        # 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} "
@@ -548,6 +583,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
         generated: list[str] = []
         current_text = prompt_text
 
+        session_id = str(uuid.uuid4())
+        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:
             stream_emit = self._start_openai_stream(model_name)
@@ -557,25 +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:
-                payload = backend.encode_prompt(current_text)
+                if use_kv and step > 0 and last_token_id is not None:
+                    try:
+                        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
-            token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
-            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)
-            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,
@@ -594,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)
@@ -687,7 +775,21 @@ 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: 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:
@@ -699,12 +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]
-                    return active_backend.decode_tail(tensor)
+                    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 f"decode error: {exc}"
-            return "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 = 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]
@@ -734,6 +841,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 "X-Meshnet-Hop-Index": str(hop_index),
                 "X-Meshnet-Start-Layer": str(start_layer),
             }
+            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:
@@ -743,12 +855,16 @@ 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:
                         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 f"pipeline error at {node_url} via relay: status {status}", None
+                except _PipelineCacheMiss:
+                    raise
                 except Exception as exc:
                     print(
                         f"  [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
@@ -767,26 +883,33 @@ 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:
+                    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 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 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)
                     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 text, int(token_id) if token_id is not None else None
                 except json.JSONDecodeError:
-                    return 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 ""
+        return "", None
 
     def _stream_openai_response(self, token_iter, model: str) -> None:
         """Stream tokens from an iterator as SSE chunks."""