README: try-it-on-better-hardware guide + predictions + support section; drop stage-0 research scripts from repo

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-05 21:02:46 +02:00
parent 1ae22a6135
commit 20008742bf
11 changed files with 58 additions and 661 deletions
@@ -21,3 +21,14 @@ c/bench/
 *.safetensors
 stats*.txt
 *.log
 # script di ricerca stadio-0 (esperimenti locali, non parte del motore)
 /engine.py
 /run.py
 /validate_ref.py
 /s0_costmodel.py
 /s1_trace_hitrate.py
 /s2_research.py
 /quant_test.py
 /profile_run.py
 /sweep.py
@@ -60,6 +60,52 @@ COLI_MODEL=/path/to/glm52_i4 ./coli chat
 Useful knobs (env or flags): `--topp 0.7` adaptive expert top-p (30–40% less disk), `--ngen N` max tokens, `STATS=f`/`PIN=f PIN_GB=g` record expert usage then pin the hottest in RAM, `THINK=1` enable GLM-5.2's reasoning block, `DRAFT=n` MTP draft depth, `TF=1` teacher-forcing validation.
 ## Got a better machine? Try it — here's what to expect
 colibrì was built on deliberately humble hardware (12 cores, 25 GB RAM, NVMe behind a WSL2 VHDX that caps random reads at ~1 GB/s). **Every one of those constraints is a knob your machine can turn up.** The engine needs: Linux (or WSL2), gcc with OpenMP, AVX2, ≥16 GB RAM, and the ~370 GB int4 model on a local NVMe (ext4 — never a network/9p mount).
 **How to test it, in order:**
 ```bash
 cd c && ./setup.sh                 # build + architecture self-test (expects 32/32)
 # 1) measure YOUR disk the way the engine uses it (parallel 19 MB random reads):
 gcc -O2 -fopenmp iobench.c -o iobench
 ./iobench /path/to/glm52_i4/out-00069.safetensors 19 64 8 0   # buffered, 8 threads
 ./iobench /path/to/glm52_i4/out-00069.safetensors 19 64 8 1   # O_DIRECT
 # 2) chat; watch the per-turn stats line (tok/s, expert hit-rate, RSS):
 COLI_MODEL=/path/to/glm52_i4 ./coli chat
 # 3) record expert usage, then pin the hottest experts in your spare RAM:
 STATS=stats.txt ./coli chat
 PIN=stats.txt PIN_GB=20 ./coli chat        # scale PIN_GB to your free RAM
 # 4) quality benchmarks (MMLU/HellaSwag/ARC):
 ./coli bench
 ```
 **Back-of-envelope predictions** (decode is disk-bound: a cold token costs ~11.4 GB of expert reads; MTP speculation roughly halves the effective cost; RAM turns cold reads into free cache hits):
 | machine | expected |
 |---|---|
 | this dev box (WSL2 VHDX, ~1 GB/s, 25 GB RAM) | ~0.05–0.1 tok/s cold — proven baseline |
 | native Linux, PCIe4 NVMe (~3–5 GB/s random), 32 GB | ~0.5–1 tok/s |
 | PCIe5 NVMe or 2×NVMe RAID0 (~8–12 GB/s), 64 GB (PIN ~40 GB of hot experts) | ~2–4 tok/s |
 | 128–256 GB RAM workstation (hot expert set mostly cached) | ~5–15 tok/s, matmul-bound — interactive |
 These are estimates, not measurements — if you run colibrì on serious hardware, **please open an issue with your numbers**: real datapoints from better machines are exactly what this project needs next.
 ## Supporting the project
 colibrì is a one-person project, written and tested entirely on a 12-core laptop with 25 GB of RAM — the numbers above are the ceiling of what I can measure at home. If this project is useful or interesting to you and you'd like to support its development (better test hardware translates *directly* into a faster engine for everyone: real NVMe scaling data, bigger pinned caches, int2/int3 quality sweeps on real benchmarks), you can:
 - ⭐ star the repo and share it;
 - 🐛 open issues with benchmark numbers from your hardware;
 - 💬 reach out via GitHub issues if you'd like to sponsor development or donate hardware.
 Every contribution, from a datapoint to a disk, moves the ceiling.
 ## Repo layout
 ```
@@ -67,10 +113,10 @@ c/glm.c          the engine (GLM-5.2 forward, streaming MoE, MTP, serve mode)
 c/st.h           safetensors reader: pread + fadvise, no mmap (RSS stays flat)
 c/tok.h          byte-level BPE tokenizer in C
 c/coli           CLI: chat / run / bench / convert / info
 c/iobench.c      parallel disk microbenchmark (measures what the engine feels)
 c/convert_fp8_to_int4.py   disk-safe FP8 → int4 converter
 c/make_glm_oracle.py       tiny-random oracle generator for validation
 c/olmoe.c        stage-A engine (OLMoE), first validation target
 *.py             research scripts (cost model, trace analysis, py engine)
 ```
 ## Why "colibrì"
