Buckets:

gijl
/

lv-storage

Files

xet

gijl/lv-storage / llama.cpp /src /llama-batch.h

gijl

15 days ago

download

raw

11.6 kB

#pragma once

#include "llama.h"

#include "llama-arch.h"
#include "llama-cparams.h"

#include <array>
#include <vector>
#include <set>
#include <bitset>
#include <memory>
#include <unordered_map>
#include <unordered_set>

// keep this struct lightweight
struct llama_ubatch {
    bool equal_seqs() const {
        return b_equal_seqs != 0;
    }

    // typical for M-RoPE cases:
    //   0 - sequential position of the tokens/embeddings in the sequence
    //   1 - y position in the image
    //   2 - x position in the image
    //   3 - other
    bool is_pos_2d() const {
        // TODO @ngxson : we may need to check for model arch when more models use >1 positions
        return n_pos >= 3;
    }

    // mixed: type picks token or embd per row, pos has n_pos sections for all rows
    bool is_mixed() const {
        return type != nullptr;
    }

    uint32_t b_equal_seqs; // note: this is a boolean, but we use an int32_t for alignment
                           //       otherwise address sanitizer complains
    // TODO: whole_seqs for embeddings?

    uint32_t n_tokens;     // total tokens (n_seq_tokens * n_seqs)
    uint32_t n_seq_tokens; // tokens per sequence set
    uint32_t n_seqs;       // sequence sets in the ubatch
    uint32_t n_seqs_unq;   // unique sequence ids in the ubatch
    uint32_t n_pos;        // number of position inputs for each token/embedding

    // seq_id_unq: unique sequence ids in the ubatch
    // seq_idx:    indices of the unique sequence ids in the ubatch in [0, n_seqs_unq)
    //             used for extracting sequence pooled embeddings

    //                          // size                     | idx | val
    llama_token  *  token;      // [n_tokens]               | i   | id, token
    float        *  embd;       // [n_embd, n_tokens]       | i   | embd
    float        *  embd_state; // [n_embd_state, n_tokens] | i   | hidden state carried over from a previous stage (e.g. MTP)
    llama_pos    *  pos;        // [n_tokens*n_pos]         | i   | pos
    int32_t      *  n_seq_id;   // [n_tokens]         | i   | -
    llama_seq_id ** seq_id;     // [n_tokens]         | s   | s0, s1, seq_id
    llama_seq_id *  seq_id_unq; // [n_seqs_unq]       | s   | seq_id
    int32_t      *  seq_idx;    // [LLAMA_MAX_SEQ]    | -   | seq_idx
    int8_t       *  output;     // [n_tokens]         | i   | -
    int8_t       *  type;       // [n_tokens]         | i   | -     (mixed ubatch only, 0 - token, 1 - embd)
    int32_t      *  decision_order; // [n_tokens], NULL if no entry has one, see llama_batch_ext_set_decision_order()

    struct data_t {
        std::vector<llama_token>    token;
        std::vector<float>          embd;
        std::vector<float>          embd_state;
        std::vector<llama_pos>      pos;
        std::vector<int32_t>        n_seq_id;
        std::vector<llama_seq_id *> seq_id;      // these point into the seq_id_data below
        std::vector<llama_seq_id>   seq_id_unq;
        std::vector<int32_t>        seq_idx;
        std::vector<int8_t>         output;
        std::vector<int8_t>         type;
        std::vector<int32_t>        batch_idxs;  // original batch index for each token
        std::vector<int32_t>        decision_order;

        std::vector<llama_seq_id> seq_id_data;
    };

    // the llama_ubatch pointers above point to this data if set. otherwise - point to external non-owning data
    std::shared_ptr<data_t> data;
};

// crash if a mixed ubatch reaches code that expects only tokens or only embd
#define ASSERT_EMBD_OR_TOKEN(ubatch) GGML_ASSERT(!(ubatch).is_mixed() && "mixed token/embd ubatch is not supported here")

struct llama_hparams;

// DFlash batches carry the fused target features at the encoder input width (n_embd_inp_enc size).
// Other batches carry token embeddings (n_embd_inp size).
size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams);

// MTP contexts also take the target model's hidden state (n_embd_out size), 0 = no state input
size_t llama_batch_ext_select_n_embd_state(llama_context_type ctx_type, const llama_hparams & hparams);

