Commit `e8a5ec2`

2026-09-11 00:34:57

perf: add persistent KV cache to skip recomputing the full sequence

Decode no longer replays every prior token through the transformer graph each step -- per-layer K/V tensors persist across generate() calls, prefill runs once, and the output projection only computes logits for the final position instead of every prompt position. Also fixes the compute arena sizing, which was a fixed per-layer constant that crashed (GGML_ASSERT) on wider models like 7B; it's now parameterized on n_embd/n_ff with a safety margin. Decode throughput: ~5.5-7.5 tok/s -> ~46-67 tok/s on Qwen2.5-0.5B.

main

1 parent 51f8f25

Changed files (2)

ext

elelem

ggml

elelem.c

README.md

ext/elelem/ggml/elelem.c

@@ -8,7 +8,9 @@
 #include "gguf.h"
 
 #define EL_COMPUTE_MEM_FLOOR (256 * 1024 * 1024)
-#define EL_COMPUTE_MEM_PER_TOKEN_PER_LAYER (256 * 1024)
+#define EL_TENSOR_OVERHEAD (sizeof(struct ggml_tensor) + 256)
+#define EL_TENSORS_PER_LAYER 40
+#define EL_DEFAULT_N_CTX 4096
 
 struct el_hparams {
   char arch[32];
@@ -22,10 +24,18 @@ struct el_hparams {
   float rms_norm_eps;
 };
 
+struct el_kv_cache {
+  struct ggml_context *ctx;
+  struct ggml_tensor **k;
+  struct ggml_tensor **v;
+};
+
 struct el_handle {
   struct gguf_context *gguf;
   struct ggml_context *weights;
   struct el_hparams hparams;
+  struct el_kv_cache kv;
+  int32_t n_ctx;
   int n_threads;
   char version[256];
 };
@@ -58,7 +68,7 @@ static float el_kv_f32(struct gguf_context *gguf, const char *key, float fallbac
   }
 }
 