@@ -1,219 +0,0 @@
 """
 Motore di inferenza MoE con EXPERT-STREAMING dal disco.
 Idea (quella dell'utente, resa reale):
  - la parte DENSA (embedding, attenzione, router, norme, lm_head) sta in RAM;
  - gli EXPERT stanno su disco in un file safetensors mappato in memoria (mmap)
    e vengono letti SOLO quando un token li attiva;
  - una cache LRU tiene in RAM gli expert "caldi" -> meno letture da disco.
 Cosi' un modello che NON entra in RAM gira lo stesso: in RAM ci tieni solo
 densa + cache, il resto vive sul disco. Validato qui su OLMoE-1B-7B.
 NB: scritto per OLMoE (Llama-style con QK-norm). I punti specifici del modello
 (routing, norme) sono isolati cosi' che lo stesso scheletro valga per GLM/DeepSeek.
 """
 import os, json, glob, struct, time, mmap, collections
 import torch
 import torch.nn.functional as F
 ST_DTYPE = {"BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32}
 class Shards:
    """Indicizza i tensori in piu' file safetensors e li legge via mmap on-demand."""
    def __init__(self, snap_dir):
        self.index = {}          # name -> (shard_path, abs_offset, nbytes, torch_dtype, shape)
        self.mm = {}             # shard_path -> mmap
        for shard in sorted(glob.glob(os.path.join(snap_dir, "model-*.safetensors"))):
            with open(shard, "rb") as f:
                hlen = struct.unpack("<Q", f.read(8))[0]
                header = json.loads(f.read(hlen))
            data_start = 8 + hlen
            for name, meta in header.items():
                if name == "__metadata__":
                    continue
                a, b = meta["data_offsets"]
                self.index[name] = (shard, data_start + a, b - a,
                                    ST_DTYPE[meta["dtype"]], tuple(meta["shape"]))
            fd = open(shard, "rb")
            self.mm[shard] = mmap.mmap(fd.fileno(), 0, prot=mmap.PROT_READ)
    def read(self, name):
        """Legge un tensore dal disco (mmap) e ne fa una copia RESIDENTE in RAM."""
        shard, off, nbytes, dt, shape = self.index[name]
        mv = memoryview(self.mm[shard])[off:off + nbytes]
        return torch.frombuffer(mv, dtype=dt).reshape(shape).clone()
    def has(self, name):
        return name in self.index
 def quant_dequant(w, bits):
    """Quantizzazione simmetrica per-riga a `bits` bit, poi dequantizza in bf16.
    Simula numericamente cosa darebbe un expert salvato a `bits` bit sul disco."""
    if bits >= 16:
        return w
    qmax = (1 << (bits - 1)) - 1            # int8->127, int4->7, int3->3, int2->1
    wf = w.float()
    scale = wf.abs().amax(dim=1, keepdim=True) / qmax
    scale = scale.clamp_min(1e-8)
    wq = torch.round(wf / scale).clamp(-qmax - 1, qmax)
    return (wq * scale).to(torch.bfloat16)
 class ExpertCache:
    """Cache LRU degli expert. capacity = quanti expert teniamo residenti PER LAYER."""
    def __init__(self, shards, n_layers, capacity, quant_bits=16):
        self.shards = shards
        self.cap = capacity
        self.quant_bits = quant_bits
        self.caches = [collections.OrderedDict() for _ in range(n_layers)]
        self.hits = 0
        self.miss = 0
    def get(self, layer, eid):
        """Ritorna (gate_w, up_w, down_w) dell'expert, da cache o da disco."""
        c = self.caches[layer]
        if eid in c:
            self.hits += 1
            c.move_to_end(eid)
            return c[eid]
        self.miss += 1
        p = f"model.layers.{layer}.mlp.experts.{eid}."
        # tengo gli expert in bf16 (niente .float(): -24% tempo, -50% RAM, piu' fedele al riferimento)
        b = self.quant_bits
        w = (quant_dequant(self.shards.read(p + "gate_proj.weight"), b),
             quant_dequant(self.shards.read(p + "up_proj.weight"), b),
             quant_dequant(self.shards.read(p + "down_proj.weight"), b))
        c[eid] = w
        if len(c) > self.cap:
            c.popitem(last=False)
        return w
    def hitrate(self):
        t = self.hits + self.miss
        return self.hits / t if t else 0.0
 def rmsnorm(x, w, eps=1e-5):
    x = x.float()
    x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    return x * w.float()
 def rotate_half(x):
    h = x.shape[-1] // 2
    return torch.cat((-x[..., h:], x[..., :h]), dim=-1)
 class OlmoeStreaming:
    def __init__(self, snap_dir, expert_cap=16, quant_bits=16):
        self.cfg = json.load(open(os.path.join(snap_dir, "config.json")))
        self.shards = Shards(snap_dir)
        c = self.cfg
        self.L = c["num_hidden_layers"]
        self.H = c["num_attention_heads"]
        self.hd = c["hidden_size"] // self.H
        self.topk = c["num_experts_per_tok"]
        self.eps = c.get("rms_norm_eps", 1e-5)
        self.norm_topk = c.get("norm_topk_prob", False)
        theta = c.get("rope_theta", 10000.0)
        self.inv_freq = 1.0 / (theta ** (torch.arange(0, self.hd, 2).float() / self.hd))
        self.cache = ExpertCache(self.shards, self.L, expert_cap, quant_bits)
        # --- parte DENSA: residente in RAM (float32) ---
        t = time.time()
        self.embed = self.shards.read("model.embed_tokens.weight").float()
        self.lm_head = self.shards.read("lm_head.weight").float()
        self.final_norm = self.shards.read("model.norm.weight").float()
        self.layers = []
        for i in range(self.L):
            p = f"model.layers.{i}."
            self.layers.append({
                "in_ln":  self.shards.read(p + "input_layernorm.weight").float(),
                "post_ln":self.shards.read(p + "post_attention_layernorm.weight").float(),
                "q":  self.shards.read(p + "self_attn.q_proj.weight").float(),
                "k":  self.shards.read(p + "self_attn.k_proj.weight").float(),
                "v":  self.shards.read(p + "self_attn.v_proj.weight").float(),
                "o":  self.shards.read(p + "self_attn.o_proj.weight").float(),
                "qn": self.shards.read(p + "self_attn.q_norm.weight").float(),
                "kn": self.shards.read(p + "self_attn.k_norm.weight").float(),
                "gate": self.shards.read(p + "mlp.gate.weight").float(),
            })
        self.dense_load_s = time.time() - t
    def _rope(self, x, pos):
        # x: (heads, seq, hd) ; pos: (seq,)
        freqs = torch.outer(pos.float(), self.inv_freq)          # (seq, hd/2)
        emb = torch.cat((freqs, freqs), dim=-1)                  # (seq, hd)
        cos, sin = emb.cos(), emb.sin()
        return x * cos + rotate_half(x) * sin
    def _attn(self, lw, x, pos, layer, kv):
        """Attenzione con KV-cache. x = SOLO i token nuovi (S in prefill, 1 in decode).
        pos = posizioni assolute dei token nuovi. kv = lista per-layer dei (k,v) passati."""
        S = x.shape[0]
        q = rmsnorm(x @ lw["q"].T, lw["qn"], self.eps).view(S, self.H, self.hd).transpose(0, 1)
        k = rmsnorm(x @ lw["k"].T, lw["kn"], self.eps).view(S, self.H, self.hd).transpose(0, 1)
        v = (x @ lw["v"].T).view(S, self.H, self.hd).transpose(0, 1)
        q = self._rope(q, pos); k = self._rope(k, pos)
        if kv is not None and kv[layer] is not None:             # concateno il passato
            pk, pv = kv[layer]
            k = torch.cat([pk, k], dim=1); v = torch.cat([pv, v], dim=1)
        if kv is not None:
            kv[layer] = (k, v)
        Tk = k.shape[1]                                          # lunghezza totale (passato+nuovi)
        scores = (q @ k.transpose(-1, -2)) / (self.hd ** 0.5)    # (H,S,Tk)
        # mask causale: query a posizione assoluta pos[i] vede key j<=pos[i]
        kpos = torch.arange(Tk)
        mask = torch.where(kpos[None, :] > pos[:, None], float("-inf"), 0.0)  # -inf dove vietato
        a = F.softmax(scores + mask, dim=-1)
        out = (a @ v).transpose(0, 1).reshape(S, self.H * self.hd)
        return out @ lw["o"].T
    def _moe(self, layer, lw, x):
        S = x.shape[0]
        logits = x @ lw["gate"].T                                # (S,64)
        probs = F.softmax(logits.float(), dim=-1)
        w, idx = torch.topk(probs, self.topk, dim=-1)            # (S,topk)
        if self.norm_topk:
            w = w / w.sum(-1, keepdim=True)
        out = torch.zeros_like(x)
        # raggruppo per expert: per ogni expert davvero usato, processo i suoi token
        for eid in torch.unique(idx).tolist():
            sel = (idx == eid)                                   # (S,topk) bool
            rows = sel.any(dim=-1).nonzero(as_tuple=True)[0]
            if rows.numel() == 0:
                continue
            gw, uw, dw = self.cache.get(layer, eid)              # <-- streaming dal disco (bf16)
            xe = x[rows].to(torch.bfloat16)                      # calcolo expert in bf16
            h = (F.silu(xe @ gw.T) * (xe @ uw.T)) @ dw.T
            weight = (w[rows] * sel[rows].float()).sum(-1, keepdim=True)
            out[rows] += weight * h.float()
        return out
    @torch.no_grad()
    def _step(self, ids_new, pos, kv):
        """Un passo del modello sui token nuovi. Ritorna logit dell'ultimo token."""
        x = self.embed[torch.tensor(ids_new)]                   # (S,hidden)
        for i, lw in enumerate(self.layers):
            x = x + self._attn(lw, rmsnorm(x, lw["in_ln"], self.eps), pos, i, kv)
            x = x + self._moe(i, lw, rmsnorm(x, lw["post_ln"], self.eps))
        x = rmsnorm(x, self.final_norm, self.eps)
        return (x @ self.lm_head.T)[-1]
    @torch.no_grad()
    def generate(self, token_ids, n_new, greedy=True):
        kv = [None] * self.L
        ids = list(token_ids)
        # PREFILL: tutti i token del prompt in un colpo, riempie la kv-cache
        logit = self._step(ids, torch.arange(len(ids)), kv)
        for s in range(n_new):
            nxt = int(torch.argmax(logit)) if greedy else int(torch.multinomial(F.softmax(logit, -1), 1))
            ids.append(nxt)
            if s == n_new - 1:
                break
            # DECODE: un solo token nuovo, usa la kv-cache (qui la cache expert torna a funzionare)
            logit = self._step([nxt], torch.tensor([len(ids) - 1]), kv)
        return ids
@@ -1,18 +0,0 @@
 """Profila dove va il tempo: lettura expert dal disco vs attenzione vs moe vs matmul."""
 import cProfile, pstats, io, glob, json
 from engine import OlmoeStreaming
 snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
 ref = json.load(open("ref.json"))
 m = OlmoeStreaming(snap, expert_cap=16)
 pr = cProfile.Profile()
 pr.enable()
 m.generate(ref["prompt_ids"], 8, greedy=True)   # 8 token bastano per il profilo
 pr.disable()
 s = io.StringIO()
 ps = pstats.Stats(pr, stream=s).sort_stats("tottime")
 ps.print_stats(15)
 print(s.getvalue())
 print(f"Hit-rate: {m.cache.hitrate()*100:.1f}%  hit={m.cache.hits} miss={m.cache.miss}")
@@ -1,22 +0,0 @@
 """La domanda che conta: a quanti bit l'output degli expert REGGE ancora?
 Quantizzo solo gli expert (la parte densa resta bf16) e confronto col riferimento."""
 import json, glob
 from engine import OlmoeStreaming
 snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
 ref = json.load(open("ref.json"))
 exp = ref["full_ids"][len(ref["prompt_ids"]):]
 n_new = len(exp)
 print(f"{'bit':>4} {'MB/expert':>10} {'match':>7}  testo")
 for bits in (16, 8, 4, 3, 2):
    m = OlmoeStreaming(snap, expert_cap=64, quant_bits=bits)  # cap64: isola l'effetto quant
    out = m.generate(ref["prompt_ids"], n_new, greedy=True)
    gen = out[len(ref["prompt_ids"]):]
    match = sum(a == b for a, b in zip(gen, exp))
    mb = 6.29 * bits / 8 / 1.0  # ~6.29M param/expert * bit / 8 -> MB
    # decode testo per vedere se e' ancora sensato
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
    txt = tok.decode(gen)
    print(f"{bits:>4} {mb:>9.1f}M {match:>4}/{n_new:<2}  {txt!r}")
@@ -1,31 +0,0 @@
 """Lancia il motore streaming, confronta con il riferimento, misura RAM/hit-rate/velocita'."""
 import json, time, glob, sys, resource
 from engine import OlmoeStreaming
 CAP = int(sys.argv[1]) if len(sys.argv) > 1 else 16   # expert residenti per layer (su 64)
 snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
 ref = json.load(open("run_ref.json" if False else "ref.json"))
 def rss_gb():
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024**2)  # KB->GB su Linux
 print(f"== Motore streaming, cache = {CAP} expert/layer su 64 ==")
 t = time.time()
 m = OlmoeStreaming(snap, expert_cap=CAP)
 print(f"densa caricata in {m.dense_load_s:.1f}s | RSS dopo load densa: {rss_gb():.2f} GB")
 n_new = len(ref["full_ids"]) - len(ref["prompt_ids"])
 t = time.time()
 out = m.generate(ref["prompt_ids"], n_new, greedy=True)
 dt = time.time() - t
 # confronto
 gen = out[len(ref["prompt_ids"]):]
 exp = ref["full_ids"][len(ref["prompt_ids"]):]
 match = sum(a == b for a, b in zip(gen, exp))
 print(f"\nRiferimento (transformers): {exp}")
 print(f"Motore streaming          : {gen}")
 print(f"Token coincidenti: {match}/{len(exp)}")
 print(f"\nRSS PICCO: {rss_gb():.2f} GB  (modello completo in bf16 = ~13 GB)")
 print(f"Hit-rate cache expert: {m.cache.hitrate()*100:.1f}%  (hit={m.cache.hits} miss={m.cache.miss})")
 print(f"Velocita': {n_new/dt:.2f} tok/s ({dt:.1f}s per {n_new} token, no kv-cache)")
