Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

.gitignore

@@ -26,3 +26,4 @@ dist/
 logs/tracker/error.log
 logs/tracker/info.log
 logs/tracker/warning.log
+.venv*

.scratch/alpha-hardening/prd.json

@@ -529,4 +529,4 @@
   "metadata": {
     "updatedAt": "2026-07-08T23:30:00.000Z"
   }
-}
+}

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):
 
@@ -130,11 +132,116 @@ 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.12 -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.
+
+Use Python 3.12 for this env. Python 3.14 is currently a bad fit for the
+Qwen3.6/FLA path because `torch.compile` is not supported there.
+
+**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.
+- `which meshnet-node` should not point at `~/.local/bin/meshnet-node` for ROCm
+  testing. Run `.venv-rocm/bin/meshnet-node ...` so the node uses the same ROCm
+  PyTorch, `transformers`, and FLA packages you verified.
+- 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`**
@@ -356,13 +463,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.

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