char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Shakhnazar Sailaukanexamples: refactor diffusion generation (#22590)d8794eecd

char/aliceai-t5-moe
2.5 KiB57 linesraw
1#pragma once
2
3#include "llama.h"
4
5#include <cstdint>
6
7enum diffusion_algorithm {
8    DIFFUSION_ALGORITHM_ORIGIN           = 0,
9    DIFFUSION_ALGORITHM_ENTROPY_BASED    = 1,
10    DIFFUSION_ALGORITHM_MARGIN_BASED     = 2,
11    DIFFUSION_ALGORITHM_RANDOM           = 3,
12    DIFFUSION_ALGORITHM_CONFIDENCE_BASED = 4,
13};
14
15// Unified transfer scheduling methods
16enum diffusion_transfer_schedule {
17    DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED = 0,  // Dream-style: (1.0 - s/t) * remaining
18    DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED    = 1,  // LLaDA-style: process in blocks with get_num_transfer_tokens
19};
20
21typedef bool (*diffusion_step_callback_t)(int32_t             step,
22                                          int32_t             total_steps,
23                                          const llama_token * tokens,
24                                          int32_t             n_tokens,
25                                          void *              user_data);
26
27struct diffusion_params {
28    int32_t                   steps                   = 0;
29    float                     temperature             = 0;
30    llama_token               mask_token_id           = LLAMA_TOKEN_NULL;
31    diffusion_step_callback_t step_callback           = nullptr;
32    void *                    step_callback_user_data = nullptr;
33    int32_t                   seed                    = 0;
34    bool                      visual_mode             = false;
35    bool                      shift_logits            = false;  // Shift logits by -1 after decode
36
37    float   top_p = 0.;
38    int32_t top_k = 0.;
39
40    diffusion_algorithm         algorithm = DIFFUSION_ALGORITHM_CONFIDENCE_BASED;
41    diffusion_transfer_schedule schedule  = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED;
42
43    float   cfg_scale        = 0.;     // Config scale for classifier-free guidance
44    float   eps              = 0.;     // Timestep scheduling
45    int32_t block_length     = 0;      // Block size (for block scheduling)
46    float   alg_temp         = 0;      // algorithm temperature (0.0 = deterministic)
47    bool    add_gumbel_noise = false;  // Add gumbel noise to the logits if temp > 0.0
48
49    int32_t max_length = 0;            // Maximum sequence length
50};
51
52void diffusion_generate(llama_context *          ctx,
53                        const llama_token *      input_tokens,
54                        llama_token *            output_tokens,
55                        int32_t                  n_input,
56                        const diffusion_params & params,
57                        int32_t &                n_generated);