model : support LongCat-Flash-Lite (ngram embeddings)

[ngxson]

Target support: https://huggingface.co/meituan-longcat/LongCat-Flash-Lite

NOTE: I'm having quite limited time recently, but quite interested by the idea of ngram embedding. Pushing this PR in a rough state mostly for discussions.

The most complex part of the model is to integrate the idea of "ngram cache" inside llama.cpp. That means new tokens can "look back" to see N tokens that were added in the past.

For example: when I given ngram_n = 3 and given 5 input tokens t0, t1, t2, t3, t4. Each token can "see" these ngrams:

• t0: (none)
• t1: t0
• t2: t1 t0
• ...

This sounds similar to SWA logic but I can't wrap my head around how to reuse the code inside set_input_kq_mask_impl. CC @ggerganov if you have any ideas.

And to store the token ID of each position, llama_kv_cell_ext is extended to store this, quite similar to x/y position for m-rope.

The current implementation does return the ngram, what's left to do is to use this info to calculate the hash (quite simple, just some multiplications and modulo):

token[0] = 128000 : ngram = 0 0 0 0
token[1] = 13347 : ngram = 0 0 0 128000
token[2] = 10 : ngram = 0 0 128000 13347
token[3] = 4925 : ngram = 0 128000 13347 10
token[4] = 674 : ngram = 128000 13347 10 4925
token[5] = 220 : ngram = 13347 10 4925 674
token[6] = 15 : ngram = 10 4925 674 220
token[7] = 220 : ngram = 4925 674 220 15
token[8] = 16 : ngram = 674 220 15 220

[ggerganov]

I think it would be more consistent if you implement a set_input_ngrams(ggml_tensor * dst, const llama_ubatch * ubatch).

In the llama_kv_cells, instead of adding token id to llama_kv_cell_ext, consider adding a map to llama_kv_cells:

// seq_id[s][p] gives the token at position p for sequence s
std::map<llama_pos, llama_token> seq_id[LLAMA_MAX_SEQ];

It will be very similar to the existing seq_pos map and updated in very similar way. It might be a bit tricky to handle cases with multiple tokens in the same position (see the comment for seq_pos), but not very important for now.

Also, might have to extend the pos_set(uint32_t i, llama_pos p) to set the token too:

void pos_set(uint32_t i, llama_pos p, llama_token id);

Once you have this map, it's trivial to construct the ngrams. You can use this in the new set_input_ngrams method.

[ggerganov]

Looking a bit more, its probably best to extend the pos concept a little bit like so:

Currently we have this:

llama.cpp/src/llama-kv-cells.h

Lines 463 to 468 in eef375c

	std::vector<llama_pos> pos;
	// stores extra info per cell
	std::vector<llama_kv_cell_ext> ext;

New version:

struct llama_kv_cell {
    llama_pos pos;
    llama_token id;
};

...
    // cell data (replaced the old `pos` vector)
    std::vector<llama_kv_cell> data;

    // stores extra info per cell
    std::vector<llama_kv_cell_ext> ext;

And then pos_set() becomes:

    void pos_set(uint32_t i, llama_kv_cell c);

// or

    void cell_set(uint32_t i, llama_kv_cell c);

etc.

[ngxson]

Thanks! Unfortunately I just spent the whole working day to realize that it's currently not possible to implement the non-ngram version with existing GGML ops: #19182

The ngram model is just the normal one + extra embeddings, so currently I'm hitting a roadblock before I can move on with this PR. But I'll come back to this if any other models also implement ngram embeddings.

[ggerganov]

For my understanding, the normal (non-ngram) model is discussed in #19182 and I understand the difficutly.

But this sentence:

The ngram model is just the normal one + extra embeddings, so currently I'm hitting a roadblock before I can move on with this PR.

What do you mean here by roadblock for the ngram model?

Edit: ah, nvm. I understand now - it also needs the zero-compute experts.

[ngxson]

I started the PR with the assumption that the non-ngram version is already supported in llama.cpp, but turns out it was not the case 😅

[PmNz8]

I know it might be out of scope for implementing support of this model, but could you consider adding the possibility to keep the engrams stored on disk during inference? From what I have read, the engrams are a bit like the Lookup Experts / Lookup Key-Value architecture/approach. For average home system without a lot of CPU-RAM to keep them in memory, reading the engram data from SSD could also be viable?

I might be totally wrong, so correct me.

[jukofyork]

I know it might be out of scope for implementing support of this model, but could you consider adding the possibility to keep the engrams stored on disk during inference? From what I have read, the engrams are a bit like the Lookup Experts / Lookup Key-Value architecture/approach. For average home system without a lot of CPU-RAM to keep them in memory, reading the engram data from SSD could also be viable?

I might be totally wrong, so correct me.

I haven't read the paper but based on what @ngxson said it uses a "hash-like" operation (ie: https://en.wikipedia.org/wiki/Feature_hashing) which will spread the embeddings throughout the store. If so, then this will effectively end up being a "random access" pattern that won't be easy to prefetch off of disk.

Without using the "hash trick" the Cartesian product of the vocab would quickly get out of hand (and also training would be very hard / impossible due to the extreme sparsity), but this would actually be much easier to lay out on disk for prefetching.

[ngxson]

I think a normal embedding (not ngram embedding) is also categorized as "random access". If we know which embeddings to be prefetched, we should have known it's correspond to which token, and to know which token, we need to firstly... predict it?

[jukofyork]

I think a normal embedding (not ngram embedding) is also categorized as "random access". If we know which embeddings to be prefetched, we should have known it's correspond to which token, and to know which token, we need to firstly... predict it?

Sorry, I ninja edited my reply whist you were typing this. It's quite a well known trick used for extremely sparse embeddings and IIRC was first used with "Field-aware" factorization machines.

