diff options
Diffstat (limited to 'llm.c')
| -rw-r--r-- | llm.c | 344 |
1 files changed, 227 insertions, 117 deletions
@@ -1,15 +1,19 @@ +#include <stdio.h> #include <err.h> +#include <string.h> #include "llm.h" #include "mem.h" -#define MAX_TOKENS 300 +#define MAX_TOKENS 300 +#define N_CTX 1024 +#define N_THREADS 4 #define SYSTEM_PROMPT \ "You are a dictionary assistant. Define the target word in the context provided.\n" \ "DO NOT repeat target word in definition. Choose one or two of Formality levels.\n\n" \ "Output format:\n" \ - "<WORD> (<part of speech>) — <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \ + "<WORD> (<part of speech>) - <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \ "DEFINITION:\n" \ "<terse, highly accurate definition in context>\n\n" \ "GENERAL TONE:\n" \ @@ -20,181 +24,287 @@ "1. <example 1>\n" \ "2. <example 2>\n\n" -struct llama_model *llm_init(const char *model_path) -{ +struct llm_ctx { struct llama_model *model; + struct llama_context *ctx; + const struct llama_vocab *vocab; + struct llama_sampler *smpl; + + /* Token history of what is currently resident in ctx's KV + * cache for sequence 0, in order. Used to compute a + * longest-common-prefix against each new request so we only + * decode the tail that actually changed. */ + llama_token *cached_tokens; + int n_cached_tokens; + + /* Token count of the system prompt.Used as a sanity-check floor: + * if a later request's LCP against the cache falls below this, the + * system-prompt prefix itself failed to match, which would + * indicate template drift. */ + int n_system_tokens; +}; + +static void llm_log_cb(enum ggml_log_level level, const char *s, void *ud) +{ + (void)ud; + if (level < GGML_LOG_LEVEL_WARN) + return; + fputs(s, stderr); +} + +/* Tokenize `text` (already a fully rendered chat-template prompt) into + * a freshly malloc'd buffer. Returns token count via *out_n, or -1 on + * failure (*out is left untouched). Caller owns *out. */ +static int tokenize_prompt(const struct llama_vocab *vocab, + const char *text, int text_len, llama_token **out) +{ + int n; + + n = -llama_tokenize(vocab, text, text_len, NULL, 0, false, true); + if (n <= 0) + return -1; + + *out = MALLOC((size_t)n * sizeof(llama_token)); + + if (llama_tokenize(vocab, text, text_len, *out, n, false, true) < 0) { + free(*out); + return -1; + } + + return n; +} + +struct llm_ctx *llm_init(const char *model_path) +{ + char *sys_prompt; + int sys_prompt_len; + llama_token *sys_tokens; + + struct llm_ctx *llm; struct llama_model_params mparams; + struct llama_context_params cparams; + struct llama_sampler_chain_params sparams; + struct llama_chat_message sys_msg[1]; + llama_log_set(llm_log_cb, NULL); llama_backend_init(); mparams = llama_model_default_params(); - mparams.n_gpu_layers = 0; /* force all layers onto CPU */ - mparams.load_mode = LLAMA_LOAD_MODE_MMAP; + mparams.n_gpu_layers = 0; /* force all layers onto CPU */ + mparams.load_mode = LLAMA_LOAD_MODE_MLOCK; /* Force model to RAM, + * don't swap/compress */ + + llm = MALLOC(sizeof(*llm)); + llm->model = llama_model_load_from_file(model_path, mparams); + if (!llm->model) { + free(llm); + return NULL; + } + + if (!llama_model_chat_template(llm->model, NULL)) + fprintf(stderr, "WARN: model has no embedded chat template\n"); + + /* Creation of context allocates KV cache and compute buffers. + * Do this once in init, instead of per request. */ + cparams = llama_context_default_params(); + cparams.n_ctx = N_CTX; + cparams.n_threads = N_THREADS; + cparams.n_threads_batch = N_THREADS; + + llm->ctx = llama_init_from_model(llm->model, cparams); + if (!llm->ctx) { + llama_model_free(llm->model); + free(llm); + return NULL; + } - model = llama_model_load_from_file(model_path, mparams); - return model; + llm->vocab = llama_model_get_vocab(llm->model); + if (!llm->vocab) { + llama_free(llm->ctx); + llama_model_free(llm->model); + free(llm); + return NULL; + } + + /* Sampler chain, reuse across requests as well. */ + sparams = llama_sampler_chain_default_params(); + llm->smpl = llama_sampler_chain_init(sparams); + + llama_sampler_chain_add(llm->smpl, + llama_sampler_init_penalties( + 64, /* last_n: lookback window */ + 1.1f, /* repeat_penalty */ + 0.0f, /* frequency_penalty */ + 0.0f /* presence_penalty */ + ) + ); + + llama_sampler_chain_add(llm->smpl, llama_sampler_init_greedy()); + + llm->cached_tokens = MALLOC((size_t)N_CTX * sizeof(llama_token)); + llm->n_cached_tokens = 0; + + /* Measure system prompt's token count. Used to check template + * drift in llm_run(), not seed the KV cache. */ + sys_msg[0].role = "system"; + sys_msg[0].content = SYSTEM_PROMPT; + + sys_prompt_len = llama_chat_apply_template(NULL, sys_msg, 1, false, NULL, 0); + if (sys_prompt_len > 0) { + sys_prompt = MALLOC((size_t)sys_prompt_len + 1); + if (llama_chat_apply_template(NULL, sys_msg, 1, false, + sys_prompt, sys_prompt_len + 1) >= 0) { + llm->n_system_tokens = tokenize_prompt(llm->vocab, + sys_prompt, sys_prompt_len, &sys_tokens); + if (llm->n_system_tokens > 0) + free(sys_tokens); + else + llm->n_system_tokens = 0; + } else { + llm->n_system_tokens = 0; + } + free(sys_prompt); + } else { + llm->n_system_tokens = 0; + } + + return llm; } -void llm_free(struct llama_model *model) +void llm_free(struct llm_ctx *llm) { - llama_model_free(model); + if (!llm) + return; + llama_sampler_free(llm->smpl); + llama_free(llm->ctx); + llama_model_free(llm->model); llama_backend_free(); + free(llm->cached_tokens); + free(llm); } -void llm_process_request(struct llama_model *model, - const char *user_prompt, FILE *out) +void llm_run(struct llm_ctx *llm, const char *user_prompt, FILE *out) { - int n_prompt_tokens, i; + int i, n; char *prompt; int prompt_len; - - llama_token *prompt_tokens; + llama_token *new_tokens; + int n_new_tokens; + int n_common; llama_token new_token_id; - struct llama_batch batch; - struct llama_context *ctx; - struct llama_context_params cparams; - struct llama_sampler *smpl; - struct llama_sampler_chain_params sparams; struct llama_chat_message messages[] = { { "system", SYSTEM_PROMPT }, { "user", user_prompt } }; - const char *tmpl = llama_model_chat_template(model, NULL); - if (!tmpl) - fprintf(stderr, "Warning: model has no embedded chat template\n"); - + /* Retokenize full prompt. Let's expensive than decode; no + * assumptions about the model's internal chat structure. */ prompt_len = llama_chat_apply_template(NULL, messages, 2, true, NULL, 0); if (prompt_len <= 0) { - fprintf(stderr, "Error: failed to calculate chat template size\n"); + fprintf(stderr, "ERROR: Chat template size calculation failed\n"); return; } prompt = MALLOC((size_t)prompt_len + 1); - if (llama_chat_apply_template(NULL, messages, 2, true, prompt, prompt_len + 1) < 0) { - fprintf(stderr, "Error: failed to apply chat template\n"); + if (llama_chat_apply_template(NULL, messages, 2, true, prompt, + prompt_len + 1) < 0) { + fprintf(stderr, "ERROR: failed to apply chat template\n"); free(prompt); return; } - cparams = llama_context_default_params(); - cparams.n_ctx = 1024; /* context size in tokens */ - cparams.n_threads = 4; - cparams.n_threads_batch = 4; - - ctx = llama_init_from_model(model, cparams); - if (!ctx) { - fprintf(stderr, "Error: failed to create context\n"); - free(prompt); + n_new_tokens = tokenize_prompt(llm->vocab, prompt, prompt_len, &new_tokens); + free(prompt); + if (n_new_tokens <= 0) { + fprintf(stderr, "ERROR: tokenization failed\n"); return; } - const struct llama_vocab *vocab = llama_model_get_vocab(model); - if (!vocab) { - fprintf(stderr, "Error: failed to obtain model vocabulary\n"); - free(prompt); - llama_free(ctx); - return; - } - - n_prompt_tokens = -llama_tokenize(vocab, - prompt, prompt_len, NULL, 0, false, true); - - if (n_prompt_tokens <= 0) { - fprintf(stderr, "Error: tokenization sizing failed\n"); - free(prompt); - llama_free(ctx); + if (n_new_tokens + MAX_TOKENS > N_CTX) { + fprintf(stderr, "ERROR: token count exceeds context size\n"); + free(new_tokens); return; } - if (n_prompt_tokens + MAX_TOKENS > (int)cparams.n_ctx) { - fprintf(stderr, "Error: token count exceeds context size\n"); - free(prompt); - llama_free(ctx); - return; + /* llm->cached_tokens contain tokens from the previous request + * (system prompt + any tokens inserted by the model + previous + * user prompt + closing tokens inserted by the model). new_tokens + * should have a similar structure, except the new user prompt. We + * walk the two token sets to locate the first new token: presumably + * the start of the new user prompt. */ + n_common = 0; + while (n_common < llm->n_cached_tokens && n_common < n_new_tokens && + llm->cached_tokens[n_common] == new_tokens[n_common]) + n_common++; + + /* Sanity check */ + if (llm->n_cached_tokens > 0 && n_common < llm->n_system_tokens) { + fprintf(stderr, "WARN: Cached prefix is shorter than the " + "system prompt. Chat template may be rendering " + "inconsistently between requests."); } - prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token)); + /* Drop everything in the cache past the common prefix. */ + llama_memory_seq_rm(llama_get_memory(llm->ctx), + 0, /* sequence id: 0 - we only use one sequence */ + n_common, /* start at first differing token */ + -1 /* -1: clear to the end of sequence */ + ); + + if (n_common < n_new_tokens) { + /* Only decode the diff */ + batch = llama_batch_get_one(new_tokens + n_common, + n_new_tokens - n_common); + if (llama_decode(llm->ctx, batch) != 0) { + fprintf(stderr, "ERROR: prompt evaluation failed\n"); + free(new_tokens); + llm->n_cached_tokens = n_common; + return; + } - if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, - n_prompt_tokens, false, true) < 0) { - fprintf(stderr, "Error: Tokenization failed\n"); - free(prompt); - free(prompt_tokens); - llama_free(ctx); - return; + /* Copy the newly tokens to the buffer */ + memcpy(llm->cached_tokens + n_common, new_tokens + n_common, + (size_t)(n_new_tokens - n_common) * sizeof(llama_token)); + llm->n_cached_tokens = n_new_tokens; } - free(prompt); /* Prompt buffer is fully tokenized and no longer needed */ - - /* Ingest prompt tokens in one batch (parallelizes matrix - * multiplications across tokens in the batch) */ - batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens); + free(new_tokens); - if (llama_decode(ctx, batch) != 0) { - fprintf(stderr, "Error: Prompt evaluation failed\n"); - free(prompt_tokens); - llama_free(ctx); - return; - } - - free(prompt_tokens); + /* Sampler holds repetition-window state (from the penalties + * stage) across calls since we no longer rebuild it each + * request. Reset that state so one request's repeated words + * don't suppress an unrelated word in the next one. */ + llama_sampler_reset(llm->smpl); /* Generation loop */ - sparams = llama_sampler_chain_default_params(); - smpl = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl, llama_sampler_init_penalties( - 64, /* last_n: lookback window (64 is standard) */ - 1.1f, /* repeat_penalty */ - 0.0f, /* frequency_penalty */ - 0.0f /* presence_penalty */ - )); - - /* Pick the top token */ - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); - - for (i = 0; i < MAX_TOKENS; i++) { - /* Model outputs next tokens for every token in the prompt. - * We need the one after the last token in the prompt */ - new_token_id = llama_sampler_sample(smpl, ctx, -1); - - /* Tell the sampler chain which token was chosen */ - llama_sampler_accept(smpl, new_token_id); - - /* Check for end-of-generation (EOG) tokens: - * EOS: end-of-sequence - * EOT: end-of-turn */ - if (llama_vocab_is_eog(vocab, new_token_id)) + for (i = 0; i < MAX_TOKENS; i++) { + new_token_id = llama_sampler_sample(llm->smpl, llm->ctx, -1); + llama_sampler_accept(llm->smpl, new_token_id); + + if (llama_vocab_is_eog(llm->vocab, new_token_id)) break; - /* Convert numeric token id to printable text */ char buf[128]; - int n = llama_token_to_piece(vocab, new_token_id, buf, - sizeof(buf), 0, false); - + n = llama_token_to_piece(llm->vocab, new_token_id, buf, + sizeof(buf), 0, false); if (n > 0) { fwrite(buf, 1, (size_t)n, out); fflush(out); } - - /* Create batch with 1 token for the next forward pass */ + batch = llama_batch_get_one(&new_token_id, 1); - - if (llama_decode(ctx, batch) != 0) { - fprintf(stderr, "llama_decode failed!\n"); + if (llama_decode(llm->ctx, batch) != 0) { + fprintf(stderr, "ERROR: llama_decode failed!\n"); break; } + + /* New token is now in KV cache. Record it for the next + * request's LCP check. */ + llm->cached_tokens[llm->n_cached_tokens++] = new_token_id; } fprintf(out, "\n\n"); fflush(out); - - //llama_perf_context_print(ctx); - - llama_sampler_free(smpl); - llama_free(ctx); } - |