Exponent Buckets Before the Nucleus Sort

A synthetic 4,096-token distribution with one sharp peak sent one candidate to Colibri's final sort. A flat distribution sent all 4,096. Both used temperature 0.7 and top_p=0.9. Setting top_p=1.0 sent the spiky fixture through a full vocabulary sort too.

The optimization depends on the distribution, not just vocabulary size. Its useful idea is to prune the sorting input while retaining the probability calculation, nucleus walk, and draw rules. That gives a narrower correctness question than comparing generated paragraphs: do the same logits and random draw still select the same token?

I examined revision 0a5504a1dc64f86fcd0fe7db298b0448c730e7e6 and compiled the unchanged sampler headers and upstream test using GCC 14.2.0 on Linux. The upstream test passed all 1,612 sampled-token comparisons against its full-sort reference. A separate probe inspected sort sizes and boundary draws on synthetic logits. No model weights, inference engine, GPU backend, or decode benchmark was run.

Keep the Distribution, Shrink the Sorting Input

Colibri contains separate C inference engines, tensor and quantization code, CPU and optional GPU paths, and a substantial test tree. Its Qwen serving implementation includes the sampler and calls it in both ordinary generation and multiplexed request handling. The analysis here concerns that shared sampling code, rather than the project's model-loading or expert-streaming claims.

The qwen36_sample.h implementation first finds the maximum logit. For positive temperature it computes unnormalized weights:

weight[i] = expf((logit[i] - maximum_logit) / temperature)
cut = top_p * sum_of_all_weights

Each weight is a float; the global sum and retained mass use double. Normalizing every weight is unnecessary for this selection rule because the cut scales with the same total mass.

The full-sort path orders all weights descending, breaking equal-weight ties by ascending token ID. It walks that order until retained mass reaches the cut, then draws against the retained mass. The optimized path aims to provide enough candidates for the same walk.

It groups each weight by the exponent field of its 32-bit float representation. Larger positive weights enter earlier buckets. It accumulates bucket mass until it reaches the cut, includes one additional bucket as a rounding margin, and copies the qualifying tokens into a candidate array. Only that array goes to qsort.

This is a coarse filter followed by an exact ordering of the retained candidates. It does not choose a fixed number of tokens or replace top-p with a top-k rule. Tokens sharing an exponent bucket travel together even when the final nucleus needs only some of them.

The implementation assumes the float representation its exponent extraction expects. This replay used the host C toolchain's usual 32-bit float layout; it did not establish portability to an unusual floating-point ABI.

Count Candidates, Not Milliseconds

My probe wrapped qsort to record its element count, then delegated to the real library sort. It supplied explicit random values of 0, RAND_MAX / 2, and RAND_MAX. The sampler implementation itself was unchanged.

These are the measured candidate counts at temperature 0.7:

Synthetic logits, vocabulary 4,096	top_p	Elements passed to sort
Flat: every logit 0	0.9	4,096
Spike: token 1,337 has logit 20; others -100	0.9	1
Same spike	1.0	4,096
Plateau: 50 scattered tokens have logit 20; others -100	0.9	50
Repeating tail: -(i % 83) / 3	0.9	300

The plateau makes the distinction visible. Fifty candidates reached the sort; that does not mean all fifty belonged to the final nucleus. Candidate filtering and nucleus truncation are separate steps.

The flat case also matters. Every token shares a weight and exponent bucket, so bucketing cannot reduce the sort input. The implementation still performs its bucket work. These counts establish how much data reached sorting on the stated fixtures; they are not latency measurements or evidence of an end-to-end speedup.

For each of four shapes, the probe checked top_p values 0.3, 0.9, and 1.0 at all three scripted draw values. All 36 token choices matched the full-sort function compiled from the same header.

Ties Are Part of the Sampling Contract

With 4,096 equal weights and top_p=0.9, the sorted-prefix walk retained 3,687 tokens. The explicit zero draw selected token 0. The explicit RAND_MAX draw selected token 3686, the last retained token ID.

This behavior follows two choices in the code: ties use token ID order, and the draw walk accepts the first cumulative mass satisfying acc >= r. A performance refactor that changes tie order can change output while preserving every probability value.

The upstream sampler regression test covers more than the small probe. It uses a vocabulary of 248,320, 13 generated logit shapes, four temperatures, and four top-p settings. Its fixtures include quantized logits with repeated values and a plateau of equal peaks. It compares separately seeded draws and consecutive draws against an independently written full-sort reference.

Its reference uses the same token-ID tie rule as the optimized implementation. Agreement therefore means agreement with that defined reference on this host. It does not promise identical output to a historical comparator that left ties unspecified, or identical rand() sequences across C library implementations.

The Escape Paths Stay Explicit

When top_p is outside the open interval (0,1), the source selects the full-sort function directly. That includes top_p=1.0, as the spike fixture confirmed. This condition is broader than a special case for exactly one.

For the bucket path, the source checks whether the sorted candidate prefix actually reaches the cut. If it does not and candidates excluded part of the vocabulary, it discards that candidate array and calls the full-sort function. I inspected this fallback branch; the probe did not force that rounding fallback to execute. Including an extra exponent bucket is a margin, while the final mass check is the explicit fallback condition.

Nonpositive temperature takes another path entirely: argmax_first.h. It finds the largest value using 16 lanes, then locates its first position. In the probe, temperature zero made no sort or random draw. Equal maxima at positions 17 and 32 selected 17; signed zeros tied at the first index.

That greedy helper preserves the comparison behavior of the original strict-> scan, including an awkward edge: a later NaN was ignored, but a NaN at index zero retained index zero. The upstream test also checks 600 greedy fixtures with ties, NaNs, infinities, signed zeros, and short vectors. These greedy checks do not validate positive-temperature sampling on non-finite logits. The custom stochastic fixtures were all finite.

Compile the Replay Without a Model

The replay script verifies SHA-256 digests for the two headers and upstream test before compiling them. The C probe instruments sort calls and supplies synthetic draw values. The script uses Python 3.12.14; both compilations use -O3 -std=c11 -Wall -Wextra and link libm, with no fast-math option.

The actual cloud invocation was:

python public/blog/exponent-buckets-nucleus/replay.py /tmp/colibri-note-nov4

The script, both compilations, and both test executables exited 0. Selected output:

Verified 3 pinned source digests
serve sampler: 1612 of 1612 draws as the full sort, argmax at temperature 0
PASSED
flat: V=4096 T=0.7 top_p=0.9 sorted_candidates=4096/4096
spike: V=4096 T=0.7 top_p=0.9 sorted_candidates=1/4096
plateau: V=4096 T=0.7 top_p=0.9 sorted_candidates=50/4096
tail: V=4096 T=0.7 top_p=0.9 sorted_candidates=300/4096
flat_ties: draw=0 picks 0; draw=RAND_MAX picks 3686; nucleus has 3687 token IDs
PASS: 36 fast/full-sort comparisons, 2 tie endpoints, 4 greedy cases; synthetic logits only

This was a scoped sampler build, not the repository's complete make check, a model oracle run, or an installation test. The useful result is specific: candidate pruning matched the defined full-sort decisions across the executed fixtures, and its sorting input ranged from one token to the whole vocabulary. Distribution shape, tie order, and escape paths all remain part of the implementation contract.