|
| 1 | +#include "common.h" |
| 2 | +#include "llama.h" |
| 3 | + |
| 4 | +#include <algorithm> |
| 5 | +#include <cmath> |
| 6 | +#include <cstdio> |
| 7 | +#include <string> |
| 8 | +#include <vector> |
| 9 | + |
| 10 | +// mutates the input string |
| 11 | +static std::vector<int> parse_list(char * p) { |
| 12 | + std::vector<int> ret; |
| 13 | + |
| 14 | + char * q = p; |
| 15 | + |
| 16 | + while (*p) { |
| 17 | + if (*p == ',') { |
| 18 | + *p = '\0'; |
| 19 | + ret.push_back(std::atoi(q)); |
| 20 | + q = p + 1; |
| 21 | + } |
| 22 | + |
| 23 | + ++p; |
| 24 | + } |
| 25 | + |
| 26 | + ret.push_back(std::atoi(q)); |
| 27 | + |
| 28 | + return ret; |
| 29 | +} |
| 30 | + |
| 31 | +int main(int argc, char ** argv) { |
| 32 | + gpt_params params; |
| 33 | + |
| 34 | + if (argc == 1 || argv[1][0] == '-') { |
| 35 | + printf("usage: %s MODEL_PATH [N_KV_MAX] [IS_PP_SHARED] [NGL] [MMQ] <PP> <TG> <PL>\n" , argv[0]); |
| 36 | + printf(" <PP>, <TG> and PL are comma-separated lists of numbers without spaces\n\n"); |
| 37 | + printf(" example: %s ggml-model-f16.gguf 2048 0 999 0 128,256,512 128,256 1,2,4,8,16,32\n\n", argv[0]); |
| 38 | + return 1 ; |
| 39 | + } |
| 40 | + |
| 41 | + int n_kv_max = 2048; |
| 42 | + int is_pp_shared = 0; |
| 43 | + int n_gpu_layers = 0; |
| 44 | + int mmq = 0; |
| 45 | + |
| 46 | + std::vector<int> n_pp = { 128, 256, 512, 1024, 2048, 3584, 7680, }; |
| 47 | + std::vector<int> n_tg = { 128, 256, }; |
| 48 | + std::vector<int> n_pl = { 1, 2, 4, 8, 16, 32, }; |
| 49 | + //std::vector<int> n_pl = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 32, }; |
| 50 | + |
| 51 | + if (argc >= 2) { |
| 52 | + params.model = argv[1]; |
| 53 | + } |
| 54 | + |
| 55 | + if (argc >= 3) { |
| 56 | + n_kv_max = std::atoi(argv[2]); |
| 57 | + } |
| 58 | + |
| 59 | + if (argc >= 4) { |
| 60 | + is_pp_shared = std::atoi(argv[3]); |
| 61 | + } |
| 62 | + |
| 63 | + if (argc >= 5) { |
| 64 | + n_gpu_layers = std::atoi(argv[4]); |
| 65 | + } |
| 66 | + |
| 67 | + if (argc >= 6) { |
| 68 | + mmq = std::atoi(argv[5]); |
| 69 | + } |
| 70 | + |
| 71 | + if (argc >= 7) { |
| 72 | + n_pp = parse_list(argv[6]); |
| 73 | + } |
| 74 | + |
| 75 | + if (argc >= 8) { |
| 76 | + n_tg = parse_list(argv[7]); |
| 77 | + } |
| 78 | + |
| 79 | + if (argc >= 9) { |
| 80 | + n_pl = parse_list(argv[8]); |
| 81 | + } |
| 82 | + |
| 83 | + // init LLM |
| 84 | + |
| 85 | + llama_backend_init(params.numa); |
| 86 | + |
| 87 | + // initialize the model |
| 88 | + |
| 89 | + llama_model_params model_params = llama_model_default_params(); |
| 90 | + |
| 91 | + model_params.n_gpu_layers = n_gpu_layers; |
| 92 | + |
| 93 | + llama_model * model = llama_load_model_from_file(params.model.c_str(), model_params); |
| 94 | + |
| 95 | + if (model == NULL) { |
| 96 | + fprintf(stderr , "%s: error: unable to load model\n" , __func__); |
| 97 | + return 1; |
| 98 | + } |
| 99 | + |
| 100 | + llama_context_params ctx_params = llama_context_default_params(); |
| 101 | + |
| 102 | + ctx_params.seed = 1234; |
| 103 | + ctx_params.n_ctx = n_kv_max; |
| 104 | + ctx_params.n_batch = 512; |
| 105 | + ctx_params.mul_mat_q = mmq; |
| 106 | + |
| 107 | + ctx_params.n_threads = params.n_threads; |
| 108 | + ctx_params.n_threads_batch = params.n_threads_batch == -1 ? params.n_threads : params.n_threads_batch; |
| 109 | + |
| 110 | + llama_context * ctx = llama_new_context_with_model(model, ctx_params); |
| 111 | + |
| 112 | + if (ctx == NULL) { |
| 113 | + fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__); |
| 114 | + return 1; |
| 115 | + } |
| 116 | + |
| 117 | + llama_batch batch = llama_batch_init(n_kv_max, 0); |
| 118 | + |
| 119 | + // decode in batches of ctx_params.n_batch tokens |
| 120 | + auto decode_helper = [](llama_context * ctx, llama_batch & batch, int32_t n_batch) { |
| 121 | + for (int32_t i = 0; i < (int32_t) batch.n_tokens; i += n_batch) { |
| 122 | + const int32_t n_tokens = std::min(n_batch, (int32_t) (batch.n_tokens - i)); |
| 123 | + |
| 124 | + llama_batch batch_view = { |
| 125 | + n_tokens, |
| 126 | + batch.token + i, |
| 127 | + nullptr, |
| 128 | + batch.pos + i, |
| 129 | + batch.seq_id + i, |
| 130 | + batch.logits + i, |
| 131 | + 0, 0, 0, // unused |
| 132 | + }; |
| 133 | + |
| 134 | + const int ret = llama_decode(ctx, batch_view); |
| 135 | + if (ret != 0) { |
| 136 | + LOG_TEE("failed to decode the batch, n_batch = %d, ret = %d\n", n_batch, ret); |
| 137 | + return false; |
| 138 | + } |
| 139 | + } |
| 140 | + |
| 141 | + return true; |
| 142 | + }; |
| 143 | + |
| 144 | + // warm up |
| 145 | + { |
| 146 | + batch.n_tokens = 16; |
| 147 | + |
| 148 | + for (int i = 0; i < batch.n_tokens; ++i) { |
| 149 | + batch.token[i] = 0; |
| 150 | + batch.pos[i] = i; |
| 151 | + batch.seq_id[i] = 0; |
| 152 | + batch.logits[i] = false; |
| 153 | + } |
| 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 | + |
| 161 | + LOG_TEE("|%6s | %6s | %4s | %6s | %8s | %8s | %8s | %8s | %8s | %8s |\n", "PP", "TG", "B", "N_KV", "T_PP s", "S_PP t/s", "T_TG s", "S_TG t/s", "T s", "S t/s"); |
| 162 | + LOG_TEE("|%6s-|-%6s-|-%4s-|-%6s-|-%8s-|-%8s-|-%8s-|-%8s-|-%8s-|-%8s-|\n", "------", "------", "----", "------", "--------", "--------", "--------", "--------", "--------", "--------"); |
| 163 | + |
| 164 | + for ( int i_pp = 0; i_pp < (int) n_pp.size(); ++i_pp) { |
| 165 | + for ( int i_tg = 0; i_tg < (int) n_tg.size(); ++i_tg) { |
| 166 | + for (int i_pl = 0; i_pl < (int) n_pl.size(); ++i_pl) { |
| 167 | + const int pp = n_pp[i_pp]; |
| 168 | + const int tg = n_tg[i_tg]; |
| 169 | + const int pl = n_pl[i_pl]; |
| 170 | + |
| 171 | + const int n_ctx_req = is_pp_shared ? pp + pl*tg : pl*(pp + tg); |
| 172 | + |
| 173 | + if (n_ctx_req > n_kv_max) { |
| 174 | + continue; |
| 175 | + } |
| 176 | + |
| 177 | + batch.n_tokens = is_pp_shared ? pp : pl*pp; |
| 178 | + |
| 179 | + for (int i = 0; i < batch.n_tokens; ++i) { |
| 180 | + batch.token[i] = 0; |
| 181 | + batch.pos[i] = i; |
| 182 | + batch.seq_id[i] = 0; |
| 183 | + batch.logits[i] = false; |
| 184 | + } |
| 185 | + batch.logits[batch.n_tokens - 1] = true; |
| 186 | + |
| 187 | + const auto t_pp_start = ggml_time_us(); |
| 188 | + |
| 189 | + llama_kv_cache_tokens_rm(ctx, -1, -1); |
| 190 | + |
| 191 | + if (!decode_helper(ctx, batch, ctx_params.n_batch)) { |
| 192 | + LOG_TEE("%s: llama_decode() failed\n", __func__); |
| 193 | + return 1; |
| 194 | + } |
| 195 | + |
| 196 | + if (is_pp_shared) { |
| 197 | + for (int32_t i = 1; i < pl; ++i) { |
| 198 | + llama_kv_cache_seq_cp(ctx, 0, i, 0, pp); |
| 199 | + } |
| 200 | + } |
| 201 | + |
| 202 | + const auto t_pp_end = ggml_time_us(); |
| 203 | + |
| 204 | + const auto t_tg_start = ggml_time_us(); |
| 205 | + |
| 206 | + for (int i = 0; i < tg; ++i) { |
| 207 | + batch.n_tokens = pl; |
| 208 | + |
| 209 | + for (int j = 0; j < pl; ++j) { |
| 210 | + batch.token[j] = 0; |
| 211 | + batch.pos[j] = pp + i; |
| 212 | + batch.seq_id[j] = j; |
| 213 | + batch.logits[j] = true; |
| 214 | + } |
| 215 | + |
| 216 | + if (!decode_helper(ctx, batch, ctx_params.n_batch)) { |
| 217 | + LOG_TEE("%s: llama_decode() failed\n", __func__); |
| 218 | + return 1; |
| 219 | + } |
| 220 | + } |
| 221 | + |
| 222 | + const auto t_tg_end = ggml_time_us(); |
| 223 | + |
| 224 | + const int32_t n_kv = n_ctx_req; |
| 225 | + |
| 226 | + const float t_pp = (t_pp_end - t_pp_start) / 1000000.0f; |
| 227 | + const float t_tg = (t_tg_end - t_tg_start) / 1000000.0f; |
| 228 | + const float t = t_pp + t_tg; |
| 229 | + |
| 230 | + const float speed_pp = is_pp_shared ? pp / t_pp : pl*pp / t_pp; |
| 231 | + const float speed_tg = pl*tg / t_tg; |
| 232 | + const float speed = n_kv / t; |
| 233 | + |
| 234 | + LOG_TEE("|%6d | %6d | %4d | %6d | %8.3f | %8.2f | %8.3f | %8.2f | %8.3f | %8.2f |\n", pp, tg, pl, n_kv, t_pp, speed_pp, t_tg, speed_tg, t, speed); |
| 235 | + } |
| 236 | + } |
| 237 | + } |
| 238 | + |
| 239 | + llama_print_timings(ctx); |
| 240 | + |
| 241 | + llama_batch_free(batch); |
| 242 | + |
| 243 | + llama_free(ctx); |
| 244 | + llama_free_model(model); |
| 245 | + |
| 246 | + llama_backend_free(); |
| 247 | + |
| 248 | + fprintf(stderr, "\n\n"); |
| 249 | + |
| 250 | + return 0; |
| 251 | +} |
0 commit comments