diff --git a/packages/react-native-executorch/common/runner/sampler.cpp b/packages/react-native-executorch/common/runner/sampler.cpp index 250d6a83ef..8215c49be5 100644 --- a/packages/react-native-executorch/common/runner/sampler.cpp +++ b/packages/react-native-executorch/common/runner/sampler.cpp @@ -36,6 +36,9 @@ #include #include #include +#include +#include +#include #include namespace executorch { @@ -141,57 +144,59 @@ template void Sampler::mask_topk(T *logits) { } } -// Mask logits whose softmax-prob falls outside the top-p nucleus to -inf. -// Keeps the token that crosses the threshold (HuggingFace convention). +// Mask logits outside the top-p nucleus to -inf. Approximates the exact +// sort-based nucleus with a histogram over (logit - max): two O(n) passes, no +// sort. Binning in logit (not probability) space keeps uniform resolution for +// peaked and flat distributions alike. kRange=40 spans exp() down to ~4e-18. template void Sampler::mask_topp(T *logits) { if (topp_ <= 0.0f || topp_ >= 1.0f) { return; } - // Softmax into a scratch probs[] (do not mutate logits yet). - T max_val = logits[0]; - for (size_t i = 1; i < vocab_size_; i++) { - if (logits[i] > max_val) { - max_val = logits[i]; - } - } - std::unique_ptr[]> probindex = - std::make_unique[]>(vocab_size_); - T sum = 0; + constexpr int32_t kBins = 2048; + // Compute in a type at least as wide as T so converting logits never loses + // precision: double stays double, everything else (float and the narrow + // half/bf16/uint16 logit types) widens to float. Accumulating in T directly + // would be unsafe for bf16, whose mantissa saturates when summing exp() + // over the full vocab. + using acc_t = std::conditional_t, double, float>; + constexpr acc_t kRange = 40; + + std::span logit_span{logits, static_cast(vocab_size_)}; + const acc_t max_val = + static_cast(*std::ranges::max_element(logit_span)); + + std::vector bin_mass(kBins, acc_t(0)); + acc_t total = 0; for (size_t i = 0; i < vocab_size_; i++) { - T e = static_cast(std::expf(static_cast(logits[i] - max_val))); - probindex[i].prob = e; - probindex[i].index = i; - sum += e; + acc_t d = static_cast(logits[i]) - max_val; + acc_t e = std::exp(d); + total += e; + int32_t bin = static_cast((d + kRange) / kRange * kBins); + bin = std::clamp(bin, 0, kBins - 1); + bin_mass[bin] += e; } - if (sum <= T(0)) { + if (total <= acc_t(0)) { return; } - for (size_t i = 0; i < vocab_size_; i++) { - probindex[i].prob /= sum; - } - std::sort(probindex.get(), probindex.get() + vocab_size_, - [](const ProbIndex &a, const ProbIndex &b) { - return a.prob > b.prob; - }); - // Find the smallest prefix whose cumulative probability >= topp_. - T cumulative = 0; - int last_idx = vocab_size_ - 1; - for (size_t i = 0; i < vocab_size_; i++) { - cumulative += probindex[i].prob; - if (static_cast(cumulative) >= topp_) { - last_idx = i; + // Highest bin downward until the kept mass reaches topp. The crossing bin is + // kept (HuggingFace "keep the token that crosses" convention). + const acc_t target = static_cast(topp_) * total; + acc_t acc = 0; + int32_t keep_bin = 0; + for (int32_t bin = kBins - 1; bin >= 0; --bin) { + acc += bin_mass[bin]; + if (acc >= target) { + keep_bin = bin; break; } } - // Mark kept indices, then -inf the rest. - std::vector keep(vocab_size_, false); - for (size_t i = 0; i <= last_idx; i++) { - keep[probindex[i].index] = true; - } + const acc_t d_threshold = + static_cast(keep_bin) / kBins * kRange - kRange; + constexpr T neg_inf = std::numeric_limits::lowest(); for (size_t i = 0; i < vocab_size_; i++) { - if (!keep[i]) { + if (static_cast(logits[i]) - max_val < d_threshold) { logits[i] = neg_inf; } } @@ -210,22 +215,28 @@ Sampler::Sampler(int32_t vocab_size, GenerationConfig config) : Sampler(vocab_size, config, std::time(nullptr)) {} template static void softmax(T *x, int size) { - // find max value (for numerical stability) + // Runs after top-k/top-p masking, which sets rejected logits to lowest(). + // Skip exp() on those: it underflows to 0 anyway and is slow on device. + constexpr T kMasked = std::numeric_limits::lowest(); T max_val = x[0]; for (size_t i = 1; i < size; i++) { if (x[i] > max_val) { max_val = x[i]; } } - // exp and sum T sum = 0; for (size_t i = 0; i < size; i++) { + if (x[i] == kMasked) { + x[i] = T(0); + continue; + } x[i] = expf(x[i] - max_val); sum += x[i]; } - // normalize for (size_t i = 0; i < size; i++) { - x[i] /= sum; + if (x[i] != T(0)) { + x[i] /= sum; + } } }