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

[KingIcyCreamProjects]

What

Fixes #369: on the default sampling path (TEMP>0, 0<NUCLEUS<1), a single NaN/+Inf logit poisoned softmax into an all-NaN g_pbuf, after which dist_sample silently emitted token 0 forever — no error, just an unbroken run of one token. The greedy path was equally blind (argmax_v NaN-pinned to index 0).

Three small changes, all O(1)/free on the happy path:

• argmax_v: skip NaN (x==x) and seed from -INFINITY → returns the max finite/+Inf entry instead of a NaN-pinned 0. Covers greedy decode and the speculative-verify argmax.
• dist_build max scan (per @woolcoxm's review on sampling: a single NaN/+Inf logit makes dist_build emit token 0 forever, silently (default path) #369): seed mx from -INFINITY and NaN-skip the scan — the collapse actually started at the max-finding (mx=lo[0] with a NaN at 0 dooms every lo[i]-mx before the sum), not at the normalize. If nothing finite survives, fall back to mx=0 and let the post-sum guard decide.
• dist_build post-sum guard: 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. Now the second line of defense, not the only one.

Clean logits take a byte-identical path.

Placement vs #354

Rebased onto current dev (post-#354/#357). As predicted in review, the #354 heap rewrite only touched the nucleus tail — the guard sits right after the normalize, before the if(g_nuc>0 && g_nuc<1.f) block, so the heap never sees an all-NaN g_pbuf.

Tests

tests/test_logit_nan.c (wired into TEST_BINS), 6 assertion blocks: argmax_v NaN-skip / +Inf-pick / all-NaN; dist_build with a NaN mid-vocab → finite one-hot on the max finite logit and dist_sample emits it (not 0); NaN at index 0 (the mx-seed failure mode from review); all-NaN stays finite; and a clean-logits regression case. Fails on stock dev (first assert: argmax NaN-pinned to 0), passes with this change.

Also re-ran test_topp (123 cases) on this tree: green — no interaction with the #354 partial select.

@woolcoxm — rebased, mx-scan point folded in, ready for your review.

[JustVugg]

Merged, reconciled with the direct #369 fix that landed on dev in the meantime — and your PR is what closed the real gap: my fix guarded dist_build (the sampling path) but left argmax_v NaN-pinned to index 0, which is exactly the greedy path (TEMP=0). Your x==x NaN-skip in argmax_v covers it, and your test_logit_nan proves it. Resolution kept the already-shipped dist_build guard and took your argmax_v + test on top; both NaN tests run green in the suite (71 tests). Sorry for the duplicate work — the maintainer-fixes-it-directly race is on us (same thing happened to your #367, which you closed graciously). Between the two PRs and the two direct fixes, non-finite logits are now handled on every path: sampling, greedy, and the all-NaN collapse, each with a named stderr warning instead of silence. Thanks for writing it anyway — it caught what the direct fix missed.