Function sample_from_gpu_probs

server/src/common/sampler.cpp:91–102

Given probabilities the GPU already computed (penalties + softmax(temp) applied, summing to ~1) for a pure top_p (no top_k) config, find the nucleus and draw. Skips all exp()/Z bookkeeping the raw-logit path needs, since the input is already normalized. Only worth calling for top_p without top_k: top_k's CPU cost is already cheap (partial_sort scales with k, not vocab — measured ~270-300us at voc

// computed the input probabilities, so skipping the CPU-side exp() pass here
// is a net win (measured ~1.4x faster end-to-end at vocab=151936).
int sample_from_gpu_probs(std::vector<float> & probs, double top_p, double r_uniform) {
    std::vector<std::pair<float, int>> cand(probs.size());
    for (size_t i = 0; i < probs.size(); i++) cand[i] = {probs[i], (int)i};

    double Z = 0.0;
    for (auto & c : cand) Z += c.first;
    const double target = top_p * Z;
    const size_t cut = nucleus_cutoff(cand, target, [](auto & c){ return (double)c.first; });
    cand.resize(cut);

    return draw_from_weights(cand, r_uniform);
}
#endif

}  // namespace