sampling : add support for GPU sampling (wip)

This is a work in progress to add support for GPU sampling.

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

For example, the GPU 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 GPU 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 GPU sampling works in a similar manner to how pooling
works, it is a function that is called by build_graph:
```c++
    // add GPU sampling layers (if any)
    llm->build_sampling(*this, params);
```

GPU samplers can be configured by creating sampler chains, where each
sampler chain is associated with a specific sequence id:
```c++
    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_gpu_init_greedy());
    std::vector<llama_sampler_seq_config> sampler_configs = {
        { 0, gpu_sampler_chain }
    };
```
The struct is defined as:
```c++
    struct llama_sampler_seq_config {
        llama_seq_id           seq_id;
        struct llama_sampler * sampler;
    };
```

These sampler configs are then passed as context params:
```c++
        llama_context_params cparams = llama_context_default_params();
        cparams.samplers = sampler_configs.data();
        cparams.n_samplers = sampler_configs.size();
```

When the graph is built, the configured sampler's _apply function is
called which allows them to add operations/nodes to the computation
graph.

This enables the sampling to happen fully, or partially on the GPU. 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:
```c++
    llama_token id = llama_get_sampled_token_ith(test_ctx.ctx, index);
```

Is it also possible to run a GPU 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.

Building and running the tests:

Download a model for testing:
```console
$ cd models && wget https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf
```
Building the test:
```console
$ cmake --build build --target test-gpu-sampling -j8
```
Runing all tests:
```console
$ env LLAMACPP_TEST_MODELFILE=../models/stories15M-q4_0.gguf \
    ctest --test-dir build -R '^test-gpu-sampling$' -V
```

The following individual tests are available:
```console
$ ctest --test-dir build -N -R test-gpu-sampling-
  Test 35: test-gpu-sampling-greedy
  Test 36: test-gpu-sampling-temp
  Test 37: test-gpu-sampling-softmax
  Test 38: test-gpu-sampling-top_k
  Test 39: test-gpu-sampling-top_p
  Test 40: test-gpu-sampling-mul_seq

Total Tests: 6
```
These can be run individually, for example:
```console
$ env LLAMACPP_TEST_MODELFILE=../models/stories15M-q4_0.gguf \
    ctest --test-dir build -R 'test-gpu-sampling-temp' -V
```

TODO:

- [ ] Allow GPU samplers to pre-allocate state tensors
- [ ] Integrate GPU samplers with llama-server
- [ ] Implement true top-p sampler on GPU
- [ ] Add missing GPU samplers (e.g. typical, mirostat, etc)

Author: danbev

common/llguidance.cpp

@@ -106,12 +106,14 @@ static void llama_sampler_llg_free(llama_sampler * smpl) {
 }
 
 static llama_sampler_i llama_sampler_llg_i = {
-    /* .name   = */ llama_sampler_llg_name,
-    /* .accept = */ llama_sampler_llg_accept_impl,
-    /* .apply  = */ llama_sampler_llg_apply,
-    /* .reset  = */ llama_sampler_llg_reset,
-    /* .clone  = */ llama_sampler_llg_clone,
-    /* .free   = */ llama_sampler_llg_free,
+    /* .name        = */ llama_sampler_llg_name,
+    /* .accept      = */ llama_sampler_llg_accept_impl,
+    /* .apply       = */ llama_sampler_llg_apply,
+    /* .reset       = */ llama_sampler_llg_reset,
+    /* .clone       = */ llama_sampler_llg_clone,
+    /* .free        = */ llama_sampler_llg_free,
+    /* .apply_ggml  = */ NULL,
+    /* .accept_ggml = */ NULL,
 };
 
 static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len,

include/llama.h

@@ -210,6 +210,13 @@ extern "C" {
         bool sorted;      // note: do not assume the data is sorted - always check this flag
     } llama_token_data_array;
 
+    struct llama_sampler_ggml_data {
+        struct ggml_tensor * logits;        // [n_vocab]          - GGML_TYPE_F32
+        struct ggml_tensor * probs;         // [n_vocab, n_vocab] - GGML_TYPE_F32
+        struct ggml_tensor * sampled_token; // [1,       n_vocab] - GGML_TYPE_I32
+        struct ggml_tensor * filtered_ids;  // [k]                - GGML_TYPE_I32
+    };
+
     typedef bool (*llama_progress_callback)(float progress, void * user_data);
 
     // Input data for llama_encode/llama_decode
@@ -300,6 +307,11 @@ extern "C" {
         bool no_host;         // bypass host buffer allowing extra buffers to be used
     };
 