struct llama_batch_ext {
    const size_t n_tokens_max;     // max number of tokens that can be stored in the batch
    const size_t n_embd_inp;       // decoder embd row width
    const size_t n_embd_inp_enc;   // encoder embd row width (e.g. eagle3/dflash extracted features)
    const size_t n_embd_state;     // state embd row width, 0 if the context takes no state
    const llama_seq_id n_seq_max;  // max number of sequences
    llama_memory_i * mem;          // memory for position inference
    const llama_token n_vocab;     // max token ID that we accept
    const size_t n_pos_per_embd;

    // actual embd row width of this batch, set by the first set_token_embd()
    // must be either n_embd_inp or n_embd_inp_enc; encode/decode verify it against the graph input
    size_t n_embd = 0;

    struct token {
        llama_token  id = LLAMA_TOKEN_NULL;
        bool         has_embd = false; // whether embd_off is set
        size_t       embd_off = 0; // index offset in the embd array
        bool         has_state = false; // whether state_off is set
        size_t       state_off = 0; // index offset in the state array
        bool         output = false; // TODO: have dedicated output flags
        int32_t      decision_order = 0; // see llama_batch_ext_set_decision_order()
        std::unordered_set<llama_seq_id> seq_ids;
        std::array<llama_pos, GGML_MROPE_SECTIONS> pos = {0, 0, 0, 0};
    };
    std::vector<token> tokens;
    std::vector<float> embd;
    std::vector<float> state;

    llama_batch_ext(llama_context * ctx);

    // build without a llama_context, used by tests
    llama_batch_ext(
            size_t n_tokens_max,
            size_t n_embd_inp,
            size_t n_embd_inp_enc,
            llama_seq_id n_seq_max,
            llama_memory_i * mem,
            llama_token n_vocab,
            size_t n_pos_per_embd);

    void clear();

    // add an entry with an undefined position
    // the caller must set it explicitly via set_token_pos()
    int32_t add_token(llama_seq_id seq_id);

    bool add_seq(int32_t idx, llama_seq_id seq_id);
    bool set_token_id(int32_t idx, llama_token id);
    bool set_token_embd(int32_t idx, llama_embd embd_in);
    bool set_token_state(int32_t idx, llama_embd state_in);
    bool set_token_pos(int32_t idx, const llama_pos * pos_in);
    bool set_output(int32_t idx, bool output_last);
    bool set_decision_order(int32_t idx, int32_t order);
};

// a helper for sanitizing, fulfilling and splitting a batch
class llama_batch_allocr {
public:
    llama_batch_allocr(uint32_t n_pos_per_embd, bool allow_mixed = false);

    // convert a llama_batch_ext to internal llama_batch and sanitize it
    bool init(
            const llama_batch_ext & batch_inp,
            const llama_vocab & vocab,
            bool output_all);

    const llama_batch & get_batch() const;

    uint32_t get_n_tokens()  const;
    uint32_t get_n_outputs() const;
    uint32_t get_n_used()    const;

    // the array of output indices in the order they were encountered during the ubatch splitting
    std::vector<int32_t> & get_out_ids();

    // min/max positions of each sequence in the current ubatch
    llama_pos seq_pos_min(llama_seq_id seq_id) const;
    llama_pos seq_pos_max(llama_seq_id seq_id) const;

    // call once before splitting the batch to reset the internal state
    void split_reset();

    // simple split, unknown number of sequence sets of unequal lengths
    llama_ubatch split_simple(uint32_t n_ubatch);

    // make ubatches of equal-length sequences sets
    // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
    // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);

    // sequence-set-wise split - each ubatch contains a single sequence-set
    // TODO: extend to support multiple _full_ sequences in the ubatch
    llama_ubatch split_seq(uint32_t n_ubatch);

    // a helper method for creating a well-defined ubatch of tokens
    // TODO: support embeddings if needed in the future
    llama_ubatch ubatch_reserve(uint32_t n_seq_tokens, uint32_t n_seqs);

private:
    void clear();

    // create the next ubatch based on the provided batch indices (idxs) and the number of sequence sets (n_seqs)
    // return llama_ubatch.n_tokens == 0 if the entire batch was consumed
    llama_ubatch ubatch_add(const std::vector<int32_t> & idxs, uint32_t n_seqs, bool equal_seqs);

    // for debugging, start with LLAMA_BATCH_DEBUG=2
    void ubatch_print(const llama_ubatch & ubatch, int debug);

    llama_batch batch;

    // only for debugging purposes
    const llama_vocab * vocab;

    // TODO: this is more of a temporary solution until we have a better way to handle multiple positions per token/embd
    //       ref: https://github.com/ggml-org/llama.cpp/issues/13694#issuecomment-2983871762
    const uint32_t n_pos_per_embd;

    const bool allow_mixed;

