common : take the top-k straight from the logits when top-k starts the chain

[sb32445]

Overview

common_sampler_sample() builds a token array for the whole vocabulary (set_logits, 248k entries) and then lets the top-k sampler partial-sort it. When top-k is the first sampler that does anything (penalties, DRY and top-n-sigma off with the current settings, no logit bias, no mirostat, no backend sampling), no grammar is applied first and the reasoning budget is not forcing, the k largest logits (k <= 128) are now selected directly from the logits pointer. The selection repeats the heap steps of libstdc++ std::partial_sort (make_heap on the first k, replace the heap top by every larger logit in index order, sort_heap), so the entries and the order of equal logits are the same as before; a 16-wide block test skips blocks without a candidate. Everything else (the rest of the chain, cur_p) runs unchanged on the k entries.

With MTP speculative decoding the sampler runs on 3 positions per step on the host while the GPU is idle (~0.55 ms of a 24.6 ms step). RTX 4070, Bonsai 2 27B PTQ1_0 + MTP head (n-max 2), 8 interleaved A/B pairs of the same binary, outputs identical in all runs: greedy benchmark +1.92 % (CI [+1.85, +2.01] %), thinking sampling (1.0 / 0.95 / 20 / 0.05, seed 42, reasoning budget 16384, context 114688) +1.85 % (CI [+1.75, +1.94] %).

Additional information

• The branch has two commits: the first contains the environment switch LLAMA_SAMPLER_FAST_TOPK (=0 restores the old path) that was used for the measurements below, the last one removes it. To reproduce a measurement, build the first commit. Penalties, logit bias, mirostat, backend sampling, k > 128 or a grammar applied first take the old path.
• Checked against std::partial_sort on 120000 random arrays (k 1-128, n up to 248k, many ties, NaN, +-inf, sorted inputs): identical index and logit at every rank. Server outputs with the switch on and off are identical for tool calls with a grammar, thinking, presence/repeat penalty, top_k 1 and 200, and logprobs (tool-call ids ignored). The checks are not part of this PR.
• It copies the sift-down/sift-up of libstdc++'s __adjust_heap. The identical tie order is therefore guaranteed (and tested) with libstdc++ only; with libc++ or MSVC STL the order of exactly equal logits may differ from what std::partial_sort gave before (the result is still a correct top-k). If you prefer not to depend on that, the same speed-up is possible with a plain selection that gives up the tie order.
• common_sampler_clone / common_sampler_copy carry the new fast_topk member along.
• Only common/sampling.cpp changes.

Test results

