prain: slot cache with eviction, scan, prefetch and CPU policies

prain/policy.py

@@ -0,0 +1,245 @@
+"""Per-layer VRAM slot cache: which slot holds which expert, what to evict, what to copy
+first, how scans and prefetched experts are admitted. Pure Python, no torch."""
+
+import json
+import math
+import os
+from bisect import bisect_right
+
+from prain.config import Policy
+
+
+class Lru:
+    """Evict the expert used longest ago."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.last = [0.0] * num_experts
+
+    def score(self, e, now):
+        return self.last[e]
+
+    def hit(self, e, now, heat):
+        self.last[e] = now
+
+    def admit(self, e, now, heat, low):
+        self.last[e] = -math.inf if low else now
+
+    def evicted(self, e):
+        pass
+
+
+class Lfu:
+    """Evict the expert with the lowest exponentially decayed heat. Scores survive eviction,
+    so an expert that comes back keeps its history."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.decay = 0.5 ** (1 / policy.lfu_half_life)
+        self.value = [0.0] * num_experts
+        self.when = [0] * num_experts
+
+    def score(self, e, now):
+        return self.value[e] * self.decay ** (now - self.when[e])
+
+    def hit(self, e, now, heat):
+        self.value[e] = self.score(e, now) + heat
+        self.when[e] = now
+
+    def admit(self, e, now, heat, low):
+        if not low:  # a low-priority admit keeps its history: never-seen experts score 0
+            self.hit(e, now, heat)
+
+    def evicted(self, e):
+        pass
+
+
+class Gdsf:
+    """GreedyDual-Size-Frequency: score = floor + frequency x refetch cost. The floor rises
+    to each victim's score, so experts that stopped being used age out. cost[e] is 1 for a
+    RAM expert and policy.disk_cost for a disk one, so slow-to-refetch experts stay longer."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.cost = cost
+        self.floor = 0.0
+        self.freq = [0.0] * num_experts
+        self.value = [0.0] * num_experts
+
+    def score(self, e, now):
+        return self.value[e]
+
+    def hit(self, e, now, heat):
+        self.freq[e] += heat
+        self.value[e] = self.floor + self.freq[e] * self.cost[e]
+
+    def admit(self, e, now, heat, low):
+        self.freq[e] = 0.0 if low else heat
+        self.value[e] = self.floor + self.freq[e] * self.cost[e]
+
+    def evicted(self, e):
+        self.floor = max(self.floor, self.value[e])
+
+
+class Belady:
+    """The offline optimum, for the simulator only: evict the expert used furthest in the
+    future. future[e] lists the forward numbers that route to e, ascending."""
+
+    def __init__(self, future):
+        self.future = future
+
+    def score(self, e, now):
+        uses = self.future.get(e, ())
+        i = bisect_right(uses, now)
+        return -uses[i] if i < len(uses) else -math.inf
+
+    def hit(self, e, now, heat):
+        pass
+
+    def admit(self, e, now, heat, low):
+        pass
+
+    def evicted(self, e):
+        pass
+
+
+POLICIES = {"lru": Lru, "lfu": Lfu, "gdsf": Gdsf}
+
+
+class SlotCache:
+    """Slots [0, len(pinned)) hold the pinned experts for the whole run. The cache_slots
+    slots after them hold experts from RAM or disk, managed by an eviction policy."""
+
+    def __init__(self, pinned, num_experts, cache_slots, policy=Policy(), cost=None, evictor=None):
+        self.num_experts = num_experts
+        self.cache_slots = cache_slots
+        self.policy = policy
+        self.pinned = {e: slot for slot, e in enumerate(pinned)}
+        self.slots = {}  # cached expert -> slot
+        first = len(self.pinned)
+        self.free = list(range(first + cache_slots - 1, first - 1, -1))  # pop() -> lowest slot
+        cost = cost or [1.0] * num_experts
+        self.evictor = evictor or POLICIES[policy.eviction](num_experts, policy, cost)
+        self.now = 0  # forwards planned so far
+        self.recent = set()  # experts routed to in the last forward
+        self.prefetched = set()  # prefetched experts that were not used yet
+        self.uses = [0] * num_experts  # forwards that routed to each expert
+        self.hits = 0  # routed expert already in VRAM
+        self.misses = 0  # routed expert copied in
+        self.cpu = 0  # routed expert computed on the CPU instead
+        self.prefetches = 0  # experts copied in ahead of time (prefetch or admit-behind)
+        self.prefetch_hits = 0  # ... that were used before being evicted
+        self.prefetch_wasted = 0  # ... that were evicted unused
+
+    def choose_cpu(self, counts, miss):
+        """Experts the CPU should compute: counts maps expert -> routed tokens."""
+        if miss.mode == "copy":
+            return set()
+        away = (e for e in counts if e not in self.pinned and e not in self.slots)
+        if miss.mode == "cpu":
+            return set(away)
+        return {e for e in away if counts[e] <= miss.cpu_max_tokens}
+
+    def plan(self, heat, tokens=1, cpu=()):
+        """Split one forward pass into rounds.
+
+        heat: {expert: heat > 0} for every expert routed to in this pass.
+        tokens: tokens in the pass; a pass with >= policy.scan_tokens tokens is a scan.
+        cpu: routed experts the CPU computes; they get no slot.
+        Returns [(slot_of, loads), ...]: slot_of maps each expert computed in that round to
+        its slot; loads lists the (expert, slot) copies to do before the round. Experts in
+        VRAM run in round 0. The others follow coldest first, so the hottest are copied last
+        and are the ones left in VRAM afterwards.
+        """
+        self.now += 1
+        now, scan = self.now, tokens >= self.policy.scan_tokens
+        norm = {e: h / tokens for e, h in heat.items()}  # per token, so scans don't dominate
+        for e in heat:
+            self.uses[e] += 1
+        self.recent = set(heat)
+        pinned = [e for e in heat if e in self.pinned]
+        resident = [e for e in heat if e in self.slots]
+        missing = sorted((e for e in heat if e not in self.pinned and e not in self.slots
+                          and e not in cpu), key=lambda e: (heat[e], e))
+        self.hits += len(pinned) + len(resident)
+        self.misses += len(missing)
+        self.cpu += len(cpu)
+        for e in resident:
+            self.evictor.hit(e, now, norm[e])
+            if e in self.prefetched:
+                self.prefetched.discard(e)
+                self.prefetch_hits += 1
+        if missing and not self.cache_slots:
+            raise RuntimeError(f"experts {missing} are not in VRAM and there are no cache slots")
+        protected = set()
+        if scan and self.policy.scan_keep:
+            ranked = sorted(self.slots, key=lambda e: (self.evictor.score(e, now), e))
+            protected = set(ranked[-self.policy.scan_keep:])
+        busy = set(resident)  # experts computed in the round being built
+        slot_of = {e: self.pinned[e] for e in pinned} | {e: self.slots[e] for e in resident}
+        rounds, i = [], 0
+        while True:
+            room = len(self.free) + sum(1 for e in self.slots if e not in busy and e not in protected)
+            group = missing[i:i + room]
+            if i < len(missing) and not group and not busy:
+                raise RuntimeError("no evictable cache slot: lower policy.scan_keep")
+            i += len(group)
+            loads = []
+            for e in group:
+                slot = self._take_slot(now, busy | protected)
+                self.slots[e] = slot
+                self.evictor.admit(e, now, norm[e], low=scan)
+                busy.add(e)
+                loads.append((e, slot))
+            slot_of.update(loads)
+            rounds.append((slot_of, loads))
+            if i >= len(missing):
+                break
+            slot_of, busy = {}, set()
+        if scan and self.policy.promote:
+            for e in missing[-self.policy.promote:]:
+                if e in self.slots:
+                    self.evictor.hit(e, now, norm[e])
+        return rounds
+
+    def prefetch(self, predicted, budget):
+        """Copy predicted experts in ahead of time, at low priority: they are evicted first
+        unless a forward uses them. Never evicts an expert routed to in the last forward.
+        predicted: [(expert, probability)], most likely first. Returns the (expert, slot) loads."""
+        loads, keep = [], set(self.recent)
+        for e, _ in predicted:
+            if len(loads) >= budget:
+                break
+            if e in self.pinned or e in self.slots:
+                continue
+            if not self.free and all(r in keep for r in self.slots):
+                break
+            slot = self._take_slot(self.now, keep)
+            self.slots[e] = slot
+            self.evictor.admit(e, self.now, 0.0, low=True)
+            self.prefetched.add(e)
+            keep.add(e)
+            loads.append((e, slot))
+        self.prefetches += len(loads)
+        return loads
+
+    def _take_slot(self, now, keep):
+        if self.free:
+            return self.free.pop()
+        victim = min((e for e in self.slots if e not in keep),
+                     key=lambda e: (self.evictor.score(e, now), e))
+        self.evictor.evicted(victim)
+        if victim in self.prefetched:
+            self.prefetched.discard(victim)
+            self.prefetch_wasted += 1
+        return self.slots.pop(victim)
+
+
+COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wasted")
+
+
+def write_stats(path, caches):
+    """Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
+    data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
+                       for layer, c in sorted(caches.items())}}
+    tmp = f"{path}.tmp"
+    with open(tmp, "w") as f:
+        json.dump(data, f)
+    os.replace(tmp, path)

