|
| 1 | +#include "ggml.h" |
| 2 | +#include "llama.h" |
| 3 | +#include "common.h" |
| 4 | +#include "llama-vocab.h" |
| 5 | + |
| 6 | +#ifdef _WIN32 |
| 7 | +#define WIN32_LEAN_AND_MEAN |
| 8 | +#ifndef NOMINMAX |
| 9 | +# define NOMINMAX |
| 10 | +#endif |
| 11 | +#include <windows.h> |
| 12 | +#endif |
| 13 | + |
| 14 | +#include <algorithm> |
| 15 | +#include <cstdlib> |
| 16 | +#include <cstdio> |
| 17 | +#include <string> |
| 18 | +#include <vector> |
| 19 | + |
| 20 | +static void print_usage(int, char ** argv) { |
| 21 | + LOG_TEE("\nexample usage:\n"); |
| 22 | + LOG_TEE("\n %s -m model.gguf -c 8192 -b 2048 -ub 512\n", argv[0]); |
| 23 | + LOG_TEE("\n"); |
| 24 | +} |
| 25 | + |
| 26 | +int main(int argc, char ** argv) { |
| 27 | + |
| 28 | + gpt_params params; |
| 29 | + |
| 30 | + if (!gpt_params_parse(argc, argv, params)) { |
| 31 | + print_usage(argc, argv); |
| 32 | + return 1; |
| 33 | + } |
| 34 | + |
| 35 | + // init LLM |
| 36 | + |
| 37 | + llama_backend_init(); |
| 38 | + llama_numa_init(params.numa); |
| 39 | + |
| 40 | + // initialize the model |
| 41 | + |
| 42 | + llama_model_params model_params = llama_model_params_from_gpt_params(params); |
| 43 | + |
| 44 | + llama_model * model = llama_load_model_from_file(params.model.c_str(), model_params); |
| 45 | + |
| 46 | + if (model == NULL) { |
| 47 | + fprintf(stderr , "%s: error: unable to load model\n" , __func__); |
| 48 | + return 1; |
| 49 | + } |
| 50 | + |
| 51 | + llama_context_params ctx_params = llama_context_params_from_gpt_params(params); |
| 52 | + |
| 53 | + llama_context * ctx = llama_new_context_with_model(model, ctx_params); |
| 54 | + |
| 55 | + if (ctx == NULL) { |
| 56 | + fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__); |
| 57 | + return 1; |
| 58 | + } |
| 59 | + |
| 60 | + const unsigned int n_kv_max = llama_n_ctx(ctx); |
| 61 | + |
| 62 | + |
| 63 | + const llama_vocab * vocab = llama_get_vocab(ctx); |
| 64 | + llama_token bos = llama_token_bos_impl(*vocab); |
| 65 | + //llama_token eos = llama_token_eos_impl(*vocab); |
| 66 | + |
| 67 | + const unsigned int n_vocab = llama_n_vocab(model); |
| 68 | + |
| 69 | + // decode in batches of ctx_params.n_batch tokens |
| 70 | + auto decode_helper = [](llama_context * ctx, llama_batch & batch, int32_t n_batch) { |
| 71 | + for (int32_t i = 0; i < (int32_t) batch.n_tokens; i += n_batch) { |
| 72 | + const int32_t n_tokens = std::min(n_batch, (int32_t) (batch.n_tokens - i)); |
| 73 | + |
| 74 | + llama_batch batch_view = { |
| 75 | + n_tokens, |
| 76 | + batch.token + i, |
| 77 | + nullptr, |
| 78 | + batch.pos + i, |
| 79 | + batch.n_seq_id + i, |
| 80 | + batch.seq_id + i, |
| 81 | + batch.logits + i, |
| 82 | + }; |
| 83 | + |
| 84 | + const int ret = llama_decode(ctx, batch_view); |
| 85 | + if (ret != 0) { |
| 86 | + LOG_TEE("failed to decode the batch, n_batch = %d, ret = %d\n", n_batch, ret); |
| 87 | + return false; |
| 88 | + } |
| 89 | + |
| 90 | + llama_synchronize(ctx); |
| 91 | + } |
| 92 | + |
| 93 | + return true; |
| 94 | + }; |
| 95 | + |
| 96 | + const unsigned int pp = params.n_ubatch; |
| 97 | + const unsigned int tg = params.n_ubatch / 4; |
| 98 | + |
| 99 | + if (!params.sweep_bench_output_jsonl) { |
| 100 | + LOG_TEE("\n"); |
| 101 | + LOG_TEE("%s: n_kv_max = %d, n_batch = %d, n_ubatch = %d, flash_attn = %d, n_gpu_layers = %d, n_threads = %u, n_threads_batch = %u\n", __func__, n_kv_max, params.n_batch, params.n_ubatch, params.flash_attn, params.n_gpu_layers, ctx_params.n_threads, ctx_params.n_threads_batch); |
| 102 | + LOG_TEE("\n"); |
| 103 | + LOG_TEE("|%6s | %6s | %6s | %8s | %8s | %8s | %8s |\n", "PP", "TG", "N_KV", "T_PP s", "S_PP t/s", "T_TG s", "S_TG t/s"); |
| 104 | + LOG_TEE("|%6s-|-%6s-|-%6s-|-%8s-|-%8s-|-%8s-|-%8s-|\n", "------", "------", "------", "--------", "--------", "--------", "--------"); |
| 105 | + } |
| 106 | + |
| 107 | + llama_batch batch = llama_batch_init(n_kv_max, 0, 1); |
| 108 | + |
| 109 | + // warm up |
| 110 | + { |
| 111 | + llama_batch_add(batch, bos, 0, { 0 }, false); |
| 112 | + |
| 113 | + if (!decode_helper(ctx, batch, ctx_params.n_batch)) { |
| 114 | + LOG_TEE("%s: llama_decode() failed\n", __func__); |
| 115 | + return 1; |
| 116 | + } |
| 117 | + } |
| 118 | + |
| 119 | + llama_batch_clear(batch); |
| 120 | + llama_kv_cache_clear(ctx); |
| 121 | + |
| 122 | + for (unsigned int n_kv = 0; n_kv < n_kv_max; n_kv += params.n_ubatch) { |
| 123 | + // clean up KV cache before generation |
| 124 | + llama_kv_cache_seq_rm(ctx, 0, n_kv, -1); |
| 125 | + |
| 126 | + // first measure token generation performance at this context size |
| 127 | + const auto t_tg_start = ggml_time_us(); |
| 128 | + |
| 129 | + for (unsigned int i = 0; i < tg; ++i) { |
| 130 | + llama_batch_clear(batch); |
| 131 | + llama_batch_add(batch, std::rand() % n_vocab, n_kv + i, { 0 }, true); |
| 132 | + |
| 133 | + if (!decode_helper(ctx, batch, ctx_params.n_batch)) { |
| 134 | + LOG_TEE("%s: llama_decode() failed\n", __func__); |
| 135 | + return 1; |
| 136 | + } |
| 137 | + } |
| 138 | + |
| 139 | + const auto t_tg_end = ggml_time_us(); |
| 140 | + |
| 141 | + // clean up KV cache after generation |
| 142 | + llama_kv_cache_seq_rm(ctx, 0, n_kv, -1); |
| 143 | + |
| 144 | + // prepare batch of pp size for prompt processing performance measurement |
| 145 | + llama_batch_clear(batch); |
| 146 | + |
| 147 | + for (unsigned int i = 0; i < pp; ++i) { |
| 148 | + llama_batch_add(batch, std::rand() % n_vocab, n_kv + i, { 0 }, false); |
| 149 | + } |
| 150 | + batch.logits[batch.n_tokens - 1] = true; |
| 151 | + |
| 152 | + // measure prompt processing performance |
| 153 | + const auto t_pp_start = ggml_time_us(); |
| 154 | + |
| 155 | + if (!decode_helper(ctx, batch, ctx_params.n_batch)) { |
| 156 | + LOG_TEE("%s: llama_decode() failed\n", __func__); |
| 157 | + return 1; |
| 158 | + } |
| 159 | + |
| 160 | + const auto t_pp_end = ggml_time_us(); |
| 161 | + |
| 162 | + // calculate and print metrics |
| 163 | + const float t_pp = (t_pp_end - t_pp_start) / 1000000.0f; |
| 164 | + const float t_tg = (t_tg_end - t_tg_start) / 1000000.0f; |
| 165 | + |
| 166 | + const float speed_pp = pp / t_pp; |
| 167 | + const float speed_tg = tg / t_tg; |
| 168 | + |
| 169 | + if(params.sweep_bench_output_jsonl) { |
| 170 | + LOG_TEE( |
| 171 | + "{\"n_kv_max\": %d, \"n_batch\": %d, \"n_ubatch\": %d, \"flash_attn\": %d, \"n_gpu_layers\": %d, \"n_threads\": %u, \"n_threads_batch\": %u, " |
| 172 | + "\"pp\": %d, \"tg\": %d, \"n_kv\": %d, \"t_pp\": %f, \"speed_pp\": %f, \"t_tg\": %f, \"speed_tg\": %f }\n", |
| 173 | + n_kv_max, params.n_batch, params.n_ubatch, params.flash_attn, params.n_gpu_layers, ctx_params.n_threads, ctx_params.n_threads_batch, |
| 174 | + pp, tg, n_kv, t_pp, speed_pp, t_tg, speed_tg |
| 175 | + ); |
| 176 | + } else { |
| 177 | + LOG_TEE("|%6d | %6d | %6d | %8.3f | %8.2f | %8.3f | %8.2f |\n", pp, tg, n_kv, t_pp, speed_pp, t_tg, speed_tg); |
| 178 | + } |
| 179 | + } |
| 180 | + |
| 181 | + llama_batch_free(batch); |
| 182 | + |
| 183 | + llama_free(ctx); |
| 184 | + llama_free_model(model); |
| 185 | + |
| 186 | + llama_backend_free(); |
| 187 | + |
| 188 | + return 0; |
| 189 | +} |
0 commit comments