@@ -1,97 +0,0 @@
 """
 Stadio 0 - Cost model + benchmark del disco per lo streaming degli expert MoE.
 Domanda: e' FISICAMENTE possibile fare streaming degli expert da disco
 e generare a velocita' usabile, su QUESTA macchina?
 Due numeri che servono:
  1. Banda effettiva del disco in lettura RANDOM, a blocchi grossi quanto un expert.
  2. Quanti byte/token dobbiamo leggere -> da cui il tetto di token/sec.
 Nessun modello richiesto. Gira in secondi.
 """
 import os, sys, time, mmap, random, argparse
 MB = 1024 * 1024
 GB = 1024 * MB
 def bench_disk(path_dir, expert_mb=12.0, total_mb=2048, n_reads=200):
    """Crea un file, poi misura lettura sequenziale e random a chunk = un expert."""
    os.makedirs(path_dir, exist_ok=True)
    fpath = os.path.join(path_dir, "_bench.bin")
    chunk = int(expert_mb * MB)
    total = int(total_mb * MB)
    total = (total // chunk) * chunk
    n_chunks = total // chunk
    # scrittura
    t = time.time()
    with open(fpath, "wb") as f:
        buf = os.urandom(chunk)
        for _ in range(n_chunks):
            f.write(buf)
        f.flush(); os.fsync(f.fileno())
    write_bw = total / (time.time() - t) / GB
    # prova a buttare via la page cache (best effort, serve permessi su Linux nativo)
    try:
        os.system("sync")
        with open("/proc/sys/vm/drop_caches", "w") as c:
            c.write("3")
    except Exception:
        pass  # su /mnt/c o senza root non si puo': il numero sara' ottimistico
    f = open(fpath, "rb")
    mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
    # random reads a chunk di un expert
    idxs = [random.randrange(n_chunks) for _ in range(n_reads)]
    s = 0
    t = time.time()
    for i in idxs:
        off = i * chunk
        s += mm[off]          # tocca la prima pagina
        s += mm[off + chunk - 1]  # e l'ultima -> forza il caricamento del range
        _ = bytes(mm[off:off + chunk])  # legge davvero l'intero expert
    rand_bw = (n_reads * chunk) / (time.time() - t) / GB
    mm.close(); f.close()
    os.remove(fpath)
    return write_bw, rand_bw
 def cost_model(name, n_layers, n_active, expert_mb, disk_bw_gbs, ram_resident_gb):
    """Stampa il tetto di token/sec in funzione dell'hit-rate della cache."""
    bytes_cold = n_layers * n_active * expert_mb / 1024  # GB letti per token se 0 cache
    print(f"\n--- {name} ---")
    print(f"  layer={n_layers}  expert_attivi/layer={n_active}  expert={expert_mb:.1f} MB")
    print(f"  parte densa residente in RAM stimata: ~{ram_resident_gb:.1f} GB")
    print(f"  byte da streammare per token (cache fredda): {bytes_cold*1024:.0f} MB")
    print(f"  tetto token/sec @ banda {disk_bw_gbs:.2f} GB/s, al variare dell'hit-rate cache:")
    for hit in (0.0, 0.5, 0.8, 0.9, 0.95, 0.99):
        eff = bytes_cold * (1 - hit)
        tps = disk_bw_gbs / eff if eff > 0 else float("inf")
        print(f"      hit {hit*100:5.1f}%  ->  {tps:6.2f} tok/s")
 if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default=".", help="cartella su cui benchmarkare il disco")
    ap.add_argument("--expert-mb", type=float, default=12.0)
    args = ap.parse_args()
    print(f"Benchmark disco su: {os.path.abspath(args.dir)}  (chunk={args.expert_mb} MB)")
    wbw, rbw = bench_disk(args.dir, expert_mb=args.expert_mb)
    print(f"  scrittura seq : {wbw:.2f} GB/s")
    print(f"  lettura random: {rbw:.2f} GB/s   <-- numero che conta per lo streaming")
    # Scenari. expert_mb a Q4 ~ (hidden*inter*3)*0.5B.
    # OLMoE 1B-7B: 16 layer, 8 attivi, hidden 2048 inter 1024 -> ~3 MB Q4
    cost_model("OLMoE 1B-7B (piccolo, lo useremo allo Stadio 1)",
               n_layers=16, n_active=8, expert_mb=3.0,
               disk_bw_gbs=rbw, ram_resident_gb=1.0)
    # DeepSeek-V3/V4 class: ~60 layer MoE, 8 attivi, expert ~6 MB Q2
    cost_model("DeepSeek/GLM class @ Q2 (il sogno finale)",
               n_layers=60, n_active=8, expert_mb=6.0,
               disk_bw_gbs=rbw, ram_resident_gb=10.0)
