perf(sampling): bound top-logprob candidate storage by requested count

edge0-infer/src/app/sampling.cpp

@@ -90,12 +90,24 @@ TokenLogprob raw_logprob(const std::vector<float>& logits,std::int32_t token,std
     const double normalizer=std::log(sum);
     TokenLogprob result;result.token=token;
     result.logprob=(static_cast<double>(logits.at(static_cast<std::size_t>(token)))-maximum)-normalizer;
-    std::vector<std::int32_t> order(logits.size());
-    for(std::size_t i=0;i<order.size();++i) order[i]=static_cast<std::int32_t>(i);
-    count=std::min(count,order.size());
-    std::partial_sort(order.begin(),order.begin()+count,order.end(),[&](auto a,auto b) {
+    count=std::min(count,logits.size());
+    if (!count) return result;
+    const auto better=[&](auto a,auto b) {
         return logits[a]>logits[b] || (logits[a]==logits[b] && a<b);
-    });
+    };
+    std::vector<std::int32_t> order;
+    order.reserve(count);
+    for (std::size_t i=0;i<logits.size();++i) {
+        const auto id=static_cast<std::int32_t>(i);
+        if (order.size()<count) {
+            order.push_back(id);std::push_heap(order.begin(),order.end(),better);
+        } else if (better(id,order.front())) {
+            std::pop_heap(order.begin(),order.end(),better);
+            order.back()=id;std::push_heap(order.begin(),order.end(),better);
+        }
+    }
+    std::sort_heap(order.begin(),order.end(),better);
+    result.top.reserve(count);
     for(std::size_t i=0;i<count;++i)
         result.top.push_back({order[i],{},(static_cast<double>(logits[order[i]])-maximum)-normalizer});
     return result;

edge0-infer/tests/unit/app/sampling_test.cpp

@@ -3,8 +3,55 @@
 #include <cmath>
 #include <limits>
 #include <iostream>
+#include <algorithm>
+#include <cstdlib>
+#include <new>
+namespace allocation_probe {
+bool active=false;
+std::size_t maximum=0,large=0;
+constexpr std::size_t threshold=32768;
+// Keep the deallocator out of line so GCC does not diagnose an inlined free
+// against the matching replacement operator new used by this test.
+#if defined(__GNUC__)
+__attribute__((noinline))
+#elif defined(_MSC_VER)
+__declspec(noinline)
+#endif
+void release(void* p) noexcept {std::free(p);}
+}
+void* operator new(std::size_t bytes) {
+    if (allocation_probe::active) {
+        allocation_probe::maximum=std::max(allocation_probe::maximum,bytes);
+        if (bytes>=allocation_probe::threshold) ++allocation_probe::large;
+    }
+    if (auto* p=std::malloc(bytes?bytes:1)) return p;
+    throw std::bad_alloc();
+}
+void operator delete(void* p) noexcept {allocation_probe::release(p);}
+void operator delete(void* p,std::size_t) noexcept {allocation_probe::release(p);}
+void* operator new[](std::size_t bytes) {return ::operator new(bytes);}
+void operator delete[](void* p) noexcept {allocation_probe::release(p);}
+void operator delete[](void* p,std::size_t) noexcept {allocation_probe::release(p);}
+
 using namespace edge0;
 int main() {
+    {
+        std::vector<float> row(262144);
+        std::vector<std::int32_t> expected(row.size());
+        for (std::size_t i=0;i<row.size();++i) {row[i]=float(i%97);expected[i]=static_cast<std::int32_t>(i);}
+        std::partial_sort(expected.begin(),expected.begin()+20,expected.end(),[&](auto a,auto b) {
+            return row[a]>row[b] || (row[a]==row[b] && a<b);
+        });
+        for (std::size_t count:{0,1,20}) {
+            allocation_probe::maximum=0;allocation_probe::active=true;
+            const auto result=app::raw_logprob(row,12345,count);
+            allocation_probe::active=false;
+            std::cout << "count=" << count << " max_allocation=" << allocation_probe::maximum << "\n";
+            CHECK(allocation_probe::maximum<4096);
+            CHECK(result.top.size()==count && std::isfinite(result.logprob));
+            for (std::size_t i=0;i<count;++i) CHECK(result.top[i].token==expected[i]);
+        }
+    }
     {
         GenerationConfiguration c;c.temperature=0;c.logit_bias={{0,4}};c.repetition_penalty=2;
         c.presence_penalty=1;c.frequency_penalty=.5;