From 4fb5b4e9754653fd978b59d4d10f73c08291b885 Mon Sep 17 00:00:00 2001
From: JustVugg <JustVugg@users.noreply.github.com>
Date: Sat, 18 Jul 2026 01:10:33 +0200
Subject: [PATCH] sampling: survive non-finite logits instead of emitting token
 0 forever (#369)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A single NaN or +Inf logit silently broke the default sampling path. +Inf
became `mx`, then `expf((Inf-mx))`/`expf((NaN-mx))` is NaN, the softmax sum went
NaN, every probability went NaN — and dist_sample's fallback loop
`if(g_pbuf[i]>0)` is false for NaN at every index, so it returned 0. Every
subsequent token: 0. No error, no warning. @KingIcyCreamProjects found it.

dist_build now takes `mx` over finite logits only, gives a non-finite logit
probability 0, and when the distribution is unusable (no finite logit, or a
non-finite/zero sum) collapses to a delta on the finite argmax and warns ONCE
on stderr — degraded, but a valid token and a visible cause, never a silent
stream of zeros. The finite argmax uses the index found during the mx pass
(robust even when lo[0] itself is NaN, where argmax_v would wrongly return 0).

tests/test_sample_nan.c: healthy logits still sample correctly; NaN/+Inf
injected at lo[0], the middle, and the last position all pick the finite
argmax; an all-non-finite vocab leaves no NaN in the buffer and doesn't crash.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 c/Makefile                |  5 +++-
 c/glm.c                   | 23 +++++++++++++--
 c/tests/test_sample_nan.c | 62 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 c/tests/test_sample_nan.c

diff --git a/c/Makefile b/c/Makefile
index feb3045..7210cbb 100644
--- a/c/Makefile
+++ b/c/Makefile
@@ -166,7 +166,7 @@ else
 PYTHON    ?= python3
 endif
 CUDA_OBJ  =
-TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
+TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
 ifneq (,$(LINUX))
 TEST_BINS += tests/test_uring$(EXE)
 endif
@@ -328,6 +328,9 @@ tests/test_topp$(EXE): tests/test_topp.c glm.c st.h uring.h json.h tok.h tok_uni
 tests/bench_topp$(EXE): tests/bench_topp.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
+tests/test_sample_nan$(EXE): tests/test_sample_nan.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
 tests/test_kv_alloc$(EXE): tests/test_kv_alloc.c glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
diff --git a/c/glm.c b/c/glm.c
index bf8acb1..2ee12ed 100644
--- a/c/glm.c
+++ b/c/glm.c
@@ -4296,9 +4296,28 @@ static void topp_siftdown(int *h, int n, int i){
  * la coda troncata va AZZERATA in g_pbuf (dist_sample la legge direttamente per id). */
 static void dist_build(const float *lo, int V){
     if(!g_pbuf){ g_pbuf=falloc(V); g_pidx=malloc(V*sizeof(int)); }
-    float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
+    /* Un solo logit NaN/+Inf avvelenava tutto (#369): +Inf diventava mx, NaN/Inf-mx
+     * -> expf NaN -> s NaN -> ogni prob NaN -> dist_sample cade sul fallback
+     * `g_pbuf[i]>0` (NaN>0 e' falso ovunque) e ritorna 0 PER SEMPRE, in silenzio.
+     * Difesa: mx solo sui finiti; un logit non finito contribuisce prob 0;
+     * se la distribuzione degenera (tutti non finiti / somma non valida) si
+     * ripiega sull'argmax dei finiti e si avvisa UNA volta, mai in silenzio. */
+    int mxi=-1; float mx=0;
+    for(int i=0;i<V;i++) if(isfinite(lo[i]) && (mxi<0 || lo[i]>mx)){ mx=lo[i]; mxi=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]; }
+    if(mxi>=0){
+        for(int i=0;i<V;i++){ g_pbuf[i]=isfinite(lo[i])?expf((lo[i]-mx)*invt):0.f; s+=g_pbuf[i]; }
+    }
+    if(mxi<0 || !isfinite(s) || s<=0.0){          /* distribuzione inutilizzabile */
+        static int warned=0;
+        if(!warned){ warned=1; fprintf(stderr,
+            "[SAMPLE] warning: non-finite logits (NaN/Inf) — falling back to argmax; "
+            "output may be degraded. This usually means a numerical blow-up upstream.\n"); }
+        int a=(mxi>=0)?mxi:0;                       /* mxi = argmax dei logit FINITI (robusto
+                                                     * anche se lo[0] e' NaN, dove argmax_v fallirebbe) */
+        for(int i=0;i<V;i++) g_pbuf[i]=0.f; g_pbuf[a]=1.f;
+        return;                                    /* delta su un token valido, niente top-p su NaN */
+    }
     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;
diff --git a/c/tests/test_sample_nan.c b/c/tests/test_sample_nan.c
new file mode 100644
index 0000000..fc33268
--- /dev/null
+++ b/c/tests/test_sample_nan.c
@@ -0,0 +1,62 @@
+/* Regressione #369: un solo logit NaN o +Inf faceva emettere a dist_build/dist_sample
+ * il token 0 PER SEMPRE, in silenzio (il fallback `g_pbuf[i]>0` e' falso su NaN ovunque).
+ * Ora la distribuzione degenere ripiega sull'argmax dei logit FINITI e avvisa una volta.
+ *
+ * Il test verifica: (a) con logit sani il campionamento resta corretto; (b) con NaN/+Inf
+ * iniettato il token scelto e' l'argmax dei FINITI (mai 0 per default), su ogni posizione
+ * del NaN inclusa lo[0]; (c) nessun NaN sopravvive in g_pbuf. */
+#define main coli_glm_main_unused
+#include "../glm.c"
+#undef main
+#include <stdio.h>
+
+static uint32_t rs=0x1234abcd; static uint32_t xr(){rs^=rs<<13;rs^=rs>>17;rs^=rs<<5;return rs;}
+
+static int pbuf_has_nan(int V){ for(int i=0;i<V;i++) if(!isfinite(g_pbuf[i])) return 1; return 0; }
+
+int main(void){
+    int fail=0, V=2000;
+    float *lo=malloc(V*sizeof(float));
+    g_temp=0.7f; g_nuc=0.90f;                     /* il path serve di default */
+
+    /* (a) logit sani: il campionato deve avere prob > 0 e nessun NaN nel buffer */
+    for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%2000)-1000)/100.f;
+    int known=1337; lo[known]=50.f;               /* picco netto */
+    dist_build(lo,V);
+    if(pbuf_has_nan(V)){ printf("  FAIL: NaN in g_pbuf con logit sani\n"); fail=1; }
+    if(g_pbuf[known]<=0.f){ printf("  FAIL: il picco ha prob 0\n"); fail=1; }
+    if(!fail) printf("  logit sani: distribuzione valida, picco vivo            ok\n");
+
+    /* (b) NaN/+Inf iniettato in varie posizioni; il finito-argmax deve vincere */
+    float bad[3]; bad[0]=NAN; bad[1]=INFINITY; bad[2]=-INFINITY;
+    const char *bn[3]={"NaN","+Inf","-Inf"};
+    for(int b=0;b<3;b++){
+        for(int pos=0;pos<3;pos++){               /* lo[0], meta', ultimo */
+            for(int i=0;i<V;i++) lo[i]=(float)((int)(xr()%400)-200)/100.f;
+            int amax=777; lo[amax]=9.0f;           /* massimo FINITO atteso */
+            int at = pos==0?0 : pos==1?V/2 : V-1;
+            if(at==amax) amax=amax+1;              /* non sovrapporre */
+            lo[amax]=9.0f;
+            lo[at]=bad[b];                         /* veleno */
+            dist_build(lo,V);
+            if(pbuf_has_nan(V)){ printf("  FAIL: NaN sopravvive (%s @ %d)\n",bn[b],at); fail=1; continue; }
+            /* con -Inf il fallback puo' non scattare (max finito resta), ma il buffer
+             * deve restare valido e sommare ~1: con +Inf/NaN scatta il delta su amax */
+            int picked=-1; float pv=-1;
+            for(int i=0;i<V;i++) if(g_pbuf[i]>pv){pv=g_pbuf[i];picked=i;}
+            if(b<2 && picked!=amax){                /* NaN e +Inf: delta esatto su amax */
+                printf("  FAIL: %s @ %d -> picked %d, atteso argmax finito %d\n",bn[b],at,picked,amax); fail=1;
+            }
+        }
+    }
+    if(!fail) printf("  NaN/+Inf iniettato: argmax dei finiti vince, mai 0/NaN   ok\n");
+
+    /* (c) caso estremo: TUTTI non finiti -> non deve crashare, buffer valido */
+    for(int i=0;i<V;i++) lo[i]=NAN;
+    dist_build(lo,V);
+    if(pbuf_has_nan(V)){ printf("  FAIL: tutti-NaN lascia NaN nel buffer\n"); fail=1; }
+    else printf("  tutti non-finiti: nessun crash, buffer valido               ok\n");
+
+    printf(fail?"test_sample_nan: FAIL\n":"test_sample_nan: ok\n");
+    return fail;
+}