| 873 | } |
| 874 | |
| 875 | struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_t n_seq_max) { |
| 876 | llama_batch batch = { |
| 877 | /*n_tokens =*/ 0, |
| 878 | /*tokens =*/ nullptr, |
| 879 | /*embd =*/ nullptr, |
| 880 | /*pos =*/ nullptr, |
| 881 | /*n_seq_id =*/ nullptr, |
| 882 | /*seq_id =*/ nullptr, |
| 883 | /*logits =*/ nullptr, |
| 884 | }; |
| 885 | |
| 886 | if (embd) { |
| 887 | batch.embd = (float *) malloc(sizeof(float) * n_tokens_alloc * embd); |
| 888 | } else { |
| 889 | batch.token = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc); |
| 890 | } |
| 891 | |
| 892 | batch.pos = (llama_pos *) malloc(sizeof(llama_pos) * n_tokens_alloc); |
| 893 | batch.n_seq_id = (int32_t *) malloc(sizeof(int32_t) * n_tokens_alloc); |
| 894 | batch.seq_id = (llama_seq_id **) malloc(sizeof(llama_seq_id *) * (n_tokens_alloc + 1)); |
| 895 | for (int i = 0; i < n_tokens_alloc; ++i) { |
| 896 | batch.seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id) * n_seq_max); |
| 897 | } |
| 898 | batch.seq_id[n_tokens_alloc] = nullptr; |
| 899 | |
| 900 | batch.logits = (int8_t *) malloc(sizeof(int8_t) * n_tokens_alloc); |
| 901 | |
| 902 | return batch; |
| 903 | } |
| 904 | |
| 905 | void llama_batch_free(struct llama_batch batch) { |
| 906 | if (batch.token) free(batch.token); |
no outgoing calls