tiered_kv_cache.py

from collections import OrderedDict


class TieredKVCache:
    """Two-tier KV block store: hot blocks in HBM (capacity-limited, LRU),
    cold blocks in CPU DRAM reached over PCIe. Models the decode access
    pattern and the cost of HBM<->CPU transfers."""

    def __init__(self, hbm_blocks, hbm_latency=1.0, pcie_latency=200.0):
        self.cap = hbm_blocks
        self.hbm = OrderedDict()            # block_id -> True; end = most recent
        self.hbm_latency = hbm_latency      # relative cost of an HBM read
        self.pcie_latency = pcie_latency    # relative cost of a CPU->HBM fetch
        self.stalls = 0                     # synchronous fetches on critical path
        self.time = 0.0

    def access(self, block_id, prefetch_next=None):
        if block_id in self.hbm:
            self.hbm.move_to_end(block_id)  # LRU touch
            self.time += self.hbm_latency
        else:                               # miss -> stall on a PCIe fetch
            self.stalls += 1
            self.time += self.pcie_latency
            self._admit(block_id)
        # A correctly predicted prefetch is issued on a copy stream and overlaps
        # compute, so it does NOT add to critical-path time.
        if prefetch_next is not None and prefetch_next not in self.hbm:
            self._admit(prefetch_next)

    def _admit(self, block_id):
        self.hbm[block_id] = True
        if len(self.hbm) > self.cap:
            self.hbm.popitem(last=False)    # evict LRU cold block back to CPU


def run(prefetch, n_blocks=200, hbm_cap=32, steps=400):
    cache = TieredKVCache(hbm_cap)
    accesses = []
    for s in range(steps):                  # working set drifts forward slowly
        base = (s // 2) % (n_blocks - hbm_cap)
        for b in range(base, base + hbm_cap // 2):
            accesses.append((b, b + 1 if prefetch else None))
    for block_id, nxt in accesses:
        cache.access(block_id, nxt)
    return cache.stalls, cache.time


if __name__ == "__main__":
    s0, t0 = run(prefetch=False)
    s1, t1 = run(prefetch=True)
    print(f"{'policy':14s} {'critical-path stalls':>22s} {'rel. time':>12s}")
    print(f"{'no prefetch':14s} {s0:>22d} {t0:12.0f}")
    print(f"{'prefetch':14s} {s1:>22d} {t1:12.0f}")
    print(f"prefetch cuts stalls {s0 / max(1, s1):.0f}x, time {t0 / t1:.1f}x")