sampling : add support for backend sampling

[danbev]

This is a work in progress to add support for backend (like GPU) sampling.

The motivation for this feature is to enable sampling to be performed directly on the backend as part of the computation graph being executed, allowing for some or all of the sampling to be done on the backend.

For example, the backend sampler chain might select/sample a token directly in which case only the sampled token needs to be transferred from device memory to host memory.

It is also possible for the backend samplers to perform filtering of the logits, or compute and filter the probability distribution, in which case only the filtered logits or probabilites need to be transferred back to system memory for further processing by CPU samplers.

Currently the backend sampling works in a similar manner to how pooling works, it is a function that is called by build_graph and the sampler operations become part of the models computation graph.

Backend samplers can be configured by creating sampler chains, where each sampler chain is associated with a specific sequence id:

    struct llama_sampler_chain_params params = llama_sampler_chain_default_params();
    struct llama_sampler * chain = llama_sampler_chain_init(params);
    llama_sampler_chain_add(chain, llama_sampler_backend_init_greedy());
    std::vector<llama_sampler_seq_config> sampler_configs = {
        { 0, chain }
    };

The struct is defined as:

    struct llama_sampler_seq_config {
        llama_seq_id           seq_id;
        struct llama_sampler * sampler;
    };

These sampler configs are then passed as context params:

    llama_context_params cparams = llama_context_default_params();
    cparams.samplers = sampler_configs.data();
    cparams.n_samplers = sampler_configs.size();

When the model graph is built the GPU samplers will be called to enable them to add their operations to the graph:

ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
    std::unique_ptr<llm_graph_context> llm;
    ...

    // add backend sampling layers (if any)
    llm->build_sampling(*this, params);

The llama_sampler_i interface as been extended with 4 new methods in the API, and they are currently all named with a _ggml suffix to indicate that they are for backend sampling:

        void                   (*init_ggml)(struct llama_sampler      * smpl,
                                            ggml_backend_buffer_type_t  buft);

        void                   (*set_input_ggml)( struct llama_sampler * smpl,
                                                       ggml_context * ctx,
                                                        ggml_cgraph * gf);

        void                   (*apply_ggml)(  struct llama_sampler * smpl,
                                                       ggml_context * ctx,
                                                        ggml_cgraph * gf,
                                            llama_sampler_ggml_data * ggml_data);

        void                   (*accept_ggml)( struct llama_sampler * smpl,
                                                       ggml_context * ctx,
                                                        ggml_cgraph * gf,
                                               struct ggml_tensor * selected_token);

The init_ggml function allows backend samplers to create input tensors that they might need. The ggml_backend_buffer_type should be used so that the tensors are created using this backend buffer type, which is the same as the output logits backend. This avoids splits in the computation graph that would require data transfer between different backends.

The set_input_ggml function is called after the computation graph has been scheduled but before it is computed. This allows the backend sampler to set any input for the tensors it created in init_ggml.

The apply_ggml function is where the backend sampler adds its operations to the graphs. When the graph is built, the configured sampler's _apply function is called which allows them to add operations/nodes to the computation graph.

The accept_ggml functions allows backend samplers to update their tensor states if needed.

This enables the sampling to happen fully, or partially on the backend. The samplers could sample a single token in which case that is what will be transferred from the device memory to host memory after llama_decode has been called. The sampled token can then be retrieved using:

    llama_token id = llama_get_backend_sampled_token_ith(test_ctx.ctx, index);

Is it also possible to run a backend sampler that only filters the logits and then only the filtered logits are transferred back to the host and the sampling can proceed on the CPU with the normal (CPU) sampler chain. In this case the CPU samplers are configured as usual but they will now operate on already filtered logits.

Similar to the above handling of logits, it is possible for a GPU samplers to compute the full probability distribution and transfer that to the host. And the CPU samplers can then operate on the those probabilities.

Configuration

Backend sampling is enabled using --backend_sampling, and the sampler chain, either explicitly specified using --samplers or the default, is automatically analyzed to determine which samplers can run on the backend. The system finds the longest contiguous chain of backend supported samplers from the start of the sampler sequence.

For example:

• If the chain is top-k -> temperature -> top-p, and both top-k and temperature are backend-supported but top-p is not, then top-k and temperature will run on the backend, while top-p and subsequent samplers run on the CPU.

• If all configured samplers are supported, the final distribution sampling will also happen on the backend, transferring only the sampled token IDs back to the host.

• If the sampler chain starts with an unsupported sampler, and the sampler is active, all sampling runs on the CPU. Note that this is currently the case with the default sampler so to use backend sampling it is required to specify a sampler chain. See below for an example.

llama-cli

Initial support for llama-cli has been added and can be used as follows:

    $ llama-cli -m models/Qwen2.5-VL-3B-Instruct-Q8_0.gguf \
        --prompt 'What is the capital of Sweden?' \
        -n 20 \
        -no-cnv \
        --verbose-prompt \
        -ngl 40 \
        --backend-sampling \
        --samplers 'top_k;temperature'

To enable a partial backend sampling (hybrid sampling), for example running top_k and temperature on the backend and typ_p on the CPU the following sampler chain could be specified:

    $ llama-cli -m models/Qwen2.5-VL-3B-Instruct-Q8_0.gguf \
        --prompt 'What is the capital of Sweden?' \
        -n 20 \
        -no-cnv \
        --verbose-prompt \
        -ngl 40 \
        --backend-sampling \
        --samplers 'top_k;temperature;top_p'

llama-server

GPU sampling can be enabled for llama-server similar to how it was done above for llama-cli