@@ -1,126 +0,0 @@
 """
 Pilastro 2 - La domanda che decide tutto:
 quanto e' alto l'hit-rate di una cache di expert su un router MoE VERO?
 Facciamo girare OLMoE-1B-7B su testo reale, registriamo per ogni (layer, posizione)
 quali 8 expert su 64 vengono attivati, e poi SIMULIAMO diverse politiche di cache
 al variare della capacita' K (quanti expert/layer teniamo residenti in RAM).
 Output: hit-rate per policy e per K  ->  mappato su token/sec col cost model.
 """
 import sys, time, collections
 import torch
 MODEL = "allenai/OLMoE-1B-7B-0924"
 TOPK = 8
 N_EXPERTS = 64
 N_LAYERS = 16
 EXPERT_MB = 12.6   # bf16
 # Banda misurata allo Stadio 0 (random read, ext4). Aggiorna se rifai il bench.
 DISK_BW_GBS = 7.33
 PROMPTS = [
    "The history of the Roman Empire spans over a thousand years, from the founding "
    "of the city to the fall of Constantinople. Its legacy in law, language and "
    "engineering still shapes the modern world.",
    "def quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr)//2]\n"
    "    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n"
    "    right = [x for x in arr if x > pivot]\n    return quicksort(left) + middle + quicksort(right)",
    "In quantum mechanics, the wave function encodes the probability amplitude of a "
    "particle's state. Measurement collapses this superposition into a definite outcome, "
    "a process that remains philosophically contested.",
    "La politica monetaria della banca centrale influenza i tassi di interesse, "
    "l'inflazione e l'occupazione. Alzare i tassi raffredda la domanda ma rischia "
    "di rallentare la crescita economica.",
 ]
 def collect_trace():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print("Carico il modello (bf16, CPU)...", flush=True)
    t = time.time()
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
    model.eval()
    print(f"  caricato in {time.time()-t:.0f}s", flush=True)
    # trace[layer] = lista (in ordine di token) di tuple di 8 expert id
    trace = [[] for _ in range(N_LAYERS)]
    for pi, p in enumerate(PROMPTS):
        ids = tok(p, return_tensors="pt").input_ids
        with torch.no_grad():
            out = model(ids, output_router_logits=True, use_cache=False)
        # out.router_logits: tupla di N_LAYERS tensori (n_token, n_experts)
        for li, rl in enumerate(out.router_logits):
            topk = rl.topk(TOPK, dim=-1).indices  # (n_token, 8)
            for row in topk.tolist():
                trace[li].append(tuple(sorted(row)))
        print(f"  prompt {pi}: {ids.shape[1]} token", flush=True)
    return trace
 def simulate(trace, K, policy="lru"):
    """Cache per-layer di capacita' K. Ritorna hit-rate globale sugli accessi a expert."""
    hits = total = 0
    for li in range(N_LAYERS):
        if policy == "lru":
            cache = collections.OrderedDict()
        elif policy == "lfu":
            cache = {}            # eid -> freq
            freq = collections.Counter()
        for experts in trace[li]:
            for e in experts:
                total += 1
                if policy == "lru":
                    if e in cache:
                        hits += 1
                        cache.move_to_end(e)
                    else:
                        cache[e] = 1
                        if len(cache) > K:
                            cache.popitem(last=False)
                elif policy == "lfu":
                    freq[e] += 1
                    if e in cache:
                        hits += 1
                    else:
                        if len(cache) >= K:
                            victim = min(cache, key=lambda x: freq[x])
                            del cache[victim]
                        cache[e] = 1
    return hits / total if total else 0.0
 def consecutive_reuse(trace):
    """Frazione di expert al token t gia' attivi al token t-1 (stesso layer)."""
    same = tot = 0
    for li in range(N_LAYERS):
        seq = trace[li]
        for t in range(1, len(seq)):
            prev = set(seq[t-1]); cur = set(seq[t])
            same += len(prev & cur); tot += len(cur)
    return same / tot if tot else 0.0
 def tok_per_sec(hitrate):
    bytes_cold_gb = N_LAYERS * TOPK * EXPERT_MB / 1024
    eff = bytes_cold_gb * (1 - hitrate)
    return DISK_BW_GBS / eff if eff > 0 else float("inf")
 if __name__ == "__main__":
    trace = collect_trace()
    ntok = sum(len(trace[0]) for _ in [0])
    print(f"\nToken totali tracciati: {len(trace[0])}  x {N_LAYERS} layer")
    print(f"Riuso consecutivo (expert in comune t vs t-1): {consecutive_reuse(trace)*100:.1f}%")
    print("\nHit-rate cache per-layer al variare di K (expert residenti su 64):")
    print(f"{'K':>4} {'RAM/GB':>7} {'LRU':>8} {'LFU':>8} {'tok/s@LRU':>10}")
    for K in (8, 12, 16, 24, 32, 48):
        ram = K * N_LAYERS * EXPERT_MB / 1024
        hl = simulate(trace, K, "lru")
        hf = simulate(trace, K, "lfu")
        print(f"{K:>4} {ram:>6.1f}G {hl*100:>7.1f}% {hf*100:>7.1f}% {tok_per_sec(hl):>9.1f}")
    print("\nNota: K=8 e' il minimo teorico (8 attivi/token). K=64 = tutto in RAM (no streaming).")
