Avoid repeated reallocation of the TopK buffer

[Magnushst]

Summary

TopK fills a std::vector<double> with one packed entry per logit but starts from an empty vector. logits.size() is the vocabulary size, so for Gemma 3's 256K vocabulary libstdc++ reallocates 19 times and copies 2.1 MB before the buffer reaches its final size. This runs once per sampled token via FusedSoftmaxAndSampleTopK, which ChooseSampleFunc selects whenever top_k is greater than 1 or an accept_token filter is set.

This reserves the full size up front and hoists the accept_token test out of the loop, so the common case (no constrained decoding) is a straight-line pack-and-append rather than a std::function test per logit. Peak transient memory also falls, because the final doubling no longer holds the old and new buffers at the same time.

Sampling results are unchanged: the same entries are appended in the same order, so the subsequent VQSelect/VQSort see identical input.

Performance

Deterministic reduction, for a 256K vocabulary:

• Allocations per call: 19 -> 1
• Bytes copied by reallocation: 2.10 MB -> 0
• Peak transient buffer: 3 MB -> 2 MB

Wall time, interleaved baseline/patched runs, best of 30 calls (300 for the small cases), median of three runs:

Case	Baseline	Patched	Change
vocab 262144, k=50	366.7 us	284.0 us	-22.6%
vocab 32768, k=50	43.4 us	33.3 us	-23.3%
vocab 256, k=3	0.63 us	0.51 us	-19%
vocab 262144, k=50, accept_token	485.2 us	461.4 us	neutral (within noise)

The accept_token case is the control: it keeps the per-element filter call, which dominates, so it is unaffected.

Environment:

• CPU: Intel Core i9-13905H
• Compiler: GCC 14.2.0
• Build: CMake Release, Ninja, -O3 -DNDEBUG
• Highway target: AVX2 (best attainable on this CPU)
• Repetitions: best of 30 per run, three interleaved runs, pinned with taskset

This is a sampling-path improvement, not an end-to-end decode speedup; on a memory-bound decode step the saving is small next to the output-head MatMul.

Testing

• New TestTopK covers the filtered and unfiltered paths, asserting exact tokens and probabilities. Its size exceeds the vector's initial capacity, so it exercises the reserved-capacity path.
• ops_test passes 40/40 across AVX2 and EMU128; compress_test and sfp_test unaffected.
• ops_test.cc built with -fsanitize=address and run with detect_leaks=1: 40/40 pass, no errors and no leaks.

Notes

Only x86-64 (AVX2) was measured; ARM was not benchmarked. The change is allocation behaviour rather than SIMD, so it should carry over, but I have not verified that.

[jan-wassenberg]

Nice improvements, thank you!

[jan-wassenberg]

hm, we have an issue with the CLA scanner because there are too many commits. I suspect this is caused by initially targeting the main branch and then rebasing to dev. Anything we can do to get a clean set of commits?

[Magnushst]

hm, we have an issue with the CLA scanner because there are too many commits. I suspect this is caused by initially targeting the main branch and then rebasing to dev. Anything we can do to get a clean set of commits?

Yeah, that is exactly why I rebased; was a bit confused at first when I saw all the commits. I'll try force pushing again and see if that resolves it