experiment: miss-aware EXPERT_BUDGET — keep all cache hits, only drop misses
The original budget dropped experts blindly — even cached ones that cost
zero disk I/O. The miss-aware version pre-scans pin/ecache residency
before applying the budget:
- ALL cache hits are kept (free to compute, no disk I/O)
- Only misses compete for the remaining budget slots
- Miss budget = EXPERT_BUDGET - nhits (min 0)
Results (budget=4, same prompt/config as before):
Original Miss-aware
tok/s 0.33 0.36 (+9%)
hit rate 16.4% 38.2% (2.3x)
prefill 8.9s 5.6s (1.6x faster)
decode 97.4s 88.9s (9% faster)
The hit rate doubling is the key quality signal: the model now gets the
full contribution from all resident experts plus the top-4 new loads,
instead of losing some hits to the budget.
This commit is contained in:
@@ -2419,9 +2419,11 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
|
||||
if(!seen[e]){ seen[e]=1; uniq[nu++]=e; }
|
||||
}
|
||||
/* EXPERT_BUDGET: cap distinct experts per layer to reduce disk I/O on cold/low-RAM
|
||||
* hosts. Keep the highest-aggregate-gate-weight experts; drop the rest from idxs[]
|
||||
* so they're never loaded. (MoE-Spec arXiv 2602.16052: top-32 of 64 capture 93%
|
||||
* routing weight.) Complementary to TOPP (per-position) — this trims cross-position. */
|
||||
* hosts. MISS-AWARE: always keep cache hits (pin/LRU — they're free, no disk I/O),
|
||||
* only drop from misses. From the misses, keep the highest-aggregate-gate-weight
|
||||
* ones up to the budget; drop the rest from idxs[] so they're never loaded.
|
||||
* (MoE-Spec arXiv 2602.16052: top-32 of 64 capture 93% routing weight.)
|
||||
* Complementary to TOPP (per-position) — this trims cross-position. */
|
||||
if(g_expert_budget>0 && nu>g_expert_budget){
|
||||
/* compute aggregate gate weight per unique expert */
|
||||
float *wsum=falloc(nu); for(int j=0;j<nu;j++) wsum[j]=0;
|
||||
@@ -2429,14 +2431,26 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
|
||||
int e=idxs[(int64_t)s*K+kk];
|
||||
for(int j=0;j<nu;j++) if(uniq[j]==e){ wsum[j]+=ws[(int64_t)s*K+kk]; break; }
|
||||
}
|
||||
/* mark which unique experts to keep (1) or drop (0) */
|
||||
/* residency pre-scan: which experts are already in pin or ecache (hits)? */
|
||||
unsigned char *is_hit=calloc(nu,1); int nhits=0;
|
||||
for(int j=0;j<nu;j++){ int eid=uniq[j];
|
||||
int found=0;
|
||||
ESlot *P=m->pin[layer];
|
||||
for(int z=0;z<m->npin[layer];z++) if(P[z].eid==eid){ found=1; break; }
|
||||
if(!found){ ESlot *Sl=m->ecache[layer]; int nn=m->ecn[layer];
|
||||
for(int z=0;z<nn;z++) if(Sl[z].eid==eid){ found=1; break; } }
|
||||
if(found){ is_hit[j]=1; nhits++; }
|
||||
}
|
||||
/* budget for misses = total budget - hits already kept (min 0) */
|
||||
int miss_budget = g_expert_budget - nhits; if(miss_budget<0) miss_budget=0;
|
||||
/* mark which unique experts to keep (1) or drop (0): keep all hits, fill rest
|
||||
* with top-weight misses up to miss_budget */
|
||||
unsigned char *keep=calloc(nu,1); int nkeep=0;
|
||||
{ /* simple selection: find top-budget by weight */
|
||||
for(int rank=0;rank<g_expert_budget && rank<nu;rank++){
|
||||
int best=-1; float bv=-1e30f;
|
||||
for(int j=0;j<nu;j++) if(!keep[j] && wsum[j]>bv){ bv=wsum[j]; best=j; }
|
||||
if(best<0) break; keep[best]=1; nkeep++;
|
||||
}
|
||||
for(int j=0;j<nu;j++) if(is_hit[j]){ keep[j]=1; nkeep++; }
|
||||
for(int rank=0;rank<miss_budget;rank++){
|
||||
int best=-1; float bv=-1e30f;
|
||||
for(int j=0;j<nu;j++) if(!keep[j] && wsum[j]>bv){ bv=wsum[j]; best=j; }
|
||||
if(best<0) break; keep[best]=1; nkeep++;
|
||||
}
|
||||
/* build a lookup: for each expert id, is it kept? (reuse seen[]) */
|
||||
memset(seen,0,(size_t)E);
|
||||
@@ -2463,7 +2477,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
|
||||
int nu2=0;
|
||||
for(int j=0;j<nu;j++) if(keep[j]) uniq[nu2++]=uniq[j];
|
||||
nu=nu2;
|
||||
free(wsum); free(keep);
|
||||
free(wsum); free(is_hit); free(keep);
|
||||
}
|
||||
/* ---- FASE C/D: risolvi (pin/cache/disco) e calcola, a blocchi di 64 unici ---- */
|
||||
float *xg=falloc((int64_t)S*D), *gg=falloc((int64_t)S*I), *uu=falloc((int64_t)S*I), *hh=falloc((int64_t)S*D);
|
||||
|
||||
Reference in New Issue
Block a user