-static void el_load_hparams(struct gguf_context *gguf, struct el_hparams *hp) {
+static void el_load_hparams(struct gguf_context *gguf, struct el_hparams *hp, int32_t *n_ctx_train) {
   memset(hp, 0, sizeof(*hp));
 
   int64_t arch_id = gguf_find_key(gguf, "general.architecture");
@@ -90,11 +100,44 @@ static void el_load_hparams(struct gguf_context *gguf, struct el_hparams *hp) {
 
   int64_t tok_id = gguf_find_key(gguf, "tokenizer.ggml.tokens");
   hp->n_vocab = tok_id >= 0 ? (int32_t) gguf_get_arr_n(gguf, tok_id) : 0;
+
+  snprintf(key, sizeof(key), "%s.context_length", arch);
+  *n_ctx_train = el_kv_i32(gguf, key, EL_DEFAULT_N_CTX);
+}
+
+static int el_kv_cache_init(struct el_kv_cache *kv, const struct el_hparams *hp, int32_t n_ctx) {
+  const int64_t n_embd_head = hp->n_embd / hp->n_head;
+  const size_t cache_bytes_per_layer = (size_t) n_embd_head * hp->n_head_kv * n_ctx * sizeof(float);
+  const size_t mem_size = (size_t) hp->n_layer * cache_bytes_per_layer * 2 + 1024 * 1024;
+
+  kv->ctx = ggml_init((struct ggml_init_params) {
+    .mem_size = mem_size,
+    .mem_buffer = NULL,
+    .no_alloc = false,
+  });
+  if (!kv->ctx) return -1;
+
+  kv->k = calloc(hp->n_layer, sizeof(struct ggml_tensor *));
+  kv->v = calloc(hp->n_layer, sizeof(struct ggml_tensor *));
+  if (!kv->k || !kv->v) return -1;
+
+  for (int32_t il = 0; il < hp->n_layer; il++) {
+    kv->k[il] = ggml_new_tensor_3d(kv->ctx, GGML_TYPE_F32, n_embd_head, hp->n_head_kv, n_ctx);
+    kv->v[il] = ggml_new_tensor_3d(kv->ctx, GGML_TYPE_F32, n_embd_head, hp->n_head_kv, n_ctx);
+    if (!kv->k[il] || !kv->v[il]) return -1;
+  }
+
+  return 0;
+}
+
+static void el_kv_cache_free(struct el_kv_cache *kv) {
+  free(kv->k);
+  free(kv->v);
+  if (kv->ctx) ggml_free(kv->ctx);
 }
 
 void *el_open(const char *model_path, int n_gpu_layers, int n_ctx, int n_threads) {
   (void) n_gpu_layers;
-  (void) n_ctx;
 
   struct ggml_context *weights = NULL;
   struct gguf_init_params params = {
@@ -114,7 +157,20 @@ void *el_open(const char *model_path, int n_gpu_layers, int n_ctx, int n_threads
   h->gguf = gguf;
   h->weights = weights;
   h->n_threads = n_threads > 0 ? n_threads : 4;
-  el_load_hparams(gguf, &h->hparams);
+
+  int32_t n_ctx_train = EL_DEFAULT_N_CTX;
+  el_load_hparams(gguf, &h->hparams, &n_ctx_train);
+
+  int32_t n_ctx_default = n_ctx_train < EL_DEFAULT_N_CTX ? n_ctx_train : EL_DEFAULT_N_CTX;
+  h->n_ctx = n_ctx > 0 ? n_ctx : n_ctx_default;
+
+  if (el_kv_cache_init(&h->kv, &h->hparams, h->n_ctx) != 0) {
+    el_kv_cache_free(&h->kv);
+    ggml_free(weights);
+    gguf_free(gguf);
+    free(h);
+    return NULL;
+  }
 
   snprintf(h->version, sizeof(h->version),
            "elelem-ggml arch=%s n_layer=%d n_embd=%d n_head=%d n_head_kv=%d n_ff=%d n_vocab=%d",
@@ -147,14 +203,20 @@ static struct ggml_tensor *el_rms_norm(struct ggml_context *ctx, struct ggml_ten
   return cur;
 }
 
-static int32_t el_forward_argmax(struct el_handle *h, const int32_t *tokens, int32_t n_tokens) {
+static int32_t el_forward_argmax(struct el_handle *h, const int32_t *new_tokens, int32_t n_new, int32_t n_past) {
   struct el_hparams *hp = &h->hparams;
 
-  const int64_t n_embd_head = hp->n_embd / hp->n_head;
+  if (n_past + n_new > h->n_ctx) return -1;
 
-  size_t linear_term = (size_t) n_tokens * hp->n_layer * EL_COMPUTE_MEM_PER_TOKEN_PER_LAYER;
-  size_t attn_term = (size_t) n_tokens * n_tokens * hp->n_head * hp->n_layer * sizeof(float) * 2;
-  size_t mem_size = linear_term + attn_term;
+  const int64_t n_embd_head = hp->n_embd / hp->n_head;
+  const int32_t n_kv = n_past + n_new;
+
+  size_t data_per_layer = (size_t) n_new * (hp->n_embd * 4 + hp->n_ff * 3) * sizeof(float);
+  size_t overhead_per_layer = EL_TENSORS_PER_LAYER * (size_t) EL_TENSOR_OVERHEAD;
+  size_t attn_term = (size_t) n_new * n_kv * hp->n_head * sizeof(float) * 2;
+  size_t output_term = (size_t) hp->n_vocab * sizeof(float) + EL_TENSOR_OVERHEAD * 8;
+  size_t estimate = (size_t) hp->n_layer * (data_per_layer + overhead_per_layer + attn_term) + output_term;
+  size_t mem_size = estimate * 2 + (16 * 1024 * 1024);
   if (mem_size < EL_COMPUTE_MEM_FLOOR) mem_size = EL_COMPUTE_MEM_FLOOR;
 
   struct ggml_context *ctx = ggml_init((struct ggml_init_params) {
@@ -172,24 +234,26 @@ static int32_t el_forward_argmax(struct el_handle *h, const int32_t *tokens, int
     return -1;
   }
 
-  struct ggml_tensor *inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
-  memcpy(inp_tokens->data, tokens, n_tokens * sizeof(int32_t));
+  struct ggml_tensor *inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new);
+  memcpy(inp_tokens->data, new_tokens, n_new * sizeof(int32_t));
 
-  struct ggml_tensor *inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
-  for (int32_t i = 0; i < n_tokens; i++) {
-    ((int32_t *) inp_pos->data)[i] = i;
+  struct ggml_tensor *inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new);
+  for (int32_t i = 0; i < n_new; i++) {
+    ((int32_t *) inp_pos->data)[i] = n_past + i;
   }
 
-  struct ggml_tensor *kq_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_tokens, n_tokens);
+  struct ggml_tensor *kq_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, n_new);
   float *mask = (float *) kq_mask->data;
-  for (int32_t q = 0; q < n_tokens; q++) {
-    for (int32_t k = 0; k < n_tokens; k++) {
-      mask[q * n_tokens + k] = k <= q ? 0.0f : -INFINITY;
+  for (int32_t q = 0; q < n_new; q++) {
+    for (int32_t k = 0; k < n_kv; k++) {
+      mask[q * n_kv + k] = k <= n_past + q ? 0.0f : -INFINITY;
     }
   }
 
   struct ggml_tensor *cur = ggml_get_rows(ctx, tok_embd, inp_tokens);
 
+  struct ggml_cgraph *gf = ggml_new_graph(ctx);
+
   for (int32_t il = 0; il < hp->n_layer; il++) {
     struct ggml_tensor *inp_sa = cur;
 
@@ -218,26 +282,46 @@ static int32_t el_forward_argmax(struct el_handle *h, const int32_t *tokens, int
     struct ggml_tensor *k = ggml_add(ctx, ggml_mul_mat(ctx, wk, cur), bk);
     struct ggml_tensor *v = ggml_add(ctx, ggml_mul_mat(ctx, wv, cur), bv);
 
-    q = ggml_reshape_3d(ctx, q, n_embd_head, hp->n_head, n_tokens);
-    k = ggml_reshape_3d(ctx, k, n_embd_head, hp->n_head_kv, n_tokens);
-    v = ggml_reshape_3d(ctx, v, n_embd_head, hp->n_head_kv, n_tokens);
+    q = ggml_reshape_3d(ctx, q, n_embd_head, hp->n_head, n_new);
+    k = ggml_reshape_3d(ctx, k, n_embd_head, hp->n_head_kv, n_new);
+    v = ggml_reshape_3d(ctx, v, n_embd_head, hp->n_head_kv, n_new);
 
     q = ggml_rope_ext(ctx, q, inp_pos, NULL, n_embd_head, GGML_ROPE_TYPE_NEOX, 0,
                        hp->rope_freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
     k = ggml_rope_ext(ctx, k, inp_pos, NULL, n_embd_head, GGML_ROPE_TYPE_NEOX, 0,
                        hp->rope_freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
 
+    struct ggml_tensor *k_cache = h->kv.k[il];
+    struct ggml_tensor *v_cache = h->kv.v[il];
+
+    struct ggml_tensor *k_cache_view = ggml_view_3d(
+      ctx, k_cache, n_embd_head, hp->n_head_kv, n_new,
+      k_cache->nb[1], k_cache->nb[2], (size_t) n_past * k_cache->nb[2]);
+    struct ggml_tensor *v_cache_view = ggml_view_3d(
+      ctx, v_cache, n_embd_head, hp->n_head_kv, n_new,
+      v_cache->nb[1], v_cache->nb[2], (size_t) n_past * v_cache->nb[2]);
+
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, k, k_cache_view));
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, v, v_cache_view));
+
+    struct ggml_tensor *k_view = ggml_view_3d(
+      ctx, k_cache, n_embd_head, hp->n_head_kv, n_kv,
+      k_cache->nb[1], k_cache->nb[2], 0);
+    struct ggml_tensor *v_view = ggml_view_3d(
+      ctx, v_cache, n_embd_head, hp->n_head_kv, n_kv,
+      v_cache->nb[1], v_cache->nb[2], 0);
+
     q = ggml_permute(ctx, q, 0, 2, 1, 3);
-    k = ggml_permute(ctx, k, 0, 2, 1, 3);
-    v = ggml_permute(ctx, v, 0, 2, 1, 3);
-    v = ggml_cont(ctx, ggml_transpose(ctx, v));
+    struct ggml_tensor *k_perm = ggml_permute(ctx, k_view, 0, 2, 1, 3);
+    struct ggml_tensor *v_perm = ggml_permute(ctx, v_view, 0, 2, 1, 3);
+    v_perm = ggml_cont(ctx, ggml_transpose(ctx, v_perm));
 
-    struct ggml_tensor *kq = ggml_mul_mat(ctx, k, q);
+    struct ggml_tensor *kq = ggml_mul_mat(ctx, k_perm, q);
     kq = ggml_soft_max_ext(ctx, kq, kq_mask, 1.0f / sqrtf((float) n_embd_head), 0.0f);
 
-    struct ggml_tensor *kqv = ggml_mul_mat(ctx, v, kq);
+    struct ggml_tensor *kqv = ggml_mul_mat(ctx, v_perm, kq);
     kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3);
-    kqv = ggml_cont_2d(ctx, kqv, n_embd_head * hp->n_head, n_tokens);
+    kqv = ggml_cont_2d(ctx, kqv, n_embd_head * hp->n_head, n_new);
 
     cur = ggml_mul_mat(ctx, wo, kqv);
     cur = ggml_add(ctx, cur, inp_sa);
@@ -255,21 +339,22 @@ static int32_t el_forward_argmax(struct el_handle *h, const int32_t *tokens, int
   }
 
   cur = el_rms_norm(ctx, cur, output_norm, hp->rms_norm_eps);
-  cur = ggml_mul_mat(ctx, output, cur);
 
-  struct ggml_cgraph *gf = ggml_new_graph(ctx);
+  struct ggml_tensor *last = ggml_view_2d(ctx, cur, hp->n_embd, 1, cur->nb[1],
+                                           (size_t)(n_new - 1) * cur->nb[1]);
+  cur = ggml_mul_mat(ctx, output, last);
+
   ggml_build_forward_expand(gf, cur);
 
   ggml_graph_compute_with_ctx(ctx, gf, h->n_threads);
 
   const float *logits = (const float *) cur->data;
-  const float *last_row = logits + (int64_t)(n_tokens - 1) * hp->n_vocab;
 
   int32_t best_id = 0;
-  float best_logit = last_row[0];
+  float best_logit = logits[0];
   for (int32_t i = 1; i < hp->n_vocab; i++) {
-    if (last_row[i] > best_logit) {
-      best_logit = last_row[i];
+    if (logits[i] > best_logit) {
+      best_logit = logits[i];
       best_id = i;
     }
   }
@@ -282,37 +367,35 @@ int32_t el_generate(void *handle, const int32_t *prompt, int32_t n_prompt,
                      int32_t max_tokens, int32_t eos_token_id, int32_t *out_ids) {
   struct el_handle *h = handle;
 
-  int32_t capacity = n_prompt + max_tokens;
-  int32_t *tokens = malloc(sizeof(int32_t) * capacity);
-  if (!tokens) return -1;
+  if (n_prompt + max_tokens > h->n_ctx) return -1;
+  if (max_tokens <= 0) return 0;
 
-  memcpy(tokens, prompt, sizeof(int32_t) * n_prompt);
-
-  int32_t n_tokens = n_prompt;
+  int32_t n_past = 0;
   int32_t n_generated = 0;
 
-  while (n_generated < max_tokens) {
-    int32_t next = el_forward_argmax(h, tokens, n_tokens);
-    if (next < 0) {
-      free(tokens);
-      return -1;
-    }
+  int32_t next = el_forward_argmax(h, prompt, n_prompt, n_past);
+  if (next < 0) return -1;
+  n_past += n_prompt;
 
-    tokens[n_tokens] = next;
-    n_tokens++;
-    out_ids[n_generated] = next;
-    n_generated++;
+  out_ids[n_generated++] = next;
+  if (next != eos_token_id) {
+    while (n_generated < max_tokens) {
+      next = el_forward_argmax(h, &next, 1, n_past);
+      if (next < 0) return -1;
+      n_past += 1;
 
-    if (next == eos_token_id) break;
+      out_ids[n_generated++] = next;
+      if (next == eos_token_id) break;
+    }
   }
 
-  free(tokens);
   return n_generated;
 }
 
 void el_close(void *handle) {
   if (!handle) return;
   struct el_handle *h = handle;
+  el_kv_cache_free(&h->kv);
   if (h->weights) ggml_free(h->weights);
   if (h->gguf) gguf_free(h->gguf);
   free(h);

README.md

@@ -9,14 +9,17 @@ library llama.cpp itself is built on -- vendored as a git submodule at
 
 The C extension (`ext/elelem/ggml/elelem.c`) loads a GGUF file directly via
 `gguf.h`, builds a Qwen2 transformer graph by hand (RMSNorm, GQA attention
-with NEOX RoPE, SwiGLU FFN), and runs greedy generation with no KV cache. ggml
-itself has no tokenizer, so BPE tokenization and the ChatML template live in
-Ruby (`lib/elelem/ggml/tokenizer.rb`, `lib/elelem/ggml/chat_ml.rb`), parsing
-vocab and merges straight out of the GGUF metadata.
-
-Other architectures, a KV cache, real sampling (top-k/top-p/temp), and
-tool-calling are not implemented yet -- see `spec/` for what's currently
-verified.
+with NEOX RoPE, SwiGLU FFN), and runs greedy generation with a persistent
+per-layer KV cache sized off `context_length` (or an explicit `n_ctx`).
+Prefill computes hidden states for the whole prompt but only projects the
+final position through the output layer; decode reuses cached K/V and only
+appends the new token. ggml itself has no tokenizer, so BPE tokenization and the
+ChatML template live in Ruby (`lib/elelem/ggml/tokenizer.rb`,
+`lib/elelem/ggml/chat_ml.rb`), parsing vocab and merges straight out of the
+GGUF metadata.
+
+Other architectures, real sampling (top-k/top-p/temp), and tool-calling are
+not implemented yet -- see `spec/` for what's currently verified.
 
 ## Setup

Commit e8a5ec2

Commit `e8a5ec2`