Optimize KV cache and compute buffers, bench results from Pi5.

1 files changed, 227 insertions, 117 deletions

diff --git a/llm.c b/llm.c
index e7e5853..23264db 100644
--- a/llm.c
+++ b/llm.c
@@ -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);
 }
-