misc

.scratch/alpha-hardening/prd.json

@@ -520,13 +520,13 @@
         "Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction"
       ],
       "priority": 25,
-      "passes": false,
+      "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": ""
+      "completionNotes": "Completed by agent"
     }
   ],
   "metadata": {
-    "updatedAt": "2026-07-08T19:15:00.000Z"
+    "updatedAt": "2026-07-08T20:09:33.742Z"
   }
 }

QUICKSTART.md

@@ -3,7 +3,9 @@
 Get from zero to a live inference request in **three terminals**: install once, start
 the tracker, start a node, send a request.
 
-Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux, CPU inference.
+Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux CPU inference.
+ROCm GPU setup is covered below, but must be verified on the host because ROCm
+support depends on the exact AMD GPU/APU, kernel, driver, and ROCm runtime.
 
 **Active development models** (what we run day-to-day):
 
@@ -129,11 +131,110 @@ Install **one** torch line into the same env as `meshnet-node`:
 |----------|---------|
 | NVIDIA CUDA | `pip install torch` (default index) |
 | CPU only | `pip install torch --index-url https://download.pytorch.org/whl/cpu` |
-| AMD ROCm | `pip install torch --index-url https://download.pytorch.org/whl/rocm6.2` |
+| AMD ROCm | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3` |
 
 On Windows `.venv`, prefix with `.\.venv\Scripts\pip.exe`. Conda users with CUDA
 torch already installed can skip this step.
 
+### Linux AMD ROCm GPU install
+
+Use this when the node machine has an AMD GPU/APU and you want PyTorch to run on
+ROCm instead of CPU. The Python wheel is not enough by itself: the host must have
+working AMD GPU device access and a compatible ROCm runtime.
+
+**Host prerequisites:**
+
+1. Confirm the AMD GPU is visible:
+
+```bash
+lspci | grep -Ei 'vga|3d|display|amd|ati'
+ls -l /dev/kfd /dev/dri/renderD* 2>/dev/null
+```
+
+2. Make sure the node user can access GPU devices. AMD ROCm documents the normal
+   Linux permission path as membership in both `video` and `render`:
+
+```bash
+groups
+sudo usermod -a -G video,render "$LOGNAME"
+# log out and back in before continuing
+```
+
+3. Confirm the ROCm runtime tools work if they are installed:
+
+```bash
+rocminfo | head
+```
+
+If `rocminfo` is missing or cannot see the GPU, fix the host ROCm install first.
+Do not debug `meshnet-node` until this works.
+
+**Install ROCm PyTorch into the node env:**
+
+```bash
+cd /path/to/neuron-tai
+python3 -m venv .venv-rocm
+source .venv-rocm/bin/activate
+python -m pip install --upgrade pip setuptools wheel
+python -m pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay
+python -m pip install "transformers>=5.12" accelerate safetensors
+python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3
+```
+
+Keep this separate from a known-good CPU `.venv` until ROCm is verified on that
+machine. ROCm wheels are large and host-runtime-sensitive; a failed ROCm install
+should not break the CPU fallback environment.
+
+**Verify PyTorch sees ROCm:**
+
+```bash
+python - <<'PY'
+import torch
+print("torch", torch.__version__)
+print("hip", torch.version.hip)
+print("cuda api available", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("device", torch.cuda.get_device_name(0))
+    x = torch.ones((1,), device="cuda")
+    torch.cuda.synchronize()
+    print("tensor", x)
+PY
+```
+
+Expected: `torch.version.hip` is not `None`, `torch.cuda.is_available()` is
+`True`, and the tensor allocation succeeds. PyTorch intentionally exposes ROCm
+through the `torch.cuda` API.
+
+**Start an AMD ROCm node:**
+
+```bash
+HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start \
+  --tracker <tracker-url> \
+  --model Qwen/Qwen2.5-0.5B-Instruct \
+  --quantization bfloat16
+```
+
+For the Qwen3.6 alpha model on Linux ROCm, install the optional FLA ROCm fast
+path in the same env:
+
+```bash
+.venv-rocm/bin/pip install 'flash-linear-attention[rocm]'
+HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start \
+  --tracker <tracker-url> \
+  --model qwen3.6-35b-a3b \
+  --quantization bfloat16
+```
+
+**Troubleshooting notes:**
+
+- `torch.version.hip is None` means you installed a CPU/CUDA torch build, not ROCm.
+- `torch.cuda.is_available() == False` with a ROCm build usually means host driver,
+  permissions, unsupported hardware, or missing runtime libraries.
+- Missing libraries such as `libamdhip64.so`, `libMIOpen.so`, `librocsolver.so`,
+  or `libroctx64.so` are host ROCm runtime problems, not meshnet-node problems.
+- Some AMD APUs and consumer GPUs require newer ROCm/Radeon support than server
+  Instinct cards. Check AMD's ROCm Radeon/Ryzen support matrix for the exact model.
+
 ### Qwen3.5/3.6-MoE notes
 
 Applies to **`qwen3.6-35b-a3b`** and other hybrid linear-attention models. **`Qwen2.5-0.5B`**
@@ -355,13 +456,20 @@ meshnet-node start --tracker http://192.168.0.179:8080 --model qwen3.6-35b-a3b -
 
 Do not add `causal-conv1d` or `flash-linear-attention[cuda]` on Windows (see Qwen3.5/3.6 notes).
 
-**Alpha model (Qwen3.6, Linux GPU — with fast path):**
+**Alpha model (Qwen3.6, Linux NVIDIA GPU — with fast path):**
 
 ```bash
 HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker <tracker-url> --model qwen3.6-35b-a3b --quantization bfloat16
 # Install once on that machine: pip install flash-linear-attention[cuda]
 ```
 
+**Alpha model (Qwen3.6, Linux AMD ROCm GPU — with fast path):**
+
+```bash
+HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start --tracker <tracker-url> --model qwen3.6-35b-a3b --quantization bfloat16
+# Install once on that machine: .venv-rocm/bin/pip install 'flash-linear-attention[rocm]'
+```
+
 After the first node registers a model, later nodes can join with only the tracker
 URL (shard auto-assigned):
 

docs/adr/0022-sharded-per-node-generation-cache.md

@@ -0,0 +1,63 @@
+# ADR-0022: Sharded per-node generation cache for distributed PyTorch routes
+
+## Status: Accepted
+
+## Context
+
+The distributed PyTorch chat path previously recomputed the full prompt-so-far for
+every generated token. The head shard embedded the entire sequence each step, forwarded
+full-sequence activations through every downstream shard, and every shard called its
+decoder layers with `use_cache=False`. On a two-node Qwen2.5-0.5B route this produced
+the expected quadratic slowdown as output length grew.
+
+ADR-0020 and ADR-0021 fixed route construction and `start_layer` semantics. They did not
+define the per-request cache lifecycle needed for efficient decode.
+
+## Decision
+
+Distributed PyTorch generation now uses one stable route session id for an entire chat
+request. The wire protocol marks each activation hop with:
+
+- `X-Meshnet-Session`: stable per generation.
+- `X-Meshnet-Cache-Mode`: `prefill`, `decode`, or `stateless`.
+- `X-Meshnet-Seq-Len`: the total sequence length represented by the step.
+
+Step 0 is prefill: the head sends the full prompt activation through the planned route.
+Each shard stores only the opaque cache state returned by its own executed layer range.
+No shard receives or stores another shard's cache.
+
+Later cached decode steps send only the newest token activation (`[1, 1, hidden]`) with
+the full sequence length and newest position id. The backend deliberately treats layer
+cache state as opaque. Standard K/V tuples, HuggingFace cache objects, and hybrid
+linear-attention recurrent state are stored without shape assumptions.
+
+## Cache lifecycle
+
+Each `TorchModelShard` owns an in-memory LRU map keyed by
+`(session_id, effective_start_layer, shard_end)`. Entries expire by TTL and by a maximum
+session count (`MESHNET_SHARD_CACHE_TTL_SECONDS`, default 600;
+`MESHNET_SHARD_CACHE_MAX_SESSIONS`, default 16).
+
+If a decode step reaches a node after restart, eviction, TTL expiry, or route mismatch,
+the node returns an explicit `cache_miss` response. The head falls back to full prefill
+for the current prompt-so-far using the same session id, rebuilding the shard-local
+caches before continuing. Alpha route repair still does not migrate cache state across
+nodes; a true route change is treated as cache loss and recovered by re-prefill.
+
+## Consequences
+
+- Healthy decode sends O(1) activation payloads per token between nodes instead of
+  O(sequence length).
+- Cache internals stay behind the model backend boundary, which keeps Qwen3.6-style
+  hybrid recurrent cache state compatible with the same route protocol.
+- Restart and eviction degrade to slower stateless/full-prefill work rather than silent
+  output corruption.
+- Cross-node cache migration, batching cache state across sessions, and speculative
+  decoding remain future work.
+
+## Verification
+
+Unit coverage in `tests/test_real_model_backend.py` verifies opaque per-layer cache
+storage, cached one-token decode, explicit cache-miss errors, and LRU eviction. Live
+two-node Qwen2.5-0.5B TPS measurement still requires the physical two-machine topology
+used to observe the regression.

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)
+        hidden_states = self._embed_input_ids(input_ids, position_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)
         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 int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
-        return self.tokenizer.decode([token_id], skip_special_tokens=True)
 
     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]:
+            next_layer_states: list[Any] = []
-                hidden_states = _call_layer(
+            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
+    batch_size, key_len = attention_mask.shape
-    if seq_len <= 1:
+    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 = torch.triu(causal, diagonal=1 + key_len - query_len)
-    causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()
+    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},
+        {"attention_mask": attention_mask, "past_key_value": past_key_value, "use_cache": use_cache},
-        {"use_cache": False},
+        {"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:

packages/node/meshnet_node/torch_server.py

@@ -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
+        # Step 0 prefills the full prompt and creates shard-local caches. Later
-        # but correct. Each step: head encodes current sequence → forwards through route
+        # cached steps send only the previous token's activation through the route.
-        # → tail returns the next token string → append → repeat.
         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 _PipelineResult(f"decode error: {exc}")
-            return "no downstream route available for non-tail shard"
+            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."""

tests/test_real_model_backend.py

@@ -1,5 +1,6 @@
 """US-012 tests for the real PyTorch node backend."""
 
+from collections import OrderedDict
 import json
 import os
 from pathlib import Path
@@ -14,6 +15,7 @@ import pytest
 from meshnet_node.model_backend import (
     InsufficientVRAMError,
     PartialModelLoadUnsupported,
+    ShardCacheMiss,
     TensorPayload,
     TorchModelShard,
     _call_layer,
@@ -43,7 +45,15 @@ class _FakeBackend:
             position_ids_header=None,
         )
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         assert shape == [1, 6, 8]
         return TensorPayload(
             body=body,
@@ -57,7 +67,15 @@ class _FakeTailBackend(_FakeBackend):
     is_head = False
     is_tail = True
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         assert len(body) == 1 * 6 * 8 * 2
         return " Paris"
 
@@ -114,7 +132,15 @@ class _FakePipelineTailBackend(_FakeTailBackend):
     def __init__(self) -> None:
         self.start_layers: list[int | None] = []
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         self.start_layers.append(start_layer)
         assert len(body) == 1 * 6 * 8 * 2
         return " token"
@@ -125,7 +151,15 @@ class _BlockingStreamingTailBackend(_FakeTailBackend):
         self._release = second_token_release
         self.calls = 0
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         self.calls += 1
         if self.calls == 1:
             return " first"
@@ -488,13 +522,118 @@ def test_call_layer_passes_rotary_position_embeddings():
             assert kwargs["position_embeddings"] == "rotary"
             return hidden_states
 
-    assert _call_layer(
+    hidden, cache_state = _call_layer(
         NeedsPositionEmbeddings(),
         "hidden",
         attention_mask=None,
         position_ids="positions",
         position_embeddings="rotary",
-    ) == "hidden"
+    )
+
+    assert hidden == "hidden"
+    assert cache_state is None
+
+
+def _fake_cache_shard(torch, *, max_sessions=16, ttl=600.0):
+    class RecordingLayer:
+        def __init__(self, index):
+            self.index = index
+            self.calls = []
+
+        def __call__(self, hidden_states, **kwargs):
+            self.calls.append({
+                "shape": tuple(hidden_states.shape),
+                "use_cache": kwargs.get("use_cache"),
+                "past_key_value": kwargs.get("past_key_value"),
+            })
+            present = {
+                "layer": self.index,
+                "shape": tuple(hidden_states.shape),
+                "opaque": object(),
+            }
+            return hidden_states + (self.index + 1), present
+
+    shard = object.__new__(TorchModelShard)
+    shard.shard_start = 0
+    shard.shard_end = 1
+    shard.torch = torch
+    shard.model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[]))
+    shard.layers = [RecordingLayer(0), RecordingLayer(1)]
+    shard._session_cache = OrderedDict()
+    shard._cache_max_sessions = max_sessions
+    shard._cache_ttl_seconds = ttl
+    return shard
+
+
+def test_shard_cache_prefill_then_decode_reuses_opaque_layer_state():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch)
+
+    prefill_hidden = torch.zeros((1, 4, 2), dtype=torch.bfloat16)
+    prefill_mask = torch.ones((1, 4), dtype=torch.long)
+    prefill_positions = torch.arange(4, dtype=torch.long).reshape(1, 4)
+    shard._run_layers(
+        prefill_hidden,
+        prefill_mask,
+        prefill_positions,
+        session_id="session-1",
+        cache_mode="prefill",
+        seq_len=4,
+    )
+
+    assert len(shard._session_cache) == 1
+    cached_states = next(iter(shard._session_cache.values())).layer_states
+    assert len(cached_states) == 2
+    assert cached_states[0]["shape"] == (1, 4, 2)
+
+    decode_hidden = torch.zeros((1, 1, 2), dtype=torch.bfloat16)
+    decode_mask = torch.ones((1, 5), dtype=torch.long)
+    decode_positions = torch.tensor([[4]], dtype=torch.long)
+    shard._run_layers(
+        decode_hidden,
+        decode_mask,
+        decode_positions,
+        session_id="session-1",
+        cache_mode="decode",
+        seq_len=5,
+    )
+
+    assert shard.layers[0].calls[-1]["shape"] == (1, 1, 2)
+    assert shard.layers[0].calls[-1]["past_key_value"] is cached_states[0]
+    assert shard.layers[1].calls[-1]["past_key_value"] is cached_states[1]
+    assert next(iter(shard._session_cache.values())).seq_len == 5
+
+
+def test_shard_cache_decode_miss_is_explicit():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch)
+
+    with pytest.raises(ShardCacheMiss):
+        shard._run_layers(
+            torch.zeros((1, 1, 2), dtype=torch.bfloat16),
+            torch.ones((1, 5), dtype=torch.long),
+            torch.tensor([[4]], dtype=torch.long),
+            session_id="missing",
+            cache_mode="decode",
+            seq_len=5,
+        )
+
+
+def test_shard_cache_lru_bounds_sessions():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch, max_sessions=1)
+
+    for session in ("old", "new"):
+        shard._run_layers(
+            torch.zeros((1, 2, 2), dtype=torch.bfloat16),
+            torch.ones((1, 2), dtype=torch.long),
+            torch.arange(2, dtype=torch.long).reshape(1, 2),
+            session_id=session,
+            cache_mode="prefill",
+            seq_len=2,
+        )
+
+    assert list(shard._session_cache.keys()) == [("new", 0, 1)]
 
 
 def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):