Skip to content

Commit c8d6a1f

Browse files
authored
simple : fix batch handling (#3803)
1 parent 2f9ec7e commit c8d6a1f

File tree

1 file changed

+4
-14
lines changed

1 file changed

+4
-14
lines changed

examples/simple/simple.cpp

+4-14
Original file line numberDiff line numberDiff line change
@@ -95,13 +95,8 @@ int main(int argc, char ** argv) {
9595
llama_batch batch = llama_batch_init(512, 0, 1);
9696

9797
// evaluate the initial prompt
98-
batch.n_tokens = tokens_list.size();
99-
100-
for (int32_t i = 0; i < batch.n_tokens; i++) {
101-
batch.token[i] = tokens_list[i];
102-
batch.pos[i] = i;
103-
batch.seq_id[i] = 0;
104-
batch.logits[i] = false;
98+
for (size_t i = 0; i < tokens_list.size(); i++) {
99+
llama_batch_add(batch, tokens_list[i], i, { 0 }, false);
105100
}
106101

107102
// llama_decode will output logits only for the last token of the prompt
@@ -148,15 +143,10 @@ int main(int argc, char ** argv) {
148143
fflush(stdout);
149144

150145
// prepare the next batch
151-
batch.n_tokens = 0;
146+
llama_batch_clear(batch);
152147

153148
// push this new token for next evaluation
154-
batch.token [batch.n_tokens] = new_token_id;
155-
batch.pos [batch.n_tokens] = n_cur;
156-
batch.seq_id[batch.n_tokens] = 0;
157-
batch.logits[batch.n_tokens] = true;
158-
159-
batch.n_tokens += 1;
149+
llama_batch_add(batch, new_token_id, n_cur, { 0 }, true);
160150

161151
n_decode += 1;
162152
}

0 commit comments

Comments
 (0)