echodict/llama.cpp

1// A basic application simulating a server with multiple clients.2// The clients submit requests to the server and they are processed in parallel.3 4#include "arg.h"5#include "common.h"6#include "sampling.h"7#include "log.h"8#include "llama.h"9 10#include <algorithm>11#include <clocale>12#include <cmath>13#include <cstdio>14#include <string>15#include <vector>16#include <ctime>17 18// trim whitespace from the beginning and end of a string19static std::string trim(const std::string & str) {20    size_t start = 0;21    size_t end = str.size();22 23    while (start < end && isspace(str[start])) {24        start += 1;25    }26 27    while (end > start && isspace(str[end - 1])) {28        end -= 1;29    }30 31    return str.substr(start, end - start);32}33 34static std::string k_system =35R"(Transcript of a never ending dialog, where the User interacts with an Assistant.36The Assistant is helpful, kind, honest, good at writing, and never fails to answer the User's requests immediately and with precision.37 38User:39Recommend a nice restaurant in the area.40Assistant:41I recommend the restaurant "The Golden Duck". It is a 5 star restaurant with a great view of the city. The food is delicious and the service is excellent. The prices are reasonable and the portions are generous. The restaurant is located at 123 Main Street, New York, NY 10001. The phone number is (212) 555-1234. The hours are Monday through Friday from 11:00 am to 10:00 pm. The restaurant is closed on Saturdays and Sundays.42User:43Who is Richard Feynman?44Assistant:45Richard Feynman was an American physicist who is best known for his work in quantum mechanics and particle physics. He was awarded the Nobel Prize in Physics in 1965 for his contributions to the development of quantum electrodynamics. He was a popular lecturer and author, and he wrote several books, including "Surely You're Joking, Mr. Feynman!" and "What Do You Care What Other People Think?".46)";47 48static std::vector<std::string> k_questions = {49    "What is the tallest mountain in the world?",50    "Who was the first person to win two Nobel Prizes?",51    "Which country invented paper?",52    "What organ is primarily responsible for pumping blood throughout the body?",53    "Which planet is known for its prominent ring system?",54    "Who directed the movie 'Inception'?",55    "What is the freezing point of water in Fahrenheit?",56    "Which animal is known to have the longest lifespan?",57    "What language has the most native speakers worldwide?",58    "What is the capital city of Canada?",59    "Who is credited with inventing the World Wide Web?",60    "Which metal is liquid at room temperature?",61    "What is the term for an animal that eats both plants and meat?",62    "Who painted 'The Starry Night'?",63    "What gas do humans exhale that plants use for photosynthesis?",64    "What year did World War II end?",65    "Which continent has the most countries?",66    "Who wrote the novel 'Frankenstein'?",67    "What does DNA stand for?",68    "What is the main ingredient in traditional Japanese miso soup?"69};70 71static std::vector<std::string> k_answers = {72    "The tallest mountain in the world is Mount Everest.",73    "Marie Curie was the first person to win two Nobel Prizes.",74    "Paper was invented in China.",75    "The heart is the organ responsible for pumping blood.",76    "Saturn is known for its prominent ring system.",77    "Christopher Nolan directed the movie 'Inception'.",78    "The freezing point of water in Fahrenheit is 32°F.",79    "The bowhead whale is known to have the longest lifespan among mammals.",80    "Mandarin Chinese has the most native speakers in the world.",81    "The capital city of Canada is Ottawa.",82    "Tim Berners-Lee is credited with inventing the World Wide Web.",83    "Mercury is the metal that is liquid at room temperature.",84    "An animal that eats both plants and meat is called an omnivore.",85    "'The Starry Night' was painted by Vincent van Gogh.",86    "Humans exhale carbon dioxide, which plants use in photosynthesis.",87    "World War II ended in 1945.",88    "Africa is the continent with the most countries.",89    "The novel 'Frankenstein' was written by Mary Shelley.",90    "DNA stands for Deoxyribonucleic Acid.",91    "The main ingredient in traditional Japanese miso soup is fermented soybean paste."92};93 94static std::vector<std::string> k_prompts = {95    "What is the meaning of life?",96    "Tell me an interesting fact about llamas.",97    "What is the best way to cook a steak?",98    "Are you familiar with the Special Theory of Relativity and can you explain it to me?",99    "Recommend some interesting books to read.",100    "What is the best way to learn a new language?",101    "How to get a job at Google?",102    "If you could have any superpower, what would it be?",103    "I want to learn how to play the piano. What would be the best way to do it?",104};105 106struct client {107    ~client() {108        if (smpl) {109            common_sampler_free(smpl);110        }111    }112 113    int32_t id = 0;114 115    llama_seq_id seq_id = -1;116 117    llama_token sampled;118 119    int64_t t_start_prompt;120    int64_t t_start_gen;121 122    int32_t n_past    = 0;123    int32_t n_prompt  = 0;124    int32_t n_decoded = 0;125    int32_t i_batch   = -1;126 127    std::string input;128    std::string prompt;129    std::string response;130 131    struct common_sampler * smpl = nullptr;132};133 134static void print_date_time() {135    std::time_t current_time = std::time(nullptr);136    std::tm* local_time = std::localtime(&current_time);137    char buffer[80];138    strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", local_time);139 140    LOG_INF("\n");141    LOG_INF("\033[35mrun parameters as of %s\033[0m\n", buffer);142    LOG_INF("\n");143}144 145// Define a split string function to ...146static std::vector<std::string> split_string(const std::string& input, char delimiter) {147    std::vector<std::string> tokens;148    std::istringstream stream(input);149    std::string token;150    while (std::getline(stream, token, delimiter)) {151        tokens.push_back(token);152    }153    return tokens;154}155 156int main(int argc, char ** argv) {157    std::setlocale(LC_NUMERIC, "C");158 159    srand(1234);160 161    common_params params;162 163    params.n_predict = 128;164    params.n_junk = 1;165 166    common_init();167 168    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PARALLEL)) {169        return 1;170    }171 172    // number of simultaneous "clients" to simulate173    const int32_t n_clients = params.n_parallel;174 175    // dedicate one sequence to the system prompt176    params.n_parallel += 1;177 178    // requests to simulate179    const int32_t n_seq = params.n_sequences;180 181    // insert new requests as soon as the previous one is done182    const bool cont_batching = params.cont_batching;183 184    // is the system prompt shared in the cache185    const bool is_sp_shared = params.is_pp_shared;186 187    // extra text to insert in each client's prompt in order to make it larger188    const int32_t n_junk = std::max(1, params.n_junk);189 190    // signed seed, use negative values to indicate different seeds for the different clients191    const int32_t & sseed = params.sampling.seed;192 193    // init llama.cpp194    llama_backend_init();195    llama_numa_init(params.numa);196 197    // load the target model198    auto llama_init = common_init_from_params(params);199 200    auto * model = llama_init->model();201    auto * ctx   = llama_init->context();202 203    auto * mem = llama_get_memory(ctx);204 205    const llama_vocab * vocab = llama_model_get_vocab(model);206 207    // load the prompts from an external file if there are any208    if (params.prompt.empty()) {209        LOG_INF("\033[32mNo new questions so proceed with build-in defaults.\033[0m\n");210    } else {211        // Output each line of the input params.prompts vector and copy to k_prompts212        int index = 0;213        LOG_INF("\033[32mNow printing the external prompt file %s\033[0m\n\n", params.prompt_file.c_str());214 215        std::vector<std::string> prompts = split_string(params.prompt, '\n');216        for (const auto& prompt : prompts) {217            k_prompts.resize(index + 1);218            k_prompts[index] = prompt;219            index++;220            LOG_INF("%3d prompt: %s\n", index, prompt.c_str());221        }222    }223 224    LOG_INF("\n\n");225 226    const int n_ctx = llama_n_ctx(ctx);227 228    if (sseed >= 0) {229        LOG_INF("%s: initializing all samplers with the same RNG seed: %d (use a negative seed to have different seeds)\n", __func__, sseed);230    } else {231        LOG_INF("%s: initializing samplers with different RNG seeds, starting from %d\n", __func__, sseed);232    }233 234    std::vector<client> clients(n_clients);235    for (size_t i = 0; i < clients.size(); ++i) {236        auto & client = clients[i];237        client.id = i;238        client.smpl = common_sampler_init(model, params.sampling);239 240        if (sseed < 0) {241            params.sampling.seed--;242        }243    }244 245    std::vector<llama_token> tokens_system;246 247    tokens_system = common_tokenize(ctx, k_system, true);248    const int32_t n_tokens_system = tokens_system.size();249 250    llama_seq_id g_seq_id = 0;251 252    // the max batch size is as large as the context to handle cases where we get very long input prompt from multiple253    // users. regardless of the size, the main loop will chunk the batch into a maximum of params.n_batch tokens at a time254    llama_batch batch = llama_batch_init(n_ctx, 0, 1);255 256    int32_t n_total_prompt = 0;257    int32_t n_total_gen    = 0;258    int32_t n_cache_miss   = 0;259 260    const auto t_main_start = ggml_time_us();261 262    LOG_INF("%s: Simulating parallel requests from clients:\n", __func__);263    LOG_INF("%s: n_parallel = %d, n_sequences = %d, cont_batching = %d, system tokens = %d\n", __func__, n_clients, n_seq, cont_batching, n_tokens_system);264    LOG_INF("\n");265 266    if (is_sp_shared) {267        LOG_INF("%s: Evaluating the system prompt ...\n", __func__);268 269        for (int32_t i = 0; i < n_tokens_system; ++i) {270            common_batch_add(batch, tokens_system[i], i, { 0 }, false);271        }272 273        if (llama_decode(ctx, batch) != 0) {274            LOG_ERR("%s: llama_decode() failed\n", __func__);275            return 1;276        }277 278        // assign the system KV cache to all parallel sequences279        for (int32_t i = 1; i <= n_clients; ++i) {280            llama_memory_seq_cp(mem, 0, i, -1, -1);281        }282 283        LOG_INF("\n");284    }285 286    LOG_INF("Processing requests ...\n\n");287 288    while (true) {289        common_batch_clear(batch);290 291        // decode any currently ongoing sequences292        for (auto & client : clients) {293            if (client.seq_id == -1) {294                continue;295            }296 297            client.i_batch = batch.n_tokens;298 299            common_batch_add(batch, client.sampled, client.n_past++, { client.id + 1 }, true);300 301            client.n_decoded += 1;302        }303 304        if (batch.n_tokens == 0) {305            // all sequences have ended - clear the entire KV cache306            for (int i = 1; i <= n_clients; ++i) {307                llama_memory_seq_rm(mem, i, -1, -1);308                // but keep the system prompt309                llama_memory_seq_cp(mem, 0, i, -1, -1);310            }311 312            LOG_INF("%s: clearing the KV cache\n", __func__);313        }314 315        // insert new sequences for decoding316        if (cont_batching || batch.n_tokens == 0) {317            for (auto & client : clients) {318                if (client.seq_id == -1 && g_seq_id < n_seq) {319                    client.seq_id = g_seq_id;320 321                    client.t_start_prompt = ggml_time_us();322                    client.t_start_gen    = 0;323 324                    client.input    = k_prompts[rand() % k_prompts.size()];325                    client.response = "";326 327                    // construct the prompt:328                    // [system prompt] + [junk] + [user prompt]329                    client.n_past = 0;330                    client.prompt = "";331                    if (is_sp_shared) {332                        client.n_past = n_tokens_system;333                    } else {334                        client.prompt += k_system;335                    }336 337                    const int n_junk_cur = rand() % n_junk;338 339                    for (int i = 0; i < n_junk_cur; ++i) {340                        const int r = rand() % k_questions.size();341                        client.prompt += "User:\n" + k_questions[r] + "\nAssistant:\n " + k_answers[r] + "\n";342                    }343                    client.prompt += "User:\n" + client.input + "\nAssistant:\n";344 345                    common_sampler_reset(client.smpl);346 347                    // do not prepend BOS because we have a system prompt!348                    std::vector<llama_token> tokens_prompt;349                    tokens_prompt = common_tokenize(ctx, client.prompt, false);350 351                    for (size_t i = 0; i < tokens_prompt.size(); ++i) {352                        common_batch_add(batch, tokens_prompt[i], client.n_past++, { client.id + 1 }, false);353                    }354 355                    // extract the logits only for the last token356                    if (batch.n_tokens > 0) {357                        batch.logits[batch.n_tokens - 1] = true;358                    }359 360                    client.n_prompt  = tokens_prompt.size();361                    client.n_decoded = 0;362                    client.i_batch   = batch.n_tokens - 1;363 364                    LOG_INF("\033[31mClient %3d, seq %4d, junk = %4d, prompt = %d, started decoding ...\033[0m\n", client.id, client.seq_id, n_junk_cur, client.n_prompt);365 366                    g_seq_id += 1;367 368                    // insert new requests one-by-one369                    //if (cont_batching) {370                    //    break;371                    //}372                }373            }374        }375 376        if (batch.n_tokens == 0) {377            break;378        }379 380        // process in chunks of params.n_batch381        int32_t n_batch = params.n_batch;382 383        int32_t i_next = 0;384 385        for (int32_t i = 0; i < batch.n_tokens; i = i_next) {386            // experiment: process in powers of 2387            //if (i + n_batch > (int32_t) batch.n_tokens && n_batch > 32) {388            //    n_batch /= 2;389            //    i -= n_batch;390            //    continue;391            //}392 393            const int32_t n_tokens = std::min(n_batch, batch.n_tokens - i);394 395            llama_batch batch_view = {396                n_tokens,397                batch.token    + i,398                nullptr,399                batch.pos      + i,400                batch.n_seq_id + i,401                batch.seq_id   + i,402                batch.logits   + i,403            };404 405            const int ret = llama_decode(ctx, batch_view);406            if (ret != 0) {407                if (n_batch == 1 || ret < 0) {408                    // if you get here, it means the KV cache is full - try increasing it via the context size409                    LOG_ERR("%s : failed to decode the batch, n_batch = %d, ret = %d\n", __func__, n_batch, ret);410                    return 1;411                }412 413                LOG_WRN("%s : failed to decode the batch, retrying with n_batch = %d\n", __func__, n_batch / 2);414 415                n_cache_miss += 1;416 417                // retry with half the batch size to try to find a free slot in the KV cache418                n_batch /= 2;419 420                continue;421            }422 423            LOG_DBG("%s : decoded batch of %d tokens\n", __func__, n_tokens);424 425            // move the head of the batch forward with the number of tokens we just processed426            i_next = i + n_tokens;427 428            // on successful decode, restore the original batch size429            n_batch = params.n_batch;430 431            for (auto & client : clients) {432                if (client.i_batch < (int) i || client.i_batch >= (int) (i + n_tokens)) {433                    continue;434                }435 436                //printf("client %d, seq %d, token %d, pos %d, batch %d\n",437                //        client.id, client.seq_id, client.sampled, client.n_decoded, client.i_batch);438 439                const llama_token id = common_sampler_sample(client.smpl, ctx, client.i_batch - i);440 441                common_sampler_accept(client.smpl, id, true);442 443                if (client.n_decoded == 1) {444                    // start measuring generation time after the first token to make sure all concurrent clients445                    // have their prompt already processed446                    client.t_start_gen = ggml_time_us();447                }448 449                const std::string token_str = common_token_to_piece(ctx, id);450 451                client.response += token_str;452                client.sampled = id;453 454                //printf("client %d, seq %d, token %d, pos %d, batch %d: %s\n",455                //        client.id, client.seq_id, id, client.n_decoded, client.i_batch, token_str.c_str());456 457                if (client.n_decoded > 2 &&458                    (llama_vocab_is_eog(vocab, id) ||459                     (params.n_predict > 0 && client.n_decoded >= params.n_predict) ||460                     client.response.find("User:") != std::string::npos)) {461                    // basic reverse prompt462                    const size_t pos = client.response.find("User:");463                    if (pos != std::string::npos) {464                        client.response = client.response.substr(0, pos);465                    }466 467                    // delete only the generated part of the sequence, i.e. keep the system prompt in the cache468                    llama_memory_seq_rm(mem,    client.id + 1, -1, -1);469                    llama_memory_seq_cp(mem, 0, client.id + 1, -1, -1);470 471                    const auto t_main_end = ggml_time_us();472 473                    LOG_INF("\033[31mClient %3d, seq %3d/%3d, prompt %4d t, response %4d t, time %5.2f s, speed %5.2f t/s, cache miss %d \033[0m \n\nInput:    %s\n\033[35mResponse: %s\033[0m\n\n",474                            client.id, client.seq_id, n_seq, client.n_prompt, client.n_decoded,475                            (t_main_end - client.t_start_prompt) / 1e6,476                            (double) (client.n_prompt + client.n_decoded) / (t_main_end - client.t_start_prompt) * 1e6,477                            n_cache_miss,478                            ::trim(client.input).c_str(),479                            ::trim(client.response).c_str());480 481                    n_total_prompt += client.n_prompt;482                    n_total_gen    += client.n_decoded;483 484                    client.seq_id = -1;485                }486 487                client.i_batch = -1;488            }489        }490    }491 492    const auto t_main_end = ggml_time_us();493 494    print_date_time();495 496    LOG_INF("%s: n_parallel = %d, n_sequences = %d, cont_batching = %d, system tokens = %d\n", __func__, n_clients, n_seq, cont_batching, n_tokens_system);497    if (params.prompt_file.empty()) {498        params.prompt_file = "used built-in defaults";499    }500    LOG_INF("External prompt file: \033[32m%s\033[0m\n", params.prompt_file.c_str());501    LOG_INF("Model and path used:  \033[32m%s\033[0m\n\n", params.model.path.c_str());502 503    LOG_INF("Total prompt tokens: %6d, speed: %5.2f t/s\n", n_total_prompt, (double) (n_total_prompt              ) / (t_main_end - t_main_start) * 1e6);504    LOG_INF("Total gen tokens:    %6d, speed: %5.2f t/s\n", n_total_gen,    (double) (n_total_gen                 ) / (t_main_end - t_main_start) * 1e6);505    LOG_INF("Total speed (AVG):   %6s  speed: %5.2f t/s\n", "",             (double) (n_total_prompt + n_total_gen) / (t_main_end - t_main_start) * 1e6);506    LOG_INF("Cache misses:        %6d\n", n_cache_miss);507 508    LOG_INF("\n");509 510    // TODO: print sampling/grammar timings for all clients511    llama_perf_context_print(ctx);512 513    llama_batch_free(batch);514 515    llama_backend_free();516 517    LOG("\n\n");518 519    return 0;520}521