popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-09 00:06:01 +03:00
parent 599aa44d97 5feb5b96f8
commit 08826f6ace
8 changed files with 952 additions and 323 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.scratch/alpha-hardening/README.md
View File

@@ -11,7 +11,7 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati
**Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`. **Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`.
**Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is `ready-for-agent`. The ADR-0020 mixed-topology `start_layer` bug is fixed, but the distributed generation loop still has no KV cache at all — every step re-encodes the full sequence and re-runs every layer on every node, causing quadratic tps decay observed live (22.3 → 12.6 tps over one generation on a 2-node Qwen2.5-0.5B pipeline). Must be architecture-aware: Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V. **Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is **implemented** ([ADR-0022](../../docs/adr/0022-sharded-per-node-kv-cache.md)): per-generation session ids, prefill/decode wire protocol (`X-Meshnet-Cache`/`X-Meshnet-Past-Len`), per-node sharded `DynamicCache(config=…)` (hybrid-attention-aware), TTL+LRU eviction with 409 cache-miss → full re-prefill fallback. Golden test proves token-identical output vs the stateless path; CPU two-shard measurement: 7.05 tps decaying 32% → 18.93 tps flat (2.68×). Remaining: re-measure on the live 2-node GPU topology and the Qwen3.6-35B-A3B mixed topology.
## Artifacts ## Artifacts

30
.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md
View File

@@ -1,4 +1,14 @@
Status: ready-for-agent Status: implemented 2026-07-08 — pending live 2-node GPU verification
Implemented in `packages/node/meshnet_node/model_backend.py` + `torch_server.py`; design in
[ADR-0022](../../../docs/adr/0022-sharded-per-node-kv-cache.md); tests in
`tests/test_kv_cache_distributed.py` (11 fast tests + env-gated golden test,
`MESHNET_REAL_MODEL_TESTS=1`).
**Measured (two-shard Qwen2.5-0.5B 0-11/12-23, CPU, 44-token prompt, 40 steps):**
stateless 7.05 tps decaying 32% (8.09 → 5.50 first-10 vs last-10); cached 18.93 tps and
FLAT (17.21 → 19.28) — 2.68× overall, gap grows quadratically with length. Remaining
acceptance item: re-measure on the live 2-node GPU topology (needs both machines).
Scoped 2026-07-08 from a live two-machine distributed-inference debugging session (Qwen2.5-0.5B GPU+GPU pipeline, and Qwen3.6-35B-A3B mixed GPU/CPU). The ADR-0020 mixed-topology `start_layer` bug is fixed (`518c259`, `e44abc9`, `1ecc599`); this issue is the next performance blocker in the same code path. Scoped 2026-07-08 from a live two-machine distributed-inference debugging session (Qwen2.5-0.5B GPU+GPU pipeline, and Qwen3.6-35B-A3B mixed GPU/CPU). The ADR-0020 mixed-topology `start_layer` bug is fixed (`518c259`, `e44abc9`, `1ecc599`); this issue is the next performance blocker in the same code path.
@@ -33,15 +43,15 @@ Observed cost of this on a live 2-node Qwen2.5-0.5B GPU pipeline (layers 0-20 /
## Acceptance criteria ## Acceptance criteria
- [ ] A session ID is stable across all steps of one chat generation (not re-minted per token) - [x] A session ID is stable across all steps of one chat generation (not re-minted per token) — minted once in `_do_chat_completions`, asserted in `test_session_is_stable_and_decode_payloads_are_single_token`
- [ ] Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes - [x] Steps after the first prefill send only the new token's activation (`[1, 1, hidden]` via `encode_next_token`) with `X-Meshnet-Cache: decode` + `X-Meshnet-Past-Len`
- [ ] Each node caches `past_key_values`/recurrent state only for its own shard's layer range; no node ever holds another node's cache - [x] Each node caches state only for its own shard's layer range (`TorchModelShard.kv_sessions`; sharding falls out of per-node layer execution)
- [ ] Cache works correctly for both standard-attention shards and Qwen3.6-style hybrid linear-attention/recurrent shards (cache abstraction is not K/V-shaped-only) - [x] Cache abstraction is not K/V-shaped-only: `DynamicCache(config=model.config)` — the same construction Qwen3.6-Next's own forward uses for hybrid linear-attention conv/delta state; store treats it as opaque; `TypeError` fallback disables caching per-backend
- [ ] Bounded memory: TTL + LRU eviction; eviction/restart triggers a documented cache-miss response, not silent corruption or an unhandled exception - [x] Bounded memory: TTL (600 s, `MESHNET_KV_TTL_SECONDS`) + LRU (8, `MESHNET_KV_MAX_SESSIONS`); miss → HTTP 409 `{"error": "cache_miss"}` → head re-prefills (tested)
- [ ] Golden-output regression test proves cached and uncached distributed generation produce equivalent output for a fixed prompt - [x] Golden-output test: cached and stateless produce identical token ids on real two-shard Qwen2.5-0.5B (`test_cached_distributed_generation_matches_stateless_golden`, passed)
- [ ] Measured tps improvement recorded on the same 2-node Qwen2.5-0.5B topology used to observe the regression (target: flat tps across generation length, not decaying) - [x] Measured (CPU two-shard proxy, 40 steps): stateless 7.05 tps w/ 32% decay → cached 18.93 tps flat, 2.68×. ⚠️ still to run on the live 2-node GPU topology
- [ ] `tests/test_two_node_pipeline.py` and `tests/test_dynamic_routing.py` still pass - [x] `tests/test_two_node_pipeline.py` and `tests/test_dynamic_routing.py` pass (30 passed; 6 tmp-dir fixture errors are a pre-existing Windows temp-permission env issue, identical on clean tree)
- [ ] Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction - [x] Design captured in [ADR-0022](../../../docs/adr/0022-sharded-per-node-kv-cache.md) incl. cache-miss/route-change interaction with ADR-0021
## Notes ## Notes

6
.scratch/alpha-hardening/prd.json
View File

@@ -507,7 +507,7 @@
{ {
"id": "AH-025", "id": "AH-025",
"title": "25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)", "title": "25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)",
"description": "Status: ready-for-agent\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.", "description": "Status: implemented 2026-07-08 — pending live 2-node GPU verification\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.",
"acceptanceCriteria": [ "acceptanceCriteria": [
"A session ID is stable across all steps of one chat generation (not re-minted per token)", "A session ID is stable across all steps of one chat generation (not re-minted per token)",
"Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes", "Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes",
@@ -523,10 +523,10 @@
"passes": true, "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.", "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": [], "dependsOn": [],
"completionNotes": "Completed by agent" "completionNotes": "Implemented 2026-07-08 (ADR-0022, docs/adr/0022-sharded-per-node-kv-cache.md). Per-generation session id; X-Meshnet-Cache prefill/decode + X-Meshnet-Past-Len wire headers; decode steps send [1,1,hidden] via encode_next_token (tail now returns token_id so the head never re-tokenizes); per-node SessionCacheStore holds DynamicCache(config=model.config) — hybrid-attention/recurrent-state aware, sharded naturally by each node's own layer range; TTL (600s) + LRU (8) eviction; 409 {\"error\":\"cache_miss\"} -> head re-prefills full sequence under the same session (stateless path kept as recovery mode; legacy nodes without the protocol degrade to per-step prefill). Tests: tests/test_kv_cache_distributed.py — 11 fast tests + env-gated golden test (MESHNET_REAL_MODEL_TESTS=1) proving token-identical cached vs stateless output on a real two-shard Qwen2.5-0.5B split. Measured (CPU two-shard, 40 steps): stateless 7.05 tps decaying 32% -> cached 18.93 tps flat, 2.68x overall. Remaining: re-measure on the live 2-node GPU topology and Qwen3.6-35B-A3B mixed topology (needs both machines)."
} }
], ],
"metadata": { "metadata": {
"updatedAt": "2026-07-08T20:09:33.742Z" "updatedAt": "2026-07-08T23:30:00.000Z"
} }
} }

