url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

sampling: survive non-finite logits instead of emitting token 0 forever (#369)

Browse Source 
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>
This commit is contained in:
JustVugg
2026-07-18 01:10:33 +02:00
parent 6ce6cc4380
commit 4fb5b4e975
3 changed files with 87 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/Makefile
+4 -1
View File
@@ -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)