Spaces:
Sleeping
Sleeping
File size: 7,714 Bytes
d9c2197 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """Tests for TokenDanceStorage — Master-Mirror diff storage.
Covers:
- register_master + register_mirror happy path
- compression_ratio() ≥ 10x on typical 5-agent shared context
- reconstruct() recovers the original within tolerance
- collective_reuse_step() updates all mirrors in O(1) per agent
- diff threshold drops near-identical blocks
"""
from __future__ import annotations
import numpy as np
import pytest
from apohara_context_forge.storage.token_dance import (
SparseKVDiff,
TokenDanceStorage,
)
# -----------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------
def _make_master_kv(n_blocks: int = 64, hidden_dim: int = 128) -> np.ndarray:
"""Synthetic master KV cache: deterministic, FP32."""
rng = np.random.default_rng(42)
return rng.standard_normal((n_blocks, hidden_dim), dtype=np.float32)
def _make_near_master(master: np.ndarray, n_diff_blocks: int) -> np.ndarray:
"""Near-master KV: identical except for n_diff_blocks tail blocks."""
out = master.copy()
rng = np.random.default_rng(7)
if n_diff_blocks > 0:
idx = np.arange(out.shape[0] - n_diff_blocks, out.shape[0])
out[idx] = rng.standard_normal(out[idx].shape, dtype=np.float32)
return out
# -----------------------------------------------------------------------
# Tests
# -----------------------------------------------------------------------
class TestTokenDanceBasics:
def test_register_master_sets_state(self):
store = TokenDanceStorage()
master = _make_master_kv()
store.register_master("retriever", master)
assert store.master_id == "retriever"
assert store.master_cache["retriever"].shape == master.shape
def test_register_master_rejects_1d(self):
store = TokenDanceStorage()
with pytest.raises(ValueError, match="at least 2D"):
store.register_master("retriever", np.zeros(8))
def test_register_mirror_requires_master(self):
store = TokenDanceStorage()
with pytest.raises(RuntimeError, match="register_master"):
store.register_mirror("reranker", _make_master_kv())
def test_register_mirror_rejects_shape_mismatch(self):
store = TokenDanceStorage()
store.register_master("retriever", _make_master_kv(64, 128))
with pytest.raises(ValueError, match="must match master shape"):
store.register_mirror("reranker", _make_master_kv(64, 64))
class TestTokenDanceCompression:
def test_compression_ratio_5_agents_realistic(self):
"""5 agents sharing 97% of blocks: ~4-5x is the upper bound by construction.
With N agents the upper bound is N (zero-diff mirrors). 11-17x in the
TokenDance paper assumes a 11-17 agent committee — see the next test.
"""
store = TokenDanceStorage()
master = _make_master_kv(n_blocks=128, hidden_dim=256)
store.register_master("retriever", master)
for aid in ("reranker", "summarizer", "critic", "responder"):
store.register_mirror(aid, _make_near_master(master, n_diff_blocks=4))
ratio = store.compression_ratio()
# 5 * 128 = 640 full vs 128 + 4*4 = 144 stored → ~4.4x
assert ratio >= 4.0
assert ratio <= 5.0 # bounded above by N
def test_compression_ratio_paper_target(self):
"""11–17x compression target from arXiv:2604.03143 — needs 11+ agents."""
store = TokenDanceStorage(diff_threshold=1e-4)
master = _make_master_kv(n_blocks=200, hidden_dim=128)
store.register_master("retriever", master)
# 11 mirrors with zero diff → 12 agents × 200 / 200 = 12x.
for i in range(11):
store.register_mirror(f"agent_{i}", master.copy())
ratio = store.compression_ratio()
assert ratio >= 10.0
assert ratio <= 17.0 # paper upper bound
def test_diff_threshold_drops_negligible_blocks(self):
store = TokenDanceStorage(diff_threshold=1.0)
master = _make_master_kv(n_blocks=32, hidden_dim=16)
store.register_master("a", master)
# Tiny perturbations should be dropped.
rng = np.random.default_rng(1)
near = master + rng.standard_normal(master.shape, dtype=np.float32) * 1e-5
diff = store.register_mirror("b", near)
assert diff.n_diff_blocks == 0
assert diff.sparsity == pytest.approx(1.0)
class TestTokenDanceReconstruction:
def test_reconstruct_master_returns_master_copy(self):
store = TokenDanceStorage()
master = _make_master_kv()
store.register_master("retriever", master)
out = store.reconstruct("retriever")
np.testing.assert_array_equal(out, master)
# Mutating the output must not poison the stored master.
out[0] = 999
np.testing.assert_array_equal(store.master_cache["retriever"], master)
def test_reconstruct_mirror_within_tolerance(self):
store = TokenDanceStorage(diff_threshold=1e-4)
master = _make_master_kv(n_blocks=64, hidden_dim=64)
store.register_master("retriever", master)
original = _make_near_master(master, n_diff_blocks=8)
store.register_mirror("critic", original)
recovered = store.reconstruct("critic")
# Reconstruction is exact for blocks above threshold (we keep their full
# delta) and exactly master for blocks below threshold. Tolerance = the
# threshold scaled by sqrt(hidden_dim) at most.
np.testing.assert_allclose(recovered, original, atol=1e-4)
def test_reconstruct_unknown_agent_raises(self):
store = TokenDanceStorage()
store.register_master("a", _make_master_kv())
with pytest.raises(KeyError):
store.reconstruct("ghost")
class TestTokenDanceCollective:
def test_collective_reuse_step_one_pass(self):
store = TokenDanceStorage()
master = _make_master_kv(n_blocks=32, hidden_dim=64)
store.register_master("retriever", master)
for aid in ("reranker", "summarizer", "critic", "responder"):
store.register_mirror(aid, master.copy())
rng = np.random.default_rng(99)
new_blocks = rng.standard_normal((4, 64), dtype=np.float32)
diff_counts = store.collective_reuse_step(
["retriever", "reranker", "summarizer", "critic", "responder"],
new_blocks,
)
# All agents covered.
assert set(diff_counts.keys()) == {
"retriever",
"reranker",
"summarizer",
"critic",
"responder",
}
# Master grew by 4 blocks; mirrors still zero-diff.
assert store.master_cache["retriever"].shape == (36, 64)
for mirror_id in ("reranker", "summarizer", "critic", "responder"):
assert store.mirrors[mirror_id].total_blocks == 36
assert store.mirrors[mirror_id].n_diff_blocks == 0
def test_collective_reuse_step_requires_master(self):
store = TokenDanceStorage()
with pytest.raises(RuntimeError):
store.collective_reuse_step(["a"], np.zeros((1, 4)))
class TestTokenDanceStats:
def test_stats_tracks_cache(self):
store = TokenDanceStorage(diff_threshold=1e-4)
master = _make_master_kv(n_blocks=16, hidden_dim=8)
store.register_master("a", master)
store.register_mirror("b", master.copy())
s = store.stats()
assert s["master_id"] == "a"
assert s["master_blocks"] == 16
assert s["n_mirrors"] == 1
assert s["diff_blocks_total"] == 0
assert s["compression_ratio"] >= 2.0
|