sampling: guard against non-finite logits (was silent token-0 spew)

On the default serve path (TEMP>0, 0<NUCLEUS<1) a single NaN or +Inf in
the logits — a bad streamed expert tile, or an fp overflow in the matmul
at a low-RAM eviction boundary — poisoned softmax: g_pbuf became all-NaN,
dist_sample never satisfied cum>=u, and the fallback returned token 0. The
engine then emitted an unbroken run of token 0 with NO error. The greedy
path was equally blind: argmax_v started bv=lo[0] and `lo[i]>NaN` is always
false, so a NaN at index 0 pinned the argmax to 0.

- argmax_v: skip NaN (x==x) and seed from -inf, so it returns the max
  finite/+Inf entry instead of being NaN-pinned to 0. Covers greedy decode
  and the speculative-verify argmax path.
- dist_build: after the softmax sum, if s is non-finite or <=0, collapse
  g_pbuf to a one-hot over the finite argmax and warn once, instead of
  dividing every entry into NaN. Covers the nucleus and verify paths.

Both are O(1)/free on the happy path (one branch after the existing loop;
one extra comparison inside the existing argmax loop). Degrade + diagnose,
never silently corrupt.

test_logit_nan (wired into TEST_BINS): asserts argmax_v skips NaN/picks
+Inf, dist_build yields a finite normalized one-hot on the max finite
logit, dist_sample emits that token (not 0), and clean logits still give a
valid distribution. Fails on stock dev, passes with this change.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
KingIcyCreamProjects
2026-07-17 14:08:50 -05:00
parent 8b36736e5d
commit 7f70a8db5d
3 changed files with 73 additions and 2 deletions
+14 -1
View File
@@ -4151,7 +4151,11 @@ static void mtp_absorb(Model *m, const int *next_ids, const float *x, int S, int
}
static inline int argmax_v(const float *lo, int V){
int b=0; float bv=lo[0]; for(int i=1;i<V;i++) if(lo[i]>bv){bv=lo[i];b=i;} return b;
/* skip NaN (x==x is false for NaN) so a poisoned logit can't pin the argmax
* to index 0 pick the max finite/+Inf entry instead. */
int b=-1; float bv=-INFINITY;
for(int i=0;i<V;i++){ float x=lo[i]; if(x==x && x>bv){ bv=x; b=i; } }
return b<0?0:b;
}
/* ---- METODO F: draft grammaticale (#48) ----
@@ -4258,6 +4262,15 @@ static void dist_build(const float *lo, int V){
float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
double s=0; float invt=1.f/(g_temp>1e-4f?g_temp:1e-4f);
for(int i=0;i<V;i++){ g_pbuf[i]=expf((lo[i]-mx)*invt); s+=g_pbuf[i]; }
/* A single NaN/+Inf logit (bad streamed expert tile, matmul overflow at an
* eviction boundary) makes s non-finite; the old code then divided every entry
* into NaN and dist_sample silently returned token 0 forever. Collapse to a
* one-hot over the max finite logit and warn once degrade, don't corrupt. */
if(!(s==s) || s<=0.0 || s>1e300){
static int warned=0; if(!warned){ warned=1; fprintf(stderr,"[sample] non-finite logits; emitting argmax of finite entries\n"); }
int b=argmax_v(lo,V); for(int i=0;i<V;i++) g_pbuf[i]=(i==b)?1.f:0.f;
return;
}
for(int i=0;i<V;i++) g_pbuf[i]/=(float)s;
if(g_nuc>0 && g_nuc<1.f){
for(int i=0;i<V;i++) g_pidx[i]=i;