#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 ints
std::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;
}