770 lines
29 KiB
Python
770 lines
29 KiB
Python
"""Isolated concurrent local Hot KV State (DGR-007).
|
|
|
|
These tests prove the KV/session manager with a *pure-numpy* KV-cached dense-Llama
|
|
reference: no download, no GPU, no torch, no API credit. The reference implements
|
|
the DGR-006 ``ShardComputation`` duck type plus ``run_layers_cached`` so cached
|
|
prefill/decode over a per-session KV context reproduces the stateless whole-model
|
|
tokens bit-for-bit. On top of that correctness core, the tests exercise the
|
|
manager's lifecycle: owned-layer allocation, prefill/decode append, truncate,
|
|
release, TTL/LRU eviction, explicit cache-miss responses, stale-epoch and
|
|
incompatible-recipe rejection, four concurrent cross-talk-free sessions, and
|
|
budget-bounded cancellation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from meshnet_node.boundary_adapter import BoundaryBundle, TailOutput
|
|
from meshnet_node.hot_kv_state import (
|
|
CacheMiss,
|
|
CacheMissReason,
|
|
HotKvStateConfig,
|
|
HotKvStateManager,
|
|
IncompatibleCacheRecipeError,
|
|
KvBoundaryAdapter,
|
|
KvBudgetExceededError,
|
|
KvCacheMissError,
|
|
KvCacheRecipe,
|
|
LayerKvCache,
|
|
StaleRouteEpochError,
|
|
kv_recipe_for,
|
|
)
|
|
|
|
PARITY_ATOL = 1e-6
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Pure-numpy KV-cached dense-Llama reference (test fixture, not production).
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class _KvDenseLlama:
|
|
"""A tiny deterministic dense-Llama with both stateless and cached runners."""
|
|
|
|
architecture_adapter = "dense-llama"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
vocab: int = 48,
|
|
hidden: int = 32,
|
|
n_layers: int = 6,
|
|
n_heads: int = 4,
|
|
intermediate: int = 64,
|
|
rms_eps: float = 1e-6,
|
|
rope_theta: float = 10000.0,
|
|
seed: int = 20260716,
|
|
) -> None:
|
|
assert hidden % n_heads == 0
|
|
self.vocab = vocab
|
|
self.hidden = hidden
|
|
self.n_layers = n_layers
|
|
self.n_heads = n_heads
|
|
self.head_dim = hidden // n_heads
|
|
assert self.head_dim % 2 == 0
|
|
self.rms_eps = rms_eps
|
|
self.rope_theta = rope_theta
|
|
|
|
rng = np.random.default_rng(seed)
|
|
|
|
def w(*shape: int) -> np.ndarray:
|
|
return (rng.standard_normal(shape) * 0.08).astype(np.float32)
|
|
|
|
self.embed = w(vocab, hidden)
|
|
self.layers = []
|
|
for _ in range(n_layers):
|
|
self.layers.append(
|
|
{
|
|
"in_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
|
"q": w(hidden, hidden),
|
|
"k": w(hidden, hidden),
|
|
"v": w(hidden, hidden),
|
|
"o": w(hidden, hidden),
|
|
"post_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
|
"gate": w(intermediate, hidden),
|
|
"up": w(intermediate, hidden),
|
|
"down": w(hidden, intermediate),
|
|
}
|
|
)
|
|
self.final_ln = (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32)
|
|
self.lm_head_w = w(vocab, hidden)
|
|
|
|
inv_freq = 1.0 / (
|
|
rope_theta ** (np.arange(0, self.head_dim, 2, dtype=np.float32) / self.head_dim)
|
|
)
|
|
self.inv_freq = inv_freq.astype(np.float32)
|
|
|
|
# -- primitive ops -----------------------------------------------------
|
|
def _rmsnorm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
|
|
variance = np.mean(x.astype(np.float32) ** 2, axis=-1, keepdims=True)
|
|
normed = x / np.sqrt(variance + self.rms_eps)
|
|
return (normed * weight).astype(np.float32)
|
|
|
|
def _rope(self, positions: np.ndarray):
|
|
angles = positions[..., None].astype(np.float32) * self.inv_freq[None, None, :]
|
|
emb = np.concatenate([angles, angles], axis=-1)
|
|
return np.cos(emb).astype(np.float32), np.sin(emb).astype(np.float32)
|
|
|
|
@staticmethod
|
|
def _rotate_half(x: np.ndarray) -> np.ndarray:
|
|
half = x.shape[-1] // 2
|
|
return np.concatenate([-x[..., half:], x[..., :half]], axis=-1)
|
|
|
|
def _apply_rope(self, t: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
|
|
cos = cos[:, None, :, :]
|
|
sin = sin[:, None, :, :]
|
|
return t * cos + self._rotate_half(t) * sin
|
|
|
|
def _project_qkv(self, normed: np.ndarray, layer: dict, positions: np.ndarray):
|
|
batch, seq, _ = normed.shape
|
|
q = (normed @ layer["q"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
k = (normed @ layer["k"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
v = (normed @ layer["v"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
|
q = q.transpose(0, 2, 1, 3)
|
|
k = k.transpose(0, 2, 1, 3)
|
|
v = v.transpose(0, 2, 1, 3)
|
|
cos, sin = self._rope(positions)
|
|
q = self._apply_rope(q, cos, sin)
|
|
k = self._apply_rope(k, cos, sin)
|
|
return q, k, v
|
|
|
|
def _attend(
|
|
self,
|
|
q: np.ndarray,
|
|
k_all: np.ndarray,
|
|
v_all: np.ndarray,
|
|
layer: dict,
|
|
q_positions: np.ndarray,
|
|
) -> np.ndarray:
|
|
batch, _, seq_new, _ = q.shape
|
|
total = k_all.shape[2]
|
|
scores = (q @ k_all.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
|
|
# Causal mask by absolute position: keys are stored in absolute order
|
|
# 0..total-1; query row i lives at absolute position q_positions[i].
|
|
key_abs = np.arange(total, dtype=np.int64)
|
|
q_abs = np.asarray(q_positions).reshape(seq_new).astype(np.int64)
|
|
mask = np.where(key_abs[None, :] <= q_abs[:, None], 0.0, -1e30).astype(np.float32)
|
|
scores = scores + mask[None, None, :, :]
|
|
scores = scores - scores.max(axis=-1, keepdims=True)
|
|
weights = np.exp(scores)
|
|
weights = weights / weights.sum(axis=-1, keepdims=True)
|
|
out = weights @ v_all
|
|
out = out.transpose(0, 2, 1, 3).reshape(batch, seq_new, self.hidden)
|
|
return (out @ layer["o"].T).astype(np.float32)
|
|
|
|
def _mlp(self, x: np.ndarray, layer: dict) -> np.ndarray:
|
|
gate = x @ layer["gate"].T
|
|
up = x @ layer["up"].T
|
|
silu = gate * (1.0 / (1.0 + np.exp(-gate)))
|
|
return ((silu * up) @ layer["down"].T).astype(np.float32)
|
|
|
|
# -- stateless whole-sequence layer (ground truth) ---------------------
|
|
def _run_layer_stateless(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
|
normed = self._rmsnorm(x, layer["in_ln"])
|
|
q, k, v = self._project_qkv(normed, layer, positions)
|
|
attn = self._attend(q, k, v, layer, positions[0])
|
|
h = x + attn
|
|
h = h + self._mlp(self._rmsnorm(h, layer["post_ln"]), layer)
|
|
return h.astype(np.float32)
|
|
|
|
def whole_model_next_token(self, token_ids: list[int]) -> int:
|
|
positions = np.arange(len(token_ids))[None, :]
|
|
h = self.embed[np.asarray(token_ids)][None, :]
|
|
for idx in range(self.n_layers):
|
|
h = self._run_layer_stateless(h, self.layers[idx], positions)
|
|
h = self._rmsnorm(h[:, -1:, :], self.final_ln)
|
|
logits = h @ self.lm_head_w.T
|
|
return int(np.argmax(logits[0, -1]))
|
|
|
|
def stateless_greedy(self, prompt: list[int], n_new: int) -> list[int]:
|
|
tokens = list(prompt)
|
|
out: list[int] = []
|
|
for _ in range(n_new):
|
|
tok = self.whole_model_next_token(tokens)
|
|
tokens.append(tok)
|
|
out.append(tok)
|
|
return out
|
|
|
|
|
|
class _KvReferenceShard:
|
|
"""A contiguous inclusive layer range with a KV-cached runner.
|
|
|
|
Satisfies the KV-aware ``ShardComputation`` duck type used by
|
|
``KvBoundaryAdapter``: DGR-006 methods plus ``run_layers_cached`` and the KV
|
|
geometry (``n_kv_heads`` / ``head_dim`` / ``kv_dtype``).
|
|
"""
|
|
|
|
kv_dtype = "float32"
|
|
|
|
def __init__(
|
|
self,
|
|
model: _KvDenseLlama,
|
|
start_layer: int,
|
|
end_layer: int,
|
|
*,
|
|
architecture_adapter: str | None = None,
|
|
) -> None:
|
|
self._model = model
|
|
self.start_layer = start_layer
|
|
self.end_layer = end_layer
|
|
self.total_layers = model.n_layers
|
|
self.n_kv_heads = model.n_heads
|
|
self.head_dim = model.head_dim
|
|
self.architecture_adapter = architecture_adapter or model.architecture_adapter
|
|
|
|
def embed_tokens(self, token_ids: np.ndarray) -> np.ndarray:
|
|
return self._model.embed[np.asarray(token_ids)]
|
|
|
|
def final_norm(self, hidden: np.ndarray) -> np.ndarray:
|
|
return self._model._rmsnorm(np.asarray(hidden, dtype=np.float32), self._model.final_ln)
|
|
|
|
def lm_head(self, hidden: np.ndarray) -> np.ndarray:
|
|
return np.asarray(hidden, dtype=np.float32) @ self._model.lm_head_w.T
|
|
|
|
def run_layers_cached(self, hidden, *, positions, past_kv):
|
|
m = self._model
|
|
x = np.asarray(hidden, dtype=np.float32)
|
|
positions = np.asarray(positions)
|
|
new_kv: dict[int, tuple[np.ndarray, np.ndarray]] = {}
|
|
for idx in range(self.start_layer, self.end_layer + 1):
|
|
layer = m.layers[idx]
|
|
normed = m._rmsnorm(x, layer["in_ln"])
|
|
q, k, v = m._project_qkv(normed, layer, positions)
|
|
# Post-RoPE new K/V stored as (seq_new, n_heads, head_dim).
|
|
new_k = k[0].transpose(1, 0, 2).copy()
|
|
new_v = v[0].transpose(1, 0, 2).copy()
|
|
cache = past_kv.get(idx)
|
|
if cache is not None and cache.length > 0:
|
|
past_k = cache.keys[None].transpose(0, 2, 1, 3)
|
|
past_v = cache.values[None].transpose(0, 2, 1, 3)
|
|
k_all = np.concatenate([past_k, k], axis=2)
|
|
v_all = np.concatenate([past_v, v], axis=2)
|
|
else:
|
|
k_all, v_all = k, v
|
|
attn = m._attend(q, k_all, v_all, layer, positions[0])
|
|
h = x + attn
|
|
x = h + m._mlp(m._rmsnorm(h, layer["post_ln"]), layer)
|
|
x = x.astype(np.float32)
|
|
new_kv[idx] = (new_k, new_v)
|
|
return x, new_kv
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Helpers.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class _FakeClock:
|
|
def __init__(self) -> None:
|
|
self.now = 0.0
|
|
|
|
def __call__(self) -> float:
|
|
return self.now
|
|
|
|
def advance(self, delta: float) -> None:
|
|
self.now += delta
|
|
|
|
|
|
def _full_shard(model: _KvDenseLlama):
|
|
return _KvReferenceShard(model, 0, model.n_layers - 1)
|
|
|
|
|
|
def _manager_for(shard, config: HotKvStateConfig | None = None, clock=None) -> HotKvStateManager:
|
|
return HotKvStateManager(kv_recipe_for(shard), config=config, clock=clock)
|
|
|
|
|
|
def _cached_greedy(
|
|
adapter: KvBoundaryAdapter,
|
|
manager: HotKvStateManager,
|
|
session_id: str,
|
|
epoch: int,
|
|
prompt: list[int],
|
|
n_new: int,
|
|
) -> list[int]:
|
|
"""Greedy decode one full-model session through the KV manager."""
|
|
out = adapter.prefill(session_id, epoch, token_ids=np.asarray(prompt))
|
|
assert isinstance(out, TailOutput)
|
|
tokens = [out.token_id]
|
|
for _ in range(n_new - 1):
|
|
step = adapter.decode(session_id, epoch, token_ids=[out.token_id])
|
|
assert isinstance(step, TailOutput)
|
|
out = step
|
|
tokens.append(out.token_id)
|
|
return tokens
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Recipe identity.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_recipe_owned_layers_and_fingerprint_aliasing():
|
|
"The KV recipe covers only owned layers and canonicalizes architecture aliases.\n\nTags: node, kv"
|
|
recipe = KvCacheRecipe(
|
|
architecture_adapter="LlamaForCausalLM",
|
|
kv_dtype="float32",
|
|
n_kv_heads=4,
|
|
head_dim=8,
|
|
total_layers=6,
|
|
start_layer=2,
|
|
end_layer=3,
|
|
)
|
|
assert recipe.owned_layers == (2, 3)
|
|
alias = KvCacheRecipe(
|
|
architecture_adapter="llama",
|
|
kv_dtype="float32",
|
|
n_kv_heads=4,
|
|
head_dim=8,
|
|
total_layers=6,
|
|
start_layer=2,
|
|
end_layer=3,
|
|
)
|
|
assert recipe.is_compatible(alias)
|
|
# A different owned range is not compatible.
|
|
other = KvCacheRecipe(
|
|
architecture_adapter="llama",
|
|
kv_dtype="float32",
|
|
n_kv_heads=4,
|
|
head_dim=8,
|
|
total_layers=6,
|
|
start_layer=0,
|
|
end_layer=1,
|
|
)
|
|
assert not recipe.is_compatible(other)
|
|
|
|
|
|
def test_recipe_bytes_per_token_scales_with_owned_layers():
|
|
"KV bytes-per-token counts keys+values across owned layers only.\n\nTags: node, kv"
|
|
base = dict(
|
|
architecture_adapter="dense-llama",
|
|
kv_dtype="float32",
|
|
n_kv_heads=4,
|
|
head_dim=8,
|
|
total_layers=6,
|
|
)
|
|
one = KvCacheRecipe(**base, start_layer=0, end_layer=0)
|
|
two = KvCacheRecipe(**base, start_layer=0, end_layer=1)
|
|
# 2 (k+v) * heads * dim * 4 bytes per layer.
|
|
assert one.bytes_per_token() == 2 * 4 * 8 * 4
|
|
assert two.bytes_per_token() == 2 * one.bytes_per_token()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Owned-layer allocation.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_manager_allocates_kv_only_for_owned_layers():
|
|
"A middle shard allocates KV state only for its owned layer range.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _KvReferenceShard(model, 2, 3)
|
|
manager = _manager_for(shard)
|
|
session = manager.open("sess-mid", 0)
|
|
assert session.owned_layers == (2, 3)
|
|
assert set(session.layers) == {2, 3}
|
|
with pytest.raises(KeyError):
|
|
session.layer(0)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Prefill / decode / truncate.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_prefill_then_decode_append_grows_owned_layers():
|
|
"Prefill and decode append advance every owned layer in lockstep.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompt = [5, 12, 3, 41]
|
|
out = adapter.prefill("s", 0, token_ids=np.asarray(prompt))
|
|
assert isinstance(out, TailOutput)
|
|
session = manager.get("s", 0)
|
|
assert session.seq_len == len(prompt)
|
|
for cache in session.layers.values():
|
|
assert cache.length == len(prompt)
|
|
|
|
step = adapter.decode("s", 0, token_ids=[out.token_id])
|
|
assert isinstance(step, TailOutput)
|
|
assert manager.get("s", 0).seq_len == len(prompt) + 1
|
|
|
|
|
|
def test_truncate_rolls_back_all_owned_layers():
|
|
"Truncate drops cached positions beyond a length across owned layers.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3, 4, 5]))
|
|
assert manager.get("s", 0).seq_len == 5
|
|
manager.truncate("s", 0, 2)
|
|
session = manager.get("s", 0)
|
|
assert session.seq_len == 2
|
|
for cache in session.layers.values():
|
|
assert cache.length == 2
|
|
|
|
|
|
def test_layer_kv_cache_rejects_wrong_shape():
|
|
"LayerKvCache rejects K/V that do not match its head geometry.\n\nTags: node, kv"
|
|
cache = LayerKvCache(0, n_kv_heads=4, head_dim=8, dtype="float32")
|
|
with pytest.raises(ValueError):
|
|
cache.append(np.zeros((1, 3, 8), dtype=np.float32), np.zeros((1, 3, 8), dtype=np.float32))
|
|
cache.append(np.zeros((2, 4, 8), dtype=np.float32), np.zeros((2, 4, 8), dtype=np.float32))
|
|
assert cache.length == 2
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Cached vs stateless parity (correctness core).
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_cached_full_shard_decode_matches_stateless_whole_model():
|
|
"Cached full-model greedy decode reproduces stateless whole-model tokens.\n\nTags: node, kv, parity"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompt = [2, 17, 8, 25, 6]
|
|
n_new = 12
|
|
reference = model.stateless_greedy(prompt, n_new)
|
|
cached = _cached_greedy(adapter, manager, "s", 0, prompt, n_new)
|
|
assert cached == reference
|
|
assert len(cached) == n_new
|
|
|
|
|
|
def test_cached_prefill_next_token_matches_whole_model_logits():
|
|
"Cached prefill produces the same next-token logits as the whole model.\n\nTags: node, kv, parity"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompt = [9, 1, 44, 6, 30, 11]
|
|
out = adapter.prefill("s", 0, token_ids=np.asarray(prompt))
|
|
assert isinstance(out, TailOutput)
|
|
assert out.token_id == model.whole_model_next_token(prompt)
|
|
|
|
|
|
def test_multi_range_cached_decode_parity_across_a_seam():
|
|
"A head/tail split with independent per-range KV reproduces whole-model decode.\n\nTags: node, kv, parity"
|
|
model = _KvDenseLlama()
|
|
head_shard = _KvReferenceShard(model, 0, 2)
|
|
tail_shard = _KvReferenceShard(model, 3, 5)
|
|
head_mgr = _manager_for(head_shard)
|
|
tail_mgr = _manager_for(tail_shard)
|
|
head = KvBoundaryAdapter(head_shard, head_mgr)
|
|
tail = KvBoundaryAdapter(tail_shard, tail_mgr)
|
|
|
|
prompt = [7, 3, 22, 5, 9]
|
|
n_new = 8
|
|
|
|
# Each range only allocates its owned layers.
|
|
def step(token_ids, is_prefill):
|
|
if is_prefill:
|
|
bundle = head.prefill("s", 0, token_ids=np.asarray(token_ids))
|
|
out = tail.prefill("s", 0, boundary=bundle)
|
|
else:
|
|
bundle = head.decode("s", 0, token_ids=[token_ids])
|
|
assert isinstance(bundle, BoundaryBundle)
|
|
out = tail.decode("s", 0, boundary=bundle)
|
|
assert isinstance(out, TailOutput)
|
|
return out.token_id
|
|
|
|
tokens = [step(prompt, True)]
|
|
for _ in range(n_new - 1):
|
|
tokens.append(step(tokens[-1], False))
|
|
|
|
assert head_mgr.get("s", 0).owned_layers == (0, 1, 2)
|
|
assert tail_mgr.get("s", 0).owned_layers == (3, 4, 5)
|
|
assert tokens == model.stateless_greedy(prompt, n_new)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Four concurrent sessions with no cross-talk.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_four_interleaved_sessions_have_no_kv_cross_talk():
|
|
"Four interleaved sessions each decode their own tokens without cross-talk.\n\nTags: node, kv, concurrency"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompts = {
|
|
"alpha": [1, 2, 3, 4],
|
|
"bravo": [40, 39, 2, 15],
|
|
"charlie": [7, 7, 7, 7],
|
|
"delta": [31, 5, 18, 22],
|
|
}
|
|
n_new = 10
|
|
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
|
# The four prompts must actually diverge, else "no cross-talk" is vacuous.
|
|
assert len({tuple(v) for v in references.values()}) == 4
|
|
|
|
generated: dict[str, list[int]] = {}
|
|
for sid, prompt in prompts.items():
|
|
out = adapter.prefill(sid, 0, token_ids=np.asarray(prompt))
|
|
assert isinstance(out, TailOutput)
|
|
generated[sid] = [out.token_id]
|
|
|
|
# Round-robin decode: every session takes one step per round, interleaved.
|
|
for _ in range(n_new - 1):
|
|
for sid in prompts:
|
|
step = adapter.decode(sid, 0, token_ids=[generated[sid][-1]])
|
|
assert isinstance(step, TailOutput)
|
|
generated[sid].append(step.token_id)
|
|
|
|
for sid in prompts:
|
|
assert generated[sid] == references[sid], sid
|
|
assert manager.session_count == 4
|
|
|
|
|
|
def test_four_sessions_on_real_threads_stay_isolated():
|
|
"Four sessions decoding on real threads produce their own reference tokens.\n\nTags: node, kv, concurrency"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard, HotKvStateConfig(max_sessions=8))
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompts = {
|
|
"t-alpha": [3, 14, 1, 5],
|
|
"t-bravo": [2, 27, 18, 4],
|
|
"t-charlie": [9, 9, 1, 2],
|
|
"t-delta": [44, 6, 30, 11],
|
|
}
|
|
n_new = 8
|
|
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
|
results: dict[str, list[int]] = {}
|
|
errors: list[Exception] = []
|
|
|
|
def run(sid: str, prompt: list[int]) -> None:
|
|
try:
|
|
results[sid] = _cached_greedy(adapter, manager, sid, 0, prompt, n_new)
|
|
except Exception as exc: # pragma: no cover - surfaced via assert below
|
|
errors.append(exc)
|
|
|
|
threads = [threading.Thread(target=run, args=(sid, p)) for sid, p in prompts.items()]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors
|
|
for sid in prompts:
|
|
assert results[sid] == references[sid], sid
|
|
|
|
|
|
def test_release_one_session_leaves_others_intact_and_returns_memory():
|
|
"Releasing one session frees its budget and does not disturb the others.\n\nTags: node, kv, concurrency"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
prompts = {"keep-1": [1, 2, 3], "drop": [10, 11, 12, 13], "keep-2": [5, 6, 7]}
|
|
n_new = 6
|
|
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
|
|
|
gen: dict[str, list[int]] = {}
|
|
for sid, prompt in prompts.items():
|
|
out = adapter.prefill(sid, 0, token_ids=np.asarray(prompt))
|
|
gen[sid] = [out.token_id]
|
|
|
|
bytes_before = manager.total_bytes
|
|
assert manager.release("drop", 0) is True
|
|
assert manager.total_bytes < bytes_before
|
|
|
|
# A decode on the released session is an explicit cache miss, not corruption.
|
|
miss = adapter.decode("drop", 0, token_ids=[gen["drop"][-1]])
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.RELEASED
|
|
|
|
# The survivors keep decoding to their own references.
|
|
for _ in range(n_new - 1):
|
|
for sid in ("keep-1", "keep-2"):
|
|
step = adapter.decode(sid, 0, token_ids=[gen[sid][-1]])
|
|
assert isinstance(step, TailOutput)
|
|
gen[sid].append(step.token_id)
|
|
for sid in ("keep-1", "keep-2"):
|
|
assert gen[sid] == references[sid], sid
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Stale epoch / incompatible recipe rejection.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_stale_route_epoch_is_rejected():
|
|
"A request for an older route epoch than the current one is rejected.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
manager = _manager_for(_full_shard(model))
|
|
manager.open("s", 5)
|
|
with pytest.raises(StaleRouteEpochError):
|
|
manager.open("s", 4)
|
|
with pytest.raises(StaleRouteEpochError):
|
|
manager.resolve("s", 4)
|
|
with pytest.raises(StaleRouteEpochError):
|
|
manager.append("s", 4, {})
|
|
|
|
|
|
def test_new_route_epoch_supersedes_and_frees_old_epoch():
|
|
"A newer route epoch supersedes the old one, freeing its KV and reporting a miss.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
adapter.prefill("s", 1, token_ids=np.asarray([1, 2, 3, 4]))
|
|
bytes_epoch1 = manager.total_bytes
|
|
assert bytes_epoch1 > 0
|
|
|
|
# Re-planned route: epoch 2 starts a fresh isolated context.
|
|
adapter.prefill("s", 2, token_ids=np.asarray([9, 8]))
|
|
assert manager.session_keys() == [("s", 2)]
|
|
# Old epoch is gone; a lookup for it is now stale (epoch < current).
|
|
with pytest.raises(StaleRouteEpochError):
|
|
manager.resolve("s", 1)
|
|
|
|
|
|
def test_incompatible_cache_recipe_is_rejected():
|
|
"A request carrying a different KV recipe is rejected, not silently reused.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
manager.open("s", 0)
|
|
|
|
incompatible = KvCacheRecipe(
|
|
architecture_adapter="dense-llama",
|
|
kv_dtype="float16", # different KV dtype
|
|
n_kv_heads=model.n_heads,
|
|
head_dim=model.head_dim,
|
|
total_layers=model.n_layers,
|
|
start_layer=0,
|
|
end_layer=model.n_layers - 1,
|
|
)
|
|
with pytest.raises(IncompatibleCacheRecipeError):
|
|
manager.resolve("s", 0, recipe=incompatible)
|
|
with pytest.raises(IncompatibleCacheRecipeError):
|
|
manager.open("s2", 0, recipe=incompatible)
|
|
|
|
|
|
def test_uncertified_architecture_recipe_fails_closed():
|
|
"A KV recipe for an uncertified architecture fails closed at construction.\n\nTags: node, kv"
|
|
from meshnet_node.boundary_adapter import UncertifiedArchitectureError
|
|
|
|
with pytest.raises(UncertifiedArchitectureError):
|
|
KvCacheRecipe(
|
|
architecture_adapter="qwen3-moe",
|
|
kv_dtype="float32",
|
|
n_kv_heads=4,
|
|
head_dim=8,
|
|
total_layers=6,
|
|
start_layer=0,
|
|
end_layer=5,
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Explicit cache-miss responses.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_unknown_session_is_an_explicit_cache_miss():
|
|
"Resolving an unknown session returns an explicit unknown-session miss.\n\nTags: node, kv"
|
|
manager = _manager_for(_full_shard(_KvDenseLlama()))
|
|
miss = manager.resolve("nope", 0)
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.UNKNOWN_SESSION
|
|
with pytest.raises(KvCacheMissError):
|
|
manager.get("nope", 0)
|
|
|
|
|
|
def test_seq_len_mismatch_is_an_explicit_cache_miss():
|
|
"A decode whose expected length disagrees with the cache is an explicit miss.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
out = adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3]))
|
|
# Cache holds 3 tokens; claim it holds 99.
|
|
miss = adapter.decode("s", 0, token_ids=[out.token_id], expected_seq_len=99)
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.SEQ_LEN_MISMATCH
|
|
|
|
|
|
def test_ttl_eviction_yields_an_explicit_cache_miss():
|
|
"A session idle past its TTL is evicted and reported as a TTL cache miss.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
clock = _FakeClock()
|
|
manager = _manager_for(shard, HotKvStateConfig(ttl_seconds=10.0), clock=clock)
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3]))
|
|
clock.advance(11.0)
|
|
miss = manager.resolve("s", 0)
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.EVICTED_TTL
|
|
assert manager.total_bytes == 0
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Eviction and budget.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_lru_eviction_by_session_cap_reports_a_miss():
|
|
"Exceeding the session cap evicts the least-recently-used session.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
manager = _manager_for(shard, HotKvStateConfig(max_sessions=2))
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
adapter.prefill("a", 0, token_ids=np.asarray([1, 2]))
|
|
adapter.prefill("b", 0, token_ids=np.asarray([3, 4]))
|
|
# Touch 'a' so 'b' becomes the LRU victim.
|
|
adapter.decode("a", 0, token_ids=[1])
|
|
adapter.prefill("c", 0, token_ids=np.asarray([5, 6]))
|
|
|
|
miss = manager.resolve("b", 0)
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.EVICTED_LRU
|
|
assert set(k[0] for k in manager.session_keys()) == {"a", "c"}
|
|
|
|
|
|
def test_budget_eviction_keeps_total_within_budget():
|
|
"Byte-budget pressure evicts LRU sessions so the store stays within budget.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
recipe = kv_recipe_for(shard)
|
|
# Budget for ~5 tokens of one session; a second big session forces eviction.
|
|
budget = recipe.bytes_per_token() * 5
|
|
manager = _manager_for(shard, HotKvStateConfig(budget_bytes=budget, max_sessions=8))
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
|
|
adapter.prefill("a", 0, token_ids=np.asarray([1, 2, 3]))
|
|
adapter.prefill("b", 0, token_ids=np.asarray([4, 5, 6, 7]))
|
|
assert manager.total_bytes <= budget
|
|
# 'a' (older, LRU) was evicted to make room for 'b'.
|
|
miss = manager.resolve("a", 0)
|
|
assert isinstance(miss, CacheMiss)
|
|
assert miss.reason is CacheMissReason.EVICTED_LRU
|
|
assert manager.get("b", 0).seq_len == 4
|
|
|
|
|
|
def test_single_session_exceeding_budget_raises():
|
|
"A single session that cannot fit the budget raises instead of evicting itself.\n\nTags: node, kv"
|
|
model = _KvDenseLlama()
|
|
shard = _full_shard(model)
|
|
recipe = kv_recipe_for(shard)
|
|
budget = recipe.bytes_per_token() * 2 # only 2 tokens fit
|
|
manager = _manager_for(shard, HotKvStateConfig(budget_bytes=budget))
|
|
adapter = KvBoundaryAdapter(shard, manager)
|
|
with pytest.raises(KvBudgetExceededError):
|
|
adapter.prefill("a", 0, token_ids=np.asarray([1, 2, 3, 4, 5]))
|