predictor: draw the sampling tail in vocabulary order again

[localai-bot]

Fixes #19.

What broke

26dd8ad ("predictor: unroll the frame into one cgraph and sample in standard ops") moved the sub-talker draw into src/sampling-graph.h. The graph tail keeps the top_k candidates in descending probability order and crosses the cdf of that list against u.

The host sampler it replaced walks the cdf in vocabulary order:

// src/sampling.h, sample_top_k_p
float r   = u * sum;
float acc = 0.0f;
for (int i = 0; i < V; i++) {
    acc += logits[i];
    if (acc >= r) {
        return i;
    }
}

and tests/cossim_common.py::patched_multinomial replicates exactly that arithmetic ("The C++ sampler draws u in [0, 1) and compares against acc/sum implicitly via acc >= u*sum. We replicate that exact arithmetic.").

Both rules are valid multinomial samplers with the same marginal distribution, so nothing looks wrong in isolation. But they map a given u to different tokens, so every stochastic acoustic code changed. Greedy was unaffected, which is why the greedy parity claim in 26dd8ad held.

Measured with a standalone harness that drives sampler_tail_build on random [2048, N] logit rows and compares against sample_top_k_p fed the same u (temperature 0.9, top_k 50):

tail	stochastic draws differing from sample_top_k_p	greedy draws differing from argmax
abab6b3	1950 / 2000	0 / 2000
this branch	0 / 2000	0 / 2000

Why it looks like a hang

An off-reference acoustic trajectory feeds the talker off-manifold and it stops emitting codec_eos, so generation runs to max_new_tokens. In the LocalAI e2e suite the clone spec went from 28 frames to the 2048 frame cap, which from the C ABI caller's side is a synthesis call that does not come back inside the 20 minute test timeout. Their CI log ends with [Pipeline] Generation done : 2048 frames immediately before the Go panic, so it is a runaway, not a deadlock.

I want to be straight about the causal claim: a different-but-valid draw is not by itself proof of a defect, and I am not claiming the descending walk produces an invalid distribution. What I am claiming is that it silently stopped reproducing this project's own reference sampler, and that the observed regression disappears exactly when parity is restored.

The fix

Keep the fused frame graph, keep the u = 0 greedy shortcut, put the draw back on the reference rule:

• argsort still supplies the rank order, but only two ranks are read out of it: rank 0 (the row max) and rank k - 1 (the top_k cutoff).
• The mask keeps everything at or above the cutoff, matching the host sampler's logits[i] < threshold -> -inf, ties included.
• Weights are the unnormalised exp(logit - max) the host sampler builds, the cumsum crossing is at u * total, and the walk runs over the vocabulary, so the crossing index is the token id and the candidate remap disappears.
• Greedy slots upload u = 0 as before; step(u) pins the cutoff at rank 0, so only the argmax survives the mask and any draw lands on it.

Cost is a few extra elementwise ops plus a cumsum over n_vocab (2048) instead of over top_k (50), against a 5 layer transformer pass per step. One file, no ABI change, no change to SamplerInputs or to the upload path.

Verified

Linux x86_64, CPU backend, qwen-talker-0.6b-base-Q4_K_M + qwen-tokenizer-12hz-Q4_K_M, built both with the default flags and with LocalAI's fallback flags (-DBUILD_SHARED_LIBS=OFF -DGGML_NATIVE=OFF -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off):

• Reproduced the regression first. A C harness replaying LocalAI's three e2e specs over the flat C ABI in one process gives, on abab6b3: 11 / 20 / 2048 frames, with [Sample] traces byte identical to their CI log (step=0 c0=1995, step=1 c0=215 on the clone spec). On 35ebe537: 17 / 21 / 28 frames.
• After this patch: 17 / 21 / 28 frames, [Sample] traces identical to 35ebe537.
• Bit identical output vs 35ebe537 on the clone path for seeds 1..12 and 42, --max-new 300: 13 / 13 WAVs match by md5.
• Greedy unchanged (--greedy): md5 identical across 35ebe537, abab6b3 and this branch.
• Sub-talker greedy only (--sub-temp 0, talker stochastic): md5 identical between 35ebe537 and abab6b3, which is how I established that the predictor logits themselves never changed and the difference is purely the draw.
• Batched shapes: the standalone tail harness at N = 1, 2, 3, 5 with slot 0 greedy and the rest stochastic, sampling from step index 1 so the strided state offset is exercised: 0 mismatches vs the host sampler in every configuration.
• --codec-fused streaming still runs and agrees with the non-fused streaming path on frame count and output size.

Not verified

• No GPU run. CUDA, Vulkan, SYCL and Metal are untested by me. The tail uses ggml_exp, ggml_neg, ggml_step, ggml_scale_bias, ggml_cumsum, ggml_argsort, ggml_get_rows, ggml_sum_rows and ggml_cast; all of those have kernels in the vendored ggml for CPU, CUDA, Vulkan, SYCL and Metal, but "has a kernel" is not the same as "I ran it". ggml_exp, ggml_neg and ggml_add are the ops that are new to this file.
• No cossim harness run. tests/debug-*-cossim.py needs the HF checkpoints, torch and the 1.7B VoiceDesign GGUFs, none of which I have here. The tests/*.log grid refreshed in abab6b3 was captured with the descending draw, so the stochastic entries will move again with this patch. Regenerating that grid is yours to do; I did not touch those files.
• I did not test max_batch > 1 end to end through tts-server, only the tail at N > 1 in isolation.

Alternative you may prefer

If you would rather keep the descending walk (it is cheaper on the cumsum), the equivalent fix is to update sampling.h and tests/cossim_common.py::patched_multinomial to the same convention so the three stay in sync. I went with restoring the graph tail because it is the one that moved, and because it puts the released behaviour back byte for byte.