gdb --args ./build-gpu-sampler/bin/llama-server \
      -m models/Qwen2.5-VL-3B-Instruct-Q8_0.gguf \
      --backend-sampling \
      --samplers 'top_k;temperature' \
      --temp 0.8 \
      --top-k 40 \
      -ngl 50 \
      -v

It is then possible to specify send GPU request parameters as follows:

curl --request POST \
    --url http://localhost:8080/completion \
    --header "Content-Type: application/json" \
    --data '{"prompt": "What is the capital of Sweden?","n_predict": 20, "top_k": 40, "backend_sampling": true}'

Building and running the tests

Download a model for testing:

$ cd models && wget https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf

Building the test:

$ cmake --build build --target test-backend-sampler -j8

Runing all tests:

$ env LLAMACPP_TEST_MODELFILE=../models/stories15M-q4_0.gguf \
    ctest --test-dir build -R '^test-backen-sampler$' -V

The following individual tests are available:

$ ctest --test-dir build-gpu-sampler/ -N -R test-backend-sampler-
Internal ctest changing into directory: /home/danbev/work/ai/llama.cpp-debug/build-gpu-sampler
Test project /home/danbev/work/ai/llama.cpp-debug/build-gpu-sampler
  Test #36: test-backend-sampler-greedy
  Test #37: test-backend-sampler-temp
  Test #38: test-backend-sampler-top_k
  Test #39: test-backend-sampler-dist
  Test #40: test-backend-sampler-dist-and-cpu
  Test #41: test-backend-sampler-logit-bias
  Test #42: test-backend-sampler-mul_seq
  Test #43: test-backend-sampler-set-sampler

Total Tests: 8

These can be run individually, for example:

$ env LLAMACPP_TEST_MODELFILE=../models/stories15M-q4_0.gguf \
    ctest --test-dir build -R 'test-backend-sampler-temp' -V

TODO

• Allocate backend sampler tensors on the same backend as the logits (dev_output.dev)
• Allow backend samplers to pre-allocate state tensors
• Integrate backend samplers with llama-cli
• Set/unset backend samplers
• Integrate backend samplers with llama-server
• Add more tests/assertions for the backend samplers to check more cases
• Rename from sampling to sampler.
• Consistent and clearer naming of backend (backend sampling) functions and data types.
• penalties samplers (to figure out/verify how accept_ggml should work) Will be done in a follow up PR.
• Add ggml_cumsum operation to CUDA backend. This operation exists for Metal and CPU already.

Implemented backend samplers

• temp
• logit_bias
• top_k
• greedy
• dist sampler
• min_p

Remaining backend samplers

The list below are the current CPU sampler that exist. All of these might not be appropriate as GPU samplers. These will be implemented separate follow up PRs.

• top_p
• typical
• temp_ext
• xtc
• top_n_sigma
• mirostat/mirostat_v2
• penalties
• dry
• infill

[am17an]

One place this would be useful immediately is the diffusion-cli. I'm happy to test this when it's ready

[ORippler]

Not sure if I have a strong opinion on this but removing hybrid sampling would reduce the complexity a bit I think (basically if we always set --gpu-dist we only have two states (either full gpu sampling or full cpu sampling, and no in-between).

[danbev]

Not sure if I have a strong opinion on this but removing hybrid sampling would reduce the complexity a bit I think (basically if we always set --gpu-dist we only have two states (either full gpu sampling or full cpu sampling, and no in-between).

My thoughts are that I think we should keep the hybrid approach even though it does come with some additional complexity like you say. I think there could be use cases where one might want to perform some sampling like temp/logit_bias/top-k sampling on the device, and then only have a smaller set of logits copied to the host memory, and still enable other CPU samplers, including grammars, to be able to process the logits.

This might turn out to be an incorrect assumption and not something anyone wants to use, but it feels safer to have the ability do hybrid sampling to play it safe.

[ggerganov]

@danbev Let's rebase on latest master to pick up the recent changes.

[ggerganov]

This separation between "backend sampling" and "backend dist" is not really necessary.

I think a more generic approach that does not require to use a separate --backend-dist argument is like this:

• Define a sampler sequence as we normally do on master: A->B->C->D->E. For example, the default chain that we have is:

sampler chain: logits -> logit-bias -> penalties -> dry -> top-n-sigma -> top-k -> typical -> top-p -> min-p -> xtc -> temp-ext -> dist 

• Note that we don't treat the dist sampler specially - it's just a regular sampler
• If the --backend-sampling parameter is not passed, we do normal CPU-based sampling as usual
• If the --backend-sampling parameter is passed, then we iterate over the sampler sequence from the start and find the longest chain of samplers that are supported by the backend. For example, if A, B, C, E are supported by the backend but D is not, then the backend sampling chain would be A->B->C and the remaining CPU chain will be D->E.
• By "supported by the backend" we mean that there is a corresponding llama_sampler_backend_ function to create the backend sampler. It does not mean to necessarily have all the ggml operators implemented by the backend.
• The struct common_sampler can be extended to maintain references to both the backend and the CPU sampling chains. The logic for the construction of the 2 chains will be implemented within common_sampler_init as described in the previous point. It's allowed for the CPU sampling chain to be empty - this means that all the sampling is done on the GPU using the backend chain.

[danbev]

I've added 9e5e09d to address this. The commit contains some notes about this and I've updated this PRs main comment with a configuration section and examples using llama-cli and llama-server.

[ORippler]

Why can't we use the get_rows to return only the values where mask == 1?

[ORippler]

I feel those are outdated as top_k has been implemented now?