feat: checkpoint distributed gguf runtime stories
This commit is contained in:
Dobromir Popov
918
packages/node/meshnet_node/hot_kv_state.py
Normal file
918
packages/node/meshnet_node/hot_kv_state.py
Normal file
@@ -0,0 +1,918 @@
|
||||
"""Isolated concurrent local Hot KV State for distributed Shards (DGR-007).
|
||||
|
||||
Hot KV State stays local to the node serving a Shard (RALPH runtime decision #7).
|
||||
A concurrent server must map each ``(Route Session ID, route epoch)`` to an
|
||||
isolated bounded KV context (decision #8) so that one request can never clear or
|
||||
corrupt another's cache.
|
||||
|
||||
This module owns the *lifecycle and storage* of that state and is deliberately
|
||||
backend-agnostic:
|
||||
|
||||
* :class:`HotKvStateManager` is the single mutation entry point. It maps
|
||||
``(session_id, route_epoch)`` to a :class:`SessionCache`, allocates KV **only
|
||||
for the owned layer range**, and enforces a byte budget, a session cap, and a
|
||||
TTL through LRU/TTL eviction. It rejects stale route epochs and incompatible
|
||||
cache recipes, and returns an **explicit** :class:`CacheMiss` when state the
|
||||
caller expected is gone (evicted, released, desynchronised, or never held) so
|
||||
the head degrades to a from-token-zero re-prefill instead of corrupting output
|
||||
(RALPH decision #14: unverified KV is never migrated silently).
|
||||
* :class:`LayerKvCache` / :class:`SessionCache` are the per-owned-layer K/V
|
||||
containers. They are plain ``numpy`` arrays so the default deterministic test
|
||||
suite needs no torch, GPU, download, or API credit; the pinned llama.cpp worker
|
||||
(DGR-008) maps a llama sequence onto the same container contract.
|
||||
* :class:`KvBoundaryAdapter` wraps a KV-aware ``ShardComputation`` (the DGR-006
|
||||
duck type plus ``run_layers_cached``) so a Shard can run cached prefill/decode
|
||||
through the manager while honouring the architecture-defined boundary contract
|
||||
(head embeds tokens, middle/tail bypass embedding, non-tail emits the
|
||||
unnormalized residual, tail samples).
|
||||
|
||||
The manager owns *all* cache mutation: a computation reads the existing cache and
|
||||
returns the new K/V for the appended positions, and the manager decides whether
|
||||
that append fits the budget. That keeps eviction, accounting, and isolation in one
|
||||
place instead of scattered across backends.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
import numpy as np
|
||||
|
||||
from meshnet_node.boundary_adapter import (
|
||||
BOUNDARY_SCHEMA_VERSION,
|
||||
BoundaryBundle,
|
||||
BoundaryContractError,
|
||||
SamplingContract,
|
||||
ShardRole,
|
||||
TailOutput,
|
||||
certified_architecture,
|
||||
role_for_range,
|
||||
)
|
||||
from meshnet_node.runtime_recipe import compatibility_fingerprint
|
||||
|
||||
|
||||
class HotKvStateError(RuntimeError):
|
||||
"""Base class for Hot KV State errors."""
|
||||
|
||||
|
||||
class StaleRouteEpochError(HotKvStateError):
|
||||
"""Raised when a request references a route epoch older than the current one.
|
||||
|
||||
A newer route epoch means the route was re-planned; the old epoch's KV is
|
||||
unverified against the new plan and must never be silently reused.
|
||||
"""
|
||||
|
||||
|
||||
class IncompatibleCacheRecipeError(HotKvStateError):
|
||||
"""Raised when a request's cache recipe does not match the loaded shard.
|
||||
|
||||
A different quantization / dtype / owned range / architecture produces a KV
|
||||
layout this node cannot reuse without corrupting output.
|
||||
"""
|
||||
|
||||
|
||||
class KvBudgetExceededError(HotKvStateError):
|
||||
"""Raised when a single session cannot fit the configured byte budget.
|
||||
|
||||
Other sessions are evicted first (LRU); this fires only when even one session
|
||||
alone exceeds the budget, which is a misconfiguration rather than pressure.
|
||||
"""
|
||||
|
||||
|
||||
class KvCacheMissError(HotKvStateError):
|
||||
"""Raised by the strict accessor when expected session state is absent.
|
||||
|
||||
Prefer :meth:`HotKvStateManager.resolve`, which returns a structured
|
||||
:class:`CacheMiss` instead of raising, when the caller wants to fall back to a
|
||||
stateless re-prefill.
|
||||
"""
|
||||
|
||||
def __init__(self, miss: "CacheMiss") -> None:
|
||||
super().__init__(str(miss))
|
||||
self.miss = miss
|
||||
|
||||
|
||||
class CacheMissReason(str, Enum):
|
||||
"""Why a lookup produced a cache miss (all benign; retry from token zero)."""
|
||||
|
||||
UNKNOWN_SESSION = "unknown-session"
|
||||
EVICTED_TTL = "evicted-ttl"
|
||||
EVICTED_LRU = "evicted-lru"
|
||||
RELEASED = "released"
|
||||
SUPERSEDED_EPOCH = "superseded-epoch"
|
||||
SEQ_LEN_MISMATCH = "seq-len-mismatch"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CacheMiss:
|
||||
"""Explicit cache-miss response the head can act on (re-prefill).
|
||||
|
||||
This is a value, not an exception: the native protocol carries a cache
|
||||
expectation/result, and a miss is a normal, expected outcome under eviction.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
reason: CacheMissReason
|
||||
detail: str = ""
|
||||
|
||||
def __str__(self) -> str:
|
||||
suffix = f": {self.detail}" if self.detail else ""
|
||||
return (
|
||||
f"cache miss for session {self.session_id[:8]} epoch "
|
||||
f"{self.route_epoch} ({self.reason.value}){suffix}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KvCacheRecipe:
|
||||
"""The identity of a Shard's KV layout, used to reject incompatible reuse.
|
||||
|
||||
Two recipes are compatible iff their fingerprints match — same certified
|
||||
architecture, KV dtype, head geometry, and owned layer range within the same
|
||||
whole-model layer count.
|
||||
"""
|
||||
|
||||
architecture_adapter: str
|
||||
kv_dtype: str
|
||||
n_kv_heads: int
|
||||
head_dim: int
|
||||
total_layers: int
|
||||
start_layer: int
|
||||
end_layer: int
|
||||
boundary_schema_version: int = BOUNDARY_SCHEMA_VERSION
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Fail closed on architecture identity (shared with the boundary adapter).
|
||||
certified_architecture(self.architecture_adapter)
|
||||
if self.n_kv_heads <= 0:
|
||||
raise ValueError("n_kv_heads must be positive")
|
||||
if self.head_dim <= 0:
|
||||
raise ValueError("head_dim must be positive")
|
||||
try:
|
||||
np.dtype(self.kv_dtype)
|
||||
except TypeError as exc: # pragma: no cover - defensive
|
||||
raise ValueError(f"invalid kv_dtype {self.kv_dtype!r}") from exc
|
||||
# role_for_range validates 0 <= start <= end <= total_layers - 1.
|
||||
role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
if self.boundary_schema_version < 1:
|
||||
raise ValueError("boundary_schema_version must be >= 1")
|
||||
|
||||
@property
|
||||
def owned_layers(self) -> tuple[int, ...]:
|
||||
return tuple(range(self.start_layer, self.end_layer + 1))
|
||||
|
||||
@property
|
||||
def role(self) -> ShardRole:
|
||||
return role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
|
||||
def bytes_per_token(self) -> int:
|
||||
"""Bytes of KV one token adds across *owned* layers (keys + values)."""
|
||||
itemsize = np.dtype(self.kv_dtype).itemsize
|
||||
per_layer = 2 * self.n_kv_heads * self.head_dim * itemsize
|
||||
return per_layer * len(self.owned_layers)
|
||||
|
||||
def fingerprint(self) -> str:
|
||||
return compatibility_fingerprint(
|
||||
{
|
||||
"kind": "hot-kv-recipe",
|
||||
# Canonicalize the architecture so 'llama' / 'LlamaForCausalLM'
|
||||
# map to the same fingerprint (they are the same layout).
|
||||
"architecture_adapter": certified_architecture(
|
||||
self.architecture_adapter
|
||||
).adapter,
|
||||
"kv_dtype": np.dtype(self.kv_dtype).name,
|
||||
"n_kv_heads": self.n_kv_heads,
|
||||
"head_dim": self.head_dim,
|
||||
"total_layers": self.total_layers,
|
||||
"start_layer": self.start_layer,
|
||||
"end_layer": self.end_layer,
|
||||
"boundary_schema_version": self.boundary_schema_version,
|
||||
}
|
||||
)
|
||||
|
||||
def is_compatible(self, other: "KvCacheRecipe") -> bool:
|
||||
return self.fingerprint() == other.fingerprint()
|
||||
|
||||
|
||||
class LayerKvCache:
|
||||
"""K/V storage for a single owned layer; sequence axis is 0.
|
||||
|
||||
Keys and values are ``(seq, n_kv_heads, head_dim)``. Backends store the
|
||||
position-encoded (post-RoPE) keys so a decode step only appends the new rows.
|
||||
"""
|
||||
|
||||
__slots__ = ("layer_index", "n_kv_heads", "head_dim", "dtype", "keys", "values")
|
||||
|
||||
def __init__(
|
||||
self, layer_index: int, n_kv_heads: int, head_dim: int, dtype: Any
|
||||
) -> None:
|
||||
self.layer_index = int(layer_index)
|
||||
self.n_kv_heads = int(n_kv_heads)
|
||||
self.head_dim = int(head_dim)
|
||||
self.dtype = np.dtype(dtype)
|
||||
self.keys = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)
|
||||
self.values = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)
|
||||
|
||||
@property
|
||||
def length(self) -> int:
|
||||
return int(self.keys.shape[0])
|
||||
|
||||
def _validate(self, array: np.ndarray, name: str) -> np.ndarray:
|
||||
arr = np.asarray(array, dtype=self.dtype)
|
||||
if arr.ndim != 3 or arr.shape[1:] != (self.n_kv_heads, self.head_dim):
|
||||
raise ValueError(
|
||||
f"layer {self.layer_index} {name} must be "
|
||||
f"(seq, {self.n_kv_heads}, {self.head_dim}), got {arr.shape}"
|
||||
)
|
||||
return arr
|
||||
|
||||
def append(self, keys: np.ndarray, values: np.ndarray) -> int:
|
||||
k = self._validate(keys, "keys")
|
||||
v = self._validate(values, "values")
|
||||
if k.shape[0] != v.shape[0]:
|
||||
raise ValueError(
|
||||
f"layer {self.layer_index} keys/values disagree on token count "
|
||||
f"({k.shape[0]} vs {v.shape[0]})"
|
||||
)
|
||||
self.keys = np.concatenate([self.keys, k], axis=0)
|
||||
self.values = np.concatenate([self.values, v], axis=0)
|
||||
return self.length
|
||||
|
||||
def truncate(self, length: int) -> None:
|
||||
length = max(0, int(length))
|
||||
self.keys = self.keys[:length]
|
||||
self.values = self.values[:length]
|
||||
|
||||
@property
|
||||
def nbytes(self) -> int:
|
||||
return int(self.keys.nbytes + self.values.nbytes)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionCache:
|
||||
"""Isolated per-``(session_id, epoch)`` KV context over the owned layers only."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
recipe: KvCacheRecipe
|
||||
layers: "OrderedDict[int, LayerKvCache]"
|
||||
created_tick: float
|
||||
last_tick: float
|
||||
released: bool = False
|
||||
|
||||
@property
|
||||
def seq_len(self) -> int:
|
||||
if not self.layers:
|
||||
return 0
|
||||
# All owned layers advance in lockstep; report the first owned layer.
|
||||
return next(iter(self.layers.values())).length
|
||||
|
||||
@property
|
||||
def owned_layers(self) -> tuple[int, ...]:
|
||||
return tuple(self.layers.keys())
|
||||
|
||||
def layer(self, index: int) -> LayerKvCache:
|
||||
try:
|
||||
return self.layers[index]
|
||||
except KeyError:
|
||||
raise KeyError(
|
||||
f"layer {index} is not owned by this shard "
|
||||
f"(owned {list(self.layers)})"
|
||||
) from None
|
||||
|
||||
def read_only_layers(self) -> Mapping[int, LayerKvCache]:
|
||||
"""The current per-layer caches a computation reads to attend over."""
|
||||
return dict(self.layers)
|
||||
|
||||
def _append(self, kv_by_layer: Mapping[int, Any]) -> int:
|
||||
provided = set(kv_by_layer)
|
||||
owned = set(self.layers)
|
||||
if provided != owned:
|
||||
raise ValueError(
|
||||
f"append must cover exactly the owned layers {sorted(owned)}, "
|
||||
f"got {sorted(provided)}"
|
||||
)
|
||||
# Pre-validate token counts so a partial append never desynchronises the
|
||||
# owned layers (append is all-or-nothing).
|
||||
new_counts = set()
|
||||
for idx, (keys, _values) in kv_by_layer.items():
|
||||
new_counts.add(int(np.asarray(keys).shape[0]))
|
||||
if len(new_counts) != 1:
|
||||
raise ValueError(
|
||||
f"append token counts disagree across layers: {sorted(new_counts)}"
|
||||
)
|
||||
for idx, (keys, values) in kv_by_layer.items():
|
||||
self.layers[idx].append(keys, values)
|
||||
return self.seq_len
|
||||
|
||||
def _truncate(self, length: int) -> None:
|
||||
for cache in self.layers.values():
|
||||
cache.truncate(length)
|
||||
|
||||
@property
|
||||
def nbytes(self) -> int:
|
||||
return sum(cache.nbytes for cache in self.layers.values())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HotKvStateConfig:
|
||||
"""Bounds for the manager: memory budget, session cap, and idle TTL."""
|
||||
|
||||
budget_bytes: int = 64 * 1024 * 1024
|
||||
max_sessions: int = 8
|
||||
ttl_seconds: float = 600.0
|
||||
miss_history: int = 256
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.budget_bytes <= 0:
|
||||
raise ValueError("budget_bytes must be positive")
|
||||
if self.max_sessions < 1:
|
||||
raise ValueError("max_sessions must be >= 1")
|
||||
if self.ttl_seconds < 0:
|
||||
raise ValueError("ttl_seconds must be >= 0")
|
||||
if self.miss_history < 0:
|
||||
raise ValueError("miss_history must be >= 0")
|
||||
|
||||
|
||||
class HotKvStateManager:
|
||||
"""Concurrent, bounded map of ``(session_id, epoch)`` to an isolated KV context."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
recipe: KvCacheRecipe,
|
||||
config: HotKvStateConfig | None = None,
|
||||
*,
|
||||
clock: Callable[[], float] | None = None,
|
||||
) -> None:
|
||||
self.recipe = recipe
|
||||
self.config = config or HotKvStateConfig()
|
||||
self._clock = clock or time.monotonic
|
||||
self._sessions: "OrderedDict[tuple[str, int], SessionCache]" = OrderedDict()
|
||||
self._latest_epoch: dict[str, int] = {}
|
||||
self._misses: "OrderedDict[tuple[str, int], CacheMiss]" = OrderedDict()
|
||||
self._lock = threading.RLock()
|
||||
|
||||
# -- introspection --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def total_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return sum(s.nbytes for s in self._sessions.values())
|
||||
|
||||
@property
|
||||
def session_count(self) -> int:
|
||||
with self._lock:
|
||||
self._evict_expired_locked(self._clock())
|
||||
return len(self._sessions)
|
||||
|
||||
def session_keys(self) -> list[tuple[str, int]]:
|
||||
with self._lock:
|
||||
return list(self._sessions.keys())
|
||||
|
||||
# -- lifecycle ------------------------------------------------------------
|
||||
|
||||
def open(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
) -> SessionCache:
|
||||
"""Create (or replace) a fresh, empty isolated context for the session.
|
||||
|
||||
A higher route epoch supersedes and frees any earlier epoch for the same
|
||||
session id; an older epoch is rejected as stale.
|
||||
"""
|
||||
self._require_text(session_id, "session_id")
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
self._supersede_older_epochs_locked(session_id, route_epoch)
|
||||
key = (session_id, route_epoch)
|
||||
# A re-open at the same epoch replaces the prior context entirely.
|
||||
self._sessions.pop(key, None)
|
||||
layers: "OrderedDict[int, LayerKvCache]" = OrderedDict(
|
||||
(
|
||||
idx,
|
||||
LayerKvCache(
|
||||
idx,
|
||||
self.recipe.n_kv_heads,
|
||||
self.recipe.head_dim,
|
||||
self.recipe.kv_dtype,
|
||||
),
|
||||
)
|
||||
for idx in self.recipe.owned_layers
|
||||
)
|
||||
session = SessionCache(
|
||||
session_id=session_id,
|
||||
route_epoch=route_epoch,
|
||||
recipe=self.recipe,
|
||||
layers=layers,
|
||||
created_tick=now,
|
||||
last_tick=now,
|
||||
)
|
||||
self._sessions[key] = session
|
||||
self._latest_epoch[session_id] = route_epoch
|
||||
self._misses.pop(key, None)
|
||||
self._enforce_capacity_locked(protect=key, incoming_bytes=0)
|
||||
return session
|
||||
|
||||
def append(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
kv_by_layer: Mapping[int, Any],
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache:
|
||||
"""Append new K/V (prefill or decode) to an existing isolated context.
|
||||
|
||||
The computation supplies exactly the owned layers' new keys/values. The
|
||||
manager evicts other sessions (LRU) to fit the byte budget before growing
|
||||
this one, and raises :class:`KvBudgetExceededError` only if this session
|
||||
alone cannot fit.
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
session = self._require_live_locked(session_id, route_epoch)
|
||||
if expected_seq_len is not None and session.seq_len != expected_seq_len:
|
||||
miss = self._drop_and_record_locked(
|
||||
(session_id, route_epoch),
|
||||
CacheMissReason.SEQ_LEN_MISMATCH,
|
||||
detail=f"cache holds {session.seq_len}, caller expected "
|
||||
f"{expected_seq_len}",
|
||||
)
|
||||
raise KvCacheMissError(miss)
|
||||
n_new = self._new_token_count(kv_by_layer)
|
||||
incoming = n_new * self.recipe.bytes_per_token()
|
||||
self._enforce_capacity_locked(
|
||||
protect=(session_id, route_epoch), incoming_bytes=incoming
|
||||
)
|
||||
session._append(kv_by_layer)
|
||||
session.last_tick = self._clock()
|
||||
self._sessions.move_to_end((session_id, route_epoch))
|
||||
return session
|
||||
|
||||
def truncate(
|
||||
self, session_id: str, route_epoch: int, length: int
|
||||
) -> SessionCache:
|
||||
"""Drop cached positions beyond ``length`` (rollback) for one session."""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
session = self._require_live_locked(session_id, route_epoch)
|
||||
if length < 0:
|
||||
raise ValueError("truncate length must be >= 0")
|
||||
session._truncate(length)
|
||||
session.last_tick = self._clock()
|
||||
self._sessions.move_to_end((session_id, route_epoch))
|
||||
return session
|
||||
|
||||
def release(self, session_id: str, route_epoch: int) -> bool:
|
||||
"""Free one session's context; other sessions are untouched.
|
||||
|
||||
Returns True if a live context was freed. A later lookup for the released
|
||||
key yields an explicit :class:`CacheMiss`.
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
key = (session_id, route_epoch)
|
||||
existed = key in self._sessions
|
||||
self._drop_and_record_locked(key, CacheMissReason.RELEASED)
|
||||
return existed
|
||||
|
||||
# -- lookup ---------------------------------------------------------------
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache | CacheMiss:
|
||||
"""Return the live context or an explicit :class:`CacheMiss`.
|
||||
|
||||
Rejects stale epochs and incompatible recipes (both are protocol
|
||||
violations, not benign misses).
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
key = (session_id, route_epoch)
|
||||
session = self._sessions.get(key)
|
||||
if session is None:
|
||||
return self._recorded_miss_locked(key)
|
||||
if expected_seq_len is not None and session.seq_len != expected_seq_len:
|
||||
return self._drop_and_record_locked(
|
||||
key,
|
||||
CacheMissReason.SEQ_LEN_MISMATCH,
|
||||
detail=f"cache holds {session.seq_len}, caller expected "
|
||||
f"{expected_seq_len}",
|
||||
)
|
||||
session.last_tick = now
|
||||
self._sessions.move_to_end(key)
|
||||
return session
|
||||
|
||||
def get(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache:
|
||||
"""Strict accessor: raises :class:`KvCacheMissError` on a miss."""
|
||||
result = self.resolve(
|
||||
session_id,
|
||||
route_epoch,
|
||||
recipe=recipe,
|
||||
expected_seq_len=expected_seq_len,
|
||||
)
|
||||
if isinstance(result, CacheMiss):
|
||||
raise KvCacheMissError(result)
|
||||
return result
|
||||
|
||||
# -- internals ------------------------------------------------------------
|
||||
|
||||
def _check_recipe(self, recipe: KvCacheRecipe | None) -> None:
|
||||
if recipe is not None and not self.recipe.is_compatible(recipe):
|
||||
raise IncompatibleCacheRecipeError(
|
||||
"request cache recipe does not match this shard's loaded recipe "
|
||||
f"(request {recipe.fingerprint()} vs shard {self.recipe.fingerprint()})"
|
||||
)
|
||||
|
||||
def _validate_epoch_locked(self, session_id: str, route_epoch: int) -> None:
|
||||
latest = self._latest_epoch.get(session_id)
|
||||
if latest is not None and route_epoch < latest:
|
||||
raise StaleRouteEpochError(
|
||||
f"session {session_id[:8]} route epoch {route_epoch} is stale; "
|
||||
f"current epoch is {latest}"
|
||||
)
|
||||
|
||||
def _supersede_older_epochs_locked(
|
||||
self, session_id: str, route_epoch: int
|
||||
) -> None:
|
||||
stale_keys = [
|
||||
key
|
||||
for key in self._sessions
|
||||
if key[0] == session_id and key[1] < route_epoch
|
||||
]
|
||||
for key in stale_keys:
|
||||
self._drop_and_record_locked(key, CacheMissReason.SUPERSEDED_EPOCH)
|
||||
|
||||
def _require_live_locked(
|
||||
self, session_id: str, route_epoch: int
|
||||
) -> SessionCache:
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
key = (session_id, route_epoch)
|
||||
session = self._sessions.get(key)
|
||||
if session is None:
|
||||
raise KvCacheMissError(self._recorded_miss_locked(key))
|
||||
return session
|
||||
|
||||
def _new_token_count(self, kv_by_layer: Mapping[int, Any]) -> int:
|
||||
owned = set(self.recipe.owned_layers)
|
||||
if set(kv_by_layer) != owned:
|
||||
raise ValueError(
|
||||
f"append must cover exactly the owned layers {sorted(owned)}, "
|
||||
f"got {sorted(kv_by_layer)}"
|
||||
)
|
||||
counts = {int(np.asarray(k).shape[0]) for k, _ in kv_by_layer.values()}
|
||||
if len(counts) != 1:
|
||||
raise ValueError(
|
||||
f"append token counts disagree across layers: {sorted(counts)}"
|
||||
)
|
||||
return counts.pop()
|
||||
|
||||
def _enforce_capacity_locked(
|
||||
self, *, protect: tuple[str, int], incoming_bytes: int
|
||||
) -> None:
|
||||
# Session cap: evict LRU sessions other than the protected one.
|
||||
while len(self._sessions) > self.config.max_sessions:
|
||||
victim = self._lru_victim_locked(protect)
|
||||
if victim is None:
|
||||
break
|
||||
self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)
|
||||
|
||||
# Byte budget: the protected session's own footprint after the append.
|
||||
protected = self._sessions.get(protect)
|
||||
protected_bytes = (protected.nbytes if protected is not None else 0) + incoming_bytes
|
||||
if protected_bytes > self.config.budget_bytes:
|
||||
raise KvBudgetExceededError(
|
||||
f"session {protect[0][:8]} needs {protected_bytes} bytes which "
|
||||
f"exceeds the KV budget {self.config.budget_bytes}"
|
||||
)
|
||||
# Evict other LRU sessions until the whole store fits with the append.
|
||||
while self._total_bytes_locked() + incoming_bytes > self.config.budget_bytes:
|
||||
victim = self._lru_victim_locked(protect)
|
||||
if victim is None:
|
||||
break
|
||||
self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)
|
||||
|
||||
def _lru_victim_locked(self, protect: tuple[str, int]) -> tuple[str, int] | None:
|
||||
for key in self._sessions: # OrderedDict iterates oldest-first.
|
||||
if key != protect:
|
||||
return key
|
||||
return None
|
||||
|
||||
def _total_bytes_locked(self) -> int:
|
||||
return sum(s.nbytes for s in self._sessions.values())
|
||||
|
||||
def _evict_expired_locked(self, now: float) -> None:
|
||||
ttl = self.config.ttl_seconds
|
||||
if ttl <= 0:
|
||||
return
|
||||
expired = [
|
||||
key
|
||||
for key, session in self._sessions.items()
|
||||
if now - session.last_tick > ttl
|
||||
]
|
||||
for key in expired:
|
||||
self._drop_and_record_locked(key, CacheMissReason.EVICTED_TTL)
|
||||
|
||||
def _drop_and_record_locked(
|
||||
self,
|
||||
key: tuple[str, int],
|
||||
reason: CacheMissReason,
|
||||
*,
|
||||
detail: str = "",
|
||||
) -> CacheMiss:
|
||||
session = self._sessions.pop(key, None)
|
||||
if session is not None:
|
||||
session.released = True
|
||||
miss = CacheMiss(
|
||||
session_id=key[0], route_epoch=key[1], reason=reason, detail=detail
|
||||
)
|
||||
self._record_miss_locked(key, miss)
|
||||
return miss
|
||||
|
||||
def _record_miss_locked(self, key: tuple[str, int], miss: CacheMiss) -> None:
|
||||
if self.config.miss_history <= 0:
|
||||
return
|
||||
self._misses.pop(key, None)
|
||||
self._misses[key] = miss
|
||||
while len(self._misses) > self.config.miss_history:
|
||||
self._misses.popitem(last=False)
|
||||
|
||||
def _recorded_miss_locked(self, key: tuple[str, int]) -> CacheMiss:
|
||||
recorded = self._misses.get(key)
|
||||
if recorded is not None:
|
||||
return recorded
|
||||
return CacheMiss(
|
||||
session_id=key[0],
|
||||
route_epoch=key[1],
|
||||
reason=CacheMissReason.UNKNOWN_SESSION,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _require_text(value: Any, name: str) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{name} must be a non-empty string")
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _require_epoch(value: Any) -> int:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError("route_epoch must be an integer")
|
||||
if value < 0:
|
||||
raise ValueError("route_epoch must be >= 0")
|
||||
return value
|
||||
|
||||
|
||||
def kv_recipe_for(computation: Any) -> KvCacheRecipe:
|
||||
"""Build a :class:`KvCacheRecipe` from a KV-aware ``ShardComputation``.
|
||||
|
||||
The computation exposes the DGR-006 duck type plus KV geometry
|
||||
(``n_kv_heads``, ``head_dim``, ``kv_dtype``).
|
||||
"""
|
||||
return KvCacheRecipe(
|
||||
architecture_adapter=str(getattr(computation, "architecture_adapter")),
|
||||
kv_dtype=str(getattr(computation, "kv_dtype", "float32")),
|
||||
n_kv_heads=int(getattr(computation, "n_kv_heads")),
|
||||
head_dim=int(getattr(computation, "head_dim")),
|
||||
total_layers=int(getattr(computation, "total_layers")),
|
||||
start_layer=int(getattr(computation, "start_layer")),
|
||||
end_layer=int(getattr(computation, "end_layer")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KvBoundaryAdapter:
|
||||
"""KV-aware boundary driver: cached prefill/decode through the manager.
|
||||
|
||||
Mirrors the DGR-006 :class:`~meshnet_node.boundary_adapter.BoundaryAdapter`
|
||||
contract (head embeds tokens, middle/tail bypass embedding and consume the
|
||||
unnormalized residual bundle, non-tail emits the unnormalized residual, tail
|
||||
normalizes + heads + prunes + samples) but threads a per-session KV context.
|
||||
|
||||
The wrapped computation must additionally expose::
|
||||
|
||||
run_layers_cached(hidden, *, positions, past_kv)
|
||||
-> (hidden_out, {layer_index: (new_keys, new_values)})
|
||||
|
||||
reading ``past_kv`` (the current per-owned-layer caches) and returning the new
|
||||
position-encoded K/V for the appended positions only. The manager, not the
|
||||
computation, commits those K/V so eviction and budget stay centralized.
|
||||
"""
|
||||
|
||||
computation: Any
|
||||
manager: HotKvStateManager
|
||||
sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
|
||||
architecture: Any = field(init=False)
|
||||
role: ShardRole = field(init=False)
|
||||
start_layer: int = field(init=False)
|
||||
end_layer: int = field(init=False)
|
||||
total_layers: int = field(init=False)
|
||||
recipe: KvCacheRecipe = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
arch_name = getattr(self.computation, "architecture_adapter", None)
|
||||
self.architecture = certified_architecture(arch_name)
|
||||
self.start_layer = int(getattr(self.computation, "start_layer"))
|
||||
self.end_layer = int(getattr(self.computation, "end_layer"))
|
||||
self.total_layers = int(getattr(self.computation, "total_layers"))
|
||||
self.role = role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
self.recipe = kv_recipe_for(self.computation)
|
||||
if not self.manager.recipe.is_compatible(self.recipe):
|
||||
raise IncompatibleCacheRecipeError(
|
||||
"manager recipe does not match this computation's KV recipe"
|
||||
)
|
||||
|
||||
@property
|
||||
def is_head(self) -> bool:
|
||||
return self.role.owns_embedding
|
||||
|
||||
@property
|
||||
def is_tail(self) -> bool:
|
||||
return self.role.owns_final_head
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
token_ids: Any | None = None,
|
||||
boundary: BoundaryBundle | None = None,
|
||||
) -> BoundaryBundle | TailOutput:
|
||||
"""Open a fresh isolated context and run the prompt through this range."""
|
||||
session = self.manager.open(session_id, route_epoch, recipe=self.recipe)
|
||||
return self._run_step(session, token_ids, boundary)
|
||||
|
||||
def decode(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
token_ids: Any | None = None,
|
||||
boundary: BoundaryBundle | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> BoundaryBundle | TailOutput | CacheMiss:
|
||||
"""Append one (or more) decode positions to an existing context.
|
||||
|
||||
Returns an explicit :class:`CacheMiss` if the context is gone so the head
|
||||
can re-prefill from token zero instead of corrupting output.
|
||||
"""
|
||||
resolved = self.manager.resolve(
|
||||
session_id,
|
||||
route_epoch,
|
||||
recipe=self.recipe,
|
||||
expected_seq_len=expected_seq_len,
|
||||
)
|
||||
if isinstance(resolved, CacheMiss):
|
||||
return resolved
|
||||
return self._run_step(resolved, token_ids, boundary)
|
||||
|
||||
# -- internals ------------------------------------------------------------
|
||||
|
||||
def _run_step(
|
||||
self,
|
||||
session: SessionCache,
|
||||
token_ids: Any | None,
|
||||
boundary: BoundaryBundle | None,
|
||||
) -> BoundaryBundle | TailOutput:
|
||||
prev_len = session.seq_len
|
||||
hidden, positions = self._ingest(prev_len, token_ids, boundary)
|
||||
hidden_out, new_kv = self.computation.run_layers_cached(
|
||||
hidden, positions=positions, past_kv=session.read_only_layers()
|
||||
)
|
||||
self.manager.append(
|
||||
session.session_id,
|
||||
session.route_epoch,
|
||||
new_kv,
|
||||
recipe=self.recipe,
|
||||
expected_seq_len=prev_len,
|
||||
)
|
||||
if self.is_tail:
|
||||
return self._emit_tail(hidden_out)
|
||||
return self._emit_boundary(hidden_out, positions)
|
||||
|
||||
def _ingest(
|
||||
self,
|
||||
prev_len: int,
|
||||
token_ids: Any | None,
|
||||
boundary: BoundaryBundle | None,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if self.role.owns_embedding:
|
||||
if token_ids is None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding and must receive token IDs"
|
||||
)
|
||||
if boundary is not None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding; it must not receive a boundary "
|
||||
"bundle from an upstream range"
|
||||
)
|
||||
ids = np.asarray(token_ids)
|
||||
if ids.ndim == 1:
|
||||
ids = ids[None, :]
|
||||
if ids.ndim != 2:
|
||||
raise BoundaryContractError("token IDs must be (seq,) or (batch, seq)")
|
||||
hidden = np.asarray(self.computation.embed_tokens(ids))
|
||||
n_new = ids.shape[1]
|
||||
positions = np.broadcast_to(
|
||||
np.arange(prev_len, prev_len + n_new, dtype=np.int64),
|
||||
ids.shape,
|
||||
).copy()
|
||||
return hidden, positions
|
||||
# Middle / tail: consume the boundary bundle (the unnormalized residual).
|
||||
if token_ids is not None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards bypass token embedding; they must not receive "
|
||||
"token IDs"
|
||||
)
|
||||
if boundary is None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards must receive the named boundary bundle"
|
||||
)
|
||||
self._check_boundary(boundary)
|
||||
return np.asarray(boundary.residual), np.asarray(boundary.positions)
|
||||
|
||||
def _check_boundary(self, boundary: BoundaryBundle) -> None:
|
||||
if certified_architecture(boundary.architecture_adapter) is not self.architecture:
|
||||
raise BoundaryContractError(
|
||||
f"boundary bundle architecture {boundary.architecture_adapter!r} "
|
||||
f"does not match this Shard's adapter {self.architecture.adapter!r}"
|
||||
)
|
||||
if boundary.schema_version != self.architecture.boundary_schema_version:
|
||||
raise BoundaryContractError(
|
||||
f"boundary schema v{boundary.schema_version} is not supported by "
|
||||
f"this Shard (expects v{self.architecture.boundary_schema_version})"
|
||||
)
|
||||
if boundary.tensor_name != self.architecture.boundary_tensor_name:
|
||||
raise BoundaryContractError(
|
||||
f"boundary tensor {boundary.tensor_name!r} is not the "
|
||||
f"architecture-defined {self.architecture.boundary_tensor_name!r}"
|
||||
)
|
||||
if boundary.normalized:
|
||||
raise BoundaryContractError(
|
||||
"boundary bundle is normalized; a Shard range must receive the "
|
||||
"UNNORMALIZED architecture-defined residual"
|
||||
)
|
||||
if boundary.next_layer != self.start_layer:
|
||||
raise BoundaryContractError(
|
||||
f"boundary hands over at layer {boundary.next_layer} but this "
|
||||
f"Shard starts at layer {self.start_layer}"
|
||||
)
|
||||
|
||||
def _emit_boundary(
|
||||
self, hidden: np.ndarray, positions: np.ndarray
|
||||
) -> BoundaryBundle:
|
||||
return BoundaryBundle(
|
||||
architecture_adapter=self.architecture.adapter,
|
||||
schema_version=self.architecture.boundary_schema_version,
|
||||
tensor_name=self.architecture.boundary_tensor_name,
|
||||
residual=np.asarray(hidden),
|
||||
positions=np.asarray(positions),
|
||||
next_layer=self.end_layer + 1,
|
||||
normalized=False,
|
||||
)
|
||||
|
||||
def _emit_tail(self, hidden: np.ndarray) -> TailOutput:
|
||||
hidden = np.asarray(hidden)
|
||||
if self.architecture.prunes_rows_at_tail:
|
||||
last_hidden = hidden[:, -1:, :]
|
||||
else: # pragma: no cover - no certified architecture takes this path yet
|
||||
last_hidden = hidden
|
||||
if self.architecture.normalizes_before_head:
|
||||
last_hidden = np.asarray(self.computation.final_norm(last_hidden))
|
||||
logits = np.asarray(self.computation.lm_head(last_hidden))
|
||||
last_logits = logits[:, -1, :]
|
||||
token_id = self.sampling.sample(last_logits)
|
||||
return TailOutput(token_id=token_id, logits=last_logits, sampling=self.sampling)
|
||||
Reference in New Issue
Block a user