url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: MTP-head probe uses the last expert by index, not a hardcoded 255

Browse Source 
The has_mtp completeness probe checked for `mlp.experts.255.down_proj.weight`,
which only exists when n_routed_experts == 256. REAP-pruned checkpoints (and any
MoE with a different expert count) have fewer experts, so the probe spuriously
reported has_mtp=0 and disabled MTP speculative decode even though the head was
present and complete. Probe `mlp.experts.<n_experts-1>` instead.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Nicholas Beerbower
2026-07-19 10:54:01 -04:00
parent 61004dcb84
commit 4586d33c60
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/colibri.c
+6 -1
View File
@@ -860,12 +860,17 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits
"self_attn.q_a_proj.weight","self_attn.q_b_proj.weight","self_attn.kv_a_proj_with_mqa.weight",
"self_attn.kv_b_proj.weight","self_attn.o_proj.weight","mlp.gate.weight",
"mlp.shared_experts.gate_proj.weight","mlp.shared_experts.down_proj.weight",
"mlp.experts.0.gate_proj.weight","mlp.experts.255.down_proj.weight"};
"mlp.experts.0.gate_proj.weight"};
char mn[256]; m->has_mtp=1;
for(unsigned q=0;q<sizeof(req)/sizeof(req[0]);q++){
snprintf(mn,sizeof(mn),"model.layers.%d.%s",c->n_layers,req[q]);
if(!st_has(&m->S,mn)){ m->has_mtp=0; break; }
}
/* probe the LAST expert by index, not a fixed 255: REAP-pruned
* checkpoints have n_routed_experts < 256 and the MTP set stays complete,
* so a hardcoded expert.255 would spuriously report has_mtp=0 on them. */
snprintf(mn,sizeof(mn),"model.layers.%d.mlp.experts.%d.down_proj.weight",c->n_layers,c->n_experts-1);
if(!st_has(&m->S,mn)) m->has_mtp=0;
if(getenv("MTP") && atoi(getenv("MTP"))==0) m->has_mtp=0;
if(m->has_mtp){
int i=c->n_layers; Layer *l=&m->mtpL;