Commit 7640081

@@ -4,12 +4,12 @@
 #include <string.h>
 
 #include "ggml.h"
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
 #include "ggml-cpu.h"
 #include "gguf.h"
 
-#define EL_COMPUTE_MEM_FLOOR (256 * 1024 * 1024)
-#define EL_TENSOR_OVERHEAD (sizeof(struct ggml_tensor) + 256)
-#define EL_TENSORS_PER_LAYER 40
+#define EL_TENSORS_PER_LAYER 64
 #define EL_DEFAULT_N_CTX 4096
 
 struct el_hparams {
@@ -36,11 +36,20 @@ struct el_handle {
   struct ggml_context *weights;
   struct el_hparams hparams;
   struct el_kv_cache kv;
+  ggml_gallocr_t galloc;
   int32_t n_ctx;
   int n_threads;
   char version[256];
 };
 
+struct el_graph {
+  struct ggml_cgraph *gf;
+  struct ggml_tensor *inp_tokens;
+  struct ggml_tensor *inp_pos;
+  struct ggml_tensor *kq_mask;
+  struct ggml_tensor *logits;
+};
+
 static int32_t el_kv_i32(struct gguf_context *gguf, const char *key, int32_t fallback) {
   int64_t id = gguf_find_key(gguf, key);
   if (id < 0) return fallback;
@@ -176,6 +185,15 @@ void *el_open(const char *model_path, int n_gpu_layers, int n_ctx, int n_threads
     return NULL;
   }
 
+  h->galloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type());
+  if (!h->galloc) {
+    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",
            h->hparams.arch, h->hparams.n_layer, h->hparams.n_embd,
@@ -207,57 +225,33 @@ 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 *new_tokens, int32_t n_new, int32_t n_past) {
-  struct el_hparams *hp = &h->hparams;
-
-  if (n_past + n_new > h->n_ctx) return -1;
+static size_t el_graph_ctx_size(int32_t n_layer) {
+  size_t n_nodes = (size_t) n_layer * EL_TENSORS_PER_LAYER + 16;
+  return ggml_tensor_overhead() * n_nodes + ggml_graph_overhead_custom(n_nodes, false);
+}
 
+static int el_build_graph(struct el_handle *h, struct ggml_context *ctx, struct el_graph *g,
+                           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;
   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) {
-    .mem_size = mem_size,
-    .mem_buffer = NULL,
-    .no_alloc = false,
-  });
-  if (!ctx) return -1;
-
   struct ggml_tensor *tok_embd = el_weight(h->weights, "token_embd.weight", -1);
   struct ggml_tensor *output_norm = el_weight(h->weights, "output_norm.weight", -1);
   struct ggml_tensor *output = el_weight(h->weights, "output.weight", -1);
   if (!output) output = tok_embd;
-  if (!tok_embd || !output_norm || !output) {
-    ggml_free(ctx);
-    return -1;
-  }
+  if (!tok_embd || !output_norm || !output) return -1;
 
   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_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_kv, n_new);
-  float *mask = (float *) kq_mask->data;
-  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;
-    }
-  }
+  ggml_set_input(inp_tokens);
+  ggml_set_input(inp_pos);
+  ggml_set_input(kq_mask);
 
   struct ggml_tensor *cur = ggml_get_rows(ctx, tok_embd, inp_tokens);
 
-  struct ggml_cgraph *gf = ggml_new_graph(ctx);
+  struct ggml_cgraph *gf = ggml_new_graph_custom(ctx, (size_t) hp->n_layer * EL_TENSORS_PER_LAYER + 16, false);
 
   for (int32_t il = 0; il < hp->n_layer; il++) {
     struct ggml_tensor *inp_sa = cur;
@@ -354,9 +348,69 @@ static int32_t el_forward_argmax(struct el_handle *h, const int32_t *new_tokens,
 
   ggml_build_forward_expand(gf, cur);
 
-  ggml_graph_compute_with_ctx(ctx, gf, h->n_threads);
+  g->gf = gf;
+  g->inp_tokens = inp_tokens;
+  g->inp_pos = inp_pos;
+  g->kq_mask = kq_mask;
+  g->logits = cur;
+
+  return 0;
+}
+
+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;
+
+  if (n_past + n_new > h->n_ctx) return -1;
+
+  const int32_t n_kv = n_past + n_new;
+  size_t ctx_size = el_graph_ctx_size(hp->n_layer);
+
+  struct ggml_context *ctx = ggml_init((struct ggml_init_params) {
+    .mem_size = ctx_size,
+    .mem_buffer = NULL,
+    .no_alloc = true,
+  });
+  if (!ctx) return -1;
+
+  struct el_graph g;
+  if (el_build_graph(h, ctx, &g, n_new, n_past) != 0) {
+    ggml_free(ctx);
+    return -1;
+  }
+
+  if (!ggml_gallocr_alloc_graph(h->galloc, g.gf)) {
+    ggml_free(ctx);
+    return -1;
+  }
+
+  ggml_backend_tensor_set(g.inp_tokens, new_tokens, 0, n_new * sizeof(int32_t));
+
+  int32_t *pos = malloc(n_new * sizeof(int32_t));
+  for (int32_t i = 0; i < n_new; i++) pos[i] = n_past + i;
+  ggml_backend_tensor_set(g.inp_pos, pos, 0, n_new * sizeof(int32_t));
+  free(pos);
+
+  float *mask = malloc((size_t) n_new * n_kv * sizeof(float));
+  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;
+    }
+  }
+  ggml_backend_tensor_set(g.kq_mask, mask, 0, (size_t) n_new * n_kv * sizeof(float));
+  free(mask);
+
+  struct ggml_cplan plan = ggml_graph_plan(g.gf, h->n_threads, NULL);
+  void *work_data = plan.work_size > 0 ? malloc(plan.work_size) : NULL;
+  plan.work_data = work_data;
+  enum ggml_status status = ggml_graph_compute(g.gf, &plan);
+  free(work_data);
+  if (status != GGML_STATUS_SUCCESS) {
+    ggml_free(ctx);
+    return -1;
+  }
 
-  const float *logits = (const float *) cur->data;
+  float *logits = malloc((size_t) hp->n_vocab * sizeof(float));
+  ggml_backend_tensor_get(g.logits, logits, 0, (size_t) hp->n_vocab * sizeof(float));
 
   int32_t best_id = 0;
   float best_logit = logits[0];
@@ -367,6 +421,7 @@ static int32_t el_forward_argmax(struct el_handle *h, const int32_t *new_tokens,
     }
   }
 
+  free(logits);
   ggml_free(ctx);
   return best_id;
 }
@@ -403,6 +458,7 @@ int32_t el_generate(void *handle, const int32_t *prompt, int32_t n_prompt,
 void el_close(void *handle) {
   if (!handle) return;
   struct el_handle *h = handle;
+  if (h->galloc) ggml_gallocr_free(h->galloc);
   el_kv_cache_free(&h->kv);
   if (h->weights) ggml_free(h->weights);
   if (h->gguf) gguf_free(h->gguf);