@@ -1,110 +0,0 @@
 """
 RICERCA - Il cardine del metodo: lo SKEW degli expert e' sfruttabile?
 Se pochi expert "caldi" coprono gran parte delle attivazioni, allora la strategia
 giusta per un modello che NON entra in RAM e':
   - PIN dei caldi (residenti per sempre in RAM, profilati offline)
   - STREAM dei freddi dal disco
 invece di una LRU dinamica (che su RAM piccola va in pressione, l'abbiamo visto).
 Test onesto: determino il "set caldo" dalla PRIMA meta' dei token, e misuro la
 copertura sulla SECONDA meta' (mai vista). Confronto PIN-caldi statico vs LRU a parita' di K.
 """
 import json, glob, collections, time
 import torch
 MODEL = "allenai/OLMoE-1B-7B-0924"
 N_EXP, TOPK, N_LAYERS = 64, 8, 16
 # testo piu' lungo e vario per statistiche decenti
 PROMPTS = [
    "The Roman Empire was one of the largest empires in history. At its height under "
    "Trajan, it covered five million square kilometres and held seventy million people, "
    "about a fifth of the world's population at the time. The empire's longevity and vast "
    "extent ensured a lasting influence on language, religion, architecture, philosophy, law "
    "and forms of government across the territory it once governed. ",
    "Photosynthesis is a biological process used by plants, algae and some bacteria to "
    "convert light energy into chemical energy stored in glucose. It occurs in the chloroplasts, "
    "specifically using the green pigment chlorophyll. The process consumes carbon dioxide and "
    "water and releases oxygen as a by-product, sustaining most life on Earth. ",
    "def fibonacci(n):\n    a, b = 0, 1\n    result = []\n    for _ in range(n):\n        "
    "result.append(a)\n        a, b = b, a + b\n    return result\n\n"
    "class Stack:\n    def __init__(self):\n        self.items = []\n    def push(self, x):\n"
    "        self.items.append(x)\n    def pop(self):\n        return self.items.pop()\n",
    "L'economia mondiale nel ventunesimo secolo e' caratterizzata da una crescente "
    "globalizzazione, dall'integrazione dei mercati finanziari e dalla rapida diffusione "
    "delle tecnologie digitali. Le banche centrali giocano un ruolo cruciale nel mantenere "
    "la stabilita' dei prezzi attraverso la politica monetaria. ",
 ]
 def collect():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print("carico modello...", flush=True)
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
                                                 low_cpu_mem_usage=True).eval()
    trace = [[] for _ in range(N_LAYERS)]
    for p in PROMPTS:
        ids = tok(p, return_tensors="pt").input_ids
        with torch.no_grad():
            out = model(ids, output_router_logits=True, use_cache=False)
        for li, rl in enumerate(out.router_logits):
            for row in rl.topk(TOPK, -1).indices.tolist():
                trace[li].append(tuple(row))
        print(f"  +{ids.shape[1]} token", flush=True)
    return trace
 def lru_hit(seq, K):
    c = collections.OrderedDict(); hit = tot = 0
    for experts in seq:
        for e in experts:
            tot += 1
            if e in c: hit += 1; c.move_to_end(e)
            else:
                c[e] = 1
                if len(c) > K: c.popitem(last=False)
    return hit / tot
 def static_hot_hit(train, test, K):
    """Set caldo = K piu' frequenti nel train; copertura misurata sul test."""
    freq = collections.Counter(e for experts in train for e in experts)
    hot = set(e for e, _ in freq.most_common(K))
    hit = tot = 0
    for experts in test:
        for e in experts:
            tot += 1
            if e in hot: hit += 1
    return hit / tot
 if __name__ == "__main__":
    trace = collect()
    ntok = len(trace[0])
    print(f"\nToken totali: {ntok} x {N_LAYERS} layer = {ntok*N_LAYERS*TOPK} accessi expert\n")
    # skew: distribuzione di frequenza (media sui layer), e curva di copertura top-K
    print("COPERTURA del set caldo (statico, profilato su prima meta', testato su seconda):")
    print(f"{'K':>4} {'RAM':>7} {'pin-caldo':>10} {'LRU':>8}  (uniforme=K/64)")
    for K in (8, 12, 16, 24, 32):
        cov_static, cov_lru = [], []
        for li in range(N_LAYERS):
            seq = trace[li]; h = len(seq) // 2
            cov_static.append(static_hot_hit(seq[:h], seq[h:], K))
            cov_lru.append(lru_hit(seq, K))
        cs = sum(cov_static)/N_LAYERS; cl = sum(cov_lru)/N_LAYERS
        ram = K * N_LAYERS * 12.6 / 1024
        print(f"{K:>4} {ram:>5.1f}GB {cs*100:>9.1f}% {cl*100:>7.1f}%   {K/64*100:>4.0f}%")
    # quanto e' skewata la distribuzione? entropia normalizzata e top-8 share
    import math
    shares = []
    for li in range(N_LAYERS):
        freq = collections.Counter(e for ex in trace[li] for e in ex)
        tot = sum(freq.values())
        top8 = sum(c for _, c in freq.most_common(8)) / tot
        shares.append(top8)
    print(f"\nSkew: gli 8 expert piu' caldi (su 64) coprono in media "
          f"{sum(shares)/len(shares)*100:.1f}% delle attivazioni (uniforme sarebbe 12.5%).")
