llama : refactor sampling v2

[ggerganov]

alt: #8643

ref: #5214

Overview

• Replace llama_sampling_ and llama_grammar_ with new llama_sampler_
• Overhaul common: struct llama_sampling_context in common: struct gpt_sampler
• Support for user-defined samplers via struct llama_sampler_i interface

API Changes

• Add struct llama_sampler and struct llama_sampler_i
• Add llama_sampler_ API
• Add llama_sampler_chain_ API for chaining multiple samplers
• Remove LLAMA_API_INTERNAL
• Remove Classifier-Free Guidance related stuff
• Remove Prompt Penalty support
• Add llama_perf_ API and remove old llama_print_timings and llama_reset_timings

Implementation details

• Move common/grammar-parser in src/llama-grammar
• The llama_context no longer comes with a built-in sampling context. The user code is responsible for creating, using, saving and loading the llama_sampler objects as needed. As a consequence, the llama_state no longer serializes the RNG
• The grammar code has been refactored, hopefully it is a bit easier to read. No functional changes.
• The samplers implemented in llama-sampling.cpp can be used as examples for implementing custom samplers in user code

Example

Comparison of user sampling code before and after:

• before

// decoding loop:
auto   n_vocab = llama_n_vocab(model);
auto * logits  = llama_get_logits_ith(ctx, i_batch[i]);

std::vector<llama_token_data> candidates;
candidates.reserve(n_vocab);

for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
    candidates.emplace_back(llama_token_data{ token_id, logits[token_id], 0.0f });
}

llama_token_data_array candidates_p = { candidates.data(), candidates.size(), false };

llama_sample_top_k(ctx, &candidates_p, top_k, 1);
llama_sample_top_p(ctx, &candidates_p, top_p, 1);
llama_sample_temp (ctx, &candidates_p, temp);

const llama_token new_token_id = llama_sample_token(ctx, &candidates_p);

• after

// prepare the sampling chain at the start
auto sparams = llama_sampler_chain_default_params();

llama_sampler * smpl = llama_sampler_chain_init(sparams);

llama_sampler_chain_add(smpl, llama_sampler_init_top_k(top_k));
llama_sampler_chain_add(smpl, llama_sampler_init_top_p(top_p, 1));
llama_sampler_chain_add(smpl, llama_sampler_init_temp (temp));
llama_sampler_chain_add(smpl, llama_sampler_init_dist (seed));

...

// decoding loop:
const llama_token new_token_id = llama_sampler_sample(smpl, ctx, i_batch[i]);

Future plan and ideas

• Attach samplers to the decoding graphs for better performance (e.g. llama_decode_with_sampler())
• Extend llama_decode_ API to support multiple decoding runs (e.g. llama_decode_n())
• Existing samplers implementation in llama-sampling.cpp could be split into separate source files
• Expose struct llama_vocab through the public API and change calls that currently use struct llama_model to use it when appropriate
• Deduplicate the ring_buffer code by implementing ggml_ring_buffer for fixed-size objects
• Measure and report the performance of the grammar
• Simplify llama_token_data llama : refactor sampling v2 #9294 (review)

---

• I have read the contributing guidelines
• Self-reported review complexity:
  • Low
  • Medium
  • High

[ExtReMLapin]

Not sure it's the right place or time to talk about this but in another issue a guy had the idea of, "if the grammar says the character/word can only be "xxx" and nothing else, don't bother asking the LLM what to say for the X next tokens".

As there is a refactoring going on, maybe it's the right time to implement it.

[ggerganov]

It's not in the scope of this change, but also the grammar never fits only one token. For example, all three tokens "x", "xx" and "xxx" would fit the grammar in that case. One way would be to use the longest token. Another way would be to tokenize "xxx" and use the resulting tokens. Not sure

[ggerganov]

This is getting close to ready. Later today will add detailed description of the changes, some comments in the code and do a bit more testing.

@slaren PTAL - any comments and suggestions are welcome.

[slaren]

I don't know what's the history of the min_keep parameter in all of these samplers, from what I can tell parameter is not used in the examples except by the server, but it seems very suspect to me.

Edit: looks like it has been there since the beginning (#1126), and there was never any explanation of why it is needed.

[ggerganov]

Removed min_keep from the top_k sampler as it didn't make sense. For the p-based samplers, I think it makes sense to guarantee minimum number of candidate results, regardless of the p value.

[ggerganov]

Thanks for the review. I got side tracked a bit with a bug in the speculative example (should be fixed now). Will apply the review tomorrow and prepare this for merging.

[clarismiranda]

Hello! Why does llama_grammar_accept return the stacks? It was previously passed by reference

[ggerganov]

Thanks for noticing. I changed it because I thought it improves the signature of the function, but I missed that it would lead to extra memory allocations. So I restored the original signature. f9762c6

Edit: though after one more look, I think it does not matter since we move stacks_new either way.

[slaren]

Looks good.

In the future we should probably simplify llama_token_data (or remove it entirely) to keep only one value per token and add a flag to llama_token_data_array to indicate whether the current values are probabilities (ie. normalized to sum 1) or not, so that samplers that can only operate with probabilities can know if they need to call softmax. Having two values per token is very confusing to me because some samplers operate on one or other, and this can lead to situations where a sampler modifies the probabilities, and the next one calls softmax which discards all the changes to the probabilities and computes them from the logits again. I cannot tell if there are already situations like that which end with some samplers that operate on probabilities being ignored.

[ggerganov]

Compiler may not though any error or warning because values set to llama_sampling_sample can also be cast to new values required by gpt_sampler_sample

Hm, it will generate an error. Don't think there is an issue.

[zhaoyinglia]

@ggerganov Hi, I'm a bit confused. Why was Classifier-Free Guidance removed? Are there any issues with it?

[ggerganov]

The implementation does not need to be part of libllama because it works on the logits and is trivial to implement in user code. I removed it also from the examples, because it was much simpler this way and refactoring the sampling was with higher priority than keeping CFG functional. We can reintroduce this functionality in one of the examples or in a new dedicated example if there is interest. PRs welcome.