sampling: a single NaN/+Inf logit makes dist_build emit token 0 forever, silently (default path)

[KingIcyCreamProjects]

Summary

On the default serve/chat sampling path (TEMP>0, 0<NUCLEUS<1), a single NaN or +Inf logit makes dist_build() produce an all-NaN distribution, after which the sampler silently emits token 0 forever — an unbroken run of one token, with no error printed. A diagnosable upstream numerical fault becomes silent corruption.

Mechanism (traced against glm.c)

dist_build() softmax:

float mx=lo[0]; for(int i=1;i<V;i++) if(lo[i]>mx) mx=lo[i];
double s=0; ...
for(int i=0;i<V;i++){ g_pbuf[i]=expf((lo[i]-mx)*invt); s+=g_pbuf[i]; }
for(int i=0;i<V;i++) g_pbuf[i]/=(float)s;

• A NaN logit → g_pbuf[that]=expf(NaN)=NaN → s=NaN → every g_pbuf[i]/=(float)s is NaN.
• (+Inf → mx=+Inf → expf(Inf-Inf)=NaN → same collapse.)
• Nucleus keep-loop cum+=g_pbuf[...]; if(cum>=g_nuc) never triggers (NaN>=x is false) → keep=V, s2=NaN, still all-NaN.
• dist_sample: cum>=u never true → fallback for(i=V-1..0) if(g_pbuf[i]>0) return i is false for all (NaN>0 false) → returns 0.
• Greedy path is equally blind: argmax_v seeds bv=lo[0] and lo[i]>NaN is always false, so a NaN at index 0 pins the argmax to 0.

There is no isfinite/isnan guard anywhere between the logits and the sampler.

Reachability

The precondition is the default, not an edge case (TEMP auto≈0.7, NUCLEUS default 0.90). A non-finite logit is plausible in this NVMe-streaming engine: a garbage/short-loaded expert tile or an fp overflow in matmul_qt at a low-RAM eviction boundary. Both dist_build call sites are affected (the sampler and the speculative-verify path).

Impact

Not a fault originator — it needs an upstream NaN/+Inf. It's a fault amplifier: it turns a catchable numeric fault into silent, unbroken token-0 output. Medium severity: it defeats loud failure on the default path.

Fix (ready)

Minimal, O(1)/free on the happy path, on fix/logit-nan-guard (off dev) on my fork:

• argmax_v: skip NaN (x==x) and seed from -inf → returns the max finite/+Inf entry, not a NaN-pinned 0.
• 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 — degrade + diagnose, never silently corrupt.
• test_logit_nan (wired into TEST_BINS): fails on stock dev, passes with the guard.

Heads-up: since #354 rewrites dist_build's internals (heap partial-select), the dist_build half of this guard should be re-placed on top of that once it lands — happy to open the PR against whichever base you prefer, or fold it into the #354 work. @woolcoxm

[nolanb-launchdarkly]

Looks like 354 was merged already

[woolcoxm]

Thanks for the clean trace — the mechanism is right, and I checked it against dev after #354 + #357 landed. Good news on the base question, plus one detail worth adding to your patch.

Base: open it against dev

#354 (top-p partial select) and #357 (DSA partial select in attention_rows) are both merged to dev now. Rebase fix/logit-nan-guard onto current dev and the placement question answers itself.

The #354 rewrite doesn't move your guard

You wrote "the dist_build half of this guard should be re-placed on top of #354." I checked — #354 only rewrote the nucleus truncation tail (the heapify + partial pop). The three lines where the NaN collapse actually happens are byte-identical to pre-#354:

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]; }
for(int i=0;i<V;i++) g_pbuf[i]/=(float)s;

So your guard lands exactly where you had it — right after the normalize, before the if(g_nuc>0 && g_nuc<1.f) nucleus block. No reshaping needed. The heap tail will then never see an all-NaN g_pbuf because you'll have collapsed it to a one-hot first.

One detail to add: the max-finding loop is also an entry point

Your writeup focuses on s going NaN and argmax_v being blind to NaN. But the collapse actually starts one line earlier — mx=lo[0]. If lo[0] itself is the NaN (plausible: token 0's logit is the one that gets pinned, and a garbage expert tile at position 0 is exactly the kind of upstream fault you describe), then mx is NaN, every lo[i] > mx is false, mx stays NaN, and (lo[i]-mx) is NaN for every i — so every expf is NaN and s is NaN before you ever reach the normalize. The softmax is doomed at the max-finding, not at the sum.

So I'd mirror the argmax_v fix here too: seed mx from -INFINITY (or the first finite entry) rather than lo[0], and skip NaNs in the max scan. Then the post-sum isfinite(s) guard is a belt-and-suspenders second line rather than the only one.

Suggested shape, consistent with your argmax_v change:

float mx=-INFINITY;
for(int i=0;i<V;i++) if(lo[i]==lo[i] && lo[i]>mx) mx=lo[i];   /* lo[i]==lo[i] rules out NaN */
if(!isfinite(mx)) mx=0.f;   /* all-NaN logits: uniform-ish, let the guard below handle */

Out of scope for your patch, but worth a note

#357 put a partial-select in attention_rows (DSA path). That's attention, not sampling, so a non-finite logit there has different reachability — but if your fault-injection test turns up anything in that neighborhood, flag it and I'll take a look on the sampling-adjacent side.

Open the PR against dev and ping me — I'll review.

[JustVugg]

Fixed on dev. You were exactly right, mechanism and all: a single NaN or +Inf logit poisoned dist_build — +Inf became mx, the softmax sum went NaN, every probability went NaN, and dist_sample's fallback (if(g_pbuf[i]>0)) is false for NaN at every index, so it returned token 0, silently, forever.

dist_build now computes the max 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) it collapses to a delta on the finite argmax and warns once on stderr — degraded output with a visible cause, never a silent stream of zeros. The argmax uses the index found during the max pass, so it's correct even when lo[0] itself is the NaN (where the plain argmax_v would wrongly return 0).

Shipped with tests/test_sample_nan.c: healthy logits still sample correctly; NaN/+Inf injected at position 0, the middle, and the last index all pick the finite argmax; an all-non-finite vocab leaves no NaN in the buffer and doesn't crash. Thanks for the clean bisect — a silent corruption on the default path is the worst kind, and you nailed it. Closing.

[JustVugg]

Fix is on dev — if you can, pull and re-trigger the case that produced the endless token 0. Instead of a silent stream of zeros you should now either get sane output (if the NaN was transient) or a single [SAMPLE] warning: non-finite logits (NaN/Inf) — falling back to argmax on stderr with a valid token, never the silent zero-loop. tests/test_sample_nan.c covers it, but a confirmation from the path that originally hit it would be the real seal. Thanks again for the clean bisect.