@@ -1,20 +0,0 @@
 """Sweep della cache: per ogni capacita' misura correttezza, hit-rate, RAM cache, velocita'."""
 import json, time, glob
 from engine import OlmoeStreaming
 snap = glob.glob("/home/vincenzo/.cache/huggingface/hub/models--allenai--OLMoE-1B-7B-0924/snapshots/*")[0]
 ref = json.load(open("ref.json"))
 exp = ref["full_ids"][len(ref["prompt_ids"]):]
 n_new = len(exp)
 EXPERT_MB_BF16 = 12.6
 print(f"{'cap':>4} {'RAMcache':>9} {'match':>6} {'hit%':>6} {'tok/s':>7} {'sec':>6}")
 for cap in (16, 32, 48, 64):
    m = OlmoeStreaming(snap, expert_cap=cap)
    t = time.time()
    out = m.generate(ref["prompt_ids"], n_new, greedy=True)
    dt = time.time() - t
    gen = out[len(ref["prompt_ids"]):]
    match = sum(a == b for a, b in zip(gen, exp))
    ram = cap * m.L * EXPERT_MB_BF16 / 1024
    print(f"{cap:>4} {ram:>7.1f}GB {match:>3}/{n_new:<2} {m.cache.hitrate()*100:>5.1f}% {n_new/dt:>7.2f} {dt:>6.1f}")
@@ -1,17 +0,0 @@
 """Genera il riferimento con transformers (greedy) e lo salva. Va lanciato da solo (usa ~13GB)."""
 import json, torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 PROMPT = "The capital of France is"
 N = 12
 tok = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
 ids = tok(PROMPT, return_tensors="pt").input_ids
 model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924",
                                             torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).eval()
 with torch.no_grad():
    out = model.generate(ids, max_new_tokens=N, do_sample=False)
 full = out[0].tolist()
 json.dump({"prompt": PROMPT, "prompt_ids": ids[0].tolist(), "full_ids": full,
           "text": tok.decode(full)}, open("ref.json", "w"))
 print("RIFERIMENTO salvato:", repr(tok.decode(full)))
 print("full_ids:", full)