KBaba7/llama.cpp

1#include "speculative.h"2 3#include "log.h"4#include "common.h"5#include "sampling.h"6 7#include <cstring>8 9#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  12810#define SPEC_VOCAB_CHECK_START_TOKEN_ID 511 12struct common_speculative {13    struct llama_context * ctx;14    struct common_sampler * smpl;15 16    llama_batch batch;17    llama_tokens prompt;18};19 20struct common_speculative * common_speculative_init(21        struct llama_context * ctx_dft) {22    auto * result = new common_speculative {23        /* .ctx    = */ ctx_dft,24        /* .smpl   = */ nullptr,25        /* .batch  = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1),26        /* .prompt = */ {},27    };28 29    // TODO: optimize or pass from outside?30#if 031    {32        common_params_sampling params;33        params.no_perf = false;34 35        params.top_k = 40;36        params.top_p = 0.9;37 38        params.samplers = {39            COMMON_SAMPLER_TYPE_TOP_K,40            COMMON_SAMPLER_TYPE_TOP_P,41            COMMON_SAMPLER_TYPE_INFILL,42        };43 44        result->smpl = common_sampler_init(llama_get_model(ctx_dft), params);45    }46#else47    {48        common_params_sampling params;49        params.no_perf = false;50 51        params.top_k = 10;52 53        params.samplers = {54            COMMON_SAMPLER_TYPE_TOP_K,55        };56 57        result->smpl = common_sampler_init(llama_get_model(ctx_dft), params);58    }59#endif60 61    return result;62}63 64void common_speculative_free(struct common_speculative * spec) {65    if (spec == nullptr) {66        return;67    }68 69    common_sampler_free(spec->smpl);70 71    llama_batch_free(spec->batch);72 73    delete spec;74}75 76bool common_speculative_are_compatible(77        const struct llama_context * ctx_tgt,78        const struct llama_context * ctx_dft) {79    const struct llama_model * model_tgt = llama_get_model(ctx_tgt);80    const struct llama_model * model_dft = llama_get_model(ctx_dft);81 82    const struct llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt);83    const struct llama_vocab * vocab_dft = llama_model_get_vocab(model_dft);84 85    const bool vocab_type_tgt = llama_vocab_type(vocab_tgt);86    LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt);87 88    const bool vocab_type_dft = llama_vocab_type(vocab_dft);89    LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft);90 91    if (vocab_type_tgt != vocab_type_dft) {92        LOG_ERR("%s: draft model vocab type must match target model to use speculation but "93                     "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt);94        return false;95    }96 97    if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||98        llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||99        llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) ||100        llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft)) {101        LOG_ERR("%s: draft vocab special tokens must match target vocab to use speculation\n", __func__);102        LOG_ERR("%s: tgt: bos = %d (%d), eos = %d (%d)\n", __func__, llama_vocab_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_tgt), llama_vocab_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_tgt));103        LOG_ERR("%s: dft: bos = %d (%d), eos = %d (%d)\n", __func__, llama_vocab_bos(vocab_dft), llama_vocab_get_add_bos(vocab_dft), llama_vocab_eos(vocab_dft), llama_vocab_get_add_eos(vocab_dft));104        return false;105    }106 107    {108        const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt);109        const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft);110 111        const int vocab_diff = std::abs(n_vocab_tgt - n_vocab_dft);112 113        if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {114            LOG_ERR("%s: draft model vocab must closely match target model to use speculation but "115                         "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",116                    __func__, n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);117            return false;118        }119 120        for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) {121            const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i);122            const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);123            if (std::strcmp(token_text_tgt, token_text_dft) != 0) {124                LOG_ERR("%s: draft vocab vocab must match target vocab to use speculation but "125                             "token %d content differs - target '%s', draft '%s'\n", __func__, i,126                        common_token_to_piece(ctx_tgt, i).c_str(),127                        common_token_to_piece(ctx_dft, i).c_str());128                return false;129            }130        }131    }132 133    return true;134}135 136llama_tokens common_speculative_gen_draft(137        struct common_speculative * spec,138        struct common_speculative_params params,139        const llama_tokens & prompt_tgt,140        llama_token id_last) {141    auto & batch  = spec->batch;142    auto & ctx    = spec->ctx;143    auto & smpl   = spec->smpl;144    auto & prompt = spec->prompt;145 146    int reuse_i = 0;147    int reuse_n = 0;148 149    const int n_ctx = llama_n_ctx(ctx) - params.n_draft;150 151    const int i_start = std::max<int>(0, (int) prompt_tgt.size() - n_ctx);152 153    // reuse as much as possible from the old draft context154    // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt155    for (int i = 0; i < (int) prompt.size(); ++i) {156        int cur = 0;157        while (i_start + cur < (int) prompt_tgt.size() &&158               i       + cur < (int) prompt.size() &&159               prompt_tgt[i_start + cur] == prompt[i + cur]) {160            cur++;161        }162 163        if ((cur >= params.n_reuse || n_ctx >= (int) prompt_tgt.size()) && cur > reuse_n) {164            reuse_i = i;165            reuse_n = cur;166        }167    }168 169    LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt.size());170 171    llama_tokens result;172    result.reserve(params.n_draft);173 174    if (reuse_n == 0) {175        llama_kv_cache_clear(ctx);176 177        prompt.clear();178    } else {179        // this happens when a previous draft has been discarded (for example, due to being too small), but the180        // target model agreed with it. in this case, we simply pass back the previous results to save compute181        if (reuse_i + reuse_n < (int) prompt.size() && prompt[reuse_i + reuse_n] == id_last) {182            for (int i = reuse_i + reuse_n + 1; i < (int) prompt.size(); ++i) {183                result.push_back(prompt[i]);184 185                if (params.n_draft <= (int) result.size()) {186                    break;187                }188            }189 190            return result;191        }192 193        if (reuse_i > 0) {194            llama_kv_cache_seq_rm (ctx, 0, 0, reuse_i);195            llama_kv_cache_seq_add(ctx, 0, reuse_i, -1, -reuse_i);196 197            prompt.erase(prompt.begin(), prompt.begin() + reuse_i);198        }199 200        if (reuse_n < (int) prompt.size()) {201            llama_kv_cache_seq_rm (ctx, 0, reuse_n, -1);202 203            prompt.erase(prompt.begin() + reuse_n, prompt.end());204        }205    }206 207    // prepare a batch to evaluate any new tokens in the prompt208    common_batch_clear(batch);209 210    for (size_t i = i_start + reuse_n; i < prompt_tgt.size(); ++i) {211        //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_tgt[i]);212        common_batch_add(batch, prompt_tgt[i], i - i_start, { 0 }, false);213 214        prompt.push_back(prompt_tgt[i]);215    }216 217    // we should rarely end-up here during normal decoding218    if (batch.n_tokens > 0) {219        //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str());220 221        llama_decode(ctx, batch);222    }223 224    const llama_pos n_past = prompt.size();225 226    LOG_DBG("%s: n_past = %d\n", __func__, n_past);227 228    common_batch_clear(batch);229    common_batch_add  (batch, id_last, n_past, { 0 }, true);230 231    prompt.push_back(id_last);232 233    //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx, prompt).c_str());234 235    llama_decode(ctx, batch);236 237    common_sampler_reset(smpl);238 239    // sample n_draft tokens from the draft model240    for (int i = 0; i < params.n_draft; ++i) {241        common_batch_clear(batch);242 243        common_sampler_sample(smpl, ctx, 0, true);244 245        const auto * cur_p = common_sampler_get_candidates(smpl);246 247        for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {248            LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",249                    k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx, cur_p->data[k].id).c_str());250        }251 252        // add drafted token for each sequence253        const llama_token id = cur_p->data[0].id;254 255        // only collect very high-confidence draft tokens256        if (cur_p->data[0].p < params.p_min) {257            break;258        }259 260        common_sampler_accept(smpl, id, true);261 262        result.push_back(id);263 264        if (params.n_draft <= (int) result.size()) {265            break;266        }267 268        common_batch_add(batch, id, n_past + i + 1, { 0 }, true);269 270        // evaluate the drafted tokens on the draft model271        llama_decode(ctx, batch);272 273        prompt.push_back(id);274    }275 276    return result;277}278