new llama_batch_init/llama_batch_free functions a bit difficult to use and may possibly leak memory

[leok7v]

1. llama_batch_init does not indicate success of failure on malloc() which makes it difficult (but not impossible) to have graceful exit (or continue) on OOM.
2. llama_batch_init always initializes .n_tokens to zero not passed n_tokens_alloc which forces caller to do batch.n_tokens = n_tokens_alloc;
3. In multithreaded environment it is possible that some malloc() succeed in the loop and some fail and return nullptr in random order. It results in llama_batch_free leaking memory stopping free too early.

    // Return batch for single sequence of tokens
    // The sequence ID will be fixed to 0
    // The position of the tokens will be tracked automatically by llama_decode
    //
    // NOTE: this is a helper function to facilitate transition to the new batch API - avoid using it
    //
    LLAMA_API struct llama_batch llama_batch_get_one(
                  llama_token * tokens,
                      int32_t   n_tokens);

    // Allocates a batch of tokens on the heap that can hold a maximum of n_tokens
    // Each token can be assigned up to n_seq_max sequence ids
    // The batch has to be freed with llama_batch_free()
    // If embd != 0, llama_batch.embd will be allocated with size of n_tokens * embd * sizeof(float)
    // Otherwise, llama_batch.token will be allocated to store n_tokens llama_token
    // The rest of the llama_batch members are allocated with size n_tokens
    // All members are left uninitialized
    LLAMA_API struct llama_batch llama_batch_init(
            int32_t n_tokens,
            int32_t embd,
            int32_t n_seq_max);

    // Frees a batch of tokens allocated with llama_batch_init()
    LLAMA_API void llama_batch_free(struct llama_batch batch);

struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_t n_seq_max) {
    llama_batch batch = {
        /*n_tokens =*/ 0,
        /*tokens   =*/ nullptr,
        /*embd     =*/ nullptr,
        /*pos      =*/ nullptr,
        /*n_seq_id =*/ nullptr,
        /*seq_id   =*/ nullptr,
        /*logits   =*/ nullptr,
    };

    if (embd) {
        batch.embd = (float *) malloc(sizeof(float) * n_tokens_alloc * embd);
    } else {
        batch.token = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc);
    }

    batch.pos      = (llama_pos *)     malloc(sizeof(llama_pos)      * n_tokens_alloc);
    batch.n_seq_id = (int32_t *)       malloc(sizeof(int32_t)        * n_tokens_alloc);
    batch.seq_id   = (llama_seq_id **) malloc(sizeof(llama_seq_id *) * (n_tokens_alloc + 1));
    for (int i = 0; i < n_tokens_alloc; ++i) {
        // some malloc() may return nullptr and some not nullptr in random order here:
        batch.seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id) * n_seq_max);
    }
    batch.seq_id[n_tokens_alloc] = nullptr;

    batch.logits   = (int8_t *)        malloc(sizeof(int8_t)         * n_tokens_alloc);

    return batch;
}

void llama_batch_free(struct llama_batch batch) {
    if (batch.token)    free(batch.token);
    if (batch.embd)     free(batch.embd);
    if (batch.pos)      free(batch.pos);
    if (batch.n_seq_id) free(batch.n_seq_id);
    if (batch.seq_id) {
        // loop exits on first nullptr while other elements of array may not be nullptr (see above)
        for (int i = 0; batch.seq_id[i] != nullptr; ++i) {
            free(batch.seq_id[i]);
        }
        free(batch.seq_id);
    }
    if (batch.logits)   free(batch.logits);
}

The code above forces called who want to check llama_batch before using it to do something like

static bool is_batch_allocated(struct llama_batch * b, int32_t n_tokens) {
    bool allocated = // everything is malloc() != NULL?
        b->n_tokens == n_tokens && (b->embd || b->token) &&
        b->pos && b->n_seq_id && b->seq_id && b->logits;
    if (allocated) {
        for (int i = 0; i < n_tokens && allocated; i++) {
            allocated = b->seq_id[i];
        }
    }
    return allocated;
}

...

    struct llama_batch b = llama_batch_init(n_tokens, /*embd:*/0, /*max_seq:*/1);
    b.n_tokens = n_tokens;
    if (!is_batch_allocated(&b, n_tokens)) { llama_batch_free(b); return false; } // it may still leak some memory

[leok7v]

Also

            llama_batch batch_view = {
                n_tokens,
                batch.token    + i,
                nullptr,
                batch.pos      + i,
                batch.n_seq_id + i,
                batch.seq_id   + i,
                batch.logits   + i,
            };

is a bit fragile (if order of fields changes). Don't remember is C++ allows C99 { .field = value, ...} initialization.

[ggerganov]

llama_batch_init always initializes .n_tokens to zero not passed n_tokens_alloc which forces caller to do batch.n_tokens = n_tokens_alloc;

Maybe this is what you misunderstand - the n_tokens are the current number of tokens that you have in the batch. When you create the batch, you don't have any tokens in it, so n_tokens == 0. Setting it to n_tokens_alloc is a mistake.

See the examples for more info.

[leok7v]

more

You underestimate me. I misunderstood much more. Had to read llama-batch.cpp line by line and step thru the whole day. The problem was not n_tokens. Problem was pos[] referring to memory. I think I finally made it clear for myself:

llama_test_batch_decode.c

But the still a bit difficult to use (I see how complicated the usage in in the server context - maybe it is just complex by nature).