diff --git a/.scratch/alpha-hardening/prd.json b/.scratch/alpha-hardening/prd.json index f6a20e6..74bb7aa 100644 --- a/.scratch/alpha-hardening/prd.json +++ b/.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" } } \ No newline at end of file diff --git a/QUICKSTART.md b/QUICKSTART.md index e7959a6..f11417f 100644 --- a/QUICKSTART.md +++ b/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 \ + --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 \ + --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 --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 --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): diff --git a/docs/adr/0022-sharded-per-node-generation-cache.md b/docs/adr/0022-sharded-per-node-generation-cache.md new file mode 100644 index 0000000..5b2df20 --- /dev/null +++ b/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. diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 49d18a2..d62136e 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/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: diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 3bc8605..38c6704 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/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 - # 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.""" diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index 38815fe..64640c0 100644 --- a/tests/test_real_model_backend.py +++ b/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):