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

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

Browse Source 
On the default serve path (TEMP>0, 0<NUCLEUS<1) a single NaN or +Inf in
the logits — a bad streamed expert tile, or an fp overflow in the matmul
at a low-RAM eviction boundary — poisoned softmax: g_pbuf became all-NaN,
dist_sample never satisfied cum>=u, and the fallback returned token 0. The
engine then emitted an unbroken run of token 0 with NO error. The greedy
path was equally blind: argmax_v started bv=lo[0] and `lo[i]>NaN` is always
false, so a NaN at index 0 pinned the argmax to 0.

- argmax_v: skip NaN (x==x) and seed from -inf, so it returns the max
  finite/+Inf entry instead of being NaN-pinned to 0. Covers greedy decode
  and the speculative-verify argmax path.
- 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, instead of
  dividing every entry into NaN. Covers the nucleus and verify paths.

Both are O(1)/free on the happy path (one branch after the existing loop;
one extra comparison inside the existing argmax loop). Degrade + diagnose,
never silently corrupt.

test_logit_nan (wired into TEST_BINS): asserts argmax_v skips NaN/picks
+Inf, dist_build yields a finite normalized one-hot on the max finite
logit, dist_sample emits that token (not 0), and clean logits still give a
valid distribution. Fails on stock dev, passes with this change.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
KingIcyCreamProjects
2026-07-17 14:08:50 -05:00
parent 8b36736e5d
commit 7f70a8db5d
3 changed files with 73 additions and 2 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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE)
ifneq (,$(LINUX))
TEST_BINS += tests/test_uring$(EXE)
endif
@@ -331,6 +331,9 @@ tests/bench_topp$(EXE): tests/bench_topp.c glm.c st.h uring.h json.h tok.h tok_u
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)
tests/test_logit_nan$(EXE): tests/test_logit_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_i4_acc512$(EXE): tests/test_i4_acc512.c
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)