sampling: partial top-p select in dist_build — O(V) heapify + k pops, not O(V log V) qsort (#335)

[woolcoxm]

What

Fixes #335.

dist_build() qsort-ed the entire 151936-entry vocab by probability on every sampled token whenever 0 < g_nuc < 1 (the serve default — temp resolves to 0.7, NUCLEUS=0.90), and again per draft position under rejection sampling. Measured on a Ryzen 9950X3D: 5.6–8.0 ms/call. The actual work is finding the few-hundred-token head whose cumulative mass reaches g_nuc.

Change

Replace the full qsort(g_pidx, V, ...) + linear scan with a max-heap partial select:

• Floyd heapify g_pidx over V by descending g_pbuf prob — O(V), cache-friendly, no recursion, no indirect comparator
• Pop winners to the array's high end until cum >= g_nuc — k · O(log V), where k is the head size (typically tens–hundreds at p=0.9)
• The remaining heap prefix g_pidx[0..out-1] is the tail → zero it, renormalize the head over g_pidx[out..V-1]

Winners land in g_pidx[out..V-1] in descending order, so s2 accumulates in the same order as the old code → head is unchanged on tie-free shapes (boundary ties were already unspecified under the unstable qsort and stay so — called out in the issue).

No API change, no caller change, no new globals.

Correctness contract (preserved)

dist_sample() reads g_pbuf[i] directly by token id over all V entries, so the four invariants from the issue all hold:

1. g_pbuf stays id-indexed (never reordered)
2. g_pidx stays internal to dist_build
3. The truncated tail is fully zeroed in g_pbuf (the heap prefix is exactly the tail)
4. The head renormalizes to 1.0

Test

c/tests/test_topp.c (new) drives the real dist_build via the #include "../glm.c" pattern (same as test_stops.c / test_i4_grouped.c) against an independent double-precision reimplementation of the old algorithm. 123 cases: 6 sizes (1, 2, 8, 64, 257, 1519 ≈ V/100 of GLM-5.2) × 5 shapes (uniform / peaked / geometric / plateau-ties / sharptail) × 4 g_nuc values (0.001, 0.5, 0.9, 0.999), plus the g_nuc ≥ 1 and g_nuc ≤ 0 guard-off paths and a V=1 edge case.

• Tie-free shapes: head values within 1e-6 rel (engine uses float, reference uses double → sub-ULP renorm noise; matches test_i4_grouped.c's tolerance convention)
• Tie shapes (incl. boundary ties at the head/tail cutoff, where membership is interchangeable): value-multiset equality
• Every shape also asserts: keep-count matches, head sums to 1.0

No scratch files → builds clean on Windows MinGW without the unmerged mkdtemp compat shim (#352). test_stops.c currently does not (pre-existing on dev, separate concern).

test_topp: 123 cases run, 0 failure(s)
test_topp: ok

Files

• c/glm.c — dist_build rewrite + topp_siftdown helper (replaces cmp_pdesc)
• c/tests/test_topp.c — new
• c/Makefile — TEST_BINS + build rule

---

Credit / co-author: the max-heap partial-select algorithm and the original 9950X3D ~7× measurement are @KingIcyCreamProjects's work from issue #335 — this PR implements their idea. Attribution is recorded here in-thread rather than as a Co-Authored-By: commit trailer only because these commits are already merged into origin/dev; amending them for the trailer would rewrite shared integration history (every contributor's dev diverges). If proper git-native co-authorship is wanted, that needs a maintainer to file it as a tracked issue/decision on the repo.

[woolcoxm]

Sibling of this PR on the attention side, same algorithmic class and the same partial-select idea — opened as #356 + branch fix/dsa-partial-select-356 (based off dev, like this one):

attention_rows() DSA indexer was full-qsorting all nk context scores per layer per token (O(nk log nk)) just to read the threshold tmp[keep-1]. Replaced with a partial_select_desc quickselect (O(nk) average) + thr = min(selected block). The position-order scans are unchanged so the kept-position set is bit-identical — actually a stronger contract than this PR's sampling heap, which could only promise multiset equality (heap unstable, accumulation order changed).

Measured (bench_dsa_select, keep=2048, median of 2000 reps):

nk (ctx)	old qsort	new quickselect	speedup
2049	119 µs	5.7 µs	21×
8192	626 µs	43 µs	15×
32768	2.8 ms	0.28 ms	10×
65536	6.6 ms	0.47 ms	14×

The gap widens with context (linear vs n-log-n), and DSA only fires past index_topk=2048 — so it bites exactly in the long-conversation regime. Ships with test_dsa_select (129 cases, element-wise identical kept-set vs an independent qsort reference) and bench_dsa_select (on-demand, not a gate), mirroring this PR's test/bench pair.

Left a note in #356 recording the three other attention-side candidates I measured and did not file (MoE top-K O(K·E) scan, 3-pass softmax, tok_encode) — all already optimal at GLM-5.2 scale, so nobody re-investigates.

[KingIcyCreamProjects]

Independent 9950X3D repro (the box from #335) — the win is bigger here than on Meteor Lake.

Ran make tests/bench_topp && ./tests/bench_topp on a Ryzen 9 9950X3D (Zen 5 — the machine that produced the original 7× in #335). GCC 16.1.0 / UCRT, -O3 -march=native, V=151936, temp=0.70, 2000 reps/cell:

shape           nuc    old ns/call    new ns/call   speedup      keep
realistic      0.50       17133800         984600    17.40x         1
realistic      0.90       16297300        1171400    13.91x      1370   <- serve default
realistic      0.95       16696600        1392000    11.99x      3955
realistic      0.99       16092000        2271400     7.08x     12612
uniform        0.50         837500         636200     1.32x     75969
uniform        0.90         922100         752000     1.23x    136743
uniform        0.95         935300         759400     1.23x    144340
uniform        0.99         962200         781900     1.23x    150417
plateau        0.50        2316100        2343200     0.99x        33
plateau        0.90        2339300        2270600     1.03x        89
plateau        0.95        2219900        2255600     0.98x       109
plateau        0.99        2233600        2231000     1.00x       168

Reading it:

• Realistic shape at the serve default (nuc=0.90): 13.91× — 16.3 ms → 1.17 ms/call, on the per-token critical path.
• The ratio grows on Zen 5 vs your Meteor Lake (13.9× vs 3.95× @0.90): old qsort is only ~1.4× faster here (16.3 vs 23.3 ms), but the new heapify is ~5× faster (1.17 vs 5.9 ms) — the cache-friendly heap loves the big V-Cache + wide Zen 5 engine. So sampling: dist_build full-qsorts all 151K vocab entries per sampled token — partial top-p select is ~7x cheaper (measured) #335's 7× was conservative for this box; at nuc=0.99 (keep=12612) it lands right at 7.08×.
• uniform (1.2×) and plateau (~1.0×) match your read exactly — theoretical floors, not shapes real logits hit.

So: provably correct (test_topp) + reproducibly 7–17× on realistic serving shapes on the 9950X3D. The "plausibly faster" caveat is closed.

Two small notes:

1. Attribution — the algorithm and the original 9950X3D measurement are from sampling: dist_build full-qsorts all 151K vocab entries per sampled token — partial top-p select is ~7x cheaper (measured) #335; mind crediting it (co-author / "implements @KingIcyCreamProjects's sampling: dist_build full-qsorts all 151K vocab entries per sampled token — partial top-p select is ~7x cheaper (measured) #335") beyond Fixes #335? Just keeping the analysis trail intact.
2. Coverage nit: test_topp's bit-equivalence sweep caps at V=1519 and leaves full V=151936 to the (non-gating) bench. One V=151936 case in the gated test_topp would lock correctness at the real vocab — cheap to add.

Serve-decode tok/s confirmation is still the open item — happy to run that next on this box.

[woolcoxm]

doing all this now including CO-AUTHOR :D thanks for the bug report, the gains were massive !!!

EDIT: is there anything special i need to do for co author?? or just credits in the texts?

[woolcoxm]

i havent been co authoring at all, i didnt really think about it, made a note to co author from now on and give credit when credit is due :D

the ai controls 100% of my comments, i automate the github chain entirely.

[KingIcyCreamProjects]

Nothing special — it's just a commit-message trailer. Last line(s) of the commit message, with a blank line before it:

fix the thing

longer description here...

Co-authored-by: KingIcyCreamProjects <KingIcyCreamProjects@users.noreply.github.com>

That's the whole mechanism. GitHub picks it up automatically: both avatars show on the commit and it counts as a contribution for the co-author. The users.noreply.github.com form always works without needing anyone's real email. If you squash-merge via the GitHub UI, it even pre-fills trailers for you.

And no need to do anything retroactive here — it's merged, the text mentions are plenty. Trailer on future commits when it applies, done. 🎉 on the merge — those bench numbers were fun to watch land.

[woolcoxm]

thanks!!!!

there seems to be an issue with doing this the proper way, i should have done it before hand and now it looks like it may be too late, can i give credit in the comments?? fixing this the proper way will break everyones code apparently.

[KingIcyCreamProjects]

Yes — comments are 100% fine, and you already did it. Never rewrite merged history for a trailer; that cure is worse than the disease. The trailer is only for future commits, before they merge. You're all square here — consider it settled 👍

[woolcoxm]

will make a note in the ai to give credit for finding bugs and such, after reviewing some of his posts it seems he is a credit thief....

[woolcoxm]

hes correcting it now, i spanked him and overloaded his ram on purpose.

[KingIcyCreamProjects]

lmaooo good, mine writes half my comments too so i cant even talk. bots crediting bots, welcome to 2026 🍻