The problem with the Cartesian product is that some of the combinations never appear in training (eg: due to being essentially ransom sequences of tokens here or random interactions for FFM), so there is a good reason for using the "hashing trick" to ensure the embedding space is activated approximately equally often...

If instead it used Cartesian product n-grams, we could definitely lay it out on disk (as say an complete m-ary tree) so that only n_vocab tokens would need to be accessible like standard embeddings (even though the full size would be n_vocab^n).

[jukofyork]

I think a normal embedding (not ngram embedding) is also categorized as "random access". If we know which embeddings to be prefetched, we should have known it's correspond to which token, and to know which token, we need to firstly... predict it?

The difference here is we do actually know the previous n-1 tokens that came before, so with say a 128K n_vocab and 8192 hidden_dim we would "only" need to prefetch 500MB-1GB of embeddings (eg: for 4-8bit quant respectively) if it were layed out "nicely" on disk.

But if you hash the embeddings, we can't really do this as even knowing the previous n-1 tokens; the hashing will essentially make all locations equally likely if the hash function is doing it's job properly.

[ngxson]

The difference here is we do actually know the previous n-1 tokens that came before, so with say a 128K n_vocab and 8192 hidden_dim we would "only" need to prefetch 500MB-1GB of embeddings

If I calculate it correctly, you mean we are loading the whole embedding tensor of dimension [8192, 128000] onto memory. Unless I misunderstood something, this is a simple "fully loaded", not "prefetched". I expect "prefetch" to be "partially load or cache something (in advanced) based on an usage pattern"

The math that I did: 128000 * 8192 * 4 bits = 4194304000 bits = 524288000 bytes = 500MB

But if you hash the embeddings, we can't really do this as even knowing the previous n-1 tokens; the hashing will essentially make all locations equally likely if the hash function is doing it's job properly.

If you know the n and n-1 tokens, the hash function is simply a pure function: f(token[n-1], token[n]) = ngram_index, so we do know for sure what is ngram index based on the past tokens.

I think by saying that you're assuming the entropy of ngram_index to be higher than entropy of input tokens. But I don't think entropy of input tokens is something predictable.

[jukofyork]

The difference here is we do actually know the previous n-1 tokens that came before, so with say a 128K n_vocab and 8192 hidden_dim we would "only" need to prefetch 500MB-1GB of embeddings

If I calculate it correctly, you mean we are loading the whole embedding tensor of dimension [8192, 128000] onto memory. Unless I misunderstood something, this is a simple "fully loaded", not "prefetched". I expect "prefetch" to be "partially load or cache something (in advanced) based on an usage pattern"

Let's say we use n=3 and n_vocab=128k and were to use the Cartesian product instead of a hash:

As soon as we finish generating the 1st token we can narrow this down to (2^17)^2 possible n-gram embeddings (that can be used for the 4th token), then after we generate the 2nd token 2^17 possible embeddings, then after we generate the 3rd token only 1 possible embedding.

But whilst the 3rd token is being generated we could (in theory) "prefetch" the 2^17 possible embeddings that can be used for the 4th token into RAM and then as soon as we start to generate the 4th token we could (in theory) avoid the added latency of disk access and get the same token generation speed as if the full set of (2^17)^3 n-gram embeddings were stored in RAM.

This isn't a great example as the numbers are all too massive anyway, and as as I linked above; there is a very good reason not to use the Cartesian product anyway! But to compare with hashing:

If you know the n and n-1 tokens, the hash function is simply a pure function: f(token[n-1], token[n]) = ngram_index, so we do know for sure what is ngram index based on the past tokens.

I think by saying that you're assuming the entropy of ngram_index to be higher than entropy of input tokens. But I don't think entropy of input tokens is something predictable.

Let's say we use n=3 and n_vocab=128k and a hash of size n_hash:

As soon as we finish generating the 1st token we can't narrow this down and all of the n_hash n-gram embeddings (that can be used for the 4th token) are still valid, the same after the 2nd token, and only after we generate the 3rd token do we know which exact embedding of the posible n_hash n-gram embeddings we are going to use (and by then we are just about to start generating the 4th token anyway, so no possibility to "prefetch" it!).

[jukofyork]

It's not a great example as the numbers are all way too big using this vocab this size and full hidden state dimensions, but with a smaller vocab like 32k, the use of a projection matrix (to upscale to the hidden dimension), and stored on a HDD instead of SSD, it would potentially be worthwhile doing this sort of prefetching (to answer the other guy's question), but only if you can narrow down the possibilities after each previous token (which I don't think we can here, so all this is irrelevant anyway).

[PmNz8]

@jukofyork so if I understand correctly, if it is not possible to narrow down the required engrams, the software would need to read the whole (or large part) of the engram data (like in tens of GBs)?
I asked about it because for some size-ranges of traditional MoE models I got acceptable tps even if the experts size > RAM size, as the most used experts were cached by system in RAM (with occasional hiccpus if experts weight were retrieved from disk).
If the engrams would display similar behavior, then for system with like 16-32 GB of RAM it also could go like this: active quantized layers in GPU, some of ngrams cached by system in RAM and the rest of them on disk, accessed less often (zipfs law). But if the access is truly random and would need scanning the (almost)/whole ngram data then SSD would be to slow when ngrams are tens or hundreds of GBs.

But this is just my layman understanding, so probably wrong. Thank you for the attention you put into my question - also @ngxson !

[kroaton]

https://huggingface.co/meituan-longcat/LongCat-Flash-Lite-Sparse new model using this architecture got released

[engrtipusultan]

https://huggingface.co/meituan-longcat/LongCat-Flash-Lite-Sparse new model using this architecture got released
@ngxson
Hi is the support of mentioned model planned in pipeline?