From e2d39abd2d4951606ca5c65fe025dc5cf8ea0fec Mon Sep 17 00:00:00 2001
From: KingIcyCreamProjects <safereesekapellen@gmail.com>
Date: Fri, 17 Jul 2026 14:39:14 -0500
Subject: [PATCH] sampling: NaN-skip the mx scan too (review follow-up on #369)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per review: the collapse starts one line before the sum — seeding
mx=lo[0] means a NaN at index 0 makes mx NaN, every (lo[i]-mx) NaN, and
the softmax is doomed at the max-finding, not the normalize. Seed mx
from -INFINITY and skip NaNs (x==x), mirroring the argmax_v change; if
nothing finite survives, fall back to mx=0 and let the post-sum guard
decide. The isfinite(s) guard is now the second line of defense rather
than the only one.

Clean logits take a byte-identical path (the extra x==x compare is
noise next to V expf calls). test_logit_nan gains the NaN-at-index-0
and all-NaN dist_build cases; test_topp's 123-case sweep still passes
on this tree, confirming no interaction with the #354 heap select.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 c/glm.c                  |  8 +++++++-
 c/tests/test_logit_nan.c | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/c/glm.c b/c/glm.c
index 9125587..54dd7fa 100644
--- a/c/glm.c
+++ b/c/glm.c
@@ -4259,7 +4259,13 @@ 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];
+    /* NaN-skip the max scan (x==x rules out NaN): seeding mx=lo[0] let a NaN at
+     * index 0 poison every (lo[i]-mx) before the sum — the softmax was doomed at
+     * the max-finding, not the normalize. The post-sum guard below stays as the
+     * second line of defense. */
+    float mx=-INFINITY;
+    for(int i=0;i<V;i++) if(lo[i]==lo[i] && lo[i]>mx) mx=lo[i];
+    if(!isfinite(mx)) mx=0.f;   /* all-NaN (or +Inf) logits: let the guard below decide */
     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
diff --git a/c/tests/test_logit_nan.c b/c/tests/test_logit_nan.c
index 036c3ea..981d74d 100644
--- a/c/tests/test_logit_nan.c
+++ b/c/tests/test_logit_nan.c
@@ -41,6 +41,21 @@ int main(void){
       assert(approx1(g_pbuf[3]) && "mass must land on the max finite logit (idx3)");
       assert(dist_sample(8,-1)==3 && "sampler emits the finite argmax, not token 0"); }
 
+    /* --- NaN at index 0: poisons the max scan itself (the old mx=lo[0] seed),
+     *     the failure mode that starts before the sum (review note on #369) --- */
+    { float lo[8]={NAN,1.f,0.5f,6.f,0.2f,-1.f,0.f,0.3f};   /* max finite = idx3 (6.0) */
+      dist_build(lo,8);
+      double sum=0; int nan=0;
+      for(int i=0;i<8;i++){ if(!(g_pbuf[i]==g_pbuf[i])) nan=1; sum+=g_pbuf[i]; }
+      assert(!nan && "NaN at lo[0] must not poison via the mx seed");
+      assert(approx1(sum) && "still normalizes to 1");
+      assert(dist_sample(8,-1)==3 && "emits the max finite logit, not token 0"); }
+
+    /* --- all-NaN logits: worst case — must stay finite, no crash --- */
+    { float lo[8]; for(int i=0;i<8;i++) lo[i]=NAN;
+      dist_build(lo,8);
+      for(int i=0;i<8;i++) assert(g_pbuf[i]==g_pbuf[i] && "all-NaN: g_pbuf stays finite"); }
+
     /* --- regression: clean logits still produce a valid distribution --- */
     { float lo[8]={0.1f,0.2f,3.0f,0.4f,0.5f,0.6f,0.7f,0.8f};   /* peak = idx2 */
       dist_build(lo,8);