1
QUICKSTART.md
View File

@@ -82,6 +82,7 @@ python -m venv .venv
.\.venv\Scripts\meshnet-node.exe --help .\.venv\Scripts\meshnet-node.exe --help
``` ```
</details> </details>
<details> <details>

102
docs/adr/0022-sharded-per-node-kv-cache.md Normal file
View File

@@ -0,0 +1,102 @@
# ADR-0022: Sharded per-node KV cache for distributed generation
## Status: Accepted, implemented (alpha-hardening issue 25)
## Context
The distributed generation loop (`torch_server.py`, `_do_chat_completions` distributed
path) had **no KV cache**: every layer-forward call passed `use_cache: False`, and each
autoregressive step re-encoded the entire prompt-so-far from scratch, re-running every
layer on every node in the route for every generated token. Measured on a live 2-node
Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 to 12.6 within a single generation —
the quadratic-cost signature. On Qwen3.6-35B-A3B mixed GPU/CPU topology this collapsed
to ~0.07 tps even after the ADR-0020 routing fix.
`X-Meshnet-Session` existed on the wire but was minted fresh **per token** and keyed no
state.
## Decision
### Session lifecycle
The head mints one session id per chat generation (not per token) and reuses it across
every step. Two new request headers extend the `/forward` wire protocol:
- `X-Meshnet-Cache: prefill | decode` — absent means legacy stateless (unchanged
behavior, and what old nodes send/understand).
- `X-Meshnet-Past-Len: N` — decode only: the number of tokens the node's session cache
must already hold. A mismatch is a cache miss, never silent corruption.
Step 0 (`prefill`) sends the full prompt activation as before; each node creates fresh
session state for its own layer range. Steps 1+ (`decode`) send only the newest token's
hidden state — `[1, 1, hidden]`, cutting per-step compute and wire payload from
O(seq_len) to O(1). The head embeds the next token directly from the `token_id` the tail
now returns alongside text (`{"text": …, "token_id": …}`), avoiding text
re-tokenization drift; EOS is detected by id against tokenizer + generation-config eos
sets.
### Per-node sharded cache
`TorchModelShard.kv_sessions` is a `SessionCacheStore`: `session_id → SessionCacheEntry`
holding cache state **only for that shard's layer range** — sharding falls out naturally
because each node only executes (and therefore only caches) its own layers. No node ever
holds another node's state.
### MoE / hybrid-attention awareness
The cached object is whatever `use_cache=True` produces: a transformers
`DynamicCache(config=model.config)` — the same construction the model's own `forward()`
uses. With the config, transformers picks the right per-layer state: K/V tensors for
standard attention, conv/recurrent delta state for Qwen3.6-style hybrid linear-attention
layers, sliding-window variants, etc. The store treats it as opaque; nothing assumes a
K/V tensor shape. Cache slots are indexed by absolute `layer_idx`, so a shard updating
only layers 12–23 leaves 0–11 empty (verified: sparse `DynamicCache.update` works).
MoE expert routing is layer-local per token and needs no cross-token state.
Layers are invoked with `past_key_values=<cache>, use_cache=True, cache_position=…`
(transformers 5.x layer API; the cache is mutated in place). If a model's layers reject
those kwargs, the backend logs once, sets `supports_kv_cache = False`, and stays on the
stateless path permanently — exotic architectures degrade to today's behavior instead of
failing.
### Cache miss and route-change interaction (ADR-0021)
Any decode-mode request that cannot be served — unknown session (evicted, node
restarted), `past_len` mismatch, `start_layer` mismatch (the route or shard overlap
changed mid-generation), or caching disabled — raises `KVCacheMiss`, answered as
**HTTP 409 `{"error": "cache_miss"}`**. The head catches it and falls back to one full
re-prefill of the accumulated sequence under the same session id, which atomically
replaces every node's session state, then continues cached. The fully-stateless path is
therefore still the recovery mode: eviction and restarts cost one prefill, never
corruption or a failed generation. A decode request against a node whose caching is
disabled is also a 409 — running a single-token payload statelessly would silently
produce garbage.
Mixed fleets degrade the same way: if the tail predates the protocol and returns no
`token_id`, the head simply prefills every step (exactly the old cost).
### Bounded memory
`SessionCacheStore` enforces TTL (default 600 s, `MESHNET_KV_TTL_SECONDS`) plus LRU cap
(default 8 sessions, `MESHNET_KV_MAX_SESSIONS`), evaluated on every access. The head
additionally drops its own session explicitly when a generation completes; downstream
nodes rely on TTL/LRU (an explicit cross-node release RPC was judged not worth the
failure modes — misses are cheap).
### Non-goals (first landing)
Cross-node cache migration on route change (evict + re-prefill is acceptable),
speculative decoding, cross-session batching.
## Consequences
- Per-token cost drops from O(seq_len) layer re-execution + O(seq_len) wire transfer per
hop to O(1) of both; tps stays flat across generation length instead of decaying.
- Golden test (`tests/test_kv_cache_distributed.py`, env-gated by
`MESHNET_REAL_MODEL_TESTS=1`) proves cached and stateless distributed generation emit
identical token ids on a real two-shard Qwen2.5-0.5B split.
- Nodes now hold per-session GPU/CPU memory between requests (bounded above); operators
sizing `max_loaded_shards` should account for ~`sessions × seq_len × kv_bytes_per_token`
per resident model.
- The wire protocol is backward- and forward-compatible: headers are additive, absent
headers mean stateless, and 409 is only sent in reply to explicit decode-mode requests.

469
packages/node/meshnet_node/model_backend.py
View File

