Codeprocastinator/optimized-tinyllama-covalent

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