919 lines
35 KiB
Python
919 lines
35 KiB
Python
"""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)
|