@@ -7,8 +7,9 @@ from collections import OrderedDict
from dataclasses import dataclass from dataclasses import dataclass
import json import json
import os import os
from pathlib import Path import threading
import time import time
from pathlib import Path
from typing import Any, Literal from typing import Any, Literal
Quantization = Literal["auto", "bfloat16", "int8", "nf4"] Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
@@ -30,8 +31,12 @@ class PartialModelLoadUnsupported(ModelBackendError):
"""Raised when a shard cannot be materialized from a local snapshot subset.""" """Raised when a shard cannot be materialized from a local snapshot subset."""
class ShardCacheMiss(ModelBackendError): class KVCacheMiss(ModelBackendError):
"""Raised when a decode step arrives after the shard-local cache was evicted.""" """Raised when a decode step references session state this node no longer holds.
The head recovers by re-prefilling the full sequence (the stateless path),
so eviction or a node restart degrades throughput instead of corrupting output.
"""
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -40,15 +45,116 @@ class TensorPayload:
shape: list[int] shape: list[int]
attention_mask_header: str | None attention_mask_header: str | None
position_ids_header: str | None position_ids_header: str | None
# Number of tokens already cached before this payload's tokens (decode steps).
past_len: int | None = None
@dataclass(frozen=True)
class TailTokenResult:
"""Tail-shard decode result: decoded text plus the raw token id.
The token id lets the head feed the next decode step (and detect EOS)
without re-tokenizing text, which is not guaranteed to round-trip.
"""
text: str
token_id: int
@dataclass @dataclass
class _ShardCacheEntry: class SessionCacheEntry:
layer_states: list[Any] """Per-session cached state for one shard's layer range.
`cache` is whatever `use_cache=True` produces for these layers — a
transformers Cache holding K/V tensors for standard attention, or
recurrent conv/delta state for hybrid linear-attention layers. The store
treats it as opaque.
"""
cache: Any
seq_len: int seq_len: int
effective_start: int
last_used: float last_used: float
class SessionCacheStore:
"""TTL + LRU bounded map of session_id → SessionCacheEntry.
Each node caches state only for its own layer range; no node ever holds
another node's cache. Stale or mismatched entries raise KVCacheMiss so the
head falls back to a full re-prefill instead of producing corrupt output.
"""
def __init__(
self,
max_sessions: int = 8,
ttl_seconds: float = 600.0,
clock: Any = None,
) -> None:
self.max_sessions = max(1, int(max_sessions))
self.ttl_seconds = float(ttl_seconds)
self._clock = clock or time.monotonic
self._entries: OrderedDict[str, SessionCacheEntry] = OrderedDict()
self._lock = threading.Lock()
def __len__(self) -> int:
with self._lock:
return len(self._entries)
def store(self, session_id: str, cache: Any, seq_len: int, effective_start: int) -> SessionCacheEntry:
now = self._clock()
with self._lock:
self._entries.pop(session_id, None)
entry = SessionCacheEntry(cache, seq_len, effective_start, now)
self._entries[session_id] = entry
self._evict_locked(now)
return entry
def lookup(
self,
session_id: str,
*,
expected_seq_len: int | None = None,
effective_start: int | None = None,
) -> SessionCacheEntry:
now = self._clock()
with self._lock:
self._evict_locked(now)
entry = self._entries.get(session_id)
if entry is None:
raise KVCacheMiss(f"no cached state for session {session_id[:8]}")
if expected_seq_len is not None and entry.seq_len != expected_seq_len:
del self._entries[session_id]
raise KVCacheMiss(
f"session {session_id[:8]} cache holds {entry.seq_len} tokens, "
f"expected {expected_seq_len}"
)
if effective_start is not None and entry.effective_start != effective_start:
del self._entries[session_id]
raise KVCacheMiss(
f"session {session_id[:8]} cached with start_layer "
f"{entry.effective_start}, requested {effective_start}"
)
entry.last_used = now
self._entries.move_to_end(session_id)
return entry
def drop(self, session_id: str) -> None:
with self._lock:
self._entries.pop(session_id, None)
def _evict_locked(self, now: float) -> None:
if self.ttl_seconds > 0:
expired = [
sid for sid, entry in self._entries.items()
if now - entry.last_used > self.ttl_seconds
]
for sid in expired:
del self._entries[sid]
while len(self._entries) > self.max_sessions:
self._entries.popitem(last=False)
def validate_quantization(value: str) -> Quantization: def validate_quantization(value: str) -> Quantization:
if value not in {"auto", "bfloat16", "int8", "nf4"}: if value not in {"auto", "bfloat16", "int8", "nf4"}:
raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4") raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4")
@@ -177,11 +283,14 @@ class TorchModelShard:
self._position_embeddings = _position_embeddings(self.model) self._position_embeddings = _position_embeddings(self.model)
self._norm = _final_norm(self.model) if self.is_tail else None 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._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")) # Per-session KV/recurrent-state cache for this shard's layer range.
self._cache_max_sessions = max(1, int(os.environ.get("MESHNET_SHARD_CACHE_MAX_SESSIONS", "16"))) self.supports_kv_cache = True
self._session_cache: OrderedDict[tuple[str, int, int], _ShardCacheEntry] = OrderedDict() self.kv_sessions = SessionCacheStore(
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
)
def encode_prompt(self, prompt: str) -> TensorPayload: def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
if not self.is_head or self._embed_tokens is None: if not self.is_head or self._embed_tokens is None:
raise ModelBackendError("text prompts can only be accepted by the head shard") raise ModelBackendError("text prompts can only be accepted by the head shard")
encoded = self.tokenizer(prompt, return_tensors="pt") encoded = self.tokenizer(prompt, return_tensors="pt")
@@ -191,49 +300,46 @@ class TorchModelShard:
attention_mask = self.torch.ones_like(input_ids) attention_mask = self.torch.ones_like(input_ids)
attention_mask = attention_mask.to(self.device) attention_mask = attention_mask.to(self.device)
position_ids = _position_ids(attention_mask, self.torch) position_ids = _position_ids(attention_mask, self.torch)
hidden_states = self._embed_input_ids(input_ids, position_ids) hidden_states = self._embed_tokens(input_ids)
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids) if self._position_embeddings is not None:
return self._payload(hidden_states, attention_mask, position_ids) hidden_states = hidden_states + self._position_embeddings(position_ids)
hidden_states = self._run_layers_session(
def encode_prompt_cached(self, prompt: str, session_id: str) -> TensorPayload: hidden_states, attention_mask, position_ids,
if not self.is_head or self._embed_tokens is None: session_id=session_id, cache_mode="prefill" if session_id else 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) return self._payload(hidden_states, attention_mask, position_ids)
def encode_token_cached(self, token_id: int, seq_len: int, session_id: str) -> TensorPayload: def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
"""Decode step: embed one new token against this head's cached session.
Raises KVCacheMiss if the session was evicted — callers fall back to a
full re-prefill via encode_prompt.
"""
if not self.is_head or self._embed_tokens is None: if not self.is_head or self._embed_tokens is None:
raise ModelBackendError("tokens can only be accepted by the head shard") raise ModelBackendError("decode steps can only start at the head shard")
if seq_len <= 0: if not self.supports_kv_cache:
raise ValueError("seq_len must be positive") raise KVCacheMiss("kv cache disabled on this backend")
entry = self.kv_sessions.lookup(
session_id, effective_start=self._effective_start(None)
)
past_len = entry.seq_len
input_ids = self.torch.tensor([[int(token_id)]], dtype=self.torch.long, device=self.device) 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([[past_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_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 = self._run_layers(
hidden_states, hidden_states, None, position_ids,
attention_mask, cache=entry.cache, past_len=past_len,
position_ids, )
session_id=session_id, entry.seq_len = past_len + 1
cache_mode="decode", return TensorPayload(
seq_len=int(seq_len), body=_tensor_to_bytes(hidden_states.to(self.torch.bfloat16).contiguous()),
shape=list(hidden_states.shape),
attention_mask_header=None,
position_ids_header=_int_tensor_header(position_ids),
past_len=past_len,
) )
return self._payload(hidden_states, attention_mask, position_ids)
def forward_bytes( def forward_bytes(
self, self,
@@ -243,9 +349,9 @@ class TorchModelShard:
position_ids_header: str | None, position_ids_header: str | None,
start_layer: int | None = None, start_layer: int | None = None,
session_id: str | None = None, session_id: str | None = None,
cache_mode: Literal["prefill", "decode", "stateless"] = "stateless", cache_mode: str | None = None,
seq_len: int | None = None, past_len: int | None = None,
) -> TensorPayload | str: ) -> TensorPayload | TailTokenResult | str:
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to( hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
self.device self.device
) )
@@ -255,32 +361,46 @@ class TorchModelShard:
position_ids = _tensor_from_int64_header( position_ids = _tensor_from_int64_header(
position_ids_header, self.torch, self.device position_ids_header, self.torch, self.device
) )
hidden_states = self._run_layers( hidden_states = self._run_layers_session(
hidden_states, hidden_states, attention_mask, position_ids, start_layer=start_layer,
attention_mask, session_id=session_id, cache_mode=cache_mode, past_len=past_len,
position_ids,
start_layer=start_layer,
session_id=session_id,
cache_mode=cache_mode,
seq_len=seq_len,
) )
if self.is_tail: if self.is_tail:
token_id = self.decode_tail_token_id(hidden_states) return self.decode_tail_token(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) return self._payload(hidden_states, attention_mask, position_ids)
def decode_tail(self, hidden_states: Any) -> str: def decode_tail(self, hidden_states: Any) -> str:
token_id = self.decode_tail_token_id(hidden_states) return self.decode_tail_token(hidden_states).text
return self.tokenizer.decode([token_id], skip_special_tokens=True)
def decode_tail_token_id(self, hidden_states: Any) -> int: def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
if self._norm is not None: if self._norm is not None:
hidden_states = self._norm(hidden_states) hidden_states = self._norm(hidden_states)
if self._lm_head is None: if self._lm_head is None:
raise ModelBackendError("tail shard has no lm_head") raise ModelBackendError("tail shard has no lm_head")
logits = self._lm_head(hidden_states) logits = self._lm_head(hidden_states)
return int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
return TailTokenResult(
text=self.tokenizer.decode([token_id], skip_special_tokens=True),
token_id=token_id,
)
def eos_token_ids(self) -> list[int]:
"""All token ids that should terminate generation (tokenizer + generation config)."""
ids: set[int] = set()
tok_eos = getattr(self.tokenizer, "eos_token_id", None)
gen_config = getattr(self.model, "generation_config", None)
gen_eos = getattr(gen_config, "eos_token_id", None) if gen_config is not None else None
for value in (tok_eos, gen_eos):
if value is None:
continue
if isinstance(value, (list, tuple)):
ids.update(int(v) for v in value)
else:
ids.add(int(value))
return sorted(ids)
def release_session(self, session_id: str) -> None:
self.kv_sessions.drop(session_id)
def generate_text( def generate_text(
self, self,
@@ -391,38 +511,108 @@ class TorchModelShard:
) )
return dict(self.tokenizer(prompt, return_tensors="pt")) return dict(self.tokenizer(prompt, return_tensors="pt"))
def _run_layers( def _effective_start(self, start_layer: int | None) -> int:
# start_layer overrides shard_start for overlapping-shard routing
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
# indexing outside the loaded weights.
return (
max(self.shard_start, start_layer)
if start_layer is not None
else self.shard_start
)
def _new_session_cache(self) -> Any | None:
"""Build the model-appropriate cache object for one session.
DynamicCache(config=...) lets transformers pick the right per-layer
state (K/V for standard attention, conv/recurrent state for hybrid
linear-attention layers) — the same construction the model's own
forward() uses when use_cache=True.
"""
try:
from transformers import DynamicCache
except ImportError:
return None
try:
return DynamicCache(config=self.model.config)
except TypeError:
return DynamicCache()
def _run_layers_session(
self, self,
hidden_states: Any, hidden_states: Any,
attention_mask: Any, attention_mask: Any,
position_ids: Any, position_ids: Any,
start_layer: int | None = None, start_layer: int | None = None,
session_id: str | None = None, session_id: str | None = None,
cache_mode: Literal["prefill", "decode", "stateless"] = "stateless", cache_mode: str | None = None,
seq_len: int | None = None, past_len: int | None = None,
) -> Any: ) -> Any:
# start_layer overrides shard_start for overlapping-shard routing """Run this shard's layers, keying cached state by session when requested.
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
# indexing outside the loaded weights. cache_mode "prefill" creates fresh session state; "decode" requires an
effective_start = ( existing entry (KVCacheMiss otherwise). None runs fully stateless —
max(self.shard_start, start_layer) today's behavior, kept as the recovery path.
if start_layer is not None """
else self.shard_start effective_start = self._effective_start(start_layer)
) if not (session_id and cache_mode and self.supports_kv_cache):
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": if cache_mode == "decode":
entry = self._session_cache.get(cache_key) # A decode payload is one token — running it stateless would
if entry is None: # silently produce garbage. Force the head to re-prefill.
raise ShardCacheMiss( raise KVCacheMiss("kv cache disabled on this backend")
f"cache miss for session {session_id} layers {effective_start}-{self.shard_end}" return self._run_layers(
hidden_states, attention_mask, position_ids, start_layer=start_layer
) )
cached_layer_states = entry.layer_states if cache_mode == "decode":
entry.last_used = time.monotonic() entry = self.kv_sessions.lookup(
self._session_cache.move_to_end(cache_key) session_id,
expected_seq_len=past_len,
effective_start=effective_start,
)
seq_len = int(hidden_states.shape[1])
# Decode attends over cache + new token; no padding, so no mask needed.
hidden_states = self._run_layers(
hidden_states, None, position_ids,
start_layer=start_layer, cache=entry.cache, past_len=entry.seq_len,
)
entry.seq_len += seq_len
return hidden_states
# Prefill: fresh cache for this session (replaces any stale entry).
cache = self._new_session_cache()
if cache is None:
return self._run_layers(
hidden_states, attention_mask, position_ids, start_layer=start_layer
)
try:
result = self._run_layers(
hidden_states, attention_mask, position_ids,
start_layer=start_layer, cache=cache, past_len=0,
)
except TypeError as exc:
# Layers reject cache kwargs (exotic architecture) — disable caching
# for this backend and stay on the stateless path.
self.supports_kv_cache = False
print(f" [node] kv cache unsupported by {self.model_id}: {exc}", flush=True)
return self._run_layers(
hidden_states, attention_mask, position_ids, start_layer=start_layer
)
self.kv_sessions.store(
session_id, cache,
seq_len=int(hidden_states.shape[1]),
effective_start=effective_start,
)
return result
def _run_layers(
self,
hidden_states: Any,
attention_mask: Any,
position_ids: Any,
start_layer: int | None = None,
cache: Any = None,
past_len: int = 0,
) -> Any:
effective_start = self._effective_start(start_layer)
position_embeddings = _rotary_position_embeddings( position_embeddings = _rotary_position_embeddings(
self.model, self.model,
hidden_states, hidden_states,
@@ -433,29 +623,23 @@ class TorchModelShard:
hidden_states, hidden_states,
self.torch, self.torch,
) )
cache_position = None
if cache is not None:
seq_len = int(hidden_states.shape[1])
cache_position = self.torch.arange(
past_len, past_len + seq_len, device=hidden_states.device
)
with self.torch.inference_mode(): with self.torch.inference_mode():
next_layer_states: list[Any] = [] for layer in self.layers[effective_start:self.shard_end + 1]:
for index, layer in enumerate(self.layers[effective_start:self.shard_end + 1]): hidden_states = _call_layer(
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, layer,
hidden_states, hidden_states,
layer_attention_mask, layer_attention_mask,
position_ids, position_ids,
position_embeddings, position_embeddings,
use_cache=use_cache, cache=cache,
past_key_value=past_state, cache_position=cache_position,
) )
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) return hidden_states.to(self.torch.bfloat16)
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload: def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
@@ -471,30 +655,6 @@ class TorchModelShard:
else None, 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( def load_torch_shard(
model_id: str, model_id: str,
@@ -842,20 +1002,19 @@ def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any)
return None return None
if len(getattr(attention_mask, "shape", ())) != 2: if len(getattr(attention_mask, "shape", ())) != 2:
return attention_mask return attention_mask
batch_size, key_len = attention_mask.shape batch_size, seq_len = attention_mask.shape
query_len = int(hidden_states.shape[-2]) if seq_len <= 1:
if key_len <= 1:
return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype) return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype)
min_value = torch.finfo(hidden_states.dtype).min min_value = torch.finfo(hidden_states.dtype).min
causal = torch.full( causal = torch.full(
(query_len, key_len), (seq_len, seq_len),
min_value, min_value,
dtype=hidden_states.dtype, dtype=hidden_states.dtype,
device=hidden_states.device, device=hidden_states.device,
) )
causal = torch.triu(causal, diagonal=1 + key_len - query_len) causal = torch.triu(causal, diagonal=1)
causal = causal[None, None, :, :].expand(batch_size, 1, query_len, key_len).clone() causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()
padding = attention_mask.to(device=hidden_states.device) padding = attention_mask.to(device=hidden_states.device)
if not bool(padding.all()): if not bool(padding.all()):
@@ -879,56 +1038,44 @@ def _call_layer(
attention_mask: Any, attention_mask: Any,
position_ids: Any, position_ids: Any,
position_embeddings: Any | None = None, position_embeddings: Any | None = None,
*, cache: Any = None,
use_cache: bool = False, cache_position: Any = None,
past_key_value: Any | None = None, ) -> Any:
) -> tuple[Any, Any | None]:
attempts = ( attempts = (
{ {
"attention_mask": attention_mask, "attention_mask": attention_mask,
"position_ids": position_ids, "position_ids": position_ids,
"position_embeddings": position_embeddings, "position_embeddings": position_embeddings,
"past_key_value": past_key_value, "use_cache": False,
"use_cache": use_cache,
}, },
{ {
"attention_mask": attention_mask, "attention_mask": attention_mask,
"position_ids": position_ids, "position_ids": position_ids,
"past_key_value": past_key_value, "use_cache": False,
"use_cache": use_cache,
}, },
{"attention_mask": attention_mask, "past_key_value": past_key_value, "use_cache": use_cache}, {"attention_mask": attention_mask, "use_cache": False},
{"past_key_value": past_key_value, "use_cache": use_cache}, {"use_cache": False},
{"use_cache": use_cache},
{}, {},
) )
last_exc: Exception | None = None last_exc: Exception | None = None
for kwargs in attempts: for kwargs in attempts:
filtered = {key: value for key, value in kwargs.items() if value is not None} filtered = {key: value for key, value in kwargs.items() if value is not None}
if cache is not None:
# transformers 5.x layers take a Cache via past_key_values and
# mutate it in place; cache_position is required by sliding-window
# and hybrid recurrent layers.
filtered["past_key_values"] = cache
filtered["use_cache"] = True
if cache_position is not None:
filtered["cache_position"] = cache_position
try: try:
output = layer(hidden_states, **filtered) output = layer(hidden_states, **filtered)
return _layer_hidden_and_cache(output) return output[0] if isinstance(output, tuple) else output
except TypeError as exc: except TypeError as exc:
last_exc = exc last_exc = exc
if last_exc is not None: if last_exc is not None:
raise last_exc raise last_exc
return _layer_hidden_and_cache(layer(hidden_states)) return layer(hidden_states)[0]
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: def _tensor_to_bytes(tensor: Any) -> bytes:

269
packages/node/meshnet_node/torch_server.py
View File

@@ -17,13 +17,17 @@ from typing import Any
from .model_backend import ( from .model_backend import (
InsufficientVRAMError, InsufficientVRAMError,
KVCacheMiss,
MissingModelDependencyError, MissingModelDependencyError,
Quantization, Quantization,
ShardCacheMiss, TailTokenResult,
TensorPayload,
TorchModelShard, TorchModelShard,
validate_quantization, validate_quantization,
) )
class _PipelineCacheMiss(Exception):
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
from .server import ( from .server import (
_WIRE_VERSION, _WIRE_VERSION,
_compress_body, _compress_body,
@@ -33,16 +37,6 @@ 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: def _endpoint_key(url: str) -> str:
"""Normalize http(s) endpoints for host:port comparison.""" """Normalize http(s) endpoints for host:port comparison."""
parsed = urllib.parse.urlparse(url.rstrip("/")) parsed = urllib.parse.urlparse(url.rstrip("/"))
@@ -106,48 +100,6 @@ def _write_progress_line(state: list[bool], message: str, *, final: bool = False
sys.stdout.flush() 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( def _relay_hop(
relay_addr: str, relay_addr: str,
path: str, path: str,
@@ -182,6 +134,13 @@ def _relay_hop(
return status, resp_headers, resp_body return status, resp_headers, resp_body
def _is_cache_miss_body(body: bytes) -> bool:
try:
return json.loads(body).get("error") == "cache_miss"
except (json.JSONDecodeError, AttributeError, UnicodeDecodeError):
return False
class _TorchHTTPServer(http.server.HTTPServer): class _TorchHTTPServer(http.server.HTTPServer):
def __init__( def __init__(
self, self,
@@ -400,6 +359,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
start_layer_header = self.headers.get("X-Meshnet-Start-Layer") start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
start_layer = int(start_layer_header) if start_layer_header else None start_layer = int(start_layer_header) if start_layer_header else None
# Session KV-cache protocol: prefill establishes per-session state on
# this node's layer range; decode reuses it. Absent header = legacy
# stateless call (also the signature fake backends implement).
cache_mode = self.headers.get("X-Meshnet-Cache")
forward_kwargs: dict[str, object] = {}
if cache_mode in ("prefill", "decode"):
past_len_header = self.headers.get("X-Meshnet-Past-Len")
forward_kwargs = {
"session_id": session,
"cache_mode": cache_mode,
"past_len": int(past_len_header) if past_len_header else None,
}
try: try:
result = server.backend.forward_bytes( result = server.backend.forward_bytes(
raw_body, raw_body,
@@ -407,28 +379,20 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.headers.get("X-Meshnet-Attn-Mask"), self.headers.get("X-Meshnet-Attn-Mask"),
self.headers.get("X-Meshnet-Position-Ids"), self.headers.get("X-Meshnet-Position-Ids"),
start_layer=start_layer, start_layer=start_layer,
session_id=session, **forward_kwargs,
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: except KVCacheMiss as exc:
self._send_json(409, {"error": "cache_miss", "detail": str(exc)}) self._send_json(409, {"error": "cache_miss", "detail": str(exc)})
return return
except Exception as exc: except Exception as exc:
self._send_json(500, {"error": str(exc)}) self._send_json(500, {"error": str(exc)})
return return
if isinstance(result, TailTokenResult):
self._send_json(200, {"text": result.text, "token_id": result.token_id})
return
if isinstance(result, str): if isinstance(result, str):
token_id = None self._send_json(200, {"text": result})
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 return
response_body = _compress_body(result.body, encoding) response_body = _compress_body(result.body, encoding)
@@ -581,9 +545,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._send_json(500, {"error": f"generation failed: {exc}"}) self._send_json(500, {"error": f"generation failed: {exc}"})
return return
# Distributed path: autoregressive generation across shards. # Distributed path: autoregressive generation across shards with a
# Step 0 prefills the full prompt and creates shard-local caches. Later # sharded per-node KV cache. Step 0 prefills the full prompt through the
# cached steps send only the previous token's activation through the route. # route (each node caches state for its own layer range, keyed by a
# per-generation session id); steps 1+ send only the newest token's
# hidden state. A 409 cache_miss from any hop (eviction/restart/route
# change) falls back to a full re-prefill — the old stateless behavior.
remaining_route = self._get_remaining_route(model_name, backend=backend) remaining_route = self._get_remaining_route(model_name, backend=backend)
print( print(
f" [node] chat route model={model_name!r} max_tokens={max_tokens} " f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
@@ -615,9 +582,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
eos_token: str = getattr(backend.tokenizer, "eos_token", "") or "" eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
generated: list[str] = [] generated: list[str] = []
current_text = prompt_text current_text = prompt_text
session_id = str(uuid.uuid4()) session_id = str(uuid.uuid4())
last_token_id: int | None = None use_kv = bool(getattr(backend, "supports_kv_cache", False))
current_seq_len: int | None = None eos_ids: set[int] = set()
if use_kv:
try:
eos_ids = set(backend.eos_token_ids())
except Exception:
eos_ids = set()
stream_emit = None stream_emit = None
if stream: if stream:
@@ -628,63 +601,60 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
gen_started = time.monotonic() gen_started = time.monotonic()
last_gen_log = gen_started last_gen_log = gen_started
progress_line = [False] progress_line = [False]
last_token_id: int | None = None
def _prefill_step() -> tuple[str, int | None]:
"""Full-sequence prefill: initial step and cache-miss recovery."""
payload = (
backend.encode_prompt(current_text, session_id=session_id)
if use_kv
else backend.encode_prompt(current_text)
)
return self._run_downstream_pipeline(
payload, remaining_route, backend=backend,
session=session_id, cache_mode="prefill" if use_kv else None,
)
for step in range(max_tokens): for step in range(max_tokens):
try: try:
if step == 0 or last_token_id is None or current_seq_len is None: if use_kv and step > 0 and last_token_id is not 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: try:
payload = backend.encode_token_cached(last_token_id, seq_len, session_id) payload = backend.encode_next_token(last_token_id, session_id)
cache_mode = "decode" token_str, token_id = self._run_downstream_pipeline(
except ShardCacheMiss: payload, remaining_route, backend=backend,
payload = _encode_prompt_for_session(backend, current_text, session_id) session=session_id, cache_mode="decode",
current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else current_seq_len )
cache_mode = "prefill" except (KVCacheMiss, _PipelineCacheMiss) as miss:
seq_len = current_seq_len # Evicted/restarted node or head lost its own session:
# re-prefill the whole sequence once and continue cached.
print(
f" [node] kv cache miss at step {step} ({miss}); "
f"re-prefilling {len(current_text)} chars",
flush=True,
)
token_str, token_id = _prefill_step()
else:
token_str, token_id = _prefill_step()
except _PipelineCacheMiss as exc:
print(f" [node] unexpected cache miss on prefill: {exc}", flush=True)
break
except Exception as exc: except Exception as exc:
print(f" [node] distributed encode error: {exc}", flush=True) print(f" [node] distributed encode error: {exc}", flush=True)
break break
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. # Stop on error responses or EOS.
if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")): if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
break break
if token_id is not None and token_id in eos_ids:
break
if eos_token and token_str == eos_token: if eos_token and token_str == eos_token:
break break
if not token_str and token_id is None:
break
last_token_id = token_id
# token_str can be empty for a skipped special token that is not
# EOS — keep generating from its token_id without emitting text.
if token_str:
generated.append(token_str) 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: if stream_emit is not None:
stream_emit(token_str) stream_emit(token_str)
current_text = current_text + token_str current_text = current_text + token_str
@@ -706,6 +676,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
) )
last_gen_log = now last_gen_log = now
if use_kv:
try:
backend.release_session(session_id)
except Exception:
pass
if generated: if generated:
elapsed = time.monotonic() - gen_started elapsed = time.monotonic() - gen_started
token_count = len(generated) token_count = len(generated)
@@ -805,10 +781,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
route: list[dict], route: list[dict],
*, *,
backend: TorchModelShard | None = None, backend: TorchModelShard | None = None,
session_id: str | None = None, session: str | None = None,
cache_mode: str = "stateless", cache_mode: str | None = None,
seq_len: int | None = None, ) -> tuple[str, int | None]:
) -> _PipelineResult: """Forward an activation through the downstream route.
Returns (token_text, token_id) — token_id is None when a hop predates
the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds
409 cache_miss (evicted/restarted node) so the caller can re-prefill.
"""
server: _TorchHTTPServer = self.server # type: ignore[assignment] server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend active_backend = backend or server.backend
if not route: if not route:
@@ -820,14 +801,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
bytearray(payload.body), # type: ignore[union-attr] bytearray(payload.body), # type: ignore[union-attr]
dtype=active_backend.torch.bfloat16, dtype=active_backend.torch.bfloat16,
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr] ).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
token_id = active_backend.decode_tail_token_id(tensor) if hasattr(active_backend, "decode_tail_token"):
text = active_backend.tokenizer.decode([token_id], skip_special_tokens=True) tail = active_backend.decode_tail_token(tensor)
return _PipelineResult(text, token_id) return tail.text, tail.token_id
return active_backend.decode_tail(tensor), None
except Exception as exc: except Exception as exc:
return _PipelineResult(f"decode error: {exc}") return f"decode error: {exc}", None
return _PipelineResult("no downstream route available for non-tail shard") return "no downstream route available for non-tail shard", None
session = session_id or str(uuid.uuid4()) # Session is stable across all steps of one generation when the caller
# provides it (KV-cache protocol); fresh per call otherwise (legacy).
session = session or str(uuid.uuid4())
shape = payload.shape # type: ignore[union-attr] shape = payload.shape # type: ignore[union-attr]
attn_mask = payload.attention_mask_header # type: ignore[union-attr] attn_mask = payload.attention_mask_header # type: ignore[union-attr]
pos_ids = payload.position_ids_header # type: ignore[union-attr] pos_ids = payload.position_ids_header # type: ignore[union-attr]
@@ -856,10 +840,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"X-Meshnet-Chunk-Total": "1", "X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": str(hop_index), "X-Meshnet-Hop-Index": str(hop_index),
"X-Meshnet-Start-Layer": str(start_layer), "X-Meshnet-Start-Layer": str(start_layer),
"X-Meshnet-Cache-Mode": cache_mode,
} }
if seq_len is not None: if cache_mode:
headers["X-Meshnet-Seq-Len"] = str(seq_len) headers["X-Meshnet-Cache"] = cache_mode
past_len = getattr(payload, "past_len", None)
if cache_mode == "decode" and past_len is not None:
headers["X-Meshnet-Past-Len"] = str(past_len)
if current_attn: if current_attn:
headers["X-Meshnet-Attn-Mask"] = current_attn headers["X-Meshnet-Attn-Mask"] = current_attn
if current_pos: if current_pos:
@@ -869,14 +855,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
status, resp_headers, resp_body = _relay_hop( status, resp_headers, resp_body = _relay_hop(
relay_addr, "/forward", current_body, headers, timeout=120.0, relay_addr, "/forward", current_body, headers, timeout=120.0,
) )
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400: if status >= 400:
if status == 409:
raise _PipelineCacheMiss(f"cache miss at {node_url}")
print( print(
f" [node] relay hop {hop_index} returned {status} from {relay_addr}", f" [node] relay hop {hop_index} returned {status} from {relay_addr}",
flush=True, flush=True,
) )
return _PipelineResult(f"pipeline error at {node_url} via relay: status {status}") return f"pipeline error at {node_url} via relay: status {status}", None
except _PipelineCacheMiss: except _PipelineCacheMiss:
raise raise
except Exception as exc: except Exception as exc:
@@ -898,33 +884,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
resp_body = r.read() resp_body = r.read()
resp_headers = {k.lower(): v for k, v in r.headers.items()} resp_headers = {k.lower(): v for k, v in r.headers.items()}
except urllib.error.HTTPError as exc: except urllib.error.HTTPError as exc:
if exc.code == 409: body = exc.read()
raise _PipelineCacheMiss(f"cache miss at {node_url}") from exc if exc.code == 409 and _is_cache_miss_body(body):
raise _PipelineCacheMiss(node_url) from exc
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return _PipelineResult(f"pipeline error at {node_url}: {exc}") return f"pipeline error at {node_url}: {exc}", None
except Exception as exc: except Exception as exc:
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return _PipelineResult(f"pipeline error at {node_url}: {exc}") return f"pipeline error at {node_url}: {exc}", None
content_type = resp_headers.get("content-type", "") content_type = resp_headers.get("content-type", "")
if "application/json" in content_type: if "application/json" in content_type:
try: try:
data = json.loads(resp_body) data = json.loads(resp_body)
if data.get("error") == "cache_miss":
raise _PipelineCacheMiss(f"cache miss at {node_url}")
text = str(data.get("text", "")) text = str(data.get("text", ""))
token_id = data.get("token_id") token_id = data.get("token_id")
if server.debug: if server.debug:
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True) print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
return _PipelineResult(text, int(token_id) if token_id is not None else None) return text, int(token_id) if token_id is not None else None
except json.JSONDecodeError: except json.JSONDecodeError:
return _PipelineResult(resp_body.decode("utf-8", errors="replace")) return resp_body.decode("utf-8", errors="replace"), None
# Binary activation — update and forward to next node # Binary activation — update and forward to next node
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape)) shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
current_shape = _parse_shape(shape_header) current_shape = _parse_shape(shape_header)
current_body = resp_body current_body = resp_body
current_attn = resp_headers.get("x-meshnet-attn-mask") current_attn = resp_headers.get("x-meshnet-attn-mask")
current_pos = resp_headers.get("x-meshnet-position-ids") current_pos = resp_headers.get("x-meshnet-position-ids")
return _PipelineResult("") return "", None
def _stream_openai_response(self, token_iter, model: str) -> None: def _stream_openai_response(self, token_iter, model: str) -> None:
"""Stream tokens from an iterator as SSE chunks.""" """Stream tokens from an iterator as SSE chunks."""

384
tests/test_kv_cache_distributed.py Normal file
View File

@@ -0,0 +1,384 @@
"""AH-25: sharded per-node KV cache for distributed generation.
Covers the SessionCacheStore (TTL + LRU + mismatch handling), the HTTP
session protocol (stable session id, O(1) decode payloads, 409 cache-miss
fallback, legacy stateless compatibility), and an env-gated golden test that
proves cached and stateless distributed generation produce identical tokens
on a real two-shard Qwen2.5-0.5B split.
"""
import json
import os
import urllib.request
import pytest
from meshnet_node.model_backend import (
KVCacheMiss,
SessionCacheStore,
TailTokenResult,
TensorPayload,
)
from meshnet_node.torch_server import TorchNodeServer
# ---------------------------------------------------------------------------
# SessionCacheStore units
# ---------------------------------------------------------------------------
class _Clock:
def __init__(self) -> None:
self.now = 0.0
def __call__(self) -> float:
return self.now
def test_store_lookup_roundtrip_advances_lru():
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0, clock=_Clock())
store.store("s1", cache=object(), seq_len=6, effective_start=12)
entry = store.lookup("s1", expected_seq_len=6, effective_start=12)
assert entry.seq_len == 6
entry.seq_len += 1
assert store.lookup("s1", expected_seq_len=7).seq_len == 7
def test_lookup_unknown_session_raises_cache_miss():
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
with pytest.raises(KVCacheMiss):
store.lookup("nope")
def test_seq_len_mismatch_drops_entry_and_raises():
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=6, effective_start=0)
with pytest.raises(KVCacheMiss):
store.lookup("s1", expected_seq_len=9)
# Entry must be gone — a poisoned cache is never reused.
with pytest.raises(KVCacheMiss):
store.lookup("s1")
def test_effective_start_mismatch_raises():
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=6, effective_start=12)
with pytest.raises(KVCacheMiss):
store.lookup("s1", effective_start=21)
def test_ttl_expiry_evicts_stale_sessions():
clock = _Clock()
store = SessionCacheStore(max_sessions=4, ttl_seconds=60.0, clock=clock)
store.store("s1", cache=object(), seq_len=6, effective_start=0)
clock.now = 61.0
with pytest.raises(KVCacheMiss):
store.lookup("s1")
assert len(store) == 0
def test_lru_eviction_bounds_session_count():
clock = _Clock()
store = SessionCacheStore(max_sessions=2, ttl_seconds=1000.0, clock=clock)
store.store("s1", cache=object(), seq_len=1, effective_start=0)
store.store("s2", cache=object(), seq_len=1, effective_start=0)
store.lookup("s1") # s1 becomes most recent → s2 is LRU
store.store("s3", cache=object(), seq_len=1, effective_start=0)
assert len(store) == 2
with pytest.raises(KVCacheMiss):
store.lookup("s2")
store.lookup("s1")
store.lookup("s3")
def test_drop_removes_session():
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=1, effective_start=0)
store.drop("s1")
with pytest.raises(KVCacheMiss):
store.lookup("s1")
# ---------------------------------------------------------------------------
# HTTP session protocol with fake cached backends
# ---------------------------------------------------------------------------
class _ChatTokenizer:
eos_token = ""
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
return "debug prompt"
class _CachedHeadBackend:
model_id = "fake-model"
total_layers = 12
is_head = True
is_tail = False
supports_kv_cache = True
tokenizer = _ChatTokenizer()
def __init__(self) -> None:
self.prefills: list[str | None] = []
self.decode_calls: list[tuple[int, str]] = []
self.released: list[str] = []
self._seq: dict[str, int] = {}
def eos_token_ids(self) -> list[int]:
return [99]
def release_session(self, session_id: str) -> None:
self.released.append(session_id)
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
self.prefills.append(session_id)
if session_id:
self._seq[session_id] = 6
return TensorPayload(
body=b"\x00" * (1 * 6 * 8 * 2),
shape=[1, 6, 8],
attention_mask_header=None,
position_ids_header=None,
)
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
self.decode_calls.append((token_id, session_id))
past = self._seq[session_id]
self._seq[session_id] = past + 1
return TensorPayload(
body=b"\x00" * (1 * 1 * 8 * 2),
shape=[1, 1, 8],
attention_mask_header=None,
position_ids_header=None,
past_len=past,
)
class _CachedTailBackend:
model_id = "fake-model"
total_layers = 12
is_head = False
is_tail = True
supports_kv_cache = True
def __init__(self, tokens, miss_on_call: int | None = None) -> None:
self._tokens = list(tokens)
self.miss_on_call = miss_on_call
self.calls: list[dict] = []
def forward_bytes(
self,
body,
shape,
attention_mask_header,
position_ids_header,
start_layer=None,
session_id=None,
cache_mode=None,
past_len=None,
):
call_index = len(self.calls)
self.calls.append({
"session": session_id,
"mode": cache_mode,
"past_len": past_len,
"shape": list(shape),
})
if self.miss_on_call is not None and call_index == self.miss_on_call:
raise KVCacheMiss("session evicted (test)")
text, token_id = self._tokens.pop(0)
return TailTokenResult(text=text, token_id=token_id)
def _chat_once(head_port: int, tail_port: int, max_tokens: int) -> str:
payload = json.dumps({
"model": "fake-model",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": max_tokens,
}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{head_port}/v1/chat/completions",
data=payload,
headers={
"Content-Type": "application/json",
"X-Meshnet-Route": json.dumps([
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 6},
]),
},
method="POST",
)
with urllib.request.urlopen(req, timeout=10) as resp:
body = json.loads(resp.read())
return body["choices"][0]["message"]["content"]
def test_session_is_stable_and_decode_payloads_are_single_token():
head_backend = _CachedHeadBackend()
tail_backend = _CachedTailBackend([(" a", 1), (" b", 2), (" c", 3)])
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
tail = TorchNodeServer(backend=tail_backend)
head_port = head.start()
tail_port = tail.start()
try:
content = _chat_once(head_port, tail_port, max_tokens=3)
finally:
head.stop()
tail.stop()
assert content == " a b c"
assert len(tail_backend.calls) == 3
# Step 0 is a full-prompt prefill; steps 1+ carry only the new token.
assert tail_backend.calls[0]["mode"] == "prefill"
assert tail_backend.calls[0]["shape"] == [1, 6, 8]
for step, call in enumerate(tail_backend.calls[1:], start=1):
assert call["mode"] == "decode"
assert call["shape"] == [1, 1, 8]
assert call["past_len"] == 6 + (step - 1)
# One session id across every step of the generation.
sessions = {call["session"] for call in tail_backend.calls}
assert len(sessions) == 1
session_id = sessions.pop()
assert head_backend.prefills == [session_id]
assert head_backend.decode_calls == [(1, session_id), (2, session_id)]
# Head releases its own session state when the generation ends.
assert head_backend.released == [session_id]
def test_eos_token_id_stops_generation():
head_backend = _CachedHeadBackend()
tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
tail = TorchNodeServer(backend=tail_backend)
head_port = head.start()
tail_port = tail.start()
try:
content = _chat_once(head_port, tail_port, max_tokens=8)
finally:
head.stop()
tail.stop()
assert content == " a"
assert len(tail_backend.calls) == 2
def test_downstream_cache_miss_falls_back_to_full_reprefill():
head_backend = _CachedHeadBackend()
# Call 1 (the first decode) raises KVCacheMiss → node answers 409 →
# head re-prefills the full sequence and keeps generating.
tail_backend = _CachedTailBackend(
[(" a", 1), (" b", 2), (" c", 3)], miss_on_call=1,
)
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
tail = TorchNodeServer(backend=tail_backend)
head_port = head.start()
tail_port = tail.start()
try:
content = _chat_once(head_port, tail_port, max_tokens=3)
finally:
head.stop()
tail.stop()
assert content == " a b c"
modes = [call["mode"] for call in tail_backend.calls]
assert modes == ["prefill", "decode", "prefill", "decode"]
# Head re-prefilled once, with the same stable session id.
assert len(head_backend.prefills) == 2
assert len(set(head_backend.prefills)) == 1
def test_kv_head_with_legacy_tail_reprefills_every_step():
"""Mixed fleet: tail predates the protocol and returns no token_id."""
class _LegacyTailBackend:
model_id = "fake-model"
total_layers = 12
is_head = False
is_tail = True
def __init__(self) -> None:
self.calls = 0
def forward_bytes(self, body, shape, attention_mask_header,
position_ids_header, start_layer=None, **kwargs):
self.calls += 1
return " x" if self.calls < 3 else ""
head_backend = _CachedHeadBackend()
tail_backend = _LegacyTailBackend()
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
tail = TorchNodeServer(backend=tail_backend)
head_port = head.start()
tail_port = tail.start()
try:
content = _chat_once(head_port, tail_port, max_tokens=5)
finally:
head.stop()
tail.stop()
assert content == " x x"
# No token_id from the tail → every step is a full prefill (legacy cost),
# never a decode against a cache the tail doesn't keep.
assert head_backend.decode_calls == []
assert len(head_backend.prefills) == 3
# ---------------------------------------------------------------------------
# Golden test on a real two-shard split (env-gated: loads Qwen2.5-0.5B twice)
# ---------------------------------------------------------------------------
_GOLDEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
requires_real_model = pytest.mark.skipif(
os.environ.get("MESHNET_REAL_MODEL_TESTS") != "1",
reason="set MESHNET_REAL_MODEL_TESTS=1 to run the real-model golden test",
)
@requires_real_model
def test_cached_distributed_generation_matches_stateless_golden():
pytest.importorskip("torch")
from meshnet_node.model_backend import TorchModelShard
head = TorchModelShard(_GOLDEN_MODEL, 0, 11)
tail = TorchModelShard(_GOLDEN_MODEL, 12, 23)
steps = 12
prompt = head.tokenizer.apply_chat_template(
[{"role": "user", "content": "Count from 1 to 5."}],
add_generation_prompt=True,
tokenize=False,
)
# Reference: today's stateless path — re-encode the full sequence each step.
stateless_ids: list[int] = []
text = prompt
for _ in range(steps):
payload = head.encode_prompt(text)
result = tail.forward_bytes(
payload.body, payload.shape,
payload.attention_mask_header, payload.position_ids_header,
start_layer=12,
)
stateless_ids.append(result.token_id)
text += result.text
# Cached path: one prefill, then single-token decode steps.
session = "golden-session"
cached_ids: list[int] = []
payload = head.encode_prompt(prompt, session_id=session)
result = tail.forward_bytes(
payload.body, payload.shape,
payload.attention_mask_header, payload.position_ids_header,
start_layer=12, session_id=session, cache_mode="prefill",
)
cached_ids.append(result.token_id)
for _ in range(steps - 1):
payload = head.encode_next_token(cached_ids[-1], session)
assert payload.shape[1] == 1, "decode payload must be a single token"
result = tail.forward_bytes(
payload.body, payload.shape,
None, payload.position_ids_header,
start_layer=12, session_id=session, cache_mode="decode",
past_len=payload.past_len,
)
cached_ids.append(result.token_id)
assert cached_ids == stateless_ids