fix: llama.cpp warmp logic

TabbyML · Nov 4, 2023 · c7c67c2 · c7c67c2
1 parent bdd9e6c
commit c7c67c2
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/crates/llama-cpp-bindings/src/engine.cc b/crates/llama-cpp-bindings/src/engine.cc
@@ -90,15 +90,16 @@ class TextInferenceEngineImpl : public TextInferenceEngine {
       batch_ = llama_batch_init(N_CTX * N_CONCURRENT_REQUESTS, 0, 1);
       // warm up
       {
-        for (int i = 0; i < 16; ++i) {
+        batch_.n_tokens = 16;
+        for (int i = 0; i < batch_.n_tokens; ++i) {
           batch_.token[i] = 0;
           batch_.pos[i] = i;
           batch_.n_seq_id[0] = 1;
           batch_.seq_id[i][0] = 0;
           batch_.logits[i] = false;
         }
 
-        if (!llama_decode(ctx_.get(), batch_)) {
+        if (llama_decode(ctx_.get(), batch_)) {
           fprintf(stderr, "%s: warmup failed\n", __func__);
         }