Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298#include <algorithm>#include <cmath>#include <cstdint>#include <cstdio>#include <cstdlib>#include <cstring>#include <fstream>#include <iostream>#include <string>#include <vector>#include <omp.h>
struct TraceHeader { char magic[4]; uint16_t version; uint16_t n_layers; uint16_t n_experts; uint16_t top_m; uint32_t n_tokens;};
static inline float fp16_to_fp32(uint16_t h) { uint32_t sign = (uint32_t)(h & 0x8000) << 16; uint32_t exp = (h >> 10) & 0x1f; uint32_t mant = h & 0x03ff; uint32_t f; if (exp == 0) { if (mant == 0) f = sign; else { exp = 1; while ((mant & 0x0400) == 0) { mant <<= 1; exp--; } mant &= 0x03ff; f = sign | ((exp + (127 - 15)) << 23) | (mant << 13); } } else if (exp == 31) { f = sign | 0x7f800000 | (mant << 13); } else { f = sign | ((exp + (127 - 15)) << 23) | (mant << 13); } float out; memcpy(&out, &f, 4); return out;}
enum Policy { POL_LRU=0, POL_LFU=1, POL_EMA=2, POL_NEAR_PICK=3, POL_INV_ROUTER=4, POL_BELADY=5 };
struct TokenRec { int n_cand; uint16_t ids[64]; float logits[64];};
struct ReplayResult { double misses_per_tok; double hit_rate;};
ReplayResult replay_one_layer( const std::vector<TokenRec>& trace, int n_experts, int cache_cap, Policy pol, int k_used, float lambda_moce, int protect_top_j, int near_pick_n = 48, double ema_decay = 0.99, int warmup_tokens = 64) { int n_tok = (int)trace.size(); int eval_tokens = std::max(1, n_tok - warmup_tokens); if (cache_cap >= n_experts) { return {0.0, 1.0}; }
std::vector<std::vector<int>> next_use(n_experts); std::vector<int> next_ptr(n_experts, 0); if (pol == POL_BELADY) { for (int t = 0; t < n_tok; ++t) { int lim = std::min(k_used, trace[t].n_cand); for (int i = 0; i < lim; ++i) { next_use[trace[t].ids[i]].push_back(t); } } }
std::vector<uint8_t> in_cache(n_experts, 0); std::vector<int> cache_list; cache_list.reserve(cache_cap); std::vector<int> last_used(n_experts, 0); std::vector<int> freq(n_experts, 0); std::vector<double> ema_score(n_experts, 0.0); double ema_mult = 1.0; double inv_decay = 1.0 / ema_decay;
int64_t misses = 0; int64_t total_picks = 0;
uint16_t chosen[16]; uint8_t in_chosen[512] = {0};
for (int t = 0; t < n_tok; ++t) { const TokenRec& rec = trace[t]; int n_chosen = 0; if (lambda_moce <= 0.0f || rec.n_cand <= k_used) { n_chosen = std::min(k_used, rec.n_cand); for (int i = 0; i < n_chosen; ++i) chosen[i] = rec.ids[i]; } else { float delta_avg = std::max(1e-4f, rec.logits[0] - rec.logits[std::min(k_used - 1, rec.n_cand - 1)]); float bonus = lambda_moce * delta_avg; struct Cand { float score; int neg_rank; uint16_t eid; }; Cand cands[64]; for (int r = 0; r < rec.n_cand; ++r) { uint16_t eid = rec.ids[r]; float s = rec.logits[r] + ((r < protect_top_j || in_cache[eid]) ? bonus : 0.0f); cands[r] = {s, -r, eid}; } std::sort(cands, cands + rec.n_cand, [](const Cand& a, const Cand& b) { if (a.score != b.score) return a.score > b.score; if (a.neg_rank != b.neg_rank) return a.neg_rank > b.neg_rank; return a.eid > b.eid; }); n_chosen = std::min(k_used, rec.n_cand); for (int i = 0; i < n_chosen; ++i) chosen[i] = cands[i].eid; }
if (pol == POL_BELADY) { for (int i = 0; i < n_chosen; ++i) { uint16_t eid = chosen[i]; auto& lst = next_use[eid]; int& ptr = next_ptr[eid]; while (ptr < (int)lst.size() && lst[ptr] <= t) ptr++; } }
for (int i = 0; i < n_chosen; ++i) in_chosen[chosen[i]] = 1; if (t >= warmup_tokens) { int missing_slots = std::max(0, k_used - n_chosen); total_picks += n_chosen + missing_slots; misses += missing_slots; }
for (int i = 0; i < n_chosen; ++i) { uint16_t eid = chosen[i]; if (!in_cache[eid]) { if (t >= warmup_tokens) misses++; if ((int)cache_list.size() >= cache_cap) { int victim_idx = -1; for (int ci = 0; ci < (int)cache_list.size(); ++ci) { int c = cache_list[ci]; if (in_chosen[c]) continue; if (victim_idx < 0) { victim_idx = ci; continue; } int best_c = cache_list[victim_idx]; bool better = false; if (pol == POL_LRU) { better = (last_used[c] < last_used[best_c]) || (last_used[c] == last_used[best_c] && c < best_c); } else if (pol == POL_LFU) { if (freq[c] != freq[best_c]) better = freq[c] < freq[best_c]; else if (last_used[c] != last_used[best_c]) better = last_used[c] < last_used[best_c]; else better = c < best_c; } else if (pol == POL_EMA || pol == POL_NEAR_PICK || pol == POL_INV_ROUTER) { if (ema_score[c] != ema_score[best_c]) better = ema_score[c] < ema_score[best_c]; else if (last_used[c] != last_used[best_c]) better = last_used[c] < last_used[best_c]; else better = c < best_c; } else if (pol == POL_BELADY) { int nt_c = (next_ptr[c] < (int)next_use[c].size()) ? next_use[c][next_ptr[c]] : 1000000000; int nt_b = (next_ptr[best_c] < (int)next_use[best_c].size()) ? next_use[best_c][next_ptr[best_c]] : 1000000000; if (nt_c != nt_b) better = nt_c > nt_b; else better = c < best_c; } if (better) victim_idx = ci; } int victim = cache_list[victim_idx]; in_cache[victim] = 0; cache_list[victim_idx] = eid; } else { cache_list.push_back(eid); } in_cache[eid] = 1; } last_used[eid] = t + 1; freq[eid] += 1; } for (int i = 0; i < n_chosen; ++i) in_chosen[chosen[i]] = 0;
if (pol == POL_EMA || pol == POL_NEAR_PICK || pol == POL_INV_ROUTER) { ema_mult *= inv_decay; if (ema_mult > 1e100) { double sc = 1.0 / ema_mult; for (int e = 0; e < n_experts; ++e) ema_score[e] *= sc; ema_mult = 1.0; } if (pol == POL_EMA) { for (int i = 0; i < n_chosen; ++i) ema_score[chosen[i]] += ema_mult; } else if (pol == POL_NEAR_PICK) { int top_n = std::min(near_pick_n, rec.n_cand); double denom = std::max(1, top_n - k_used); for (int r = 0; r < top_n; ++r) { double cred = (r < k_used) ? 1.0 : 0.35 * (1.0 - (double)(r - k_used) / denom); ema_score[rec.ids[r]] += cred * ema_mult; } } else if (pol == POL_INV_ROUTER && rec.n_cand > 0) { float z_max = rec.logits[0]; float z_min = rec.logits[rec.n_cand - 1]; float span = std::max(1e-4f, z_max - z_min); for (int r = 0; r < rec.n_cand; ++r) { double util = (rec.logits[r] - z_min) / span + ((r < k_used) ? 1.0 : 0.0); ema_score[rec.ids[r]] += util * ema_mult; } } } }
double m_per_tok = (double)misses / (double)eval_tokens; double hit_rate = 1.0 - ((double)misses / (double)std::max<int64_t>(1, total_picks)); return {m_per_tok, hit_rate};}
// Parse simple keep_map file exported as 48 lines of space-separated expert intsstd::vector<std::vector<int>> load_keep_lines(const std::string& path) { std::vector<std::vector<int>> out(48); if (path == "NONE") return out; std::ifstream f(path); for (int l = 0; l < 48; ++l) { int k; f >> k; out[l].resize(k); for (int i = 0; i < k; ++i) f >> out[l][i]; } return out;}
int main(int argc, char** argv) { if (argc < 7) { fprintf(stderr, "Usage: %s <trace.q4rt> <keep.txt|NONE> <cache_cap> <pol 0..5> <lambda> <top_j>\n", argv[0]); return 1; } FILE* fp = fopen(argv[1], "rb"); if (!fp) { fprintf(stderr, "Failed to open trace file: %s\n", argv[1]); return 1; } TraceHeader hdr; if (fread(&hdr, sizeof(hdr), 1, fp) != 1) return 1; fseek(fp, 0, SEEK_END); long fsz = ftell(fp); fseek(fp, sizeof(hdr), SEEK_SET); int rec_bytes = hdr.top_m * 4; int actual_tokens = (fsz - 16) / (hdr.n_layers * rec_bytes); int n_tok = (hdr.n_tokens == 0 || actual_tokens < (int)hdr.n_tokens) ? actual_tokens : (int)hdr.n_tokens;
std::vector<uint8_t> raw((size_t)n_tok * hdr.n_layers * rec_bytes); if (fread(raw.data(), 1, raw.size(), fp) != raw.size()) return 1; fclose(fp);
auto keep = load_keep_lines(argv[2]); int cache_cap = atoi(argv[3]); if (cache_cap < 10) { fprintf(stderr, "cache_cap (%d) must be >= k_used (10)\n", cache_cap); return 1; } Policy pol = (Policy)atoi(argv[4]); float lam = atof(argv[5]); int top_j = atoi(argv[6]);
double tot_misses = 0.0; double tot_hit = 0.0;
#pragma omp parallel for reduction(+:tot_misses,tot_hit) schedule(dynamic) for (int l = 0; l < hdr.n_layers; ++l) { int remap[512]; for (int e = 0; e < 512; ++e) remap[e] = -1; int n_exp = hdr.n_experts; if (!keep[l].empty()) { n_exp = (int)keep[l].size(); for (int i = 0; i < n_exp; ++i) remap[keep[l][i]] = i; } else { for (int i = 0; i < 512; ++i) remap[i] = i; }
std::vector<TokenRec> l_trace(n_tok); for (int t = 0; t < n_tok; ++t) { const uint8_t* p = raw.data() + ((size_t)t * hdr.n_layers + l) * rec_bytes; const uint16_t* ids = (const uint16_t*)p; const uint16_t* lgs = (const uint16_t*)(p + hdr.top_m * 2); TokenRec& tr = l_trace[t]; tr.n_cand = 0; for (int i = 0; i < hdr.top_m; ++i) { int nid = remap[ids[i]]; if (nid >= 0) { tr.ids[tr.n_cand] = (uint16_t)nid; tr.logits[tr.n_cand] = fp16_to_fp32(lgs[i]); tr.n_cand++; } } } ReplayResult r = replay_one_layer(l_trace, n_exp, cache_cap, pol, 10, lam, top_j); tot_misses += r.misses_per_tok; tot_hit += r.hit_rate; } printf("%.6f %.6f\n", tot_misses, tot_hit / hdr.n_layers); return 0;}