+    struct llama_sampler_seq_config {
+        llama_seq_id           seq_id;
+        struct llama_sampler * sampler;
+    };
+
     // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations
     //       https://github.com/ggml-org/llama.cpp/pull/7544
     struct llama_context_params {
@@ -348,6 +360,10 @@ extern "C" {
         bool kv_unified;  // use a unified buffer across the input sequences when computing the attention
                           // try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix
                           // ref: https://github.com/ggml-org/llama.cpp/pull/14363
+
+        // GPU sampler chain configuration
+        struct llama_sampler_seq_config * samplers;
+        size_t                            n_samplers;
     };
 
     // model quantization parameters
@@ -948,6 +964,29 @@ extern "C" {
     // otherwise: float[n_embd] (1-dimensional)
     LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id);
 
+    // Get the GPU sampled token for the ith token.
+    // Returns LLAMA_TOKEN_NULL if no token was sampled.
+    LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i);
+
+    // Get the GPU sampled probabilites for the ith token
+    // The index matches llama_get_sampled_token_ith().
+    // Returns NULL if no probabilites were generated.
+    LLAMA_API float * llama_get_sampled_probs_ith(struct llama_context * ctx, int32_t i);
+
+    // Get the GPU sampled logits for the ith token
+    // Returns NULL if no logits were sampled.
+    LLAMA_API float * llama_get_sampled_logits_ith(struct llama_context * ctx, int32_t i);
+
+    // Get the GPU sampled token ids associated with the sampled logits for the ith token
+    // Returns NULL if no logits were sampled.
+    LLAMA_API llama_token * llama_get_sampled_token_ids_ith(struct llama_context * ctx, int32_t i);
+
+    // Get the number of GPU sampled logits for the ith token.
+    LLAMA_API uint32_t llama_get_sampled_logits_count_ith(struct llama_context * ctx, int32_t i);
+
+    // Get the number of GPU sampled probabilites for the ith token.
+    LLAMA_API uint32_t llama_get_sampled_probs_count_ith(struct llama_context * ctx, int32_t i);
+
     //
     // Vocab
     //
@@ -1133,6 +1172,17 @@ extern "C" {
         struct llama_sampler * (*clone) (const struct llama_sampler * smpl);                                 // can be NULL if ctx is NULL
         void                   (*free)  (      struct llama_sampler * smpl);                                 // can be NULL if ctx is NULL
 
+        void                   (*apply_ggml)(  struct llama_sampler           * smpl,
+                                               struct ggml_context            * ctx,
+                                               struct  ggml_cgraph            * gf,
+                                               struct llama_sampler_ggml_data * ggml_data);
+
+        void                   (*accept_ggml)( struct llama_sampler * smpl,
+                                               struct ggml_context  * ctx,
+                                               struct ggml_cgraph   * gf,
+                                               struct ggml_tensor   * selected_token);
+
+
         // TODO: API for internal libllama usage for appending the sampling to an existing ggml_cgraph
         //void (*apply_ggml) (struct llama_sampler * smpl, ...);
     };
@@ -1150,7 +1200,16 @@ extern "C" {
     LLAMA_API void                   llama_sampler_reset (      struct llama_sampler * smpl);
     LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl);
     // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add)
-    LLAMA_API void                   llama_sampler_free  (      struct llama_sampler * smpl);
+    LLAMA_API void                   llama_sampler_free  (      struct llama_sampler           * smpl);
+    LLAMA_API void                   llama_sampler_apply_ggml(  struct llama_sampler           * smpl,
+                                                                struct ggml_context            * ctx,
+                                                                struct ggml_cgraph             * gf,
+                                                                struct llama_sampler_ggml_data * ggml_data);
+
+    LLAMA_API void                   llama_sampler_accept_ggml( struct llama_sampler * smpl,
+                                                                struct ggml_context  * ctx,
+                                                                struct  ggml_cgraph  * gf,
+                                                                struct ggml_tensor   * selected_token);
 
     // llama_sampler_chain
     // a type of llama_sampler that can chain multiple samplers one after another
@@ -1164,6 +1223,7 @@ extern "C" {
 
     // after removing a sampler, the chain will no longer own it, and it will not be freed when the chain is freed
     LLAMA_API struct llama_sampler * llama_sampler_chain_remove(   struct llama_sampler * chain, int32_t i);
+    LLAMA_API uint64_t               llama_sampler_chain_get_version(const struct llama_sampler * chain);
 
     // available samplers:
 

src/CMakeLists.txt

@@ -31,6 +31,7 @@ add_library(llama
             llama-model.cpp
             llama-quant.cpp
             llama-sampling.cpp
+            llama-gpu-sampling.cpp
             llama-vocab.cpp
             unicode-data.cpp
             unicode.cpp