add batch.clear() and batch.n_tokens()

Author: ngxson

common/speculative.cpp

@@ -204,7 +204,7 @@ llama_tokens common_speculative_gen_draft(
     }
 
     // prepare a batch to evaluate any new tokens in the prompt
-    llama_batch_ext_clear(batch.get());
+    batch.clear();
 
     for (size_t i = i_start + reuse_n; i < prompt_tgt.size(); ++i) {
         //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_tgt[i]);
@@ -214,7 +214,7 @@ llama_tokens common_speculative_gen_draft(
     }
 
     // we should rarely end-up here during normal decoding
-    if (llama_batch_ext_get_n_tokens(batch.get()) > 0) {
+    if (batch.n_tokens() > 0) {
         //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str());
 
         llama_decode_ext(ctx, batch.get());
@@ -224,7 +224,7 @@ llama_tokens common_speculative_gen_draft(
 
     LOG_DBG("%s: n_past = %d\n", __func__, n_past);
 
-    llama_batch_ext_clear(batch.get());
+    batch.clear();
     batch.add_text(id_last, n_past, 0, true);
 
     prompt.push_back(id_last);
@@ -237,7 +237,7 @@ llama_tokens common_speculative_gen_draft(
 
     // sample n_draft tokens from the draft model
     for (int i = 0; i < params.n_draft; ++i) {
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         common_sampler_sample(smpl, ctx, 0, true);
 

examples/gritlm/gritlm.cpp

@@ -17,7 +17,7 @@ static std::vector<std::vector<float>> encode(llama_context * ctx, const std::ve
     llama_batch_ext_ptr batch(llama_batch_ext_init(llama_n_batch(ctx), 1));
 
     for (uint64_t i = 0; i < sentences.size(); i++) {
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         const std::string input_string = instruction + sentences[i];
 
@@ -111,7 +111,7 @@ static std::string generate(llama_context * ctx, llama_sampler * smpl, const std
     int32_t i_current_token = 0;
 
     while (true) {
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
         {
             const int32_t n_inputs = inputs.size();
 
@@ -123,7 +123,7 @@ static std::string generate(llama_context * ctx, llama_sampler * smpl, const std
 
         llama_decode_ext(ctx, batch.get());
 
-        llama_token token = llama_sampler_sample(smpl, ctx, llama_batch_ext_get_n_tokens(batch.get()) - 1);
+        llama_token token = llama_sampler_sample(smpl, ctx, batch.n_tokens() - 1);
 
         if (token == eos_token) {
             break;

examples/llava/gemma3-cli.cpp

@@ -90,7 +90,7 @@ struct gemma3_context {
 
 static int eval_text(gemma3_context & ctx, std::string input, bool logits_last = false) {
     llama_tokens tokens = common_tokenize(ctx.lctx, input, false, true);
-    llama_batch_ext_clear(ctx.batch.get());
+    ctx.batch.clear();
     for (llama_token & t : tokens) {
         ctx.batch.add_text(t, ctx.n_past++, 0, false);
     }
@@ -178,7 +178,7 @@ static int generate_response(gemma3_context & ctx, common_sampler * smpl, int n_
         fflush(stdout);
 
         // eval the token
-        llama_batch_ext_clear(ctx.batch.get());
+        ctx.batch.clear();
         ctx.batch.add_text(token_id, ctx.n_past++, 0, true);
         if (llama_decode_ext(ctx.lctx, ctx.batch.get())) {
             LOG_ERR("failed to decode token\n");

examples/lookup/lookup.cpp

@@ -197,7 +197,7 @@ int main(int argc, char ** argv){
         // clean the cache of draft tokens that weren't accepted
         llama_kv_self_seq_rm(ctx, 0, n_past, -1);
 
-        llama_batch_ext_clear(batch_tgt.get());
+        batch_tgt.clear();
         batch_tgt.add_text(draft[0], n_past, 0, true);
 
         // Draft already contains a single token sampled from the model:

examples/parallel/parallel.cpp

@@ -217,23 +217,23 @@ int main(int argc, char ** argv) {
             common_kv_cache_dump_view_seqs(kvc_view, 40);
         }
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         // decode any currently ongoing sequences
         for (auto & client : clients) {
             if (client.seq_id == -1) {
                 continue;
             }
 
-            client.i_batch = llama_batch_ext_get_n_tokens(batch.get());
+            client.i_batch = batch.n_tokens();
 
             llama_seq_id seq_id = client.id + 1;
             batch.add_text(client.sampled, n_tokens_system + client.n_prompt + client.n_decoded, seq_id, true);
 
             client.n_decoded += 1;
         }
 
-        if (llama_batch_ext_get_n_tokens(batch.get()) == 0) {
+        if (batch.n_tokens() == 0) {
             // all sequences have ended - clear the entire KV cache
             for (int i = 1; i <= n_clients; ++i) {
                 llama_kv_self_seq_rm(ctx, i, -1, -1);
@@ -245,7 +245,7 @@ int main(int argc, char ** argv) {
         }
 
         // insert new sequences for decoding
-        if (cont_batching || llama_batch_ext_get_n_tokens(batch.get()) == 0) {
+        if (cont_batching || batch.n_tokens() == 0) {
             for (auto & client : clients) {
                 if (client.seq_id == -1 && g_seq_id < n_seq) {
                     client.seq_id = g_seq_id;
@@ -269,13 +269,13 @@ int main(int argc, char ** argv) {
                     }
 
                     // extract the logits only for the last token
-                    if (llama_batch_ext_get_n_tokens(batch.get()) > 0) {
+                    if (batch.n_tokens() > 0) {
                         llama_batch_ext_set_output_last(batch.get());
                     }
 
                     client.n_prompt  = tokens_prompt.size();
                     client.n_decoded = 0;
-                    client.i_batch   = llama_batch_ext_get_n_tokens(batch.get()) - 1;
+                    client.i_batch   = batch.n_tokens() - 1;
 
                     LOG_INF("\033[31mClient %3d, seq %4d, started decoding ...\033[0m\n", client.id, client.seq_id);
 
@@ -289,14 +289,14 @@ int main(int argc, char ** argv) {
             }
         }
 
-        if (llama_batch_ext_get_n_tokens(batch.get()) == 0) {
+        if (batch.n_tokens() == 0) {
             break;
         }
 
         // process in chunks of params.n_batch
         int32_t n_batch = params.n_batch;
 
-        int32_t n_tokens_in_batch = llama_batch_ext_get_n_tokens(batch.get());
+        int32_t n_tokens_in_batch = batch.n_tokens();
         for (int32_t i = 0; i < (int32_t) n_tokens_in_batch; i += n_batch) {
             // experiment: process in powers of 2
             //if (i + n_batch > (int32_t) batch.n_tokens && n_batch > 32) {

examples/passkey/passkey.cpp

@@ -141,7 +141,7 @@ int main(int argc, char ** argv) {
             n_past = llama_kv_self_seq_pos_max(ctx, 0) + 1;
         }
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {
             batch.add_text(tokens_list[i + j], n_past++, 0, false);
@@ -175,7 +175,7 @@ int main(int argc, char ** argv) {
 
         n_past = llama_kv_self_seq_pos_max(ctx, 0) + 1;
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {
             batch.add_text(tokens_list[i + j], n_past++, 0, false);
@@ -224,7 +224,7 @@ int main(int argc, char ** argv) {
     while (n_cur <= n_len) {
         // sample the next token
         {
-            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, llama_batch_ext_get_n_tokens(batch.get()) - 1);
+            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens() - 1);
 
             // is it an end of generation?
             if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len) {
@@ -238,7 +238,7 @@ int main(int argc, char ** argv) {
             n_decode += 1;
 
             // prepare the next batch
-            llama_batch_ext_clear(batch.get());
+            batch.clear();
 
             // push this new token for next evaluation
             llama_seq_id seq_id = 0;

examples/perplexity/perplexity.cpp

@@ -369,7 +369,7 @@ static results_perplexity perplexity_v2(llama_context * ctx, const common_params
             const int batch_start = start + j * n_batch;
             const int batch_size  = std::min(end - batch_start, n_batch);
 
-            llama_batch_ext_clear(batch.get());
+            batch.clear();
             for (int i = 0; i < batch_size; i++) {
                 batch.add_text(tokens[batch_start + i], j*n_batch + i, 0, true);
             }
@@ -552,7 +552,7 @@ static results_perplexity perplexity(llama_context * ctx, const common_params &
 
             int n_outputs = 0;
 
-            llama_batch_ext_clear(batch.get());
+            batch.clear();
             for (int seq = 0; seq < n_seq_batch; seq++) {
                 int seq_start = batch_start + seq*n_ctx;
 
@@ -846,7 +846,7 @@ static void hellaswag_score(llama_context * ctx, const common_params & params) {
         size_t i1 = i0;
         size_t i_logits = 0; // this tells us how many logits were needed before this point in the batch
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         // batch as much tasks as possible into the available context
         // each task has 4 unique sequence ids - one for each ending
@@ -1131,7 +1131,7 @@ static void winogrande_score(llama_context * ctx, const common_params & params)
         size_t i1 = i0;
         size_t i_logits = 0;
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         while (n_cur + (int) data[i1].required_tokens <= n_ctx) {
             int n_logits = 0;
@@ -1485,7 +1485,7 @@ static void multiple_choice_score(llama_context * ctx, const common_params & par
         size_t i1 = i0;
         size_t i_logits = 0; // this tells us how many logits were needed before this point in the batch
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
 
         // batch as much tasks as possible into the available context
         // each task has 4 unique sequence ids - one for each ending
@@ -1744,7 +1744,7 @@ static void kl_divergence(llama_context * ctx, const common_params & params) {
                 tokens[batch_start] = llama_vocab_bos(vocab);
             }
 
-            llama_batch_ext_clear(batch.get());
+            batch.clear();
             for (int i = 0; i < batch_size; i++) {
                 batch.add_text(tokens[batch_start + i], j*n_batch + i, 0, true);
             }

examples/run/run.cpp

@@ -954,7 +954,7 @@ static int tokenize_prompt(const llama_vocab * vocab, const std::string & prompt
 static int check_context_size(const llama_context_ptr & ctx, const llama_batch_ext_ptr & batch) {
     const int n_ctx      = llama_n_ctx(ctx.get());
     const int n_ctx_used = llama_kv_self_used_cells(ctx.get());
-    if (n_ctx_used + llama_batch_ext_get_n_tokens(batch.get()) > n_ctx) {
+    if (n_ctx_used + batch.n_tokens() > n_ctx) {
         printf(LOG_COL_DEFAULT "\n");
         printe("context size exceeded\n");
         return 1;
@@ -1001,7 +1001,7 @@ static int generate(LlamaData & llama_data, const std::string & prompt, std::str
             return 1;
         }
 
-        llama_data.n_past += llama_batch_ext_get_n_tokens(batch.get());
+        llama_data.n_past += batch.n_tokens();
 
         // sample the next token, check is it an end of generation?
         new_token_id = llama_sampler_sample(llama_data.sampler.get(), llama_data.context.get(), -1);

examples/save-load-state/save-load-state.cpp

@@ -52,7 +52,7 @@ int main(int argc, char ** argv) {
 
     // evaluate prompt
     llama_decode_ext(ctx, batch.get());
-    n_past += llama_batch_ext_get_n_tokens(batch.get());
+    n_past += batch.n_tokens();
 
     // save state (rng, logits, embedding and kv_cache) to file
     {
@@ -79,7 +79,7 @@ int main(int argc, char ** argv) {
         printf("%s", next_token_str.c_str());
         result0 += next_token_str;
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
         batch.add_text(next_token, 0, 0, true);
 
         if (llama_decode_ext(ctx, batch.get())) {
@@ -131,7 +131,7 @@ int main(int argc, char ** argv) {
         printf("%s", next_token_str.c_str());
         result1 += next_token_str;
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
         batch.add_text(next_token, 0, 0, true);
 
         if (llama_decode_ext(ctx2, batch.get())) {
@@ -212,7 +212,7 @@ int main(int argc, char ** argv) {
         printf("%s", next_token_str.c_str());
         result2 += next_token_str;
 
-        llama_batch_ext_clear(batch.get());
+        batch.clear();
         batch.add_text(next_token, 0, 1, true);
 
         if (llama_decode_ext(ctx3, batch.get())) {