    uint32_t n_embd;
    uint32_t n_embd_state;
    uint32_t n_seq_max;
    uint32_t n_outputs;

    std::vector<llama_token>    token_vec;    // owned token IDs built from llama_batch_ext
    std::vector<float>          embd_vec;     // owned embeddings built from llama_batch_ext
    std::vector<int8_t>         is_embd_vec;  // mixed batch only (= 1 if embd, 0 if text token)
    std::vector<float>          state_vec;    // owned state embeddings built from llama_batch_ext, llama_batch has no slot for them
    std::vector<llama_seq_id>   seq_id_data;  // flat storage for seq_id pointers below

    std::vector<llama_pos>      pos;
    std::vector<int32_t>        n_seq_id;
    std::vector<llama_seq_id *> seq_id;
    std::vector<llama_seq_id>   seq_id_unq;
    std::vector<int32_t>        seq_idx;
    std::vector<int8_t>         output;
    std::vector<int32_t>        decision_order; // empty if no entry has one

    using pos_set_t = std::set<llama_pos>;
    using seq_cpl_t = std::vector<bool>;

    // helper flag to quickly determine if there are any coupled sequences in the batch
    bool has_cpl = false;

    std::vector<pos_set_t> seq_pos; // seq_pos[s]: the set of positions in sequence s
    std::vector<seq_cpl_t> seq_cpl; // seq_cpl[s0][s1]: if sequence s0 is coupled to sequence s1

    using idx_vec_t = std::vector<int32_t>;
    using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;

    std::vector<seq_set_t> seq_set; // seq_set[i]: the sequence set of token i

    std::unordered_map<seq_set_t, idx_vec_t> seq_set_map; // the indices at which the sequence set appears

    // batch indices of the output
    std::vector<int32_t> out_ids;

    uint32_t n_used;

    // used[i] indicates if token i has already been used in a previous ubatch
    std::vector<bool> used;

    int debug;
};

// RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext
struct llama_batch_compat {
    llama_batch_ext * batch_ext;

    // n_embd_row is the embd row width of batch_inp, 0 = use the decoder width
    llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0);
    ~llama_batch_compat();

    // fill an existing llama_batch_ext from a llama_batch (old API)
    // note: this is called directly by the tests, skipping llama_context creation
    static void init(llama_batch_ext & batch_ext, const llama_batch & batch_inp, size_t n_embd_row = 0);
};

Xet Storage Details

