sampling: drop the quadratic RNG rewind and the per-token vocab readback

[CryptVenture]

What this changes

Three changes. None of them alters a sampled token.

1. Quadratic RNG rewind, reachable on non-CUDA backends

hf_sampler rebuilt the Torch-CPU MT19937 on every call and discarded call_index * vocab * 2 words to reach the right stream position, so a sequential walk cost O(call_index × vocab) per token. Most call sites gate this behind cuda_fast_path, but fireredtts3:730, magpie_tts:1236, muscriptor:1504 and personaplex:410 do not — MuScriptor at 1500 steps burns on the order of 1e10 wasted RNG words.

Replaced with a 4-entry thread_local LRU generator cache keyed on (seed, draws_per_call, next_call_index).

The stream is identical by construction. torch_cpu_exponential_float is called at the top of the loop body, before the isfinite check, so every entry draws exactly two words whether or not it is -inf, and there is no break or early return. After call n the generator therefore sits at exactly the offset a rebuild for call n+1 would produce. Any key mismatch falls through to the previous construct-and-discard path verbatim.

The max_score finiteness throw moved above generator acquisition. That is unobservable in the stream — no words were drawn before it — and it guarantees no exception can occur between acquire and commit, so the cache can never record a position the generator did not reach.

ENGINE_TORCH_CPU_RNG_CACHE=0 restores the old path.

Note on existing coverage: tests/unittests/test_torch_random.cpp does not cover this. It includes only torch_random.h and exercises the CUDA Philox path, never the CPU MT19937. So the new test carries a verbatim copy of the previous algorithm as an independent oracle and compares token by token across a 64-step sequential walk, out-of-order call indices, forced cache eviction, three and seven interleaved seeds, alternating vocab sizes, and a row that is two-thirds -inf. Anti-vacuity guards assert at least 8 distinct tokens so it cannot pass by returning a constant.

2. VibeVoice copied 152,064 floats per token to read four of them

decoder.cpp:1085 pulled the full vocabulary row; generator.cpp:429-441 indexes exactly four fixed control-token ids. A 38,016× overfetch — 608,256 bytes and a 594 KiB allocation, every token.

compact_logits_readback already existed in the framework, gathering an arbitrary token subset on-device via ggml_get_rows, with exactly one user (minimax_music3). VibeVoice now drives the same mechanism across all four of its logits-producing graphs: 608,256 → 16 bytes per token. The gather body was hoisted into build_compact_logits_gather so there is one implementation and it is unit-testable; the node sequence is unchanged, so minimax_music3 is bit-identical.

The other six full-vocabulary readback consumers were audited and none can use this: four run in Hidden mode where no logits tensor is built at all, and two are genuinely unconstrained samplers that need the whole row.

3. Fish Audio allocated ~2.5 MB of heap per frame

Across six sites, including two full 155,776-float logits vectors because sample_from_logits runs twice. Adopts the run_step_into pattern that higgs_audio_tts/generator.cpp:451-453 already uses; steady state is now zero per-frame heap traffic on those paths. Call order, RNG consumption and the weights handed to std::discrete_distribution are unchanged.

Investigated and deliberately left alone

Fish Audio's top-p-before-temperature ordering. ar.cpp:487-527 is a line-by-line port of fish-speech's logits_to_probs — confirmed by the && i != 0 keep-at-least-one guard and the max(temperature, 1e-5) clamp — so it is an intentional divergence from the shared sampler, not a defect. Changing it would break reference parity and move generated audio for every existing seed. Documented in place, with test_hf_sampler_ordering pinning the shared sampler's HF order (hand-computed fixture, both orders asserted to actually differ) so the two do not get harmonised by mistake.

Validation

Build

cmake -S . -B build -DENGINE_BUILD_TESTS=ON
cmake --build build

Test

ctest -R "hf_sampler_determinism_test|compact_logits_readback_test|hf_sampler_ordering_test"

Backend tested: Metal, Apple M4 Max. A VibeVoice render is byte-identical before and after (SHA-256 d38f9ebd…) — that is the check that matters for the on-device gather, since a Metal GET_ROWS divergence would change which control token wins without erroring. Full suite 41/41.

Affects

Performance only; output asserted unchanged. If any sampled-token change is observed on FireRedTTS3 / MagpieTTS / MuScriptor / PersonaPlex, set ENGINE_TORCH_CPU_RNG_CACHE=0 first to confirm the cache is the cause.

Known limitations

Two genuinely unconstrained full-vocabulary readbacks remain and are out of scope here — fireredtts3/pipeline.cpp:718 (151,936 floats per prefill step) and qwen3_asr/thinker.cpp:710-711. Both sit on model-owned graphs that logits_readback_token_ids cannot reach; they need a device-side top-k rather than a subset gather.

🤖 Generated with Claude Code

https://claude.ai/code/session_01ATa5YkLUPMDPRL7w1gCo9p

[0xShug0]

Similar issues as other PRs. It touches shared sampler, framework logits readback, Fish Audio, and VibeVoice. This is high-risk and must be split with model-level A/B per affected model.