tests/test_policy.py

@@ -0,0 +1,121 @@
+import json
+import math
+import os
+import random
+import unittest
+
+from prain.config import Miss, Policy
+from prain.policy import SlotCache, write_stats
+from tests.helpers import TmpCase
+
+
+def lru(**kw):
+    return Policy(eviction="lru", **kw)
+
+
+class Eviction(unittest.TestCase):
+    def test_pinned_experts_stay_and_lru_evicts_the_oldest(self):
+        cache = SlotCache(pinned=[3, 5], num_experts=8, cache_slots=2, policy=lru())
+        self.assertEqual(cache.plan({0: 1, 3: 1}), [({0: 2, 3: 0}, [(0, 2)])])
+        self.assertEqual(cache.plan({0: 1, 1: 1}), [({0: 2, 1: 3}, [(1, 3)])])
+        self.assertEqual(cache.plan({1: 1, 2: 1}), [({1: 3, 2: 2}, [(2, 2)])])  # 0 was oldest
+        self.assertEqual((cache.hits, cache.misses), (3, 3))
+        self.assertEqual(cache.uses, [2, 2, 1, 1, 0, 0, 0, 0])
+
+    def test_lfu_keeps_the_frequent_expert_lru_the_recent_one(self):
+        for policy, survivor in ((lru(), 1), (Policy(eviction="lfu"), 0)):
+            cache = SlotCache([], 3, 2, policy)
+            for e in (0, 0, 0, 1, 2):
+                cache.plan({e: 1})
+            self.assertIn(survivor, cache.slots, policy.eviction)
+
+    def test_gdsf_keeps_the_expensive_disk_expert(self):
+        cost = [1.0, 4.0, 1.0]  # expert 1 lives on disk
+        for eviction, survivor in (("lfu", 0), ("gdsf", 1)):
+            cache = SlotCache([], 3, 2, Policy(eviction=eviction), cost)
+            for e in (0, 0, 1, 2):
+                cache.plan({e: 1})
+            self.assertIn(survivor, cache.slots, eviction)
+
+
+class Rounds(unittest.TestCase):
+    def test_hottest_experts_are_copied_last_and_stay(self):
+        cache = SlotCache([], 10, 2, lru())
+        rounds = cache.plan({0: 1, 1: 9, 2: 3, 3: 8, 4: 2, 5: 1})
+        self.assertEqual([sorted(slot_of) for slot_of, _ in rounds], [[0, 5], [2, 4], [1, 3]])
+        self.assertEqual(set(cache.slots), {1, 3})
+
+    def test_scan_keeps_the_protected_expert_and_promotes_the_hottest(self):
+        cache = SlotCache([], 6, 3, lru(scan_tokens=4, scan_keep=1, promote=1))
+        cache.plan({0: 1})
+        cache.plan({1: 1})
+        rounds = cache.plan({2: 1, 3: 5, 4: 2, 5: 3}, tokens=8)
+        self.assertEqual(rounds, [({2: 2, 4: 0}, [(2, 2), (4, 0)]),
+                                  ({5: 2, 3: 0}, [(5, 2), (3, 0)])])
+        self.assertEqual(set(cache.slots), {1, 3, 5})  # 1 was protected
+        self.assertEqual(cache.evictor.last[3], cache.now)  # promoted
+        self.assertEqual(cache.evictor.last[5], -math.inf)  # admitted at low priority
+
+    def test_cpu_experts_get_no_slot(self):
+        cache = SlotCache([0], 6, 2, lru())
+        cache.plan({1: 1})
+        counts = {0: 1, 1: 1, 2: 1, 3: 4}
+        self.assertEqual(cache.choose_cpu(counts, Miss()), set())
+        self.assertEqual(cache.choose_cpu(counts, Miss(mode="cpu")), {2, 3})
+        cpu = cache.choose_cpu(counts, Miss(mode="auto", cpu_max_tokens=2))
+        self.assertEqual(cpu, {2})
+        self.assertEqual(cache.plan(counts, 4, cpu), [({0: 0, 1: 1, 3: 2}, [(3, 2)])])
+        self.assertEqual((cache.cpu, cache.misses), (1, 2))
+
+    def test_random_plans_always_compute_every_expert_from_the_right_slot(self):
+        rng = random.Random(1)
+        for _ in range(300):
+            n = rng.randint(2, 12)
+            pinned = rng.sample(range(n), rng.randint(0, n // 2))
+            slots = rng.randint(1, 4)
+            policy = Policy(eviction=rng.choice(["lru", "lfu", "gdsf"]),
+                            scan_tokens=rng.randint(1, 6), scan_keep=rng.randint(0, slots - 1),
+                            promote=rng.randint(0, 3))
+            cache = SlotCache(pinned, n, slots, policy, [rng.choice([1.0, 4.0]) for _ in range(n)])
+            content = {slot: e for e, slot in cache.pinned.items()}  # what each slot really holds
+            for _ in range(12):
+                if rng.random() < 0.3:
+                    guesses = [(e, 0.5) for e in rng.sample(range(n), rng.randint(1, n))]
+                    recent = set(cache.recent)
+                    for e, slot in cache.prefetch(guesses, rng.randint(0, 3)):
+                        self.assertNotIn(content.get(slot), recent - {None})
+                        content[slot] = e
+                routed = rng.sample(range(n), rng.randint(1, n))
+                heat = {e: rng.randint(1, 9) for e in routed}
+                cpu = cache.choose_cpu(heat, Miss(mode=rng.choice(["copy", "auto"]), cpu_max_tokens=3))
+                computed = []
+                for slot_of, loads in cache.plan(heat, rng.randint(1, 8), cpu):
+                    self.assertEqual(len(set(slot_of.values())), len(slot_of))
+                    for e, slot in loads:
+                        self.assertEqual(slot_of[e], slot)
+                        content[slot] = e
+                    for e, slot in slot_of.items():
+                        self.assertEqual(content[slot], e)
+                    computed += slot_of
+                self.assertEqual(sorted(computed + sorted(cpu)), sorted(routed))
+                held = list(cache.slots.values())
+                self.assertEqual(len(set(held)), len(held))
+                self.assertEqual(len(held) + len(cache.free), slots)
+
+
+class Prefetch(TmpCase):
+    def test_prefetch_is_low_priority_and_never_evicts_the_last_forward(self):
+        cache = SlotCache([], 4, 2, Policy(eviction="lfu"))
+        cache.plan({0: 1})
+        self.assertEqual(cache.prefetch([(1, 0.5), (2, 0.4)], budget=2), [(1, 1)])  # 0 is recent
+        cache.plan({1: 1})  # the prefetch paid off
+        self.assertEqual(cache.prefetch([(3, 0.9)], budget=1), [(3, 0)])  # evicts 0, not 1
+        cache.plan({2: 1})  # 3 was never used: it goes first
+        self.assertEqual((cache.prefetches, cache.prefetch_hits, cache.prefetch_wasted), (2, 1, 1))
+        self.assertEqual(set(cache.slots), {1, 2})
+        path = os.path.join(self.dir, "stats.json")
+        write_stats(path, {7: cache})
+        with open(path) as f:
+            stats = json.load(f)["layers"]["7"]
+        self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]),
+                         ([1, 1, 1, 0], 1, 2))