cli: auto-size context window to the prompt (long-context usable end-to-end)

README.md

@@ -319,7 +319,7 @@ exceeds the 2 GB GitHub-release limit. (Runs real-time on a Pi 5 —
 
 <sub>**Baseline versions:** llama.cpp `d05fe1d` (Pi 5, M1 Max) · `b9827` (x86) — bitnet.cpp = [microsoft/BitNet](https://github.com/microsoft/BitNet) `master` (its bundled llama.cpp fork, unpinned `--depth 1` clone). Full methodology: [`benchmark/`](benchmark/README.md).</sub>
 
-<sub>**Metal backend status** (`BACKENDS="… metal"`): experimental — greedy-decode tokens verified identical to the `cpu_scalar` reference at every optimization step. Decode is within **12 %** of llama.cpp (81.2 vs 91.3 t/s) and holds up at long context (73 t/s at kv≈2100, vs 27 before the head_dim-512 flash kernels). All attention — including the head_dim-512 full-attention layers — runs simdgroup flash (8-simdgroup prefill variant + 16-chunk split-KV decode variant) on a native f16 KV cache (half the KV memory, zero per-call conversion), with fused per-layer blocks (q/k/v norm+RoPE+KV-append, gate+up GeGLU matvec, k/v pair, PLE), llama-style pipelined command buffers, and a device greedy argmax (4-byte token readback). The remaining prefill gap is the shared ~6-TF q4_K GEMM plateau plus flash-kernel efficiency — kernel-level work that needs M3+ profiler counters to attribute; the full measurement ledger lives in `docs/proposals/metal-beat-llamacpp-plan.md`. Known issue: decode after a ≥4096-token prefill no-ops (RoPE table sizing, fix in progress). Cool-state protocol (240 s idle), geist and llama.cpp measured back-to-back (Homebrew llama.cpp, `BLAS,MTL`), M1 Max 32-core.</sub>
+<sub>**Metal backend status** (`BACKENDS="… metal"`): experimental — greedy-decode tokens verified identical to the `cpu_scalar` reference at every optimization step. Decode is within **12 %** of llama.cpp (81.2 vs 91.3 t/s) and holds up at long context (73 t/s at kv≈2100, vs 27 before the head_dim-512 flash kernels). All attention — including the head_dim-512 full-attention layers — runs simdgroup flash (8-simdgroup prefill variant + 16-chunk split-KV decode variant) on a native f16 KV cache (half the KV memory, zero per-call conversion), with fused per-layer blocks (q/k/v norm+RoPE+KV-append, gate+up GeGLU matvec, k/v pair, PLE), llama-style pipelined command buffers, and a device greedy argmax (4-byte token readback). The remaining prefill gap is the shared ~6-TF q4_K GEMM plateau plus flash-kernel efficiency — kernel-level work that needs M3+ profiler counters to attribute; the full measurement ledger lives in `docs/proposals/metal-beat-llamacpp-plan.md`. Long context works past the 4096 default: the session window grows with the workload (the CLI sizes it from prompt + decode budget) and over-capacity prefills are rejected up-front with `GEIST_E_TOO_MANY_TOKENS` instead of the old silent decode no-op — pp5235 greedy on Metal matches the CPU reference across the 4096 boundary. Cool-state protocol (240 s idle), geist and llama.cpp measured back-to-back (Homebrew llama.cpp, `BLAS,MTL`), M1 Max 32-core.</sub>
 </details>
 
 ---

tools/geist.c

@@ -402,15 +402,51 @@ int main(int argc, char **argv) {
     }
     fprintf(stderr, "loaded %s (arch: %s)\n", src, geist_model_arch(model));
 
-    /* Zero-initialized opts == greedy decode (temperature 0). */
+    /* max_new is a hard cap when the user passed -n; otherwise it's a soft target:
+     * keep going past it until a sentence ends (capped at 2x), so a bare completion
+     * prompt — the base model never emits an end token for one — stops on a clean
+     * boundary instead of mid-word. */
+    int budget = n_explicit ? max_new : max_new * 2;
+
+    /* Zero-initialized opts == greedy decode (temperature 0). Size the KV
+     * window to prompt + decode budget so long prompts clear the default
+     * 4096 ceiling (otherwise prefill returns GEIST_E_TOO_MANY_TOKENS).
+     * A throwaway default session tokenizes first — tokenize touches no KV
+     * cache, so this only pays for the tokenizer pass. */
     struct geist_session_opts opts = {0};
-    struct geist_session *sess = nullptr;
+    struct geist_session     *sess = nullptr;
     if (geist_session_create(model, be, &opts, &sess) != GEIST_OK) {
         fprintf(stderr, "session_create failed\n");
         geist_model_destroy(model);
         geist_backend_destroy(be);
         return 1;
     }
+    {
+        const size_t   cap  = strlen(prompt) + 8; /* one token/byte upper bound + BOS */
+        size_t         np   = 0;
+        geist_token_t *tmp  = malloc(cap * sizeof(geist_token_t));
+        if (tmp != nullptr &&
+            geist_session_tokenize(sess, prompt, cap, tmp, &np) == GEIST_OK) {
+            const size_t need = np + (size_t) budget + 8;
+            /* 4096 = the state-default window (arch_state.c). Only rebuild
+             * when the workload exceeds it, so short/medium prompts keep
+             * the single-session fast path. */
+            if (need > 4096u) {
+                geist_session_destroy(sess);
+                opts.max_seq_len = need;
+                sess             = nullptr;
+                if (geist_session_create(model, be, &opts, &sess) != GEIST_OK) {
+                    fprintf(stderr, "session_create (ctx=%zu) failed: %s\n", need,
+                            geist_last_create_error());
+                    free(tmp);
+                    geist_model_destroy(model);
+                    geist_backend_destroy(be);
+                    return 1;
+                }
+            }
+        }
+        free(tmp);
+    }
 
     if (geist_session_set_prompt(sess, prompt) != GEIST_OK) {
         fprintf(stderr, "set_prompt failed: %s\n", geist_session_errmsg(sess));
@@ -420,12 +456,6 @@ int main(int argc, char **argv) {
 
     printf("%s", prompt);
     fflush(stdout);
-
-    /* max_new is a hard cap when the user passed -n; otherwise it's a soft target:
-     * keep going past it until a sentence ends (capped at 2x), so a bare completion
-     * prompt — the base model never emits an end token for one — stops on a clean
-     * boundary instead of mid-word. */
-    int budget = n_explicit ? max_new : max_new * 2;
     for (int i = 0; i < budget; i++) {
         geist_token_t tok = 0;
         if (geist_session_decode_step(sess, &tok) != GEIST_OK) {