Size:: 11.6 kB
Xet hash:: 381aece1e1ca342ea3545540daa6662959cdf928bb8458a4c60791773891984f

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#pragma once

	#include "llama.h"

	#include "llama-arch.h"
	#include "llama-cparams.h"

	#include <array>
	#include <vector>
	#include <set>
	#include <bitset>
	#include <memory>
	#include <unordered_map>
	#include <unordered_set>

	// keep this struct lightweight
	struct llama_ubatch {
	bool equal_seqs() const {
	return b_equal_seqs != 0;
	}

	// typical for M-RoPE cases:
	// 0 - sequential position of the tokens/embeddings in the sequence
	// 1 - y position in the image
	// 2 - x position in the image
	// 3 - other
	bool is_pos_2d() const {
	// TODO @ngxson : we may need to check for model arch when more models use >1 positions
	return n_pos >= 3;
	}

	// mixed: type picks token or embd per row, pos has n_pos sections for all rows
	bool is_mixed() const {
	return type != nullptr;
	}

	uint32_t b_equal_seqs; // note: this is a boolean, but we use an int32_t for alignment
	// otherwise address sanitizer complains
	// TODO: whole_seqs for embeddings?

	uint32_t n_tokens; // total tokens (n_seq_tokens * n_seqs)
	uint32_t n_seq_tokens; // tokens per sequence set
	uint32_t n_seqs; // sequence sets in the ubatch
	uint32_t n_seqs_unq; // unique sequence ids in the ubatch
	uint32_t n_pos; // number of position inputs for each token/embedding

	// seq_id_unq: unique sequence ids in the ubatch
	// seq_idx: indices of the unique sequence ids in the ubatch in [0, n_seqs_unq)
	// used for extracting sequence pooled embeddings

	// // size \| idx \| val
	llama_token * token; // [n_tokens] \| i \| id, token
	float * embd; // [n_embd, n_tokens] \| i \| embd
	float * embd_state; // [n_embd_state, n_tokens] \| i \| hidden state carried over from a previous stage (e.g. MTP)
	llama_pos * pos; // [n_tokens*n_pos] \| i \| pos
	int32_t * n_seq_id; // [n_tokens] \| i \| -
	llama_seq_id ** seq_id; // [n_tokens] \| s \| s0, s1, seq_id
	llama_seq_id * seq_id_unq; // [n_seqs_unq] \| s \| seq_id
	int32_t * seq_idx; // [LLAMA_MAX_SEQ] \| - \| seq_idx
	int8_t * output; // [n_tokens] \| i \| -
	int8_t * type; // [n_tokens] \| i \| - (mixed ubatch only, 0 - token, 1 - embd)
	int32_t * decision_order; // [n_tokens], NULL if no entry has one, see llama_batch_ext_set_decision_order()

	struct data_t {
	std::vector<llama_token> token;
	std::vector<float> embd;
	std::vector<float> embd_state;
	std::vector<llama_pos> pos;
	std::vector<int32_t> n_seq_id;
	std::vector<llama_seq_id *> seq_id; // these point into the seq_id_data below
	std::vector<llama_seq_id> seq_id_unq;
	std::vector<int32_t> seq_idx;
	std::vector<int8_t> output;
	std::vector<int8_t> type;
	std::vector<int32_t> batch_idxs; // original batch index for each token
	std::vector<int32_t> decision_order;

	std::vector<llama_seq_id> seq_id_data;
	};

	// the llama_ubatch pointers above point to this data if set. otherwise - point to external non-owning data
	std::shared_ptr<data_t> data;
	};

	// crash if a mixed ubatch reaches code that expects only tokens or only embd
	#define ASSERT_EMBD_OR_TOKEN(ubatch) GGML_ASSERT(!(ubatch).is_mixed() && "mixed token/embd ubatch is not supported here")

	struct llama_hparams;

	// DFlash batches carry the fused target features at the encoder input width (n_embd_inp_enc size).
	// Other batches carry token embeddings (n_embd_inp size).
	size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams);

	// MTP contexts also take the target model's hidden state (n_embd_out size), 0 = no state input
	size_t llama_batch_ext_select_n_embd_state(llama_context_type ctx_type, const llama_hparams & hparams);

	struct llama_batch_ext {
	const size_t n_tokens_max; // max number of tokens that can be stored in the batch
	const size_t n_embd_inp; // decoder embd row width
	const size_t n_embd_inp_enc; // encoder embd row width (e.g. eagle3/dflash extracted features)
	const size_t n_embd_state; // state embd row width, 0 if the context takes no state
	const llama_seq_id n_seq_max; // max number of sequences
	llama_memory_i * mem; // memory for position inference
	const llama_token n_vocab; // max token ID that we accept
	const size_t n_pos_per_embd;

	// actual embd row width of this batch, set by the first set_token_embd()
	// must be either n_embd_inp or n_embd_inp_enc; encode/decode verify it against the graph input
	size_t n_embd = 0;

	struct token {
	llama_token id = LLAMA_TOKEN_NULL;
	bool has_embd = false; // whether embd_off is set
	size_t embd_off = 0; // index offset in the embd array
	bool has_state = false; // whether state_off is set
	size_t state_off = 0; // index offset in the state array
	bool output = false; // TODO: have dedicated output flags
	int32_t decision_order = 0; // see llama_batch_ext_set_decision_order()
	std::unordered_set<llama_seq_id> seq_ids;
	std::array<llama_pos, GGML_MROPE_SECTIONS> pos = {0, 0, 0, 0};
	};
	std::vector<token> tokens;
	std::vector<float> embd;
	std::vector<float> state;

	llama_batch_ext(llama_context * ctx);

	// build without a llama_context, used by tests
	llama_batch_ext(
	size_t n_tokens_max,
	size_t n_embd_inp,
	size_t n_embd_inp_enc,
	llama_seq_id n_seq_max,
	llama_memory_i * mem,
	llama_token n_vocab,
	size_t n_pos_per_embd);

	void clear();

	// add an entry with an undefined position
	// the caller must set it explicitly via set_token_pos()
	int32_t add_token(llama_seq_id seq_id);

	bool add_seq(int32_t idx, llama_seq_id seq_id);
	bool set_token_id(int32_t idx, llama_token id);
	bool set_token_embd(int32_t idx, llama_embd embd_in);
	bool set_token_state(int32_t idx, llama_embd state_in);
	bool set_token_pos(int32_t idx, const llama_pos * pos_in);
	bool set_output(int32_t idx, bool output_last);
	bool set_decision_order(int32_t idx, int32_t order);
	};

	// a helper for sanitizing, fulfilling and splitting a batch
	class llama_batch_allocr {
	public:
	llama_batch_allocr(uint32_t n_pos_per_embd, bool allow_mixed = false);

	// convert a llama_batch_ext to internal llama_batch and sanitize it
	bool init(
	const llama_batch_ext & batch_inp,
	const llama_vocab & vocab,
	bool output_all);

	const llama_batch & get_batch() const;

	uint32_t get_n_tokens() const;
	uint32_t get_n_outputs() const;
	uint32_t get_n_used() const;

	// the array of output indices in the order they were encountered during the ubatch splitting
	std::vector<int32_t> & get_out_ids();

	// min/max positions of each sequence in the current ubatch
	llama_pos seq_pos_min(llama_seq_id seq_id) const;
	llama_pos seq_pos_max(llama_seq_id seq_id) const;

	// call once before splitting the batch to reset the internal state
	void split_reset();

	// simple split, unknown number of sequence sets of unequal lengths
	llama_ubatch split_simple(uint32_t n_ubatch);

	// make ubatches of equal-length sequences sets
	// if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
	// n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
	llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);

	// sequence-set-wise split - each ubatch contains a single sequence-set
	// TODO: extend to support multiple _full_ sequences in the ubatch
	llama_ubatch split_seq(uint32_t n_ubatch);

	// a helper method for creating a well-defined ubatch of tokens
	// TODO: support embeddings if needed in the future
	llama_ubatch ubatch_reserve(uint32_t n_seq_tokens, uint32_t n_seqs);

	private:
	void clear();

	// create the next ubatch based on the provided batch indices (idxs) and the number of sequence sets (n_seqs)
	// return llama_ubatch.n_tokens == 0 if the entire batch was consumed
	llama_ubatch ubatch_add(const std::vector<int32_t> & idxs, uint32_t n_seqs, bool equal_seqs);

	// for debugging, start with LLAMA_BATCH_DEBUG=2
	void ubatch_print(const llama_ubatch & ubatch, int debug);

	llama_batch batch;

	// only for debugging purposes
	const llama_vocab * vocab;

	// TODO: this is more of a temporary solution until we have a better way to handle multiple positions per token/embd
	// ref: https://github.com/ggml-org/llama.cpp/issues/13694#issuecomment-2983871762
	const uint32_t n_pos_per_embd;

	const bool allow_mixed;

	uint32_t n_embd;
	uint32_t n_embd_state;
	uint32_t n_seq_max;
	uint32_t n_outputs;

	std::vector<llama_token> token_vec; // owned token IDs built from llama_batch_ext
	std::vector<float> embd_vec; // owned embeddings built from llama_batch_ext
	std::vector<int8_t> is_embd_vec; // mixed batch only (= 1 if embd, 0 if text token)
	std::vector<float> state_vec; // owned state embeddings built from llama_batch_ext, llama_batch has no slot for them
	std::vector<llama_seq_id> seq_id_data; // flat storage for seq_id pointers below

	std::vector<llama_pos> pos;
	std::vector<int32_t> n_seq_id;
	std::vector<llama_seq_id *> seq_id;
	std::vector<llama_seq_id> seq_id_unq;
	std::vector<int32_t> seq_idx;
	std::vector<int8_t> output;
	std::vector<int32_t> decision_order; // empty if no entry has one

	using pos_set_t = std::set<llama_pos>;
	using seq_cpl_t = std::vector<bool>;

	// helper flag to quickly determine if there are any coupled sequences in the batch
	bool has_cpl = false;

	std::vector<pos_set_t> seq_pos; // seq_pos[s]: the set of positions in sequence s
	std::vector<seq_cpl_t> seq_cpl; // seq_cpl[s0][s1]: if sequence s0 is coupled to sequence s1

	using idx_vec_t = std::vector<int32_t>;
	using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;

	std::vector<seq_set_t> seq_set; // seq_set[i]: the sequence set of token i

	std::unordered_map<seq_set_t, idx_vec_t> seq_set_map; // the indices at which the sequence set appears

	// batch indices of the output
	std::vector<int32_t> out_ids;

	uint32_t n_used;

	// used[i] indicates if token i has already been used in a previous ubatch
	std::vector<bool> used;

	int debug;
	};

	// RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext
	struct llama_batch_compat {
	llama_batch_ext * batch_ext;

	// n_embd_row is the embd row width of batch_inp, 0 = use the decoder width
	llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0);
	~llama_batch_compat();

	// fill an existing llama_batch_ext from a llama_batch (old API)
	// note: this is called directly by the tests, skipping llama_context creation
	static void init(llama_batch_ext & batch_ext, const llama_batch & batch_inp, size_t n_embd_row = 0);
	};