• Hardware / software: RTX 4070 12 GB (AD104, cc 8.9), Linux 6.18, NVIDIA driver 615.71, CUDA 13.4, GCC 16.2 (libstdc++); Release build, -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=89.
• Base: speed numbers were measured on prism at 88c4bc60b; the four commits since (SYCL, WebGPU and cuda: fused FWHT quantizer for 64-wide warps (#303)) do not touch the sampler. The branch is rebased on 2459f68b5 and builds.
• Model: Ternary-Bonsai-2-27B (PTQ1_0) with a community MTP draft head (Q8_0), --spec-type draft-mtp --spec-draft-n-max 2, q4_0 K/V cache, one slot; vocabulary 248320 tokens.
• Method: the same binary, one switch (LLAMA_SAMPLER_FAST_TOPK, first commit of the branch) flipped through an environment variable, alternating runs A B B A, 8 pairs, paired differences with a 95 % bootstrap interval; run-to-run noise about 0.04 to 0.12 %.
• Speed, 4 greedy prompts x 256 tokens, depth 0: 107.8 -> 109.85 tok/s, +1.92 % (95 % CI [+1.85, +2.01] %), outputs identical in all 16 runs.
• Agent-like setup (context 114688, --reasoning-format deepseek --reasoning-budget 16384, thinking sampling 1.0 / 0.95 / 20 / 0.05, fixed seed 42): 100.84 -> 102.71 tok/s, +1.85 % (CI [+1.75, +1.94] %), outputs identical, acceptance unchanged (61.6 %).
• Where the time went: with MTP the sampler runs on 3 positions per step on the host while the GPU is idle, about 0.55 ms of a 24.6 ms step (set_logits 0.36 ms, top-k partial sort 0.18 ms; nsys with CPU sampling).
• Correctness of the selection: compared with std::partial_sort on 120000 random arrays (k 1 to 128, n up to 248k, many ties, NaN, +-inf, ascending and descending inputs): identical index and logit at every rank (libstdc++, GCC 16). This test is not part of the PR; I can turn it into a test case in tests/ if you want it.
• Server outputs with LLAMA_SAMPLER_FAST_TOPK=0 and =1 on the same build are identical for: tool calls with a grammar (sampling and greedy), thinking with sampling, presence_penalty 1.5 (old path), repeat_penalty 1.1 (old path), top_k 200 (old path) and top_k 1, and logprobs/top_logprobs (tool-call ids ignored).
• test-sampling passes on the rebased branch (it exercises llama_sampler_*, not common_sampler, so it does not cover the new path).
• Not tested: libc++ / MSVC STL (see above), several slots, top_k above 128 (old path), grammar-first sampling with an active grammar (old path by construction).

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: The patches were developed with Claude Code (Anthropic's coding agent): it wrote the code, the measurement scripts and the first drafts of the commit messages and PR texts. I decided what to work on (which kernels and host paths to optimise, based on profiles of my own decode setup). The measurements and checks listed in the PR texts were run in the Claude Code sessions; I did not re-run them independently. I will maintain the changes. Commits where Claude Code was used carry a Co-Authored-By trailer.

[bri-prism]

No blocking findings from this review. The fast-path eligibility checks match the existing no-op conditions, and grammar rejection correctly rebuilds the full vocabulary before resampling. The reasoning-budget FORCING guard also preserves access to a forced token outside the raw top-k.

Independent validation on AppleClang/libc++, CPU-only Release build:

• The PR builds and test-sampling passes.
• An additional probe exercises the actual common_sampler_sample implementation with controlled logits. All 422 paired samples across 23 configurations match the original path, covering top-k 1/20/128 and fallback at 0/129, active and disabled penalties/DRY/top-n-sigma, reordered chains, logit bias, greedy and stochastic sampling, adaptive-p, mirostat, grammar rejection, grammar-first sampling, and forced reasoning-budget completion. Clone/copy preserve the eligibility flag.
• 2,000 selection comparisons against the existing top-k sampler match every selected logit. Unique-logit arrays also match token IDs and order. Of 1,000 tied-logit arrays, 984 differ in token ordering on libc++, confirming the compatibility caveat already disclosed in the PR. This can change fixed-seed output, including which tied maximum greedy sampling selects; it is not equivalent to the old path on this standard library.

Non-blocking follow-up: please retain focused coverage of the new common-sampler path, especially grammar rejection/forcing guards and tie behavior. The existing test-sampling pass alone does not cover this optimization. If preserving existing fixed-seed output on libc++/MSVC is a merge requirement, gate the optimization by supported standard library until that behavior is covered.

I have not reproduced the RTX 4070 performance numbers, run MSVC, or tested backend sampling with a live backend. The comparison above uses controlled logits and a small synthetic model, so it establishes these code paths rather than real-model quality or throughput.

Review assisted by Codex.

[sb32445]

@bri-prism thanks, all three points are right.

• Bias: the PTQ1_0 2 to 4 column kernel reads the bias as bias[j*stride_col_dst + row], so it needs the layout of dst. All ADD paths already require that both ADD operands have the same shape, so a broadcast bias [n, 1] is not fused for more than one column, but the same shape does not mean the same layout. ggml_cuda_should_fuse_mul_mat_vec_q now takes the bias tensors and returns false for ncols > 1 when one of them is not contiguous; every caller that has a bias passes it. I have no test for this check: test_mul_mat_vec_fusion builds its own bias and cannot make a non-contiguous one. Tell me if you want a case for it.
• Tests: added 8 cases to MUL_MAT_VEC_FUSION: PTQ1_0 SwiGLU with 2, 3 and 4 columns at K 512 and 5120 (100 rows), the Bonsai 2 FFN shape (3 columns, 17408 x 5120) and the transposed shape (K 17408), which runs into the shared memory check in ggml_cuda_mmvq_ptq1_0_can_fuse_mc. test-backend-ops -o MUL_MAT,MUL_MAT_VEC_FUSION,GLU,SWIGLU,RMS_NORM on the RTX 4070: 2175/2175. In an nsys trace the 3 column case runs mul_mat_vec_ptq1_0_pt<3, 4, true, true>, so it takes the fused kernel and not the fallback.
• cuda: make the ADD+RMS_NORM+MUL fusion reachable #209: a trial merge has one conflict block, in ggml_cuda_check_fusion_memory_ranges: my if (nodes_overlap(dst, src)) with the skip for q8 rows that live in a pool block, against if (nodes_overlap(dst, src) && !is_exact_alias(dst, src)) in cuda: make the ADD+RMS_NORM+MUL fusion reachable #209. The two conditions combine without trouble. I will rebase cuda: gate (SwiGLU) fused PTQ1_0 mat-vec for 2-4 columns #310 on whichever lands first.

Speed is not the point of this change: on the 4070 the kTrain stack with it is +0.17 % at depth 0 (12 pairs, CI [+0.11, +0.23] %, outputs identical). An A/A control shows +0.10 % for the second run of a pair, so I count it as no change.

The checks were run with Claude Code (AI assistance).