diff --git a/CHANGELOG.md b/CHANGELOG.md index 4bcf0bd..5c04580 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,100 @@ # Changelog +## v0.8.33 + +Performance and robustness release for the batching `Server` and generation paths +(the `perf-batching-prefix-cache` plan: 22 tasks across hot-loop, prefix-caching, +API-routing, and correctness phases). Benchmarked before/after on an M1 Max +(`bench/results/v0.8.32-e6e1ef1-baseline-perf-m1max.md` vs +`bench/results/v0.8.32-perf-branch-final-m1max.md`). Headline: with 8 interleaved +conversations sharing a system prompt on 4 slots, TTFT median drops 115 → 35.5 ms +(3.2x) at a ~75% prefix-cache hit ratio; server-routed multi-turn `chat_completion` +is 1.6x faster than the stateless path. Full suite: 252 tests, 0 failures. + +### Added + +- **Cross-slot prefix sharing** — with unified KV (`kv_unified: true`, new `Server` + and `Context.create/2` option), a prefix cached by any slot (idle or still + generating) is adopted by other slots via a metadata-only `llama_memory_seq_cp`; + a shared system prompt prefills once, ever. Gated off automatically on + configurations where it is unsafe (split KV, hybrid-GDN models). +- **Session affinity** — `session:` request option pins a conversation to the slot + holding its cache; freed slots serve queued requests session-first. +- **Level-2 RAM prompt cache** — `prompt_cache_ram_mb:` server option (default 0 = + off): slot KV states about to be destroyed are serialized to RAM (new + `llama_state_seq_get_size/get_data/set_data` NIFs) and restored later instead of + re-prefilling; FIFO eviction under a byte budget checked before copying. +- **Per-request options** — `:cache_prompt`, `:session`, and all sampling params + (`:temp`, `:seed`, `:top_k`, `:top_p`, `:min_p`, penalties, `:grammar`) now work + per request on `Server.generate/stream/generate_tokens/stream_tokens`, overriding + server defaults; each request gets a fresh sampler chain. +- **Server-routed chat completions** — `LlamaCppEx.chat_completion/3` and + `stream_chat_completion/3` accept a running `Server` (pid or name); templating and + tokenization happen in the caller and generation uses continuous batching with + prefix caching. New `Server.complete_tokens/3` returns + `%{text, completion_tokens, finish_reason}`. `ModelManager` gains matching + `chat_completion/3` and `stream_chat_completion/3` routing. +- **Cancellation** — the server monitors each request's consumer and frees the slot + when it dies; halting a `Server.stream/stream_tokens` early cancels generation + (`Server.cancel/2` also public). Stateless `generate`/`stream`/MTP loops carry a + cancel-flag NIF resource installed as `llama_set_abort_callback`, so even a long + prefill aborts mid-decode instead of running to completion for a departed caller. +- **`:max_queue` backpressure** — the previously documented but unenforced option now + rejects overflow with `{:error, :queue_full}` immediately; streams surface it (and + mid-generation errors) as a single `{:error, reason}` element. +- **KV-pressure recovery** — on `llama_decode == 1` the fused tick NIF purges idle + slots' cached KV, then recursively halves the batch; a sequence that still cannot + fit one token fails alone with `{:error, :context_full}` while other slots keep + generating. New `[:llama_cpp_ex, :server, :kv_pressure]`, `[:..., :ram_cache]`, and + `[:..., :prefix_instability]` telemetry events. +- **Embeddings for Nx** — `format: :binary` on `Embedding.embed/embed_batch` returns + the raw native-endian f32 binary (`Nx.from_binary(bin, :f32)`), skipping the boxed + float list entirely. +- Tests: `test/server_test.exs` (pure slot-pick/session/donor/RAM-cache logic), + `test/server_smoke_test.exs` (cache semantics, affinity, backpressure, + cancellation, overflow isolation), `test/utf8_stream_test.exs`; two new bench + scripts (`bench/prefix_cache_concurrent.exs`, `bench/chat_completion_server.exs`). + +### Changed + +- **Server hot loop** — one fused dirty-CPU NIF per tick (`batch_eval_sample`: + decode + per-slot sampling + detokenization + EOG check) replaces the previous + 1 dirty + (2 normal-scheduler NIFs + send) × generating-slots pattern; streamed + pieces are delivered one tick earlier; a persistent `llama_batch` is reused across + decode NIFs; `sampler_sample`, `sampler_sample_at`, `chat_apply_template_jinja`, + `json_schema_to_grammar`, and `speculative_init` are dirty-flagged. +- **Defaults** — `cache_prompt` on the `Server` now defaults to `true` (llama-server + parity; per-request overridable); `kv_unified` defaults to `true` (slots share the + full `n_ctx` budget instead of a fixed `n_ctx/n_parallel` split; no measurable + throughput cost in A/B); `n_batch` defaults to `min(n_ctx, 2048)` (bounds + worst-case tick latency). +- **Error shapes** — failure replies are atoms/tagged tuples: `:context_full`, + `:queue_full`, `:empty_prompt`, `{:inference_failed, reason}` (previously + formatted strings). A batch failure now fails only the affected request instead of + every active slot. +- **Tokenization moved to the caller** — `Server.generate/stream` encode client-side + using a `:persistent_term`-cached model handle; `Server.get_model/1` no longer + round-trips the server mailbox. +- `Context.decode/2` uses an explicit batch requesting logits only for the final + token (previously the last token of every `n_batch` chunk). + +### Fixed + +- **UTF-8-safe streaming** — pieces that split a multibyte codepoint across tokens + are buffered and emitted whole at every emission point (server streams, stateless + streams, chat-completion chunks); previously consumers could receive invalid UTF-8. +- A prompt exactly matching its cached prefix hung the slot until the call timeout + (nothing left to prefill, no logits to sample); cached reuse is now capped at + `prompt_len - 1` like llama-server, so the last prompt token is always re-decoded. +- A tiny incidental prefix match (e.g. 2 tokens) could destroy a long cached prefix + via trim; own-slot, donor-slot, and RAM-cache candidates now compete by match + length with cost-ordered tie-breaking. +- Abandoned streams no longer burn batch budget to `max_tokens` (see cancellation). +- `stream_tokens` with an empty token list now returns `{:error, :empty_prompt}` + instead of hanging until the stream timeout. +- Embedding decodes clear only their own sequences instead of the whole context — + safe if pointed at a shared context. + ## v0.8.32 ### Changed diff --git a/README.md b/README.md index 19e20c2..305f851 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,10 @@ Built with C++ NIFs using [fine](https://github.com/elixir-nx/fine) for ergonomi - **Multi-model manager** — keep several models resident, route requests by id, with a placement-aware (per-GPU VRAM) memory budget - **Device introspection** — `LlamaCppEx.devices/0` lists GPUs/accelerators with per-device VRAM - **Multi-Token Prediction (MTP) speculative decoding** — ~2x token-generation speedup on Qwen 3.6 with live acceptance-rate stats -- **Prefix caching** — same-slot KV cache reuse for multi-turn chat (1.23x faster) +- **Prefix caching** — cross-slot KV reuse with session affinity and an optional RAM prompt cache (TTFT 115 → 35 ms at ~75% hit ratio under concurrency) - **Pluggable batching strategies** — DecodeMaximal, PrefillPriority, Balanced - **Pre-tokenized API** — tokenize outside the GenServer for lower contention +- **Request lifecycle controls** — per-request sampling params, `:max_queue` backpressure, and cancellation (dead or halted stream consumers free their slot immediately) - Telemetry integration for observability ## Installation @@ -262,17 +263,40 @@ Multiple callers are batched into a single forward pass per tick, improving thro ### Prefix Caching -The server caches KV state between requests on the same slot. Multi-turn chat benefits automatically — the system prompt and prior turns aren't recomputed: +The server caches KV state between requests (on by default) and shares it **across slots**: with unified KV (`kv_unified: true`, the default) a system prompt prefilled by any slot is adopted by every other slot via a metadata-only copy, so it is computed once, ever. Requests carrying a `:session` term stick to their slot under concurrency, keeping conversations on their cached prefix: ```elixir {:ok, server} = LlamaCppEx.Server.start_link( model_path: "model.gguf", n_parallel: 4, - cache_prompt: true # opt-in (default: false) + cache_prompt: true, # default: true; also overridable per request + prompt_cache_ram_mb: 1024 # optional level-2 RAM cache for evicted prefixes (default: 0 = off) ) + +{:ok, text} = LlamaCppEx.Server.generate(server, prompt, session: "conversation-42") +``` + +Benchmark (8 interleaved conversations × 4 turns on 4 slots, shared system prompt): **TTFT median 115 → 35.5 ms (3.2x)** at a ~75% prefix-cache hit ratio. + +Notes: + +- Hybrid GDN models (e.g. Qwen 3.5/3.6) only hit on exact-prefix continuations; dense-attention models additionally get partial-prefix and cross-slot hits. +- If a chat template rewrites history (e.g. stripping thinking blocks), cache hits silently degrade — the server emits `[:llama_cpp_ex, :server, :prefix_instability]` telemetry when it detects this. + +### Chat Completions via the Server + +`chat_completion/3` and `stream_chat_completion/3` accept a running server in place of a `%Model{}` — templating and tokenization happen in the caller, and the multi-turn prompt benefits from the prefix cache (**1.6x faster** than the stateless path over a 4-turn conversation): + +```elixir +{:ok, completion} = + LlamaCppEx.chat_completion(server, messages, max_tokens: 200, session: "conversation-42") ``` -Benchmark: **1.23x faster** for multi-turn conversations (487ms vs 597ms per 4-turn exchange). +### Backpressure & Cancellation + +- `max_queue: n` bounds the request queue; overflow returns `{:error, :queue_full}` immediately and streams emit a single `{:error, :queue_full}` element (default: `0`, unlimited). +- Halting a stream early (`Enum.take/2`, consumer exit) cancels generation and frees the slot right away instead of decoding to `max_tokens`; `LlamaCppEx.Server.cancel/2` is also available explicitly. +- Sampling options (`:temp`, `:seed`, `:grammar`, ...) can be set per request, overriding the server defaults. ### Batching Strategies diff --git a/bench/chat_completion_server.exs b/bench/chat_completion_server.exs new file mode 100644 index 0000000..ba54740 --- /dev/null +++ b/bench/chat_completion_server.exs @@ -0,0 +1,102 @@ +Code.require_file("helpers.exs", __DIR__) + +# Chat-completion routing + abandoned-stream benchmark (plan 5.3). +# +# Part 1 — multi-turn chat_completion: stateless %Model{} path (fresh context +# + full prefill per turn) vs Server-routed path (continuous batching + +# prefix cache). Reports wall time for a K-turn conversation. +# +# Part 2 — abandoned streams: aggregate throughput of complete requests while +# a fraction of streams is abandoned after a few chunks. With cancellation, +# abandoned slots free immediately instead of decoding to max_tokens. +# +# Run: LLAMA_MODEL_PATH=... mix run bench/chat_completion_server.exs + +defmodule Bench.ChatCompletionServer do + @moduledoc false + + @turns 4 + @max_tokens 24 + + def run do + model_path = System.get_env("LLAMA_MODEL_PATH") || raise "Set LLAMA_MODEL_PATH" + + IO.puts("Chat completion: stateless vs Server-routed (#{@turns} turns)") + + {:ok, model} = LlamaCppEx.load_model(model_path, n_gpu_layers: -1) + + stateless_ms = + time_conversation(fn msgs, opts -> LlamaCppEx.chat_completion(model, msgs, opts) end) + + IO.puts(" stateless : #{stateless_ms} ms") + + server = Bench.Helpers.start_server(n_parallel: 2, n_ctx: 8192) + + server_ms = + time_conversation(fn msgs, opts -> LlamaCppEx.chat_completion(server, msgs, opts) end) + + IO.puts(" server-routed: #{server_ms} ms") + GenServer.stop(server) + + IO.puts("") + IO.puts("Abandoned streams: 8 requests, half abandoned after 3 chunks") + + for label <- ["with cancellation"] do + ms = abandoned_stream_wall() + IO.puts(" #{label}: 4 full + 4 abandoned in #{ms} ms") + end + end + + defp time_conversation(complete_fun) do + base = [ + %{role: "system", content: "You are terse. Answer with at most one sentence."} + ] + + t0 = System.monotonic_time(:millisecond) + + Enum.reduce(1..@turns, base, fn turn, msgs -> + msgs = msgs ++ [%{role: "user", content: "Question #{turn}: name #{turn} colors."}] + + {:ok, completion} = + complete_fun.(msgs, max_tokens: @max_tokens, temp: 0.0, timeout: 120_000) + + [%{message: %{content: content}}] = completion.choices + msgs ++ [%{role: "assistant", content: content}] + end) + + System.monotonic_time(:millisecond) - t0 + end + + defp abandoned_stream_wall do + model_path = System.get_env("LLAMA_MODEL_PATH") + server = Bench.Helpers.start_server(n_parallel: 2, n_ctx: 8192) + _warm = LlamaCppEx.Server.generate(server, "hi", max_tokens: 4) + + t0 = System.monotonic_time(:millisecond) + + tasks = + for i <- 1..8 do + Task.async(fn -> + if rem(i, 2) == 0 do + # Abandon after 3 chunks — cancellation should free the slot. + LlamaCppEx.Server.stream(server, "Endless story #{i}:", max_tokens: 200) + |> Enum.take(3) + else + {:ok, _} = + LlamaCppEx.Server.generate(server, "Short answer #{i}: 2+2=", + max_tokens: 16, + timeout: 120_000 + ) + end + end) + end + + Task.await_many(tasks, 300_000) + wall = System.monotonic_time(:millisecond) - t0 + GenServer.stop(server) + _ = model_path + wall + end +end + +Bench.ChatCompletionServer.run() diff --git a/bench/helpers.exs b/bench/helpers.exs index a8b6290..1882c2e 100644 --- a/bench/helpers.exs +++ b/bench/helpers.exs @@ -27,6 +27,8 @@ defmodule Bench.Helpers do server_opts |> maybe_add(opts, :cache_prompt) |> maybe_add(opts, :batch_strategy) + |> maybe_add(opts, :kv_unified) + |> maybe_add(opts, :prompt_cache_ram_mb) {:ok, server} = LlamaCppEx.Server.start_link(server_opts) server diff --git a/bench/prefix_cache_concurrent.exs b/bench/prefix_cache_concurrent.exs new file mode 100644 index 0000000..0e2312a --- /dev/null +++ b/bench/prefix_cache_concurrent.exs @@ -0,0 +1,130 @@ +Code.require_file("helpers.exs", __DIR__) + +# Prefix caching under concurrency — the headline benchmark for the caching +# work: K interleaved multi-turn conversations sharing one system prompt, +# served by n_parallel slots. Reports per-turn TTFT and the prefix-cache hit +# ratio for cache-off vs cache-on (cross-slot sharing + session affinity). +# +# Run: LLAMA_MODEL_PATH=... mix run bench/prefix_cache_concurrent.exs +# (Use a dense-attention model for cross-slot sharing; hybrid GDN models +# degrade to exact-prefix per-slot hits.) + +defmodule Bench.PrefixCacheConcurrent do + @moduledoc false + + @n_parallel 4 + @n_conversations 8 + @n_turns 4 + @max_tokens 24 + + def run do + system_prompt = + "You are a helpful assistant. Answer concisely. " <> + String.duplicate("Follow the house style: short sentences, plain words. ", 24) + + IO.puts("Prefix Cache Under Concurrency") + + IO.puts( + " #{@n_conversations} conversations x #{@n_turns} turns, " <> + "n_parallel #{@n_parallel}, shared system prompt, #{@max_tokens} tok/turn" + ) + + IO.puts("") + + modes = [ + {"cache OFF ", cache_prompt: false}, + {"per-slot cache ", cache_prompt: true, kv_unified: false}, + {"cross-slot + affinity", cache_prompt: true, kv_unified: true} + ] + + for {label, opts} <- modes do + {ttfts, cache_ratios} = run_workload(opts, system_prompt) + avg_ratio = Enum.sum(cache_ratios) / length(cache_ratios) + + IO.puts( + " #{label}: TTFT median #{fmt(median(ttfts))} ms, p95 #{fmt(percentile(ttfts, 95))} ms" <> + ", mean hit ratio #{Float.round(avg_ratio * 100, 1)}%" + ) + end + end + + defp run_workload(opts, system_prompt) do + server = + Bench.Helpers.start_server([n_parallel: @n_parallel, n_ctx: 8192] ++ opts) + + collector = start_collector() + + # Each conversation runs its turns sequentially; conversations run + # concurrently (2x the slot count, so the queue and eviction paths are + # exercised too). + tasks = + for conv <- 1..@n_conversations do + Task.async(fn -> + opener = "\nUser: Question #{conv}.1 — name #{conv} colors.\nAssistant:" + + Enum.reduce(1..@n_turns, system_prompt <> opener, fn turn, prompt -> + {:ok, reply} = + LlamaCppEx.Server.generate(server, prompt, + max_tokens: @max_tokens, + timeout: 120_000, + session: conv + ) + + prompt <> reply <> "\nUser: Question #{conv}.#{turn + 1} — one more.\nAssistant:" + end) + end) + end + + Task.await_many(tasks, 600_000) + + measurements = stop_collector(collector) + GenServer.stop(server) + + ttfts = Enum.map(measurements, & &1.ttft_ms) + ratios = Enum.map(measurements, & &1.prefix_cache_ratio) + {ttfts, ratios} + end + + # ETS-based collector — no extra process needed for a bench script. + defp start_collector do + table = :ets.new(:prefix_cache_bench, [:public, :duplicate_bag]) + + :telemetry.attach( + {__MODULE__, table}, + [:llama_cpp_ex, :server, :request, :done], + fn _event, m, _meta, _ -> + :ets.insert(table, {:m, m.ttft_ms, m.prefix_cache_ratio}) + end, + nil + ) + + table + end + + defp stop_collector(table) do + :telemetry.detach({__MODULE__, table}) + + measurements = + for {:m, ttft, ratio} <- :ets.tab2list(table) do + %{ttft_ms: ttft, prefix_cache_ratio: ratio} + end + + :ets.delete(table) + measurements + end + + defp median(values) do + sorted = Enum.sort(values) + Enum.at(sorted, div(length(sorted), 2)) + end + + defp percentile(values, p) do + sorted = Enum.sort(values) + idx = min(length(sorted) - 1, round(p / 100 * (length(sorted) - 1))) + Enum.at(sorted, idx) + end + + defp fmt(ms), do: :erlang.float_to_binary(ms / 1, decimals: 1) +end + +Bench.PrefixCacheConcurrent.run() diff --git a/bench/results/v0.8.32-e6e1ef1-baseline-perf-m1max.md b/bench/results/v0.8.32-e6e1ef1-baseline-perf-m1max.md new file mode 100644 index 0000000..ee4e744 --- /dev/null +++ b/bench/results/v0.8.32-e6e1ef1-baseline-perf-m1max.md @@ -0,0 +1,58 @@ +# Benchmark Baseline — perf work (batching / prefix caching) + +Baseline snapshot taken BEFORE any Phase 1–4 change from the +`perf-batching-prefix-cache` plan. Re-run after each phase to attribute deltas. + +- **Repo**: e6e1ef1 (post v0.8.32) +- **llama.cpp**: 2d973636e (b9870) +- **Hardware**: Apple M1 Max, 64GB unified memory +- **Backend**: Metal (built from source, `LLAMA_BACKEND=metal MIX_ENV=bench`) +- **Model**: Qwen3.5-0.8B-UD-Q4_K_XL.gguf (hybrid GDN — reports seq_rm `:full`) +- **Date**: 2026-07-04 + +## single_generate (stateless path, Benchee ips / median) + +| Input | 32 tokens | 128 tokens | +|-------|-----------|------------| +| short | 4.00 ips (median 0.25 s) | 0.92 ips (median 1.09 s) | +| medium | 3.78 ips (median 0.25 s) | 0.92 ips (median 1.11 s) | +| long | 3.07 ips (median 0.33 s) | 0.92 ips (median 1.13 s) | + +## server_concurrent (max_tokens 32, short prompt, n_parallel 4) + +| Concurrency | Wall median (p99) | Total tok/s | Per-req tok/s | Speedup | Avg batch | +|-------------|-------------------|-------------|---------------|---------|-----------| +| 1 | 291.2 ms (320.0 ms) | 109.9 | 109.9 | 1.00x | 1.1 | +| 2 | 362.7 ms (397.3 ms) | 176.4 | 88.2 | 1.61x | 2.2 | +| 4 | 550.4 ms (573.9 ms) | 232.5 | 58.1 | 2.12x | 4.5 | + +## prefix_cache (4-turn chat, n_parallel 1, max_tokens 16/turn) + +| Scenario | ips | median | +|----------|-----|--------| +| WITH prefix cache | 1.70 | 615.08 ms | +| WITHOUT prefix cache | 1.57 | 638.11 ms | + +Cache advantage only 1.08x — expected to grow substantially after Phase 2. +Note: model is hybrid GDN (`seq_rm :full`), so partial-trim cache hits are +disabled; only exact-prefix continuations hit. + +## strategies (single request, max_tokens 32) + +| Strategy | ips | median | +|----------|-----|--------| +| prefill_priority | 4.24 | 206.51 ms | +| decode_maximal | 3.94 | 256.87 ms | +| balanced | 3.53 | 284.48 ms | + +## tokenize_overhead (Server, max_tokens 32) + +| Input | text API | pre-tokenized | +|-------|----------|---------------| +| short | 5.10 ips | 5.14 ips | +| medium | 4.50 ips | 4.54 ips | + +## Notes + +- All runs single-shot Benchee defaults from bench/*.exs (warmup 1, time 10–15 s). +- Raw outputs archived in session scratchpad (`baseline/*.txt`), not committed. diff --git a/bench/results/v0.8.32-perf-branch-final-m1max.md b/bench/results/v0.8.32-perf-branch-final-m1max.md new file mode 100644 index 0000000..f845c5f --- /dev/null +++ b/bench/results/v0.8.32-perf-branch-final-m1max.md @@ -0,0 +1,81 @@ +# Benchmark Results — perf branch final (batching / prefix caching) + +Final numbers for the `perf-batching-prefix-cache` plan (Phases 1–4 complete). +Baseline for comparison: `v0.8.32-e6e1ef1-baseline-perf-m1max.md` (same machine, +same build config, model, and scripts). + +- **Repo**: perf/batching-prefix-cache (27 commits over e6e1ef1) +- **llama.cpp**: 2d973636e (b9870) +- **Hardware**: Apple M1 Max, 64GB unified memory +- **Backend**: Metal (built from source, `LLAMA_BACKEND=metal MIX_ENV=bench`) +- **Models**: Qwen3.5-0.8B-UD-Q4_K_XL (standard suite; hybrid GDN), + Llama-3.2-1B-Instruct-Q4_K_M (feature benches; dense attention) +- **Date**: 2026-07-05 + +## Headline: prefix cache under concurrency (Llama-3.2-1B, dense) + +8 interleaved conversations × 4 turns on 4 slots, shared system prompt +(`bench/prefix_cache_concurrent.exs`): + +| Mode | TTFT median | TTFT p95 | Hit ratio | +|------|-------------|----------|-----------| +| cache OFF | 115.0 ms | 348.9 ms | 0% | +| per-slot cache (kv_unified off) | 39.0 ms | 362.0 ms | 75.6% | +| cross-slot + session affinity | **35.5 ms (3.2x)** | 341.0 ms | 74.8% | + +## Chat completions + cancellation (Llama-3.2-1B) + +`bench/chat_completion_server.exs`: + +| Scenario | Result | +|----------|--------| +| 4-turn chat_completion, stateless | 346 ms | +| 4-turn chat_completion, server-routed | **211 ms (1.6x)** | +| 8 requests, half abandoned after 3 chunks | 348 ms total (pre-cancellation each abandoned stream decoded 200 tokens to completion) | + +## server_concurrent (Qwen3.5-0.8B, max_tokens 32, short prompt) + +| Concurrency | Baseline tok/s | Final tok/s | Δ | +|-------------|----------------|-------------|---| +| 1 | 109.9 | 133.2 | +21% | +| 2 | 176.4 | 178.9 | +1% | +| 4 | 232.5 | 257.8 | +11% | + +## prefix_cache (Qwen3.5-0.8B hybrid, 4-turn chat, n_parallel 1) + +| Scenario | Baseline avg | Final avg | +|----------|--------------|-----------| +| WITH prefix cache | 588 ms (1.08x vs off) | **448 ms (1.20x vs off)** | +| WITHOUT prefix cache | 637 ms | 536 ms | + +Hybrid GDN models only hit on exact-prefix continuations (`seq_rm :full`); +dense models additionally get partial-trim and cross-slot hits. + +## single_generate (stateless path, Qwen3.5-0.8B) + +| Input | Baseline 128tok | Final 128tok | +|-------|-----------------|--------------| +| short | 0.92 ips | 1.21 ips | +| medium | 0.92 ips | 0.99 ips | +| long | 0.92 ips | 1.12 ips | + +## strategies (single request, max_tokens 32) + +| Strategy | Baseline median | Final median | +|----------|-----------------|--------------| +| prefill_priority | 206.5 ms | 181.4 ms | +| decode_maximal | 256.9 ms | 195.5 ms | +| balanced | 284.5 ms | 210.2 ms | + +## Notes + +- Machine noise on this hardware is significant for the 5-iteration + server_concurrent samples (±15–30% run to run); the multi-run Benchee + numbers and the controlled A/Bs recorded in the plan scratchpad are the + more reliable signal. +- `kv_unified: true` (new Server default) showed no regression on + zero-shared-prefix concurrent workloads in direct A/B (653 vs 672 ms dense; + 884 vs 889 ms hybrid). +- The GGML Metal `rsets` assert at BEAM shutdown (llama.cpp #17869) predates + this branch and fires after benches/tests complete; it occasionally turns a + clean run's exit code into 134. diff --git a/c_src/llama_cpp_ex/llama_nif.cpp b/c_src/llama_cpp_ex/llama_nif.cpp index 5aff7fd..31ffd67 100644 --- a/c_src/llama_cpp_ex/llama_nif.cpp +++ b/c_src/llama_cpp_ex/llama_nif.cpp @@ -20,6 +20,38 @@ FINE_RESOURCE(LlamaModel); FINE_RESOURCE(LlamaContext); FINE_RESOURCE(LlamaSampler); FINE_RESOURCE(LlamaSpeculative); +FINE_RESOURCE(CancelFlag); + +// --- Cancellation --- + +fine::ResourcePtr cancel_flag_new(ErlNifEnv* env) { + return fine::make_resource(); +} +FINE_NIF(cancel_flag_new, 0); + +fine::Ok<> request_cancel(ErlNifEnv* env, fine::ResourcePtr flag) { + flag->cancelled.store(true, std::memory_order_relaxed); + return fine::Ok(); +} +FINE_NIF(request_cancel, 0); + +static bool cancel_abort_cb(void* data) { + return static_cast(data)->cancelled.load(std::memory_order_relaxed); +} + +// Installs the flag as the context's abort callback for the duration of a +// generation loop, so a cancel interrupts even a long prefill decode +// (llama_decode returns 2). Cleared on scope exit. +struct AbortCallbackScope { + llama_context* ctx; + + AbortCallbackScope(llama_context* c, CancelFlag* flag) : ctx(c) { + llama_set_abort_callback(ctx, cancel_abort_cb, flag); + } + ~AbortCallbackScope() { + llama_set_abort_callback(ctx, nullptr, nullptr); + } +}; // --- Backend --- @@ -299,6 +331,7 @@ context_create( int64_t attention_type, bool no_perf, bool swa_full, + bool kv_unified, // Speculative decoding / MTP int64_t ctx_type, int64_t n_rs_seq) @@ -339,6 +372,10 @@ context_create( params.attention_type = static_cast(attention_type); params.no_perf = no_perf; params.swa_full = swa_full; + // Unified KV: all sequences share one buffer/stream, making cross-seq + // llama_memory_seq_cp a metadata-only tag copy for ANY position range. + // In split mode (false), partial cross-stream seq_cp aborts the process. + params.kv_unified = kv_unified; // Speculative decoding / MTP params.ctx_type = static_cast(ctx_type); @@ -440,23 +477,48 @@ fine::Ok<> sampler_reset(ErlNifEnv* env, fine::ResourcePtr sampler } FINE_NIF(sampler_reset, 0); +// Dirty: the sampler chain runs a softmax over the full vocab (100k+ entries) +// and grammar samplers can take multiple ms — too slow for a normal scheduler. int64_t sampler_sample(ErlNifEnv* env, fine::ResourcePtr sampler, fine::ResourcePtr ctx) { return llama_sampler_sample(sampler->sampler, ctx->ctx, -1); } -FINE_NIF(sampler_sample, 0); +FINE_NIF(sampler_sample, ERL_NIF_DIRTY_JOB_CPU_BOUND); // --- Decode --- std::variant, fine::Error> decode(ErlNifEnv* env, fine::ResourcePtr ctx, std::vector token_ids) { - std::vector tokens(token_ids.begin(), token_ids.end()); + int n_tokens = static_cast(token_ids.size()); + if (n_tokens == 0) { + return fine::Ok(); + } + + // Explicit batch instead of llama_batch_get_one: get_one requests logits + // for the LAST TOKEN OF EVERY CHUNK, computing (and copying) a full-vocab + // logit row per n_batch tokens of prompt for nothing. Only the very last + // token needs logits. Positions continue seq 0 from wherever the context + // is, preserving get_one's append semantics for repeated decode calls. + llama_pos start = + llama_memory_seq_pos_max(llama_get_memory(ctx->ctx), 0) + 1; - // Process in chunks of n_batch int n_batch = llama_n_batch(ctx->ctx); - for (size_t i = 0; i < tokens.size(); i += n_batch) { - int n = std::min(static_cast(tokens.size() - i), n_batch); - llama_batch batch = llama_batch_get_one(tokens.data() + i, n); + llama_batch& batch = ctx->reserve_batch(std::min(n_tokens, n_batch)); + + for (int i = 0; i < n_tokens; i += n_batch) { + int n = std::min(n_tokens - i, n_batch); + bool is_last_chunk = (i + n >= n_tokens); + + batch.n_tokens = n; + + for (int j = 0; j < n; j++) { + batch.token[j] = static_cast(token_ids[i + j]); + batch.pos[j] = start + static_cast(i + j); + batch.n_seq_id[j] = 1; + batch.seq_id[j][0] = 0; + batch.logits[j] = (is_last_chunk && j == n - 1); + } + int ret = llama_decode(ctx->ctx, batch); if (ret != 0) { return fine::Error(std::string("llama_decode failed with code: " + std::to_string(ret))); @@ -556,8 +618,10 @@ embed_decode( return fine::Error(std::string("empty token list")); } - // Clear memory for a fresh decode - llama_memory_clear(llama_get_memory(ctx->ctx), true); + // Fresh decode for THIS sequence only — clearing the whole memory would + // clobber every other sequence if the context is ever shared. + llama_memory_seq_rm(llama_get_memory(ctx->ctx), + static_cast(seq_id), -1, -1); // Build batch with explicit seq_id and position tracking llama_batch batch = llama_batch_init(n_tokens, 0, 1); @@ -582,7 +646,11 @@ embed_decode( } FINE_NIF(embed_decode, ERL_NIF_DIRTY_JOB_CPU_BOUND); -std::variant>, fine::Error> +// Returns the embedding as a raw little-endian f32 binary (native order on +// all supported targets): one 16 KB refc binary for a 4096-dim vector instead +// of ~100 KB of boxed floats + list cells per call. The Elixir wrapper +// decodes to a list by default and passes the binary through for Nx users. +std::variant, fine::Error> get_embeddings( ErlNifEnv* env, fine::ResourcePtr ctx, @@ -606,14 +674,16 @@ get_embeddings( return fine::Error(std::string("failed to get embeddings (null pointer)")); } - std::vector out(n_embd); + ERL_NIF_TERM bin; + float* out = reinterpret_cast( + enif_make_new_binary(env, static_cast(n_embd) * sizeof(float), &bin)); if (normalize == 2) { // L2 normalization double sum = 0.0; for (int i = 0; i < n_embd; i++) sum += (double)embd[i] * (double)embd[i]; - double norm = sum > 0.0 ? 1.0 / std::sqrt(sum) : 0.0; - for (int i = 0; i < n_embd; i++) out[i] = (double)embd[i] * norm; + float norm = sum > 0.0 ? static_cast(1.0 / std::sqrt(sum)) : 0.0f; + for (int i = 0; i < n_embd; i++) out[i] = embd[i] * norm; } else if (normalize == 0) { // Max-abs normalization double max_abs = 0.0; @@ -621,14 +691,14 @@ get_embeddings( double a = std::abs((double)embd[i]); if (a > max_abs) max_abs = a; } - double norm = max_abs > 0.0 ? 1.0 / max_abs : 0.0; - for (int i = 0; i < n_embd; i++) out[i] = (double)embd[i] * norm; + float norm = max_abs > 0.0 ? static_cast(1.0 / max_abs) : 0.0f; + for (int i = 0; i < n_embd; i++) out[i] = embd[i] * norm; } else { // No normalization - for (int i = 0; i < n_embd; i++) out[i] = (double)embd[i]; + std::memcpy(out, embd, static_cast(n_embd) * sizeof(float)); } - return fine::Ok(out); + return fine::Ok(fine::Term(bin)); } FINE_NIF(get_embeddings, 0); @@ -657,8 +727,12 @@ embed_batch_decode( return fine::Error(std::string("no tokens to decode")); } - // Fresh decode for this batch - llama_memory_clear(llama_get_memory(ctx->ctx), true); + // Fresh decode for the sequences in THIS batch only (successive groups + // reuse the same seq ids) — never clear the whole memory. + for (auto& [seq_id, tokens] : sequences) { + llama_memory_seq_rm(llama_get_memory(ctx->ctx), + static_cast(seq_id), -1, -1); + } llama_batch batch = llama_batch_init(total, 0, 1); batch.n_tokens = total; @@ -702,12 +776,12 @@ prefill( } int n_batch = llama_n_batch(ctx->ctx); + llama_batch& batch = ctx->reserve_batch(std::min(n_tokens, n_batch)); for (int i = 0; i < n_tokens; i += n_batch) { int n = std::min(n_tokens - i, n_batch); bool is_last_chunk = (i + n >= n_tokens); - llama_batch batch = llama_batch_init(n, 0, 1); batch.n_tokens = n; for (int j = 0; j < n; j++) { @@ -720,7 +794,6 @@ prefill( } int ret = llama_decode(ctx->ctx, batch); - llama_batch_free(batch); if (ret != 0) { return fine::Error(std::string("prefill decode failed with code: " + std::to_string(ret))); @@ -814,7 +887,7 @@ decode_token( int64_t pos, int64_t seq_id) { - llama_batch batch = llama_batch_init(1, 0, 1); + llama_batch& batch = ctx->reserve_batch(1); batch.n_tokens = 1; batch.token[0] = static_cast(token_id); batch.pos[0] = static_cast(pos); @@ -823,7 +896,6 @@ decode_token( batch.logits[0] = true; int ret = llama_decode(ctx->ctx, batch); - llama_batch_free(batch); if (ret != 0) { return fine::Error(std::string("decode_token failed with code: " + std::to_string(ret))); @@ -846,7 +918,7 @@ batch_eval( return fine::Error(std::string("empty entries list")); } - llama_batch batch = llama_batch_init(n, 0, 1); + llama_batch& batch = ctx->reserve_batch(n); batch.n_tokens = n; for (int i = 0; i < n; i++) { @@ -859,7 +931,6 @@ batch_eval( } int ret = llama_decode(ctx->ctx, batch); - llama_batch_free(batch); if (ret != 0) { return fine::Error(std::string("batch_eval failed with code: " + std::to_string(ret))); @@ -869,8 +940,197 @@ batch_eval( } FINE_NIF(batch_eval, ERL_NIF_DIRTY_JOB_CPU_BOUND); +// --- Fused batch eval + sample (Server hot loop) --- +// +// One NIF call per Server tick: builds a batch from `entries` +// ({token, pos, seq_id, wants_logits}), runs llama_decode, then samples every +// wants_logits entry whose seq_id has a sampler registered in `samplers`, +// returning {seq_id, new_token, piece, is_eog}. The piece is empty for EOG +// tokens. Unlike decode_batch above, samplers are per-sequence resources owned +// by the caller — their grammar/penalty state advances across ticks and is +// never reset or shared here. +// +// KV-pressure policy (llama_decode == 1, "no KV slot found"), mirroring +// llama-server's update_slots(): first drop whole sequences listed in +// `purgeable_seq_ids` (idle slots whose cache the caller is willing to lose), +// retrying after each purge; once the purge list is exhausted, recursively +// halve the batch — each half's logits entries are sampled right after its +// sub-decode, before the next decode invalidates the logits buffer. A +// single-token batch that still fails means THAT sequence is out of KV +// budget: it is added to `failed` and skipped for the rest of the call so the +// other sequences keep going (the caller fails just that request). Purged and +// failed seq ids plus the split count are returned so the caller can fix up +// its bookkeeping and emit telemetry. Purgeable seqs must not have entries in +// the batch. + +static int bes_decode_range( + llama_context* ctx, + const llama_vocab* vocab, + const std::vector>& entries, + size_t begin, size_t end, + const std::vector>& samplers, + const std::vector& purgeable, + size_t& purge_idx, + std::vector& purged, + int64_t& n_splits, + std::vector& failed, + std::vector>& results, + llama_batch& batch) // reserved by the caller for >= entries.size() tokens +{ + // Skip entries whose sequence already failed this call — decoding past a + // failed (missing) position would leave a hole in that sequence's KV. + std::vector idxs; + idxs.reserve(end - begin); + for (size_t i = begin; i < end; i++) { + int64_t seq_id = std::get<2>(entries[i]); + if (std::find(failed.begin(), failed.end(), seq_id) == failed.end()) { + idxs.push_back(i); + } + } + + size_t n = idxs.size(); + if (n == 0) { + return 0; + } + + batch.n_tokens = static_cast(n); + + for (size_t i = 0; i < n; i++) { + const auto& [token_id, pos, seq_id, logits] = entries[idxs[i]]; + batch.token[i] = static_cast(token_id); + batch.pos[i] = static_cast(pos); + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = static_cast(seq_id); + batch.logits[i] = logits; + } + + int ret = llama_decode(ctx, batch); + + // Purge donatable idle caches one at a time while the KV cache is full. + while (ret == 1 && purge_idx < purgeable.size()) { + auto victim = static_cast(purgeable[purge_idx++]); + llama_memory_seq_rm(llama_get_memory(ctx), victim, -1, -1); + purged.push_back(victim); + ret = llama_decode(ctx, batch); + } + + if (ret == 0) { + // Sample now: these logits belong to THIS decode call and the next + // sub-decode would overwrite them. + for (size_t i = 0; i < n; i++) { + const auto& [token_id, pos, seq_id, logits] = entries[idxs[i]]; + if (!logits) continue; + + llama_sampler* smpl = nullptr; + for (const auto& [sid, s] : samplers) { + if (sid == seq_id) { smpl = s; break; } + } + if (!smpl) continue; // logits-only entry, nothing to sample + + // llama_sampler_sample() already accepts the selected token — + // calling llama_sampler_accept() again would double-advance + // grammar state. + llama_token new_token = + llama_sampler_sample(smpl, ctx, static_cast(i)); + bool is_eog = llama_vocab_is_eog(vocab, new_token); + + std::string piece; + if (!is_eog) { + char buf[1024]; + int pn = llama_token_to_piece(vocab, new_token, buf, sizeof(buf), 0, false); + if (pn < 0) { + std::vector large_buf(-pn); + pn = llama_token_to_piece(vocab, new_token, + large_buf.data(), large_buf.size(), 0, false); + if (pn > 0) piece.assign(large_buf.data(), pn); + } else if (pn > 0) { + piece.assign(buf, pn); + } + } + + results.emplace_back(seq_id, static_cast(new_token), + std::move(piece), is_eog); + } + return 0; + } + + if (ret == 1 && n == 1) { + // A single token still can't fit: this sequence is out of KV budget. + // Fail it and let the rest of the batch proceed. + failed.push_back(std::get<2>(entries[idxs[0]])); + return 0; + } + + if (ret == 1) { + // Halve and retry — explicit positions/seq_ids make any split valid, + // and per-seq entries stay in position order across the halves. + n_splits++; + size_t mid = begin + (end - begin) / 2; + int rc = bes_decode_range(ctx, vocab, entries, begin, mid, samplers, + purgeable, purge_idx, purged, n_splits, + failed, results, batch); + if (rc != 0) return rc; + return bes_decode_range(ctx, vocab, entries, mid, end, samplers, + purgeable, purge_idx, purged, n_splits, + failed, results, batch); + } + + return ret; +} + +std::variant< + fine::Ok>, + std::vector, + int64_t, + std::vector>, + fine::Error +> +batch_eval_sample( + ErlNifEnv* env, + fine::ResourcePtr ctx, + std::vector> entries, + std::vector>> samplers, + std::vector purgeable_seq_ids) +{ + if (entries.empty()) { + return fine::Error(std::string("empty entries list")); + } + + const auto* vocab = ctx->model->vocab(); + + // The ResourcePtr arguments keep the samplers alive for the whole call. + std::vector> smpls; + smpls.reserve(samplers.size()); + for (auto& [sid, s] : samplers) { + smpls.emplace_back(sid, s->sampler); + } + + std::vector> results; + std::vector purged; + std::vector failed; + int64_t n_splits = 0; + size_t purge_idx = 0; + + llama_batch& batch = ctx->reserve_batch(static_cast(entries.size())); + + int rc = bes_decode_range(ctx->ctx, vocab, entries, 0, entries.size(), + smpls, purgeable_seq_ids, purge_idx, purged, + n_splits, failed, results, batch); + + if (rc != 0) { + return fine::Error(std::string( + "batch_eval_sample failed with code: " + std::to_string(rc))); + } + + return fine::Ok(std::move(results), std::move(purged), n_splits, + std::move(failed)); +} +FINE_NIF(batch_eval_sample, ERL_NIF_DIRTY_JOB_CPU_BOUND); + // --- Sampler sample at batch index --- +// Dirty for the same reason as sampler_sample: full-vocab softmax + optional +// grammar evaluation exceed the ~1 ms normal-scheduler guideline. int64_t sampler_sample_at( ErlNifEnv* env, fine::ResourcePtr sampler, @@ -879,7 +1139,62 @@ int64_t sampler_sample_at( { return llama_sampler_sample(sampler->sampler, ctx->ctx, static_cast(idx)); } -FINE_NIF(sampler_sample_at, 0); +FINE_NIF(sampler_sample_at, ERL_NIF_DIRTY_JOB_CPU_BOUND); + +// --- Sequence state save/restore (RAM prompt cache) --- + +// Size of one sequence's serialized KV state. Cheap metadata walk — used by +// the caller to enforce a byte budget BEFORE paying for the copy. +int64_t state_seq_get_size(ErlNifEnv* env, fine::ResourcePtr ctx, + int64_t seq_id) { + return static_cast( + llama_state_seq_get_size(ctx->ctx, static_cast(seq_id))); +} +FINE_NIF(state_seq_get_size, 0); + +// Serializes a sequence's KV state into a binary. Dirty: the state is +// KV-sized (potentially hundreds of MB on long contexts). +std::variant, fine::Error> +state_seq_get_data(ErlNifEnv* env, fine::ResourcePtr ctx, + int64_t seq_id) { + size_t size = llama_state_seq_get_size(ctx->ctx, static_cast(seq_id)); + if (size == 0) { + return fine::Error(std::string("sequence has no state")); + } + + ERL_NIF_TERM bin; + unsigned char* data = enif_make_new_binary(env, size, &bin); + size_t written = llama_state_seq_get_data( + ctx->ctx, data, size, static_cast(seq_id)); + + if (written != size) { + return fine::Error(std::string("state serialization size mismatch")); + } + + return fine::Ok(fine::Term(bin)); +} +FINE_NIF(state_seq_get_data, ERL_NIF_DIRTY_JOB_CPU_BOUND); + +// Restores a previously serialized sequence state into dest_seq_id (which +// must be empty). Returns bytes read. Dirty: KV-sized memcpy. +std::variant, fine::Error> +state_seq_set_data(ErlNifEnv* env, fine::ResourcePtr ctx, + fine::Term state_bin, int64_t dest_seq_id) { + ErlNifBinary bin; + if (!enif_inspect_binary(env, state_bin, &bin)) { + return fine::Error(std::string("state must be a binary")); + } + + size_t read = llama_state_seq_set_data( + ctx->ctx, bin.data, bin.size, static_cast(dest_seq_id)); + + if (read == 0) { + return fine::Error(std::string("state restore failed")); + } + + return fine::Ok(static_cast(read)); +} +FINE_NIF(state_seq_set_data, ERL_NIF_DIRTY_JOB_CPU_BOUND); // --- Chat template --- @@ -956,7 +1271,9 @@ std::string chat_apply_template_jinja( auto result = common_chat_templates_apply(model->chat_templates.get(), inputs); return result.prompt; } -FINE_NIF(chat_apply_template_jinja, 0); +// Dirty: minja template rendering allocates and walks a full AST per call — +// multi-ms for large templates/histories. +FINE_NIF(chat_apply_template_jinja, ERL_NIF_DIRTY_JOB_CPU_BOUND); // --- Speculative decoding (MTP) --- @@ -1008,7 +1325,9 @@ speculative_init( spec, std::move(ctx_tgt), std::move(ctx_dft), static_cast(n_draft), needs_ckpt)); } -FINE_NIF(speculative_init, 0); +// Dirty: common_speculative_init probes the contexts (KV clear + setup work) +// and can block well past the normal-scheduler budget. +FINE_NIF(speculative_init, ERL_NIF_DIRTY_JOB_CPU_BOUND); // Build the live counter snapshot as a flat map { atom => term }. Used by // speculative_stats (queried from Elixir) and by the streaming NIF when it @@ -1097,7 +1416,8 @@ fine::Ok<> generate_mtp_tokens( int64_t max_tokens, int64_t emit_stats_every, ErlNifPid caller_pid, - fine::Term ref) + fine::Term ref, + fine::ResourcePtr cancel) { auto& sp = *spec_res; auto* ctx_tgt = sp.ctx_tgt->ctx; @@ -1107,6 +1427,11 @@ fine::Ok<> generate_mtp_tokens( const llama_seq_id seq_id = 0; const int32_t n_draft = static_cast(sp.n_draft); + // Cancellation: polled per speculative iteration and installed as the + // target context's abort callback so the prompt prefill can stop + // mid-decode (ret == 2 handled by the prefill error path below). + AbortCallbackScope abort_scope(ctx_tgt, cancel.get()); + ErlNifEnv* msg_env = enif_alloc_env(); auto send_error = [&](const std::string& msg) { @@ -1292,6 +1617,12 @@ fine::Ok<> generate_mtp_tokens( // Main speculative loop. while (n_emitted < max_tokens) { + if (cancel->cancelled.load(std::memory_order_relaxed)) { + // Caller abandoned the stream — stop quietly. + enif_free_env(msg_env); + return fine::Ok(); + } + sp.n_iters.fetch_add(1, std::memory_order_relaxed); // Anchor for the "other" bucket: time between known timer ends @@ -1546,12 +1877,17 @@ fine::Ok<> generate_tokens( std::vector prompt_token_ids, int64_t max_tokens, ErlNifPid caller_pid, - fine::Term ref) + fine::Term ref, + fine::ResourcePtr cancel) { auto* ctx = ctx_res->ctx; auto* sampler = sampler_res->sampler; const auto* vocab = ctx_res->model->vocab(); + // Cancellation: polled per generated token AND installed as the abort + // callback so a long prefill decode stops mid-flight (ret == 2). + AbortCallbackScope abort_scope(ctx, cancel.get()); + std::vector prompt_tokens(prompt_token_ids.begin(), prompt_token_ids.end()); if (prompt_tokens.empty()) { @@ -1572,7 +1908,12 @@ fine::Ok<> generate_tokens( for (size_t i = 0; i < prompt_tokens.size(); i += n_batch) { int n = std::min(static_cast(prompt_tokens.size() - i), n_batch); llama_batch batch = llama_batch_get_one(prompt_tokens.data() + i, n); - if (llama_decode(ctx, batch) != 0) { + int ret = llama_decode(ctx, batch); + if (ret == 2) { + // Aborted by the cancel flag mid-prefill — stop quietly. + return fine::Ok(); + } + if (ret != 0) { ErlNifEnv* msg_env = enif_alloc_env(); ERL_NIF_TERM ref_copy = enif_make_copy(msg_env, ref); ERL_NIF_TERM msg = enif_make_tuple2(msg_env, ref_copy, @@ -1598,6 +1939,13 @@ fine::Ok<> generate_tokens( // Generation loop for (int64_t i = 0; i < max_tokens; i++) { + if (cancel->cancelled.load(std::memory_order_relaxed)) { + // Caller abandoned the stream — stop quietly; the consumer is + // not reading messages anymore. + enif_free_env(msg_env); + return fine::Ok(); + } + // llama_sampler_sample() already accepts the selected token; calling // llama_sampler_accept() again would double-advance grammar state. llama_token new_token = llama_sampler_sample(sampler, ctx, -1); @@ -1642,7 +1990,13 @@ fine::Ok<> generate_tokens( // Decode next token llama_batch batch = llama_batch_get_one(&new_token, 1); - if (llama_decode(ctx, batch) != 0) { + int ret = llama_decode(ctx, batch); + if (ret == 2) { + // Aborted by the cancel flag — stop quietly. + enif_free_env(msg_env); + return fine::Ok(); + } + if (ret != 0) { enif_clear_env(msg_env); ref_copy = enif_make_copy(msg_env, ref); ERL_NIF_TERM err_msg = enif_make_tuple2(msg_env, ref_copy, @@ -1674,12 +2028,17 @@ generate( fine::ResourcePtr ctx_res, fine::ResourcePtr sampler_res, std::vector prompt_token_ids, - int64_t max_tokens) + int64_t max_tokens, + fine::ResourcePtr cancel) { auto* ctx = ctx_res->ctx; auto* sampler = sampler_res->sampler; const auto* vocab = ctx_res->model->vocab(); + // Cancellation: polled per token; the abort callback interrupts long + // prefill decodes (ret == 2). A cancelled call returns the partial text. + AbortCallbackScope abort_scope(ctx, cancel.get()); + // Convert prompt tokens std::vector prompt_tokens(prompt_token_ids.begin(), prompt_token_ids.end()); @@ -1693,6 +2052,9 @@ generate( int n = std::min(static_cast(prompt_tokens.size() - i), n_batch); llama_batch batch = llama_batch_get_one(prompt_tokens.data() + i, n); int ret = llama_decode(ctx, batch); + if (ret == 2) { + return fine::Ok(std::string()); + } if (ret != 0) { return fine::Error(std::string("prompt decode failed with code: " + std::to_string(ret))); } @@ -1701,6 +2063,10 @@ generate( // Generation loop std::string result; for (int64_t i = 0; i < max_tokens; i++) { + if (cancel->cancelled.load(std::memory_order_relaxed)) { + break; + } + // llama_sampler_sample() applies the sampler chain, selects a token, and // already accepts it (advancing grammar state / penalties). Do NOT call // llama_sampler_accept() again — a double-accept corrupts grammar state. @@ -1725,6 +2091,9 @@ generate( // Decode the new token for next iteration llama_batch batch = llama_batch_get_one(&new_token, 1); int ret = llama_decode(ctx, batch); + if (ret == 2) { + break; + } if (ret != 0) { return fine::Error(std::string("generation decode failed with code: " + std::to_string(ret))); } @@ -1746,7 +2115,9 @@ json_schema_to_grammar_nif(ErlNifEnv* env, std::string json_str) { return fine::Error(std::string(e.what())); } } -FINE_NIF(json_schema_to_grammar_nif, 0); +// Dirty: JSON parsing + grammar construction scale with schema size and can +// run for milliseconds on real-world schemas. +FINE_NIF(json_schema_to_grammar_nif, ERL_NIF_DIRTY_JOB_CPU_BOUND); // --- Init --- diff --git a/c_src/llama_cpp_ex/llama_nif.h b/c_src/llama_cpp_ex/llama_nif.h index e3e0928..f1298f2 100644 --- a/c_src/llama_cpp_ex/llama_nif.h +++ b/c_src/llama_cpp_ex/llama_nif.h @@ -34,14 +34,39 @@ class LlamaModel { // RAII wrapper for llama_context* // Holds a ResourcePtr to the model to prevent premature GC. +// +// INVARIANT: a context is driven by a single process (the Server GenServer or +// one owning caller). The reusable `batch` below relies on that — decode-side +// NIFs must never run concurrently on the same context from multiple +// processes. (A single process can't overlap NIF calls, so no locking.) class LlamaContext { public: llama_context* ctx; fine::ResourcePtr model; + // Reusable explicit batch for the decode-side NIFs (batch_eval, + // batch_eval_sample, decode_token, prefill), allocated once on first use + // and grown on demand instead of llama_batch_init/free per call. + llama_batch batch{}; + int32_t batch_capacity = 0; + LlamaContext(llama_context* c, fine::ResourcePtr m) : ctx(c), model(std::move(m)) {} + + // Returns the reusable batch with capacity for at least n tokens + // (per-token seq-id capacity 1 — all decode builders use single-seq + // entries). Contents are stale; the caller fills 0..n-1 and n_tokens. + llama_batch& reserve_batch(int32_t n) { + if (batch_capacity < n) { + if (batch_capacity > 0) llama_batch_free(batch); + batch = llama_batch_init(n, 0, 1); + batch_capacity = n; + } + return batch; + } + ~LlamaContext() { + if (batch_capacity > 0) llama_batch_free(batch); if (ctx) llama_free(ctx); } @@ -49,6 +74,15 @@ class LlamaContext { LlamaContext& operator=(const LlamaContext&) = delete; }; +// Cooperative cancellation flag for the stateless generation loops. The +// owning Elixir process holds the resource and sets it via request_cancel/1; +// the generating NIF polls it per iteration and also installs it as the +// context's abort callback so a long prefill aborts mid-decode. +class CancelFlag { +public: + std::atomic cancelled{false}; +}; + // RAII wrapper for llama_sampler* class LlamaSampler { public: diff --git a/lib/llama_cpp_ex.ex b/lib/llama_cpp_ex.ex index 6990631..6b28a5f 100644 --- a/lib/llama_cpp_ex.ex +++ b/lib/llama_cpp_ex.ex @@ -36,8 +36,10 @@ defmodule LlamaCppEx do Grammar, Model, Sampler, + Server, Thinking, - Tokenizer + Tokenizer, + UTF8Stream } @context_opt_keys [ @@ -220,22 +222,42 @@ defmodule LlamaCppEx do # Start: tokenize, create context+sampler, spawn generator {:ok, tokens} = Tokenizer.encode(model, prompt) {:ok, ctx, sampler} = create_gen_resources(model, tokens, cfg) - {ref, gen_pid} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens) - {ref, gen_pid, cfg.timeout} + {ref, gen_pid, cancel} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens) + {ref, gen_pid, cancel, cfg.timeout, _utf8_pending = ""} end, - fn {ref, _gen_pid, timeout} = state -> - receive do - {^ref, {:token, _id, text}} -> {[text], state} - {^ref, outcome} when outcome in [:eog, :done] -> {:halt, state} - {^ref, {:error, _reason}} -> {:halt, state} - after - timeout -> {:halt, state} - end + fn + {:flushed, _ref, _gen_pid, _cancel} = state -> + {:halt, state} + + {ref, gen_pid, cancel, timeout, pending} = state -> + receive do + {^ref, {:token, _id, text}} -> + # Emit only whole codepoints; hold back a split multibyte tail. + {out, pending} = UTF8Stream.push(pending, text) + emitted = if out == "", do: [], else: [out] + {emitted, {ref, gen_pid, cancel, timeout, pending}} + + {^ref, outcome} when outcome in [:eog, :done] -> + flush_then_halt(pending, {:flushed, ref, gen_pid, cancel}) + + {^ref, {:error, _reason}} -> + {:halt, state} + after + timeout -> {:halt, state} + end end, - fn {ref, gen_pid, _timeout} -> stop_generator(ref, gen_pid) end + fn + {:flushed, ref, gen_pid, cancel} -> stop_generator(ref, gen_pid, cancel) + {ref, gen_pid, cancel, _timeout, _pending} -> stop_generator(ref, gen_pid, cancel) + end ) end + # Emits held-back trailing bytes (if any) as a final element, then halts on + # the next pull. + defp flush_then_halt("", state), do: {:halt, state} + defp flush_then_halt(pending, state), do: {[pending], state} + # --- Shared generation plumbing --- # Option handling shared by the generation entry points. @@ -263,20 +285,35 @@ defmodule LlamaCppEx do # Runs the NIF token generator in a linked process that sends # {ref, {:token, id, text} | :eog | :done | {:error, reason}} to the caller. + # The cancel flag lets stop_generator/3 halt the NIF loop cooperatively — + # killing the process alone cannot interrupt a running NIF. defp spawn_generator(ctx, sampler, tokens, max_tokens) do ref = make_ref() parent = self() + cancel = LlamaCppEx.NIF.cancel_flag_new() gen_pid = spawn_link(fn -> - LlamaCppEx.NIF.generate_tokens(ctx.ref, sampler.ref, tokens, max_tokens, parent, ref) + LlamaCppEx.NIF.generate_tokens( + ctx.ref, + sampler.ref, + tokens, + max_tokens, + parent, + ref, + cancel + ) end) - {ref, gen_pid} + {ref, gen_pid, cancel} end - # Kills a generator (if still running) and drains its remaining messages. - defp stop_generator(ref, gen_pid) do + # Cancels a generator (if still running) and drains its remaining messages. + # The flag stops the NIF loop — even mid-prefill via the abort callback — + # so the dirty scheduler is freed promptly instead of decoding to + # max_tokens for a departed consumer. + defp stop_generator(ref, gen_pid, cancel) do + LlamaCppEx.NIF.request_cancel(cancel) Process.unlink(gen_pid) Process.exit(gen_pid, :kill) flush_stream_messages(ref) @@ -290,9 +327,11 @@ defmodule LlamaCppEx do end end - # OpenAI-style finish_reason for a generator outcome message. + # OpenAI-style finish_reason. :eog/:done come from the stateless generator + # protocol; :max_tokens from the Server's completion metadata. defp finish_reason(:eog), do: "stop" defp finish_reason(:done), do: "length" + defp finish_reason(:max_tokens), do: "length" @doc """ Applies the chat template and generates a response. @@ -351,62 +390,110 @@ defmodule LlamaCppEx do completion.choices |> hd() |> Map.get(:message) |> Map.get(:content) + ## Server routing + + When given a running `LlamaCppEx.Server` (pid or name) instead of a + `%Model{}`, the request is served by the batching server: chat templating + and tokenization happen in the caller, and the multi-turn prompt benefits + from the server's prefix cache. Server-only request options like + `:session` and `:cache_prompt` are forwarded. """ - @spec chat_completion(Model.t(), [Chat.message()], keyword()) :: + @spec chat_completion(Model.t() | GenServer.server(), [Chat.message()], keyword()) :: {:ok, ChatCompletion.t()} | {:error, term()} - def chat_completion(%Model{} = model, messages, opts \\ []) when is_list(messages) do + def chat_completion(model_or_server, messages, opts \\ []) + + def chat_completion(%Model{} = model, messages, opts) when is_list(messages) do {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts() cfg = gen_config(gen_opts) with {:ok, prompt} <- Chat.apply_template(model, messages, chat_opts), {:ok, prompt_tokens} <- Tokenizer.encode(model, prompt) do {:ok, ctx, sampler} = create_gen_resources(model, prompt_tokens, cfg) - {ref, gen_pid} = spawn_generator(ctx, sampler, prompt_tokens, cfg.max_tokens) + {ref, gen_pid, cancel} = spawn_generator(ctx, sampler, prompt_tokens, cfg.max_tokens) {texts, finish_reason, completion_tokens} = collect_completion_tokens(ref, cfg.timeout) - # On timeout the generator may still be running — kill it and drain any - # stragglers so they don't pollute the caller's mailbox. - stop_generator(ref, gen_pid) - - raw_text = Enum.join(texts) - enable_thinking = Keyword.get(chat_opts, :enable_thinking, false) - - {reasoning_content, content} = - if enable_thinking do - {rc, c} = Thinking.parse(raw_text) - {if(rc == "", do: nil, else: rc), c} - else - {nil, raw_text} - end - - completion = %ChatCompletion{ - id: "chatcmpl-" <> random_hex(12), - object: "chat.completion", - created: System.os_time(:second), - model: Model.desc(model), - choices: [ - %{ - index: 0, - message: %{ - role: "assistant", - content: content, - reasoning_content: reasoning_content - }, - finish_reason: finish_reason - } - ], - usage: %{ - prompt_tokens: length(prompt_tokens), - completion_tokens: completion_tokens, - total_tokens: length(prompt_tokens) + completion_tokens - } - } + # On timeout the generator may still be running — cancel it and drain + # any stragglers so they don't pollute the caller's mailbox. + stop_generator(ref, gen_pid, cancel) + + completion = + build_chat_completion( + Model.desc(model), + length(prompt_tokens), + Enum.join(texts), + finish_reason, + completion_tokens, + chat_opts + ) + + {:ok, completion} + end + end + + def chat_completion(server, messages, opts) when is_list(messages) do + {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts() + model = Server.get_model(server) + + with {:ok, prompt} <- Chat.apply_template(model, messages, chat_opts), + {:ok, prompt_tokens} <- Tokenizer.encode(model, prompt), + {:ok, result} <- Server.complete_tokens(server, prompt_tokens, gen_opts) do + completion = + build_chat_completion( + Model.desc(model), + length(prompt_tokens), + result.text, + finish_reason(result.finish_reason), + result.completion_tokens, + chat_opts + ) {:ok, completion} end end + defp build_chat_completion( + model_desc, + prompt_tokens, + raw_text, + finish_reason, + completion_tokens, + chat_opts + ) do + enable_thinking = Keyword.get(chat_opts, :enable_thinking, false) + + {reasoning_content, content} = + if enable_thinking do + {rc, c} = Thinking.parse(raw_text) + {if(rc == "", do: nil, else: rc), c} + else + {nil, raw_text} + end + + %ChatCompletion{ + id: "chatcmpl-" <> random_hex(12), + object: "chat.completion", + created: System.os_time(:second), + model: model_desc, + choices: [ + %{ + index: 0, + message: %{ + role: "assistant", + content: content, + reasoning_content: reasoning_content + }, + finish_reason: finish_reason + } + ], + usage: %{ + prompt_tokens: prompt_tokens, + completion_tokens: completion_tokens, + total_tokens: prompt_tokens + completion_tokens + } + } + end + @doc """ Returns a lazy stream of OpenAI-compatible chat completion chunks. @@ -429,8 +516,11 @@ defmodule LlamaCppEx do end) """ - @spec stream_chat_completion(Model.t(), [Chat.message()], keyword()) :: Enumerable.t() - def stream_chat_completion(%Model{} = model, messages, opts \\ []) when is_list(messages) do + @spec stream_chat_completion(Model.t() | GenServer.server(), [Chat.message()], keyword()) :: + Enumerable.t() + def stream_chat_completion(model_or_server, messages, opts \\ []) + + def stream_chat_completion(%Model{} = model, messages, opts) when is_list(messages) do {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts() cfg = gen_config(gen_opts) @@ -439,22 +529,11 @@ defmodule LlamaCppEx do {:ok, prompt} = Chat.apply_template(model, messages, chat_opts) {:ok, tokens} = Tokenizer.encode(model, prompt) {:ok, ctx, sampler} = create_gen_resources(model, tokens, cfg) - {ref, gen_pid} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens) - - enable_thinking = Keyword.get(chat_opts, :enable_thinking, false) + {ref, gen_pid, cancel} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens) - %{ - ref: ref, - gen_pid: gen_pid, - timeout: cfg.timeout, - id: "chatcmpl-" <> random_hex(12), - created: System.os_time(:second), - model: Model.desc(model), - phase: :first, - enable_thinking: enable_thinking, - thinking_parser: - if(enable_thinking, do: Thinking.stream_parser(thinking: true), else: nil) - } + chat_chunk_state(ref, cfg.timeout, Model.desc(model), chat_opts) + |> Map.put(:gen_pid, gen_pid) + |> Map.put(:cancel, cancel) end, fn %{phase: :first} = state -> @@ -463,10 +542,56 @@ defmodule LlamaCppEx do %{phase: :streaming, ref: ref, timeout: timeout} = state -> receive do {^ref, {:token, _id, text}} -> - token_chunks(state, text) + # Whole codepoints only — the thinking parser and consumers + # must never see a split multibyte sequence. + {out, pending} = UTF8Stream.push(state.utf8_pending, text) + state = %{state | utf8_pending: pending} + if out == "", do: {[], state}, else: token_chunks(state, out) {^ref, outcome} when outcome in [:eog, :done] -> - {[chunk(state, %{}, finish_reason(outcome))], %{state | phase: :done}} + {tail, state} = flush_pending_chunks(state) + {tail ++ [chunk(state, %{}, finish_reason(outcome))], %{state | phase: :done}} + + {^ref, {:error, _reason}} -> + {:halt, state} + after + timeout -> {:halt, state} + end + + %{phase: :done} = state -> + {:halt, state} + end, + fn %{ref: ref, gen_pid: gen_pid, cancel: cancel} -> stop_generator(ref, gen_pid, cancel) end + ) + end + + # Server-routed streaming: templating/tokenization in the caller, tokens + # served by the batching server with prefix caching. + def stream_chat_completion(server, messages, opts) when is_list(messages) do + {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts() + timeout = Keyword.get(gen_opts, :timeout, 30_000) + + Stream.resource( + fn -> + model = Server.get_model(server) + {:ok, prompt} = Chat.apply_template(model, messages, chat_opts) + {:ok, tokens} = Tokenizer.encode(model, prompt) + ref = make_ref() + :ok = Server.subscribe_stream_tokens(server, tokens, ref, gen_opts) + + chat_chunk_state(ref, timeout, Model.desc(model), chat_opts) + end, + fn + %{phase: :first} = state -> + {[chunk(state, %{role: "assistant", content: ""}, nil)], %{state | phase: :streaming}} + + %{phase: :streaming, ref: ref, timeout: timeout} = state -> + receive do + {^ref, {:token, text}} -> + token_chunks(state, text) + + {^ref, {:done, reason}} -> + {[chunk(state, %{}, finish_reason(reason))], %{state | phase: :done}} {^ref, {:error, _reason}} -> {:halt, state} @@ -477,10 +602,37 @@ defmodule LlamaCppEx do %{phase: :done} = state -> {:halt, state} end, - fn %{ref: ref, gen_pid: gen_pid} -> stop_generator(ref, gen_pid) end + fn %{ref: ref, phase: phase} -> + # Halted before the final chunk — cancel server-side generation. + if phase != :done, do: Server.cancel(server, ref) + flush_stream_messages(ref) + end ) end + defp flush_pending_chunks(%{utf8_pending: ""} = state), do: {[], state} + + defp flush_pending_chunks(%{utf8_pending: pending} = state) do + token_chunks(%{state | utf8_pending: ""}, pending) + end + + # Shared per-stream state for the chat-completion chunk builders. + defp chat_chunk_state(ref, timeout, model_desc, chat_opts) do + enable_thinking = Keyword.get(chat_opts, :enable_thinking, false) + + %{ + ref: ref, + timeout: timeout, + id: "chatcmpl-" <> random_hex(12), + created: System.os_time(:second), + model: model_desc, + phase: :first, + utf8_pending: "", + enable_thinking: enable_thinking, + thinking_parser: if(enable_thinking, do: Thinking.stream_parser(thinking: true), else: nil) + } + end + # Emits the chunk(s) for one generated token, routing through the thinking # parser when enabled (one token can yield reasoning and content chunks). defp token_chunks(%{enable_thinking: true} = state, text) do diff --git a/lib/llama_cpp_ex/context.ex b/lib/llama_cpp_ex/context.ex index 5fdcdde..2a66b98 100644 --- a/lib/llama_cpp_ex/context.ex +++ b/lib/llama_cpp_ex/context.ex @@ -55,6 +55,10 @@ defmodule LlamaCppEx.Context do Use `:non_causal` for embedding models. * `:no_perf` - Disable performance timing. Defaults to `true`. * `:swa_full` - Use full-size sliding window attention cache. Defaults to `true`. + * `:kv_unified` - Share one KV buffer across all sequences instead of + splitting `n_ctx` evenly between them. Required for cheap cross-sequence + `memory_seq_cp` (prefix sharing); sequences then compete for the shared + `n_ctx` budget. Defaults to `false` (llama.cpp default). ### Speculative decoding / MTP @@ -104,6 +108,7 @@ defmodule LlamaCppEx.Context do attention_type = Keyword.get(opts, :attention_type, :unspecified) |> attention_type_to_int() no_perf = Keyword.get(opts, :no_perf, true) swa_full = Keyword.get(opts, :swa_full, true) + kv_unified = Keyword.get(opts, :kv_unified, false) # Speculative decoding / MTP ctx_type = Keyword.get(opts, :ctx_type, :default) |> ctx_type_to_int() @@ -135,6 +140,7 @@ defmodule LlamaCppEx.Context do attention_type, no_perf, swa_full, + kv_unified, ctx_type, n_rs_seq ) do @@ -193,7 +199,8 @@ defmodule LlamaCppEx.Context do opts \\ [] ) do max_tokens = Keyword.get(opts, :max_tokens, 256) - LlamaCppEx.NIF.generate(ctx_ref, sampler_ref, tokens, max_tokens) + cancel = Keyword.get_lazy(opts, :cancel, fn -> LlamaCppEx.NIF.cancel_flag_new() end) + LlamaCppEx.NIF.generate(ctx_ref, sampler_ref, tokens, max_tokens, cancel) end # --- Enum mappings --- diff --git a/lib/llama_cpp_ex/embedding.ex b/lib/llama_cpp_ex/embedding.ex index bfd0a43..ed3f122 100644 --- a/lib/llama_cpp_ex/embedding.ex +++ b/lib/llama_cpp_ex/embedding.ex @@ -18,13 +18,18 @@ defmodule LlamaCppEx.Embedding do * `:pooling_type` - Pooling type. Defaults to `:unspecified` (model's default). Values: `:unspecified`, `:none`, `:mean`, `:cls`, `:last`. * `:normalize` - Normalization mode. `2` = L2 (default), `0` = max-abs, `-1` = none. + * `:format` - `:list` (default) returns a list of floats; `:binary` + returns the raw native-endian f32 binary as produced by the NIF — + zero-copy and directly loadable via `Nx.from_binary(bin, :f32)`. """ - @spec embed(Model.t(), String.t(), keyword()) :: {:ok, t()} | {:error, String.t()} + @spec embed(Model.t(), String.t(), keyword()) :: + {:ok, t() | binary()} | {:error, String.t()} def embed(%Model{} = model, text, opts \\ []) when is_binary(text) do n_ctx = Keyword.get(opts, :n_ctx, 2048) pooling_type = Keyword.get(opts, :pooling_type, :unspecified) normalize = Keyword.get(opts, :normalize, 2) + format = Keyword.get(opts, :format, :list) {:ok, tokens} = Tokenizer.encode(model, text) ctx_size = max(n_ctx, length(tokens) + 8) @@ -36,7 +41,7 @@ defmodule LlamaCppEx.Embedding do pooling_type: pooling_type ), :ok <- embed_decode(ctx, tokens, 0) do - get_embeddings(ctx, 0, normalize) + get_embeddings(ctx, 0, normalize, format) end end @@ -69,6 +74,7 @@ defmodule LlamaCppEx.Embedding do n_ctx = Keyword.get(opts, :n_ctx, 2048) pooling_type = Keyword.get(opts, :pooling_type, :unspecified) normalize = Keyword.get(opts, :normalize, 2) + format = Keyword.get(opts, :format, :list) max_seqs = Keyword.get(opts, :max_batch_sequences, @default_max_batch_sequences) # texts is non-empty here (embed_batch handles the [] case), so tokenized @@ -86,7 +92,7 @@ defmodule LlamaCppEx.Embedding do pooling_type: pooling_type, n_seq_max: max(max_group, 1) ) do - decode_groups(ctx, groups, normalize) + decode_groups(ctx, groups, normalize, format) end end end @@ -118,17 +124,17 @@ defmodule LlamaCppEx.Embedding do # Decodes each group in one batch and extracts per-sequence embeddings, # preserving the original text order. - defp decode_groups(ctx, groups, normalize) do - with {:ok, nested} <- map_while_ok(groups, &decode_group(ctx, &1, normalize)) do + defp decode_groups(ctx, groups, normalize, format) do + with {:ok, nested} <- map_while_ok(groups, &decode_group(ctx, &1, normalize, format)) do {:ok, Enum.concat(nested)} end end - defp decode_group(ctx, group, normalize) do + defp decode_group(ctx, group, normalize, format) do sequences = Enum.with_index(group, fn tokens, i -> {i, tokens} end) with :ok <- embed_batch_decode(ctx, sequences) do - map_while_ok(0..(length(group) - 1)//1, &get_embeddings(ctx, &1, normalize)) + map_while_ok(0..(length(group) - 1)//1, &get_embeddings(ctx, &1, normalize, format)) end end @@ -162,6 +168,14 @@ defmodule LlamaCppEx.Embedding do defp embed_batch_decode(%Context{ref: ref}, sequences), do: LlamaCppEx.NIF.embed_batch_decode(ref, sequences) - defp get_embeddings(%Context{ref: ref}, seq_id, normalize), - do: LlamaCppEx.NIF.get_embeddings(ref, seq_id, normalize) + # The NIF returns a raw native-endian f32 binary; decode to a float list + # unless the caller asked for the binary itself (Nx interop). + defp get_embeddings(%Context{ref: ref}, seq_id, normalize, format) do + with {:ok, bin} <- LlamaCppEx.NIF.get_embeddings(ref, seq_id, normalize) do + case format do + :binary -> {:ok, bin} + :list -> {:ok, for(<>, do: f)} + end + end + end end diff --git a/lib/llama_cpp_ex/model_manager.ex b/lib/llama_cpp_ex/model_manager.ex index 4627eef..f66e421 100644 --- a/lib/llama_cpp_ex/model_manager.ex +++ b/lib/llama_cpp_ex/model_manager.ex @@ -235,6 +235,44 @@ defmodule LlamaCppEx.ModelManager do ) end + @doc """ + Routes an OpenAI-shaped chat completion to model `id` (or `:default`). + + Server-backed models serve it through the batching server (prefix caching, + `:session` affinity); direct models use the stateless path. + """ + @spec chat_completion(id(), [LlamaCppEx.Chat.message()], keyword()) :: + {:ok, LlamaCppEx.ChatCompletion.t()} | {:error, term()} + def chat_completion(id, messages, opts \\ []) do + with_route( + id, + &LlamaCppEx.chat_completion(&1, messages, opts), + &LlamaCppEx.chat_completion(&1, messages, opts) + ) + end + + @doc """ + Routes a streaming OpenAI-shaped chat completion to model `id` (or `:default`). + + Raises `ArgumentError` if the model is not loaded and ready (a lazy stream + cannot carry an error tuple). + """ + @spec stream_chat_completion(id(), [LlamaCppEx.Chat.message()], keyword()) :: Enumerable.t() + def stream_chat_completion(id, messages, opts \\ []) do + case with_route( + id, + &LlamaCppEx.stream_chat_completion(&1, messages, opts), + &LlamaCppEx.stream_chat_completion(&1, messages, opts) + ) do + {:error, reason} -> + raise ArgumentError, + "cannot stream chat completion from model #{inspect(id)}: #{inspect(reason)}" + + stream -> + stream + end + end + # A server-backed model exposes generate/stream only, so chat templating # happens here before handing the rendered prompt to the server. defp server_chat(pid, messages, opts) do diff --git a/lib/llama_cpp_ex/mtp.ex b/lib/llama_cpp_ex/mtp.ex index 2c473f1..6c35e30 100644 --- a/lib/llama_cpp_ex/mtp.ex +++ b/lib/llama_cpp_ex/mtp.ex @@ -179,6 +179,7 @@ defmodule LlamaCppEx.MTP do ref = make_ref() parent = self() + cancel = LlamaCppEx.NIF.cancel_flag_new() gen_pid = spawn_link(fn -> @@ -189,14 +190,22 @@ defmodule LlamaCppEx.MTP do max_tokens, emit_stats_every, parent, - ref + ref, + cancel ) end) # Keep `sampler` alive for the duration of the stream so it doesn't # get GC'd while the NIF is still using it. `done?` records that a # terminal event was already emitted, so we halt on the next pull. - %{ref: ref, gen_pid: gen_pid, sampler: sampler, timeout: timeout, done?: false} + %{ + ref: ref, + gen_pid: gen_pid, + sampler: sampler, + cancel: cancel, + timeout: timeout, + done?: false + } end, fn %{done?: true} = state -> @@ -213,7 +222,10 @@ defmodule LlamaCppEx.MTP do timeout -> {:halt, state} end end, - fn %{ref: ref, gen_pid: gen_pid} -> + fn %{ref: ref, gen_pid: gen_pid, cancel: cancel} -> + # The flag stops the NIF loop cooperatively — killing the process + # alone cannot interrupt a running NIF. + LlamaCppEx.NIF.request_cancel(cancel) Process.unlink(gen_pid) Process.exit(gen_pid, :kill) flush(ref) diff --git a/lib/llama_cpp_ex/nif.ex b/lib/llama_cpp_ex/nif.ex index 815f15b..6de46ff 100644 --- a/lib/llama_cpp_ex/nif.ex +++ b/lib/llama_cpp_ex/nif.ex @@ -84,6 +84,7 @@ defmodule LlamaCppEx.NIF do _attention_type, _no_perf, _swa_full, + _kv_unified, _ctx_type, _n_rs_seq ), @@ -124,6 +125,11 @@ defmodule LlamaCppEx.NIF do def memory_seq_pos_max(_ctx, _seq_id), do: :erlang.nif_error(:not_loaded) def context_can_seq_rm(_ctx), do: :erlang.nif_error(:not_loaded) + # Sequence state save/restore (RAM prompt cache) + def state_seq_get_size(_ctx, _seq_id), do: :erlang.nif_error(:not_loaded) + def state_seq_get_data(_ctx, _seq_id), do: :erlang.nif_error(:not_loaded) + def state_seq_set_data(_ctx, _state_bin, _dest_seq_id), do: :erlang.nif_error(:not_loaded) + # Chat template def chat_apply_template(_template, _messages, _add_assistant), do: :erlang.nif_error(:not_loaded) @@ -138,8 +144,12 @@ defmodule LlamaCppEx.NIF do ), do: :erlang.nif_error(:not_loaded) + # Cancellation (cooperative flag polled by the stateless generation loops) + def cancel_flag_new, do: :erlang.nif_error(:not_loaded) + def request_cancel(_flag), do: :erlang.nif_error(:not_loaded) + # Streaming generation (sends messages to caller_pid tagged with ref) - def generate_tokens(_ctx, _sampler, _prompt_tokens, _max_tokens, _caller_pid, _ref), + def generate_tokens(_ctx, _sampler, _prompt_tokens, _max_tokens, _caller_pid, _ref, _cancel), do: :erlang.nif_error(:not_loaded) # Speculative decoding (MTP) @@ -154,12 +164,14 @@ defmodule LlamaCppEx.NIF do _max_tokens, _emit_stats_every, _caller_pid, - _ref + _ref, + _cancel ), do: :erlang.nif_error(:not_loaded) # High-level generation - def generate(_ctx, _sampler, _prompt_tokens, _max_tokens), do: :erlang.nif_error(:not_loaded) + def generate(_ctx, _sampler, _prompt_tokens, _max_tokens, _cancel), + do: :erlang.nif_error(:not_loaded) # Embeddings def embed_decode(_ctx, _tokens, _seq_id), do: :erlang.nif_error(:not_loaded) @@ -173,6 +185,10 @@ defmodule LlamaCppEx.NIF do # Continuous batching def batch_eval(_ctx, _entries), do: :erlang.nif_error(:not_loaded) + + def batch_eval_sample(_ctx, _entries, _samplers, _purgeable_seq_ids), + do: :erlang.nif_error(:not_loaded) + def sampler_sample_at(_sampler, _ctx, _idx), do: :erlang.nif_error(:not_loaded) # JSON Schema to Grammar diff --git a/lib/llama_cpp_ex/server.ex b/lib/llama_cpp_ex/server.ex index 9f63092..16194f9 100644 --- a/lib/llama_cpp_ex/server.ex +++ b/lib/llama_cpp_ex/server.ex @@ -79,8 +79,57 @@ defmodule LlamaCppEx.Server do * `:server` - PID of the server process. * `:seq_id` - Slot sequence ID (integer). * `:mode` - `:generate` or `:stream`. - * `:stop_reason` - `:eog` (end-of-generation token sampled) or - `:max_tokens` (request `max_tokens` reached). + * `:stop_reason` - `:eog` (end-of-generation token sampled), + `:max_tokens` (request `max_tokens` reached), or `:cancelled` (consumer + died or cancelled the request). + + ### `[:llama_cpp_ex, :server, :kv_pressure]` + + Emitted when a forward pass hit KV-cache pressure (`llama_decode == 1`) and + the server recovered by purging idle slots' cached prefixes and/or splitting + the batch. + + Measurements: + + * `:purged_slots` - Number of idle slots whose cached KV was dropped. + * `:batch_splits` - Number of times the batch was halved to fit. + + Metadata: + + * `:server` - PID of the server process. + * `:purged_seq_ids` - Sequence IDs whose caches were purged. + + ### `[:llama_cpp_ex, :server, :prefix_instability]` + + Emitted when a cache-eligible request matches only 10–50% of a slot's + cached history — the signature of a chat template that rewrites earlier + turns (e.g. stripping thinking blocks), silently defeating prefix caching. + + Measurements: + + * `:matched_tokens` - Length of the common prefix actually reusable. + * `:cached_tokens` - Length of the cached history that was expected to match. + + Metadata: + + * `:server` - PID of the server process. + * `:seq_id` - Slot sequence ID. + + ### `[:llama_cpp_ex, :server, :ram_cache]` + + Emitted on level-2 RAM prompt cache activity (see `:prompt_cache_ram_mb`). + + Measurements: + + * `:bytes` - Size of the entry involved. + * `:tokens` - Cached prefix length of the entry involved. + * `:total_bytes` - Cache size after the operation. + * `:entries` - Entry count after the operation. + + Metadata: + + * `:server` - PID of the server process. + * `:op` - `:save`, `:restore`, or `:evict`. ### `[:llama_cpp_ex, :server, :request, :exception]` @@ -103,7 +152,39 @@ defmodule LlamaCppEx.Server do require Logger - alias LlamaCppEx.{Context, Model, Sampler, Tokenizer} + alias LlamaCppEx.{Context, Model, Sampler, Tokenizer, UTF8Stream} + + # Sampling options accepted both at server start (defaults) and per request + # (overrides). Keep in sync with LlamaCppEx.Sampler.create/2. + @sampler_opt_keys [ + :seed, + :temp, + :top_k, + :top_p, + :min_p, + :penalty_repeat, + :penalty_freq, + :penalty_present, + :grammar, + :grammar_root + ] + + # Per-request options accepted by generate/stream and carried through the + # queue into the slot. Server-level options provide the defaults. + @request_opt_keys [:cache_prompt, :session] ++ @sampler_opt_keys + + # Evicted caches below this many tokens are not worth a KV-sized state copy. + @ram_cache_min_tokens 32 + + # Restore a RAM cache entry only when the usable fraction clears this bar + # (llama-server's f_keep heuristic) — restoring is a KV-sized memcpy. + @ram_cache_min_keep 0.25 + + # Prefix-instability heuristic: a cached history this long that matches + # 10–50% of its tokens looks like a chat template rewriting history + # (thinking-strip etc.) rather than a brand-new conversation (~0%) or a + # clean continuation (~100%). + @prefix_instability_min_cached 64 defstruct [ :model, @@ -111,10 +192,22 @@ defmodule LlamaCppEx.Server do :sampler_opts, slots: %{}, queue: nil, + max_queue: 0, + sessions: %{}, n_parallel: 4, n_batch: 2048, chunk_size: 512, - cache_prompt: false, + cache_prompt: true, + # True when cross-slot prefix sharing is safe: unified KV (partial + # cross-stream seq_cp aborts in split mode) AND partial seq_rm support + # (hybrid GDN recurrent state cannot be partially copied). + cross_slot_sharing: false, + # Level-2 prompt cache: evicted slot KV states saved to RAM, FIFO under a + # byte budget. 0 MB = disabled. + prompt_cache_ram_mb: 0, + ram_cache: [], + ram_cache_bytes: 0, + prefix_instability_warned: false, # `:part`/`:rs` = partial seq_rm works; `:full` = whole-sequence only # (hybrid GDN models like Qwen 3.5/3.6); `:no` = no memory module. We # only do prefix-cache partial trims when this is `:part` or `:rs`. @@ -134,11 +227,33 @@ defmodule LlamaCppEx.Server do * `:n_gpu_layers` - GPU layers. Defaults to `99`. * `:n_ctx` - Total context size (shared across slots). Defaults to `8192`. * `:n_parallel` - Number of concurrent slots. Defaults to `4`. - * `:n_batch` - Batch size. Defaults to `n_ctx`. + * `:n_batch` - Max tokens per forward pass. Defaults to `min(n_ctx, 2048)`. + Bounds worst-case tick latency: one huge prompt can occupy at most + `n_batch` tokens of a tick, so decode tokens of other slots are never + delayed by more than one `n_batch`-sized pass. Raise it for pure batch + throughput (fewer, larger passes); lower it (or lower `:chunk_size`) for + smoother streaming latency under mixed load. * `:chunk_size` - Max prefill tokens per slot per tick. Defaults to `512`. - * `:max_queue` - Max queued requests. `0` for unlimited. Defaults to `0`. + * `:max_queue` - Max queued requests waiting for a slot. When the bound is + hit, calls return `{:error, :queue_full}` immediately and streams emit a + single `{:error, :queue_full}` element — no silent queueing until the + call timeout. `0` for unlimited. Defaults to `0`; `2 * n_parallel` is a + reasonable bound for latency-sensitive deployments. * `:cache_prompt` - Retain KV cache between requests on the same slot for - prefix reuse. Defaults to `false`. Set to `true` for multi-turn chat. + prefix reuse. Defaults to `true` (matching llama-server). Overridable + per request via the `:cache_prompt` option on `generate/3` and friends. + * `:kv_unified` - Share one KV buffer across all slots instead of splitting + `n_ctx` evenly (`n_ctx/n_parallel` each). Enables cross-slot prefix + sharing: a system prompt cached by any slot is adopted by every other + slot via a metadata-only copy. Slots then compete for the shared `n_ctx` + budget, and idle slots' caches are purged under KV pressure. Defaults to + `true`. Set `false` for strictly isolated per-slot budgets. + * `:prompt_cache_ram_mb` - Byte budget (in MB) for the level-2 RAM prompt + cache: when a slot's cached prefix is about to be destroyed it is + serialized to RAM and can be restored later instead of re-prefilling. + State blobs are KV-sized (up to hundreds of MB for long contexts) — + entries larger than the budget are never stored, so a small budget + degrades to "disabled" rather than OOM. Defaults to `0` (off). * `:batch_strategy` - Batch building strategy module. Defaults to `LlamaCppEx.Server.Strategy.DecodeMaximal`. See `LlamaCppEx.Server.BatchStrategy`. * Sampling options: `:temp`, `:top_k`, `:top_p`, `:min_p`, `:seed`, `:penalty_repeat`, @@ -159,6 +274,14 @@ defmodule LlamaCppEx.Server do * `:max_tokens` - Maximum tokens to generate. Defaults to `256`. * `:timeout` - Call timeout in ms. Defaults to `60_000`. + * `:cache_prompt` - Reuse/retain this request's KV prefix on the slot. + Defaults to the server-level `:cache_prompt` setting. + * `:session` - Any term identifying a conversation. Requests with the same + session are routed to the same slot whenever it is free, keeping their + cached prefix intact under concurrency. + * Sampling options (`:temp`, `:top_k`, `:top_p`, `:min_p`, `:seed`, + `:penalty_repeat`, `:penalty_freq`, `:penalty_present`, `:grammar`, + `:grammar_root`) - override the server-level defaults for this request. """ @spec generate(GenServer.server(), String.t(), keyword()) :: @@ -166,48 +289,86 @@ defmodule LlamaCppEx.Server do def generate(server, prompt, opts \\ []) do timeout = Keyword.get(opts, :timeout, 60_000) max_tokens = Keyword.get(opts, :max_tokens, 256) - GenServer.call(server, {:generate, prompt, max_tokens}, timeout) + req_opts = Keyword.take(opts, @request_opt_keys) + + # Tokenize in the caller: parallel across clients and off the server's + # mailbox. The model handle comes from a :persistent_term cache. + {:ok, token_ids} = Tokenizer.encode(get_model(server), prompt) + GenServer.call(server, {:generate_tokens, token_ids, max_tokens, req_opts}, timeout) end @doc """ Returns a stream of generated text chunks. + If the request is rejected (`:queue_full`) or fails mid-generation, the + stream emits a single `{:error, reason}` element and halts — consumers that + need to distinguish errors from text should match on it. + ## Options * `:max_tokens` - Maximum tokens to generate. Defaults to `256`. * `:timeout` - Per-token timeout. Defaults to `30_000`. + Also accepts the per-request options documented on `generate/3`. """ @spec stream(GenServer.server(), String.t(), keyword()) :: Enumerable.t() def stream(server, prompt, opts \\ []) do max_tokens = Keyword.get(opts, :max_tokens, 256) timeout = Keyword.get(opts, :timeout, 30_000) + req_opts = Keyword.take(opts, @request_opt_keys) Stream.resource( fn -> + {:ok, token_ids} = Tokenizer.encode(get_model(server), prompt) ref = make_ref() - :ok = GenServer.call(server, {:stream, prompt, max_tokens, self(), ref}) - {ref, timeout} - end, - fn {ref, timeout} = state -> - receive do - {^ref, {:token, text}} -> {[text], state} - {^ref, :done} -> {:halt, state} - {^ref, {:error, _reason}} -> {:halt, state} - after - timeout -> {:halt, state} + + case GenServer.call( + server, + {:stream_tokens, token_ids, max_tokens, self(), ref, req_opts} + ) do + :ok -> {server, ref, timeout} + {:error, reason} -> {:rejected, ref, reason} end end, - fn {ref, _timeout} -> - receive do - {^ref, _} -> :ok - after - 0 -> :ok - end - end + &stream_next/1, + &stream_cleanup/1 ) end + # Shared next/cleanup functions for the piece-streaming Stream.resources. + # A rejected request ({:error, :queue_full}) surfaces as a single error + # element instead of silence-then-timeout. Halting an ACTIVE stream early + # (cleanup before :done) cancels the request server-side — otherwise an + # abandoned stream would keep consuming batch budget to max_tokens. + defp stream_next({:rejected, ref, reason}), do: {[{:error, reason}], {:done, ref}} + defp stream_next({:done, _ref} = state), do: {:halt, state} + + defp stream_next({server, ref, timeout} = state) do + receive do + {^ref, {:token, text}} -> {[text], state} + {^ref, {:done, _reason}} -> {:halt, {:done, ref}} + {^ref, {:error, reason}} -> {[{:error, reason}], {:done, ref}} + after + timeout -> {:halt, {server, ref, timeout}} + end + end + + defp stream_cleanup({:rejected, ref, _reason}), do: drain_stream_messages(ref) + defp stream_cleanup({:done, ref}), do: drain_stream_messages(ref) + + defp stream_cleanup({server, ref, _timeout}) do + cancel(server, ref) + drain_stream_messages(ref) + end + + defp drain_stream_messages(ref) do + receive do + {^ref, _} -> drain_stream_messages(ref) + after + 0 -> :ok + end + end + @doc """ Generates text from pre-tokenized input. Blocks until generation is complete. @@ -224,7 +385,37 @@ defmodule LlamaCppEx.Server do def generate_tokens(server, token_ids, opts \\ []) when is_list(token_ids) do timeout = Keyword.get(opts, :timeout, 60_000) max_tokens = Keyword.get(opts, :max_tokens, 256) - GenServer.call(server, {:generate_tokens, token_ids, max_tokens}, timeout) + req_opts = Keyword.take(opts, @request_opt_keys) + GenServer.call(server, {:generate_tokens, token_ids, max_tokens, req_opts}, timeout) + end + + @doc """ + Like `generate_tokens/3`, but returns completion metadata alongside the text. + + Returns `{:ok, %{text: text, completion_tokens: n, finish_reason: reason}}` + where `reason` is `:eog` or `:max_tokens`. Used by the OpenAI-shaped + `LlamaCppEx.chat_completion/3` when routed through a server. + """ + @spec complete_tokens(GenServer.server(), [integer()], keyword()) :: + {:ok, %{text: String.t(), completion_tokens: non_neg_integer(), finish_reason: atom()}} + | {:error, term()} + def complete_tokens(server, token_ids, opts \\ []) when is_list(token_ids) do + timeout = Keyword.get(opts, :timeout, 60_000) + max_tokens = Keyword.get(opts, :max_tokens, 256) + req_opts = Keyword.take(opts, @request_opt_keys) ++ [reply: :full] + GenServer.call(server, {:generate_tokens, token_ids, max_tokens, req_opts}, timeout) + end + + @doc false + # Subscribes `self()` to a token stream using the raw message protocol: + # {ref, {:token, piece}} per piece, then {ref, {:done, :eog | :max_tokens}} + # or {ref, {:error, reason}}. Internal — used by the chat-completion + # streaming path; library users should call stream/3 or stream_tokens/3. + @spec subscribe_stream_tokens(GenServer.server(), [integer()], reference(), keyword()) :: :ok + def subscribe_stream_tokens(server, token_ids, ref, opts \\ []) when is_list(token_ids) do + max_tokens = Keyword.get(opts, :max_tokens, 256) + req_opts = Keyword.take(opts, @request_opt_keys) + GenServer.call(server, {:stream_tokens, token_ids, max_tokens, self(), ref, req_opts}) end @doc """ @@ -240,41 +431,56 @@ defmodule LlamaCppEx.Server do def stream_tokens(server, token_ids, opts \\ []) when is_list(token_ids) do max_tokens = Keyword.get(opts, :max_tokens, 256) timeout = Keyword.get(opts, :timeout, 30_000) + req_opts = Keyword.take(opts, @request_opt_keys) Stream.resource( fn -> ref = make_ref() - :ok = GenServer.call(server, {:stream_tokens, token_ids, max_tokens, self(), ref}) - {ref, timeout} - end, - fn {ref, timeout} = state -> - receive do - {^ref, {:token, text}} -> {[text], state} - {^ref, :done} -> {:halt, state} - {^ref, {:error, _reason}} -> {:halt, state} - after - timeout -> {:halt, state} + + case GenServer.call( + server, + {:stream_tokens, token_ids, max_tokens, self(), ref, req_opts} + ) do + :ok -> {server, ref, timeout} + {:error, reason} -> {:rejected, ref, reason} end end, - fn {ref, _timeout} -> - receive do - {^ref, _} -> :ok - after - 0 -> :ok - end - end + &stream_next/1, + &stream_cleanup/1 ) end + @doc """ + Cancels an in-flight or queued stream request by its subscription reference. + + The slot stops being scheduled immediately and is freed for other requests + (its prefix cache is retained per the request's `:cache_prompt`). Consumer + death is detected automatically via monitors — explicit cancel is for + consumers that stop reading without exiting. `Server.stream/3` and + `stream_tokens/3` call this from their cleanup, so halting those streams + early (e.g. `Enum.take/2`) cancels generation instead of burning batch + budget to `max_tokens`. + """ + @spec cancel(GenServer.server(), reference()) :: :ok + def cancel(server, ref) when is_reference(ref) do + GenServer.cast(server, {:cancel, ref}) + end + @doc """ Returns the model struct for external tokenization. The model resource is reference-counted and thread-safe for read-only - operations like tokenization. + operations like tokenization. Served from a `:persistent_term` cache — no + round-trip through the server's mailbox. """ @spec get_model(GenServer.server()) :: Model.t() def get_model(server) do - GenServer.call(server, :get_model) + with pid when is_pid(pid) <- GenServer.whereis(server), + %Model{} = model <- :persistent_term.get({__MODULE__, pid}, nil) do + model + else + _ -> GenServer.call(server, :get_model) + end end @doc """ @@ -293,24 +499,15 @@ defmodule LlamaCppEx.Server do n_gpu_layers = Keyword.get(opts, :n_gpu_layers, 99) n_parallel = Keyword.get(opts, :n_parallel, 4) n_ctx = Keyword.get(opts, :n_ctx, 8192) - n_batch = Keyword.get(opts, :n_batch, n_ctx) + n_batch = Keyword.get(opts, :n_batch, min(n_ctx, 2048)) chunk_size = Keyword.get(opts, :chunk_size, 512) - cache_prompt = Keyword.get(opts, :cache_prompt, false) + cache_prompt = Keyword.get(opts, :cache_prompt, true) + kv_unified = Keyword.get(opts, :kv_unified, true) + prompt_cache_ram_mb = Keyword.get(opts, :prompt_cache_ram_mb, 0) + max_queue = Keyword.get(opts, :max_queue, 0) batch_strategy = Keyword.get(opts, :batch_strategy, LlamaCppEx.Server.Strategy.DecodeMaximal) - sampler_opts = - Keyword.take(opts, [ - :seed, - :temp, - :top_k, - :top_p, - :min_p, - :penalty_repeat, - :penalty_freq, - :penalty_present, - :grammar, - :grammar_root - ]) + sampler_opts = Keyword.take(opts, @sampler_opt_keys) model_opts = Keyword.take(opts, [ @@ -342,13 +539,22 @@ defmodule LlamaCppEx.Server do :swa_full ]) + # Trap exits so terminate/2 reliably erases the persistent_term model + # cache on shutdown. + Process.flag(:trap_exit, true) + :ok = LlamaCppEx.init() {:ok, model} = Model.load(model_path, [n_gpu_layers: n_gpu_layers] ++ model_opts) + # Cache the model handle for callers (get_model/1): tokenization and chat + # templating happen client-side without a GenServer.call round-trip. + :persistent_term.put({__MODULE__, self()}, model) + {:ok, ctx} = Context.create( model, - [n_ctx: n_ctx, n_batch: n_batch, n_seq_max: n_parallel] ++ context_opts + [n_ctx: n_ctx, n_batch: n_batch, n_seq_max: n_parallel, kv_unified: kv_unified] ++ + context_opts ) # Probe seq_rm support BEFORE any decode work — the call has the side @@ -368,10 +574,19 @@ defmodule LlamaCppEx.Server do ) end + now = System.monotonic_time() + slots = for seq_id <- 0..(n_parallel - 1), into: %{} do {:ok, sampler} = Sampler.create(model, sampler_opts) - {seq_id, Map.put(idle_slot_fields([], 0), :sampler, sampler)} + + slot = + idle_slot_fields([], 0) + |> Map.put(:sampler, sampler) + |> Map.put(:t_last_used, now) + |> Map.put(:session, nil) + + {seq_id, slot} end state = %__MODULE__{ @@ -380,10 +595,13 @@ defmodule LlamaCppEx.Server do sampler_opts: sampler_opts, slots: slots, queue: :queue.new(), + max_queue: max_queue, n_parallel: n_parallel, n_batch: n_batch, chunk_size: chunk_size, cache_prompt: cache_prompt, + cross_slot_sharing: kv_unified and seq_rm_kind == :part, + prompt_cache_ram_mb: prompt_cache_ram_mb, seq_rm_kind: seq_rm_kind, batch_strategy: batch_strategy } @@ -392,67 +610,41 @@ defmodule LlamaCppEx.Server do end @impl true - def handle_call({:generate, prompt, max_tokens}, from, state) do - {:ok, tokens} = Tokenizer.encode(state.model, prompt) - - case acquire_slot(state, tokens) do - {:ok, seq_id, state} -> - state = init_slot(state, seq_id, tokens, max_tokens, from, nil, nil) - state = maybe_schedule_tick(state) - {:noreply, state} - - :no_slots -> - state = enqueue_request(state, {:generate, tokens, max_tokens, from, nil, nil}) - {:noreply, state} - end - end - - def handle_call({:generate_tokens, token_ids, max_tokens}, from, state) do + def handle_call({:generate_tokens, token_ids, max_tokens, req_opts}, from, state) do if token_ids == [] do - {:reply, {:error, "token list cannot be empty"}, state} + {:reply, {:error, :empty_prompt}, state} else - case acquire_slot(state, token_ids) do + case acquire_slot(state, token_ids, req_opts) do {:ok, seq_id, state} -> - state = init_slot(state, seq_id, token_ids, max_tokens, from, nil, nil) + state = init_slot(state, seq_id, token_ids, max_tokens, from, nil, nil, req_opts) state = maybe_schedule_tick(state) {:noreply, state} :no_slots -> - state = enqueue_request(state, {:generate, token_ids, max_tokens, from, nil, nil}) - {:noreply, state} + request = {:generate, token_ids, max_tokens, from, nil, nil, req_opts} + enqueue_or_reject(state, request, from, _reply_ok? = false) end end end - def handle_call({:stream, prompt, max_tokens, pid, ref}, from, state) do - {:ok, tokens} = Tokenizer.encode(state.model, prompt) - - case acquire_slot(state, tokens) do - {:ok, seq_id, state} -> - state = init_slot(state, seq_id, tokens, max_tokens, nil, pid, ref) - GenServer.reply(from, :ok) - state = maybe_schedule_tick(state) - {:noreply, state} - - :no_slots -> - GenServer.reply(from, :ok) - state = enqueue_request(state, {:stream, tokens, max_tokens, nil, pid, ref}) - {:noreply, state} - end - end - - def handle_call({:stream_tokens, token_ids, max_tokens, pid, ref}, from, state) do - case acquire_slot(state, token_ids) do - {:ok, seq_id, state} -> - state = init_slot(state, seq_id, token_ids, max_tokens, nil, pid, ref) - GenServer.reply(from, :ok) - state = maybe_schedule_tick(state) - {:noreply, state} + def handle_call({:stream_tokens, token_ids, max_tokens, pid, ref, req_opts}, from, state) do + if token_ids == [] do + # Same guard as :generate_tokens — an empty prompt would enter the tick + # with nothing to prefill and no logits to sample, hanging the consumer + # until its stream timeout. + {:reply, {:error, :empty_prompt}, state} + else + case acquire_slot(state, token_ids, req_opts) do + {:ok, seq_id, state} -> + state = init_slot(state, seq_id, token_ids, max_tokens, nil, pid, ref, req_opts) + GenServer.reply(from, :ok) + state = maybe_schedule_tick(state) + {:noreply, state} - :no_slots -> - GenServer.reply(from, :ok) - state = enqueue_request(state, {:stream, token_ids, max_tokens, nil, pid, ref}) - {:noreply, state} + :no_slots -> + request = {:stream, token_ids, max_tokens, nil, pid, ref, req_opts} + enqueue_or_reject(state, request, from, _reply_ok? = true) + end end end @@ -472,12 +664,25 @@ defmodule LlamaCppEx.Server do prefilling_slots: counts.prefilling, queue_depth: :queue.len(state.queue), n_parallel: state.n_parallel, - n_batch: state.n_batch + n_batch: state.n_batch, + ram_cache_entries: length(state.ram_cache), + ram_cache_bytes: state.ram_cache_bytes } {:reply, stats, state} end + @impl true + def handle_cast({:cancel, ref}, state) do + state = + case Enum.find(state.slots, fn {_id, slot} -> slot.stream_ref == ref end) do + {seq_id, _slot} -> cancel_slot(state, seq_id) + nil -> drop_queued_request(state, ref) + end + + {:noreply, state} + end + @impl true def handle_info(:tick, state) do state = %{state | tick_scheduled: false} @@ -485,31 +690,116 @@ defmodule LlamaCppEx.Server do {:noreply, state} end + def handle_info({:DOWN, mref, :process, pid, _reason}, state) do + # A request's consumer died — free its slot instead of generating into + # the void until max_tokens. Queued requests from the dead pid are + # dropped as well. + state = + case Enum.find(state.slots, fn {_id, slot} -> slot.monitor_ref == mref end) do + {seq_id, _slot} -> cancel_slot(state, seq_id) + nil -> state + end + + {:noreply, drop_queued_requests_from(state, pid)} + end + + def handle_info({:EXIT, _pid, reason}, state) do + # trap_exit is on (for terminate/2 cleanup) — honor exit signals. + {:stop, reason, state} + end + + @impl true + def terminate(_reason, _state) do + :persistent_term.erase({__MODULE__, self()}) + :ok + end + # --- Internal: Slot management --- - defp acquire_slot(state, tokens) do + defp request_cache_prompt?(state, req_opts) do + Keyword.get(req_opts, :cache_prompt, state.cache_prompt) + end + + # Queues a request for the next free slot, or rejects it when :max_queue is + # bound and full. Stream subscriptions reply :ok up-front (tokens arrive as + # messages); sync requests stay unanswered until finish_slot replies. + defp enqueue_or_reject(state, request, from, reply_ok?) do + if state.max_queue > 0 and :queue.len(state.queue) >= state.max_queue do + {:reply, {:error, :queue_full}, state} + else + if reply_ok?, do: GenServer.reply(from, :ok) + {:noreply, enqueue_request(state, request)} + end + end + + defp acquire_slot(state, tokens, req_opts) do idle_slots = Enum.filter(state.slots, fn {_id, slot} -> slot.state == :idle end) case idle_slots do [] -> :no_slots - slots when state.cache_prompt and tokens != [] -> - # Prefer the slot with the longest cached prefix match - {best_id, _} = - Enum.max_by(slots, fn {_id, slot} -> - common_prefix_length(tokens, slot.cached_tokens) - end) + slots -> + session_pick = session_slot_if_idle(state, Keyword.get(req_opts, :session), slots) + + cond do + session_pick != nil -> + {:ok, session_pick, state} + + tokens != [] and request_cache_prompt?(state, req_opts) -> + {:ok, pick_cached_slot(slots, tokens), state} + + true -> + {:ok, pick_lru_slot(slots), state} + end + end + end - {:ok, best_id, state} + @doc false + # Session affinity: the slot that served this session last, if it is idle. + # Overrides the similarity rule — the session's cache lives there. + def session_slot_if_idle(_state, nil, _idle_slots), do: nil + + def session_slot_if_idle(state, session, idle_slots) do + with seq_id when seq_id != nil <- Map.get(state.sessions, session), + true <- Enum.any?(idle_slots, fn {id, _} -> id == seq_id end) do + seq_id + else + _ -> nil + end + end - [{seq_id, _} | _] -> - {:ok, seq_id, state} + @doc false + # llama-server's slot-pick rule (server-context.cpp): reuse the slot with the + # best cached-prefix similarity only when it clears a threshold + # (LCP/prompt_len > 0.1); otherwise take the least-recently-used idle slot so + # a tiny unrelated request doesn't evict a valuable long cache. + # Public (like common_prefix_length/2) for direct unit testing. + def pick_cached_slot(idle_slots, tokens) do + prompt_len = length(tokens) + + {best_id, best_lcp} = + idle_slots + |> Enum.map(fn {id, slot} -> {id, common_prefix_length(tokens, slot.cached_tokens)} end) + |> Enum.max_by(fn {_id, lcp} -> lcp end) + + if best_lcp / prompt_len > 0.1 do + best_id + else + pick_lru_slot(idle_slots) end end - defp init_slot(state, seq_id, tokens, max_tokens, from, stream_pid, stream_ref) do + @doc false + def pick_lru_slot(idle_slots) do + {seq_id, _} = Enum.min_by(idle_slots, fn {_id, slot} -> slot.t_last_used end) + seq_id + end + + defp init_slot(state, seq_id, tokens, max_tokens, from, stream_pid, stream_ref, req_opts) do + state = update_session_mapping(state, seq_id, Keyword.get(req_opts, :session)) slot = state.slots[seq_id] + cache_prompt? = request_cache_prompt?(state, req_opts) # Prefix cache: find common prefix with cached KV. On models that only # support whole-sequence seq_rm (`:full`, e.g. hybrid GDN), a partial @@ -517,38 +807,45 @@ defmodule LlamaCppEx.Server do # an M-RoPE position-mismatch abort on the next decode. Disable the # cache hit in that case unless the new prompt extends the old one # exactly (no trim needed). + # + # A cached match may never cover the WHOLE prompt: at least the last + # prompt token must be decoded to produce logits for sampling the first + # generated token (llama-server does the same n_past-- adjustment). An + # uncapped full match would enter the tick with nothing to prefill and no + # logits to sample — a stuck slot. + max_reuse = length(tokens) - 1 + raw_match = - if state.cache_prompt do - common_prefix_length(tokens, slot.cached_tokens) + if cache_prompt? do + min(common_prefix_length(tokens, slot.cached_tokens), max_reuse) else 0 end needs_trim = raw_match > 0 and raw_match < slot.cached_pos - n_match = + own_match = if needs_trim and state.seq_rm_kind == :full do 0 else raw_match end - cond do - n_match > 0 and n_match < slot.cached_pos -> - # Trim KV cache beyond the matched prefix (only safe on `:part`/`:rs`). - true = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, n_match, -1) + state = maybe_warn_prefix_instability(state, seq_id, slot, raw_match, cache_prompt?) - n_match > 0 -> - # Exact-prefix continuation; nothing to trim. - :noop + {state, n_match} = resolve_prefix_cache(state, seq_id, tokens, own_match, cache_prompt?) + slot = state.slots[seq_id] - true -> - # No usable match — clear everything for this slot. - _ = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) - end + # Fresh sampler per request: request opts override server defaults, and a + # new chain means clean grammar/penalty state and a fresh seed. The old + # sampler resource is dropped and freed by GC. + sampler_opts = Keyword.merge(state.sampler_opts, Keyword.take(req_opts, @sampler_opt_keys)) + {:ok, sampler} = Sampler.create(state.model, sampler_opts) - # Reset sampler for fresh generation - Sampler.reset(slot.sampler) + # Watch the consumer (stream subscriber or sync caller) so its death + # cancels the request instead of burning batch budget to max_tokens. + watch_pid = stream_pid || caller_pid(from) + monitor_ref = if watch_pid, do: Process.monitor(watch_pid) slot = %{ slot @@ -556,11 +853,16 @@ defmodule LlamaCppEx.Server do from: from, stream_pid: stream_pid, stream_ref: stream_ref, + monitor_ref: monitor_ref, + sampler: sampler, + reply_mode: Keyword.get(req_opts, :reply, :text), + cache_prompt: cache_prompt?, prompt_tokens: tokens, prompt_tokens_tuple: List.to_tuple(tokens), prefill_pos: n_match, pos: n_match, pending_token: nil, + pending_eog: false, batch_idx: -1, tokens_generated: 0, max_tokens: max_tokens, @@ -581,23 +883,329 @@ defmodule LlamaCppEx.Server do put_in(state.slots[seq_id], slot) end + # Decides where this request's cached prefix comes from and prepares the + # slot's KV accordingly. The three sources compete by match length, with + # ties broken by cost: the slot's own cache (free) beats a donor slot + # (metadata-only seq_cp under unified KV) beats the RAM prompt cache + # (KV-sized memcpy). Whenever the slot's own cache is about to be destroyed + # or truncated, it is offered to the RAM cache first. + defp resolve_prefix_cache(state, seq_id, tokens, own_match, cache_prompt?) do + slot = state.slots[seq_id] + + donor = best_donor(state, seq_id, tokens, own_match, cache_prompt?) + donor_lcp = if donor, do: elem(donor, 1), else: 0 + + ram = if cache_prompt?, do: best_ram_candidate(state, tokens), else: nil + ram_lcp = if ram, do: elem(ram, 1), else: 0 + + cond do + donor_lcp > own_match and donor_lcp >= ram_lcp -> + adopt_donor_cache(state, seq_id, slot, donor) + + ram_lcp > own_match -> + adopt_ram_cache(state, seq_id, slot, ram) + + own_match > 0 -> + keep_own_cache(state, seq_id, slot, own_match) + + true -> + # No usable match anywhere — clear the slot. + state = maybe_save_to_ram_cache(state, seq_id, slot) + _ = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + {state, 0} + end + end + + defp adopt_donor_cache(state, seq_id, slot, {donor_id, donor_lcp}) do + state = maybe_save_to_ram_cache(state, seq_id, slot) + _ = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + :ok = LlamaCppEx.NIF.memory_seq_cp(state.ctx.ref, donor_id, seq_id, 0, donor_lcp) + {state, donor_lcp} + end + + defp adopt_ram_cache(state, seq_id, slot, {entry, ram_lcp}) do + state = maybe_save_to_ram_cache(state, seq_id, slot) + _ = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + {state, apply_ram_restore(state, seq_id, entry, ram_lcp)} + end + + defp keep_own_cache(state, seq_id, slot, own_match) do + if own_match < slot.cached_pos do + # Trim KV cache beyond the matched prefix (only safe on `:part`/`:rs`). + # The truncated tail may still be valuable to another conversation — + # offer the full state to the RAM cache before cutting it. + state = maybe_save_to_ram_cache(state, seq_id, slot) + true = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, own_match, -1) + {state, own_match} + else + # Exact-prefix continuation; nothing to trim. + {state, own_match} + end + end + + # A cache-eligible request whose prompt shares only a partial prefix with + # the slot's cached history usually means the chat template rewrote earlier + # turns (e.g. Qwen thinking-strip re-rendering) — every such request pays a + # near-full re-prefill that the caller probably believes is cached. Emit + # telemetry per occurrence and warn in the log once per server. + defp maybe_warn_prefix_instability(state, seq_id, slot, raw_match, cache_prompt?) do + suspicious? = + cache_prompt? and slot.cached_pos >= @prefix_instability_min_cached and + raw_match > slot.cached_pos * 0.1 and raw_match < slot.cached_pos * 0.5 + + if suspicious? do + :telemetry.execute( + [:llama_cpp_ex, :server, :prefix_instability], + %{matched_tokens: raw_match, cached_tokens: slot.cached_pos}, + %{server: self(), seq_id: seq_id} + ) + + warn_prefix_instability_once(state, raw_match, slot.cached_pos) + else + state + end + end + + defp warn_prefix_instability_once(%{prefix_instability_warned: true} = state, _match, _cached), + do: state + + defp warn_prefix_instability_once(state, raw_match, cached_pos) do + Logger.warning( + "LlamaCppEx.Server: request matched only #{raw_match}/#{cached_pos} cached prompt " <> + "tokens — the chat template may be rewriting history (e.g. stripping thinking " <> + "blocks), which defeats prefix caching. Emitted per-request as " <> + "[:llama_cpp_ex, :server, :prefix_instability]; this warning logs once." + ) + + %{state | prefix_instability_warned: true} + end + + # Offers a slot's about-to-be-destroyed cache to the RAM prompt cache. + # Skipped when the feature is off, the cache is too small to be worth a + # KV-sized copy, an existing entry already covers it, or the blob exceeds + # the whole budget (degrade to disabled, never OOM). + defp maybe_save_to_ram_cache(%{prompt_cache_ram_mb: 0} = state, _seq_id, _slot), do: state + + defp maybe_save_to_ram_cache(state, seq_id, slot) do + budget = state.prompt_cache_ram_mb * 1024 * 1024 + + with true <- slot.cached_pos >= @ram_cache_min_tokens, + false <- ram_cache_covers?(state.ram_cache, slot.cached_tokens, slot.cached_pos), + bytes = LlamaCppEx.NIF.state_seq_get_size(state.ctx.ref, seq_id), + true <- bytes > 0 and bytes <= budget, + {:ok, bin} <- LlamaCppEx.NIF.state_seq_get_data(state.ctx.ref, seq_id) do + entry = %{tokens: slot.cached_tokens, len: slot.cached_pos, bin: bin, bytes: bytes} + + state = %{ + state + | ram_cache: state.ram_cache ++ [entry], + ram_cache_bytes: state.ram_cache_bytes + bytes + } + + state = evict_ram_cache_to_budget(state, budget) + emit_ram_cache_telemetry(state, :save, entry) + state + else + _ -> state + end + end + + @doc false + def ram_cache_covers?(entries, cached_tokens, cached_pos) do + Enum.any?(entries, fn entry -> + entry.len >= cached_pos and + common_prefix_length(cached_tokens, entry.tokens) == cached_pos + end) + end + + @doc false + # FIFO eviction: drop oldest entries until the cache fits the budget. The + # empty-cache clause is defensive — entries larger than the whole budget + # are never stored, so over-budget implies non-empty today. + def evict_ram_cache_to_budget(state, budget) do + case state.ram_cache do + _ when state.ram_cache_bytes <= budget -> + state + + [] -> + %{state | ram_cache_bytes: 0} + + [evicted | rest] -> + state = %{state | ram_cache: rest, ram_cache_bytes: state.ram_cache_bytes - evicted.bytes} + emit_ram_cache_telemetry(state, :evict, evicted) + evict_ram_cache_to_budget(state, budget) + end + end + + # Best RAM cache entry for this prompt, if one clears the f_keep bar + # (restoring is a KV-sized memcpy — not worth it for a sliver) and its + # unusable tail can actually be trimmed on this model. + defp best_ram_candidate(%{ram_cache: []}, _tokens), do: nil + + defp best_ram_candidate(state, tokens) do + # Cap at len-1: the last prompt token must be decoded for logits. + max_reuse = length(tokens) - 1 + + {entry, lcp} = + state.ram_cache + |> Enum.map(fn entry -> + {entry, min(common_prefix_length(tokens, entry.tokens), max_reuse)} + end) + |> Enum.max_by(fn {_entry, lcp} -> lcp end) + + usable? = + lcp > 0 and lcp / entry.len >= @ram_cache_min_keep and + (lcp == entry.len or state.seq_rm_kind != :full) + + if usable?, do: {entry, lcp} + end + + # Restores a RAM cache entry into an (empty) sequence and trims the unusable + # tail. Returns the number of reusable prefix tokens. + defp apply_ram_restore(state, seq_id, entry, lcp) do + case LlamaCppEx.NIF.state_seq_set_data(state.ctx.ref, entry.bin, seq_id) do + {:ok, _bytes} -> + if lcp < entry.len do + true = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, lcp, -1) + end + + emit_ram_cache_telemetry(state, :restore, entry) + lcp + + {:error, reason} -> + # A partial restore could leave garbage in the sequence — clear it. + Logger.warning("LlamaCppEx.Server: RAM cache restore failed: #{inspect(reason)}") + _ = LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + 0 + end + end + + defp emit_ram_cache_telemetry(state, op, entry) do + :telemetry.execute( + [:llama_cpp_ex, :server, :ram_cache], + %{ + bytes: entry.bytes, + tokens: entry.len, + total_bytes: state.ram_cache_bytes, + entries: length(state.ram_cache) + }, + %{server: self(), op: op} + ) + end + + # Finds the slot whose in-KV tokens share the longest prefix with the new + # prompt, when that beats the assigned slot's own match. Active donors count + # too — only their FED tokens are in the KV, so the match is capped at the + # fed length. The pos_max probe guards against any bookkeeping drift between + # slot state and the actual KV contents. + defp best_donor(state, dst_seq_id, tokens, own_match, cache_prompt?) do + if cache_prompt? and state.cross_slot_sharing do + # Cap at len-1: the last prompt token must be decoded for logits. + max_reuse = length(tokens) - 1 + + state.slots + |> Enum.reject(fn {id, _slot} -> id == dst_seq_id end) + |> Enum.map(fn {id, slot} -> {id, min(donor_prefix_match(slot, tokens), max_reuse)} end) + |> Enum.max_by(fn {_id, lcp} -> lcp end, fn -> {nil, 0} end) + |> validate_donor(state, own_match) + else + nil + end + end + + defp validate_donor({donor_id, lcp}, state, own_match) when lcp > own_match and lcp > 0 do + if LlamaCppEx.NIF.memory_seq_pos_max(state.ctx.ref, donor_id) + 1 >= lcp do + {donor_id, lcp} + end + end + + defp validate_donor(_candidate, _state, _own_match), do: nil + + @doc false + def donor_prefix_match(%{state: :idle} = slot, tokens) do + common_prefix_length(tokens, slot.cached_tokens) + end + + def donor_prefix_match(%{state: :prefilling} = slot, tokens) do + # Only positions 0..prefill_pos-1 are in the KV so far. + min(common_prefix_length(tokens, slot.prompt_tokens), slot.prefill_pos) + end + + def donor_prefix_match(%{state: :generating} = slot, tokens) do + fed = slot.prompt_tokens ++ Enum.reverse(slot.generated_token_ids) + min(common_prefix_length(tokens, fed), slot.pos) + end + defp slot_mode(%{stream_pid: pid}) when is_pid(pid), do: :stream defp slot_mode(_slot), do: :generate + # Keeps sessions ↔ slots consistent: the slot remembers its session (so an + # idle slot can be re-picked by affinity) and the reverse map points each + # session at the slot serving it. A slot taken over by a different session + # drops its old mapping — but only if that mapping still points here (the + # old session may already have moved to another slot). + defp update_session_mapping(state, seq_id, session) do + old_session = state.slots[seq_id].session + + sessions = + if old_session != nil and old_session != session and + Map.get(state.sessions, old_session) == seq_id do + Map.delete(state.sessions, old_session) + else + state.sessions + end + + sessions = if session != nil, do: Map.put(sessions, session, seq_id), else: sessions + + state = put_in(state.slots[seq_id].session, session) + %{state | sessions: sessions} + end + defp enqueue_request(state, request) do %{state | queue: :queue.in(request, state.queue)} end + # Freed slots serve queued requests session-first, then FIFO — otherwise a + # busy queue scatters a conversation across slots and shreds its cache. defp dequeue_into_slot(state) do + state + |> dequeue_session_matches() + |> dequeue_fifo() + end + + defp dequeue_session_matches(state) do + requests = :queue.to_list(state.queue) + + {state, remaining} = + Enum.reduce(requests, {state, []}, fn request, {state, rest} -> + opts = request_opts(request) + + case session_slot_if_idle(state, Keyword.get(opts, :session), idle_slots(state)) do + nil -> + {state, [request | rest]} + + seq_id -> + {assign_queued_request(state, seq_id, request), rest} + end + end) + + %{state | queue: :queue.from_list(Enum.reverse(remaining))} + end + + defp idle_slots(state) do + Enum.filter(state.slots, fn {_id, slot} -> slot.state == :idle end) + end + + defp dequeue_fifo(state) do case :queue.out(state.queue) do {{:value, request}, queue} -> state = %{state | queue: queue} tokens = request_tokens(request) - case acquire_slot(state, tokens) do + case acquire_slot(state, tokens, request_opts(request)) do {:ok, seq_id, state} -> state = assign_queued_request(state, seq_id, request) - dequeue_into_slot(state) + dequeue_fifo(state) :no_slots -> # Put it back @@ -609,14 +1217,16 @@ defmodule LlamaCppEx.Server do end end - defp request_tokens({_type, tokens, _max, _from, _pid, _ref}), do: tokens + defp request_tokens({_type, tokens, _max, _from, _pid, _ref, _req_opts}), do: tokens - defp assign_queued_request(state, seq_id, {:generate, tokens, max_tokens, from, _, _}) do - init_slot(state, seq_id, tokens, max_tokens, from, nil, nil) + defp request_opts({_type, _tokens, _max, _from, _pid, _ref, req_opts}), do: req_opts + + defp assign_queued_request(state, seq_id, {:generate, tokens, max_tokens, from, _, _, req_opts}) do + init_slot(state, seq_id, tokens, max_tokens, from, nil, nil, req_opts) end - defp assign_queued_request(state, seq_id, {:stream, tokens, max_tokens, _, pid, ref}) do - init_slot(state, seq_id, tokens, max_tokens, nil, pid, ref) + defp assign_queued_request(state, seq_id, {:stream, tokens, max_tokens, _, pid, ref, req_opts}) do + init_slot(state, seq_id, tokens, max_tokens, nil, pid, ref, req_opts) end # --- Internal: Tick loop --- @@ -631,12 +1241,6 @@ defmodule LlamaCppEx.Server do # Phase 2: Build batch {entries, state} = build_batch(state) - # Capture TTFT immediately after the strategy ran. The strategy sends - # each decode token piece to its stream consumer before returning, so - # by now any slot whose first decode just streamed has tokens_generated - # == 1 — that's the user-perceived first-token timestamp. - state = mark_first_tokens(state) - if entries == [] do state else @@ -644,25 +1248,33 @@ defmodule LlamaCppEx.Server do end end - # Phases 3-5: forward pass, sampling, telemetry, and continuation + # Phases 3-5: fused forward pass + sampling, result handling, telemetry, + # and continuation. One dirty-CPU NIF per tick does decode, per-slot + # sampling, detokenization, and EOG checks; streaming sends happen here, + # right after the NIF returns — one tick earlier than the previous design, + # which deferred them to the next tick's batch building. defp run_forward_pass(state, entries) do - # Count decode tokens before sampling (slots may transition after) + # Count decode tokens before result handling (slots may transition after) n_decode = Enum.count(state.slots, fn {_id, s} -> s.state == :generating and s.batch_idx >= 0 end) - # Phase 3: Forward pass + samplers = active_samplers(state) + purgeable = purgeable_seq_ids(state) + + # Phase 3: Fused forward pass + sample tick_start = System.monotonic_time() - case LlamaCppEx.NIF.batch_eval(state.ctx.ref, entries) do - :ok -> + case LlamaCppEx.NIF.batch_eval_sample(state.ctx.ref, entries, samplers, purgeable) do + {:ok, results, purged, n_splits, failed} -> tick_end = System.monotonic_time() - # Phase 4: Sample - state = sample_generating_slots(state) - state = sample_completed_prefills(state) - state = advance_incomplete_prefills(state) + # Phase 4: Apply sampled results, stream pieces, clear tick markers + state = handle_kv_pressure(state, purged, n_splits) + state = apply_sample_results(state, results) + state = fail_overflowed_slots(state, failed) + state = clear_batch_indices(state) emit_tick_telemetry(state, entries, n_decode, tick_end - tick_start) @@ -675,6 +1287,146 @@ defmodule LlamaCppEx.Server do end end + # Sequences that could not fit a single further token even after purging and + # batch splitting — their KV budget is exhausted. Fail just those requests + # with a clean error; other slots keep generating. + defp fail_overflowed_slots(state, []), do: state + + defp fail_overflowed_slots(state, failed) do + Enum.reduce(failed, state, fn seq_id, state -> + Logger.warning("LlamaCppEx.Server: slot #{seq_id} out of context — failing request") + fail_slot(state, seq_id, :context_full) + end) + end + + defp fail_slot(state, seq_id, reason) do + slot = state.slots[seq_id] + demonitor_slot(slot) + + if slot.from do + GenServer.reply(slot.from, {:error, reason}) + end + + if slot.stream_pid && slot.stream_ref do + send(slot.stream_pid, {slot.stream_ref, {:error, reason}}) + end + + emit_request_exception(slot, seq_id, reason) + + # Clear KV cache — don't preserve cache on error + LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + + slot = + slot + |> Map.merge(idle_slot_fields([], 0)) + |> Map.put(:t_last_used, System.monotonic_time()) + + put_in(state.slots[seq_id], slot) + end + + # Samplers for all active slots, keyed by seq_id — the NIF samples every + # logits-requesting entry whose seq_id is registered here. + defp active_samplers(state) do + for {seq_id, slot} <- state.slots, slot.state != :idle do + {seq_id, slot.sampler.ref} + end + end + + # Idle slots whose retained prefix-cache KV may be dropped by the NIF under + # KV pressure (llama_decode == 1). Active slots are never purgeable. + defp purgeable_seq_ids(state) do + for {seq_id, slot} <- state.slots, slot.state == :idle, slot.cached_pos > 0 do + seq_id + end + end + + # The NIF purged these idle slots' KV to relieve pressure — drop the + # corresponding prefix-cache bookkeeping and surface telemetry. + defp handle_kv_pressure(state, [], 0), do: state + + defp handle_kv_pressure(state, purged, n_splits) do + Logger.warning( + "LlamaCppEx.Server: KV pressure — purged #{length(purged)} idle slot cache(s), " <> + "#{n_splits} batch split(s)" + ) + + :telemetry.execute( + [:llama_cpp_ex, :server, :kv_pressure], + %{purged_slots: length(purged), batch_splits: n_splits}, + %{server: self(), purged_seq_ids: purged} + ) + + Enum.reduce(purged, state, fn seq_id, state -> + slot = state.slots[seq_id] + put_in(state.slots[seq_id], %{slot | cached_tokens: [], cached_pos: 0}) + end) + end + + # Phase 4: apply per-slot sampled tokens from the fused NIF. Each result is + # {seq_id, token, piece, is_eog} for a slot that had a logits-requesting + # entry this tick — either a generating slot's decode token or a slot whose + # prefill just completed. Non-EOG pieces are streamed/accumulated + # immediately; the token itself is fed to the KV on the NEXT tick (unless + # the slot finishes first — see finish_completed_slots/1). + defp apply_sample_results(state, results) do + now = System.monotonic_time() + + Enum.reduce(results, state, fn {seq_id, token, piece, is_eog}, state -> + slot = state.slots[seq_id] + + slot = + case slot.state do + :prefilling -> %{slot | state: :generating, pos: slot.n_prompt_tokens} + :generating -> %{slot | pos: slot.pos + 1} + end + + slot = %{slot | pending_token: token, pending_eog: is_eog} + slot = if is_eog, do: slot, else: emit_piece(slot, piece, now) + + put_in(state.slots[seq_id], slot) + end) + end + + # Streams a sampled piece to the slot's subscriber (if any) and folds it + # into the slot's accumulated output/counters. Streamed bytes pass through + # a per-slot UTF-8 buffer so a codepoint split across tokens is never sent + # partially; the accumulated text keeps raw pieces (whole result is valid). + defp emit_piece(slot, piece, now) do + utf8_pending = + if slot.stream_pid && slot.stream_ref do + {out, pending} = UTF8Stream.push(slot.utf8_pending, piece) + + if out != "" do + send(slot.stream_pid, {slot.stream_ref, {:token, out}}) + end + + pending + else + slot.utf8_pending + end + + %{ + slot + | accumulated_pieces: [piece | slot.accumulated_pieces], + utf8_pending: utf8_pending, + tokens_generated: slot.tokens_generated + 1, + t_first_token: slot.t_first_token || now + } + end + + # batch_idx is only meaningful within a single tick (it marks slots the + # strategy fed this batch). Clear it so the next tick's bookkeeping starts + # clean even for slots that were skipped by budget limits. + defp clear_batch_indices(state) do + slots = + Map.new(state.slots, fn + {id, %{batch_idx: -1} = slot} -> {id, slot} + {id, slot} -> {id, %{slot | batch_idx: -1}} + end) + + %{state | slots: slots} + end + # Schedules the next tick while any slot is still active. defp continue_if_active(state) do if Enum.any?(state.slots, fn {_id, slot} -> slot.state != :idle end) do @@ -699,7 +1451,10 @@ defmodule LlamaCppEx.Server do ) end - # Phase 1: Check generating slots for completion + # Phase 1: Check generating slots for completion. EOG was determined by the + # fused NIF at sample time (pending_eog) — no per-token NIF call here. A slot + # whose pending token is EOG or whose streamed output hit max_tokens finishes + # without feeding that pending token to the KV. defp finish_completed_slots(state) do generating_slots = state.slots @@ -709,11 +1464,9 @@ defmodule LlamaCppEx.Server do Enum.reduce(generating_slots, state, fn {seq_id, _slot}, state -> slot = state.slots[seq_id] - token = slot.pending_token - is_eog = LlamaCppEx.NIF.vocab_is_eog(state.model.ref, token) cond do - is_eog -> + slot.pending_eog -> # pending_token is the EOG control token — finish without streaming it. finish_slot(state, seq_id, :eog) @@ -737,103 +1490,25 @@ defmodule LlamaCppEx.Server do {entries, %{state | slots: updated_slots}} end - defp mark_first_tokens(state) do - now = System.monotonic_time() - - slots = - Enum.reduce(state.slots, state.slots, fn {seq_id, slot}, acc -> - if slot.t_first_token == nil and slot.tokens_generated >= 1 do - Map.put(acc, seq_id, %{slot | t_first_token: now}) - else - acc - end - end) - - %{state | slots: slots} - end - - # Phase 4a: Sample for generating slots - defp sample_generating_slots(state) do - generating_slots = - state.slots - |> Enum.filter(fn {_id, slot} -> - slot.state == :generating and slot.batch_idx >= 0 - end) - - Enum.reduce(generating_slots, state, fn {seq_id, _slot}, state -> - slot = state.slots[seq_id] - - # sampler_sample_at/3 already accepts the selected token into the sampler - # (advancing grammar/penalty state); a second accept would double-advance. - next_token = - LlamaCppEx.NIF.sampler_sample_at(slot.sampler.ref, state.ctx.ref, slot.batch_idx) - - slot = %{ - slot - | pos: slot.pos + 1, - pending_token: next_token, - batch_idx: -1 - } - - put_in(state.slots[seq_id], slot) - end) - end - - # Phase 4b: Sample for prefilling slots that completed - defp sample_completed_prefills(state) do - completed_prefills = - state.slots - |> Enum.filter(fn {_id, slot} -> - slot.state == :prefilling and slot.batch_idx >= 0 and - slot.prefill_pos >= slot.n_prompt_tokens - end) - - Enum.reduce(completed_prefills, state, fn {seq_id, _slot}, state -> - slot = state.slots[seq_id] - - # sampler_sample_at/3 already accepts the token internally — no extra accept. - first_token = - LlamaCppEx.NIF.sampler_sample_at(slot.sampler.ref, state.ctx.ref, slot.batch_idx) - - slot = %{ - slot - | state: :generating, - pending_token: first_token, - pos: slot.n_prompt_tokens, - batch_idx: -1 - } - - put_in(state.slots[seq_id], slot) - end) - end - - # Phase 4c: Advance incomplete prefills (no sampling needed) - defp advance_incomplete_prefills(state) do - incomplete_prefills = - state.slots - |> Enum.filter(fn {_id, slot} -> - slot.state == :prefilling and slot.prefill_pos < slot.n_prompt_tokens - end) - - Enum.reduce(incomplete_prefills, state, fn {seq_id, _slot}, state -> - slot = state.slots[seq_id] - slot = %{slot | batch_idx: -1} - put_in(state.slots[seq_id], slot) - end) - end - # --- Internal: Slot completion --- defp finish_slot(state, seq_id, stop_reason) do slot = state.slots[seq_id] t_end = System.monotonic_time() + demonitor_slot(slot) if slot.from do - GenServer.reply(slot.from, {:ok, accumulated_text(slot)}) + GenServer.reply(slot.from, {:ok, completion_reply(slot, stop_reason)}) end if slot.stream_pid && slot.stream_ref do - send(slot.stream_pid, {slot.stream_ref, :done}) + # Flush any held-back trailing bytes (best effort — an incomplete + # codepoint at end-of-generation is the model's own output). + if slot.utf8_pending != "" do + send(slot.stream_pid, {slot.stream_ref, {:token, slot.utf8_pending}}) + end + + send(slot.stream_pid, {slot.stream_ref, {:done, stop_reason}}) end # Emit telemetry @@ -842,6 +1517,51 @@ defmodule LlamaCppEx.Server do reset_slot(state, seq_id) end + # Releases a slot whose consumer went away (died or explicitly cancelled). + # Nothing to reply to — the prefix cache is retained per the request's + # cache_prompt so a resumed conversation can still hit it, and the freed + # slot immediately serves the queue. + defp cancel_slot(state, seq_id) do + slot = state.slots[seq_id] + demonitor_slot(slot) + emit_request_done(slot, seq_id, System.monotonic_time(), :cancelled) + + state + |> reset_slot(seq_id) + |> dequeue_into_slot() + |> continue_if_active() + end + + defp demonitor_slot(%{monitor_ref: nil}), do: :ok + defp demonitor_slot(%{monitor_ref: mref}), do: Process.demonitor(mref, [:flush]) + + defp caller_pid({pid, _tag}) when is_pid(pid), do: pid + defp caller_pid(_from), do: nil + + # Drops a queued (not yet slotted) stream request by its subscription ref. + defp drop_queued_request(state, ref) do + queue = + :queue.filter( + fn {_type, _tokens, _max, _from, _pid, req_ref, _opts} -> req_ref != ref end, + state.queue + ) + + %{state | queue: queue} + end + + # Drops queued requests whose consumer (stream pid or sync caller) died. + defp drop_queued_requests_from(state, pid) do + queue = + :queue.filter( + fn {_type, _tokens, _max, from, stream_pid, _ref, _opts} -> + stream_pid != pid and caller_pid(from) != pid + end, + state.queue + ) + + %{state | queue: queue} + end + defp emit_request_done(slot, seq_id, t_end, stop_reason) do m = request_measurements(slot, t_end) @@ -911,43 +1631,60 @@ defmodule LlamaCppEx.Server do defp reset_slot(state, seq_id) do slot = state.slots[seq_id] - Sampler.reset(slot.sampler) - # Build full token history for prefix cache + # Build the token history matching what is actually in the KV. Retention + # follows the per-request setting: a cache_prompt: false request leaves + # nothing behind. A slot cancelled mid-prefill only has prefill_pos + # positions in KV — caching the full prompt would advertise positions + # that don't exist. (No sampler reset — every request gets a fresh + # sampler at init_slot.) {cached_tokens, cached_pos} = - if state.cache_prompt do - all_tokens = slot.prompt_tokens ++ Enum.reverse(slot.generated_token_ids) - {all_tokens, slot.pos} - else - LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) - {[], 0} + cond do + not slot.cache_prompt -> + LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) + {[], 0} + + slot.state == :prefilling -> + {Enum.take(slot.prompt_tokens, slot.prefill_pos), slot.prefill_pos} + + true -> + {slot.prompt_tokens ++ Enum.reverse(slot.generated_token_ids), slot.pos} end - slot = Map.merge(slot, idle_slot_fields(cached_tokens, cached_pos)) + slot = + slot + |> Map.merge(idle_slot_fields(cached_tokens, cached_pos)) + |> Map.put(:t_last_used, System.monotonic_time()) put_in(state.slots[seq_id], slot) end # The single source of truth for a slot's per-request fields. init/1, - # reset_slot/2, and fail_all_active_slots/2 all build from this map, so a - # new slot field cannot silently carry stale data across requests. The - # prefix-cache carry-over is the only caller-controlled part; :sampler is - # the only field that lives outside it. + # reset_slot/2, and the failure paths all build from this map, so a new + # slot field cannot silently carry stale data across requests. The + # prefix-cache carry-over is the only caller-controlled part; :sampler, + # :t_last_used, and :session are the only fields that live outside it + # (slot metadata that must survive request resets). defp idle_slot_fields(cached_tokens, cached_pos) do %{ state: :idle, from: nil, stream_pid: nil, stream_ref: nil, + monitor_ref: nil, + reply_mode: :text, + cache_prompt: false, prompt_tokens: [], prompt_tokens_tuple: {}, prefill_pos: 0, pos: 0, pending_token: nil, + pending_eog: false, batch_idx: -1, tokens_generated: 0, max_tokens: 0, accumulated_pieces: [], + utf8_pending: "", t_start: nil, t_first_token: nil, n_prompt_tokens: 0, @@ -963,26 +1700,27 @@ defmodule LlamaCppEx.Server do slot.accumulated_pieces |> Enum.reverse() |> IO.iodata_to_binary() end + # generate/generate_tokens reply with plain text; complete_tokens gets + # completion metadata for the OpenAI-shaped API. + defp completion_reply(%{reply_mode: :full} = slot, stop_reason) do + %{ + text: accumulated_text(slot), + completion_tokens: slot.tokens_generated, + finish_reason: stop_reason + } + end + + defp completion_reply(slot, _stop_reason), do: accumulated_text(slot) + + # Error replies are shaped {:error, atom | {atom, detail}} across all + # failure paths: :context_full, :queue_full, :empty_prompt, and + # {:inference_failed, reason} here. defp fail_all_active_slots(state, reason) do active_slots = Enum.filter(state.slots, fn {_id, slot} -> slot.state != :idle end) - Enum.reduce(active_slots, state, fn {seq_id, slot}, state -> - if slot.from do - GenServer.reply(slot.from, {:error, "inference failed: #{reason}"}) - end - - if slot.stream_pid && slot.stream_ref do - send(slot.stream_pid, {slot.stream_ref, {:error, reason}}) - end - - emit_request_exception(slot, seq_id, reason) - - # Clear KV cache and reset — don't preserve cache on error - LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1) - Sampler.reset(slot.sampler) - - put_in(state.slots[seq_id], Map.merge(slot, idle_slot_fields([], 0))) + Enum.reduce(active_slots, state, fn {seq_id, _slot}, state -> + fail_slot(state, seq_id, {:inference_failed, reason}) end) end diff --git a/lib/llama_cpp_ex/server/strategy/balanced.ex b/lib/llama_cpp_ex/server/strategy/balanced.ex index 57bb3ae..7e5ce86 100644 --- a/lib/llama_cpp_ex/server/strategy/balanced.ex +++ b/lib/llama_cpp_ex/server/strategy/balanced.ex @@ -15,9 +15,7 @@ defmodule LlamaCppEx.Server.Strategy.Balanced do alias LlamaCppEx.Server.Strategy.Batch @impl true - def build_batch(slots, budget, chunk_size, opts) do - model_ref = Keyword.fetch!(opts, :model_ref) - + def build_batch(slots, budget, chunk_size, _opts) do n_generating = Enum.count(slots, fn {_id, slot} -> slot.state == :generating and slot.pending_token != nil @@ -28,7 +26,7 @@ defmodule LlamaCppEx.Server.Strategy.Balanced do prefill_budget = budget - decode_budget {entries, n_entries, slots, decode_remaining} = - Batch.add_decode_tokens(slots, [], 0, decode_budget, model_ref) + Batch.add_decode_tokens(slots, [], 0, decode_budget) # Prefill gets its half plus any unused decode budget. effective_prefill_budget = prefill_budget + decode_remaining diff --git a/lib/llama_cpp_ex/server/strategy/batch.ex b/lib/llama_cpp_ex/server/strategy/batch.ex index 4d28b2c..1d77013 100644 --- a/lib/llama_cpp_ex/server/strategy/batch.ex +++ b/lib/llama_cpp_ex/server/strategy/batch.ex @@ -6,6 +6,11 @@ defmodule LlamaCppEx.Server.Strategy.Batch do the order and budget split between decode and prefill — the per-slot assembly of decode tokens and prefill chunks is identical, so it lives here. + Sampling, detokenization, and stream delivery all happen after the forward + pass (fused into the `batch_eval_sample` NIF and the server's post-tick + result handling) — these helpers only assemble batch entries and update + slot bookkeeping. + ## Performance notes These helpers run on every tick of the server loop, once per active slot, and @@ -26,9 +31,14 @@ defmodule LlamaCppEx.Server.Strategy.Batch do @doc """ Adds one decode token for each generating slot (lowest seq_id first) until the - budget is exhausted. Streams each token piece to the slot's subscriber. + budget is exhausted. + + `generated_token_ids` is updated here — at feed time — so it tracks exactly + the tokens present in the slot's KV cache (the prefix-cache bookkeeping + depends on that invariant; a sampled-but-never-fed final token must not + appear in `cached_tokens`). """ - def add_decode_tokens(slots, entries, n_entries, budget, model_ref) do + def add_decode_tokens(slots, entries, n_entries, budget) do generating_slots = slots |> Enum.filter(fn {_id, slot} -> @@ -44,17 +54,9 @@ defmodule LlamaCppEx.Server.Strategy.Batch do slot = slots[seq_id] token = slot.pending_token - piece = LlamaCppEx.NIF.token_to_piece(model_ref, token) - - if slot.stream_pid && slot.stream_ref do - send(slot.stream_pid, {slot.stream_ref, {:token, piece}}) - end - slot = %{ slot - | accumulated_pieces: [piece | slot.accumulated_pieces], - batch_idx: n_entries, - tokens_generated: slot.tokens_generated + 1, + | batch_idx: n_entries, generated_token_ids: [token | slot.generated_token_ids] } diff --git a/lib/llama_cpp_ex/server/strategy/decode_maximal.ex b/lib/llama_cpp_ex/server/strategy/decode_maximal.ex index 8d2f845..fb5f7a6 100644 --- a/lib/llama_cpp_ex/server/strategy/decode_maximal.ex +++ b/lib/llama_cpp_ex/server/strategy/decode_maximal.ex @@ -15,12 +15,10 @@ defmodule LlamaCppEx.Server.Strategy.DecodeMaximal do alias LlamaCppEx.Server.Strategy.Batch @impl true - def build_batch(slots, budget, chunk_size, opts) do - model_ref = Keyword.fetch!(opts, :model_ref) - + def build_batch(slots, budget, chunk_size, _opts) do # Decode tokens first (priority), then prefill chunks fill remaining budget. {entries, n_entries, slots, budget} = - Batch.add_decode_tokens(slots, [], 0, budget, model_ref) + Batch.add_decode_tokens(slots, [], 0, budget) {entries, _n_entries, slots, _budget} = Batch.add_prefill_chunks(slots, entries, n_entries, budget, chunk_size) diff --git a/lib/llama_cpp_ex/server/strategy/prefill_priority.ex b/lib/llama_cpp_ex/server/strategy/prefill_priority.ex index 7a301d6..50acc54 100644 --- a/lib/llama_cpp_ex/server/strategy/prefill_priority.ex +++ b/lib/llama_cpp_ex/server/strategy/prefill_priority.ex @@ -13,15 +13,13 @@ defmodule LlamaCppEx.Server.Strategy.PrefillPriority do alias LlamaCppEx.Server.Strategy.Batch @impl true - def build_batch(slots, budget, chunk_size, opts) do - model_ref = Keyword.fetch!(opts, :model_ref) - + def build_batch(slots, budget, chunk_size, _opts) do # Prefill chunks first (priority), then decode tokens fill remaining budget. {entries, n_entries, slots, budget} = Batch.add_prefill_chunks(slots, [], 0, budget, chunk_size) {entries, _n_entries, slots, _budget} = - Batch.add_decode_tokens(slots, entries, n_entries, budget, model_ref) + Batch.add_decode_tokens(slots, entries, n_entries, budget) {Enum.reverse(entries), slots} end diff --git a/lib/llama_cpp_ex/utf8_stream.ex b/lib/llama_cpp_ex/utf8_stream.ex new file mode 100644 index 0000000..a146a6b --- /dev/null +++ b/lib/llama_cpp_ex/utf8_stream.ex @@ -0,0 +1,44 @@ +defmodule LlamaCppEx.UTF8Stream do + @moduledoc false + # Pending-bytes buffering for streamed token pieces. A tokenizer piece can + # end mid-codepoint (e.g. one token per byte of a multibyte emoji), so + # emitting pieces verbatim hands consumers invalid UTF-8. `push/2` returns + # the longest emittable prefix and holds back a trailing incomplete + # codepoint until the next piece completes it. Invalid bytes that cannot + # become valid by appending more data are passed through unchanged — the + # model produced them, hiding them would corrupt offsets. + + @doc """ + Appends `chunk` to the held-back bytes and splits off what can be emitted. + + Returns `{emit, pending}` — `emit` never ends in a partial codepoint. + """ + @spec push(binary(), binary()) :: {binary(), binary()} + def push(pending, chunk) when is_binary(pending) and is_binary(chunk) do + data = pending <> chunk + size = byte_size(data) + holdback = incomplete_suffix_len(data, size) + {binary_part(data, 0, size - holdback), binary_part(data, size - holdback, holdback)} + end + + # Number of trailing bytes forming an incomplete-but-completable UTF-8 + # codepoint (0..3). Scans back at most 3 bytes: a lead byte further back + # than that heads a sequence that is already complete or already invalid. + defp incomplete_suffix_len(data, size) do + Enum.find_value(1..min(size, 3)//1, 0, fn back -> + byte = :binary.at(data, size - back) + + cond do + byte < 0x80 -> 0 + byte < 0xC0 -> nil + byte >= 0xF8 -> 0 + back < expected_len(byte) -> back + true -> 0 + end + end) + end + + defp expected_len(byte) when byte >= 0xF0, do: 4 + defp expected_len(byte) when byte >= 0xE0, do: 3 + defp expected_len(_byte), do: 2 +end diff --git a/mix.exs b/mix.exs index 89c4f2e..2b4e41c 100644 --- a/mix.exs +++ b/mix.exs @@ -37,7 +37,7 @@ end defmodule LlamaCppEx.MixProject do use Mix.Project - @version "0.8.32" + @version "0.8.33" @source_url "https://github.com/nyo16/llama_cpp_ex" def project do diff --git a/test/model_manager_test.exs b/test/model_manager_test.exs index 301eb69..a766c7f 100644 --- a/test/model_manager_test.exs +++ b/test/model_manager_test.exs @@ -5,9 +5,14 @@ defmodule LlamaCppEx.ModelManagerTest do alias LlamaCppEx.ModelManager - # A stub standing in for LlamaCppEx.Server. It answers the two messages the - # manager's dispatch sends it: {:generate, prompt, max_tokens} and :get_model. - # Started unlinked so killing it (DOWN test) doesn't take down the manager. + # A stub standing in for LlamaCppEx.Server. It answers the raw protocol + # messages a routed request produces: {:generate_tokens, ...} and + # :get_model. NOTE: full-path ModelManager.generate/3 can't be unit-tested + # against a stub anymore — Server.generate tokenizes in the CALLER via + # get_model/1, which needs a real model NIF resource; routing tests assert + # via route/1 + the raw tokens call instead (real-path coverage lives in + # the smoke tests). Started unlinked so killing it (DOWN test) doesn't take + # down the manager. defmodule StubServer do use GenServer @@ -17,12 +22,22 @@ defmodule LlamaCppEx.ModelManagerTest do def init(reply), do: {:ok, reply} @impl true - def handle_call({:generate, _prompt, _max}, _from, reply), do: {:reply, reply, reply} + def handle_call({:generate_tokens, _tokens, _max, _req_opts}, _from, reply), + do: {:reply, reply, reply} + def handle_call(:get_model, _from, reply), do: {:reply, fake_model(), reply} defp fake_model, do: %LlamaCppEx.Model{ref: make_ref()} end + # Routes id through the manager and performs the stubbed raw call — the + # unit-testable equivalent of ModelManager.generate/3 for :server entries. + defp generate_via_route(id) do + with {:ok, {:server, pid, _entry}} <- ModelManager.route(id) do + GenServer.call(pid, {:generate_tokens, [1, 2, 3], 256, []}) + end + end + # Fake Backend: no real GGUF files, no native loads. Behaviour is driven by # keys threaded through the load opts (`fake_bytes`, `fake_*_error`, ...). defmodule FakeIO do @@ -160,7 +175,7 @@ defmodule LlamaCppEx.ModelManagerTest do test "generate routes to the server-backed model" do {:ok, _} = ModelManager.load("chat", {:path, "chat.gguf"}, fake_reply: {:ok, "hello"}) - assert {:ok, "hello"} = ModelManager.generate("chat", "hi") + assert {:ok, "hello"} = generate_via_route("chat") end test "generate on a missing model returns :not_loaded" do @@ -212,7 +227,7 @@ defmodule LlamaCppEx.ModelManagerTest do ModelManager.load("chat", {:path, "chat.gguf"}, default: true, fake_reply: {:ok, "d"}) assert ModelManager.default() == "chat" - assert {:ok, "d"} = ModelManager.generate(:default, "hi") + assert {:ok, "d"} = generate_via_route(:default) end test "set_default/1 updates the default" do @@ -221,7 +236,7 @@ defmodule LlamaCppEx.ModelManagerTest do assert :ok = ModelManager.set_default("b") assert ModelManager.default() == "b" - assert {:ok, "from-b"} = ModelManager.generate(:default, "hi") + assert {:ok, "from-b"} = generate_via_route(:default) end test "set_default on a missing model errors" do diff --git a/test/server_smoke_test.exs b/test/server_smoke_test.exs new file mode 100644 index 0000000..496fe99 --- /dev/null +++ b/test/server_smoke_test.exs @@ -0,0 +1,184 @@ +defmodule LlamaCppEx.ServerSmokeTest do + # Integration tests for Server behaviors that need a real model. Run with: + # + # LLAMA_SMOKE_GEN_MODEL=/path/to/chat-model.gguf mix test --include smoke + # + # async: false — each test starts its own server against the GPU. + use ExUnit.Case, async: false + + @moduletag :smoke + @moduletag timeout: 300_000 + + @gen_model System.get_env("LLAMA_SMOKE_GEN_MODEL") + + if @gen_model && File.exists?(@gen_model) do + alias LlamaCppEx.Server + + # No explicit teardown: the server is linked and traps exits, so it stops + # itself (running terminate/2) when the test process exits. + defp start_server(opts) do + defaults = [ + model_path: @gen_model, + n_gpu_layers: -1, + n_parallel: 2, + n_ctx: 2048, + temp: 0.0 + ] + + {:ok, server} = Server.start_link(Keyword.merge(defaults, opts)) + server + end + + defp attach_collector(event) do + parent = self() + handler_id = {__MODULE__, self(), event} + + :telemetry.attach( + handler_id, + event, + fn _event, measurements, metadata, _ -> + send(parent, {:telemetry, measurements, metadata}) + end, + nil + ) + + on_exit(fn -> :telemetry.detach(handler_id) end) + end + + defp next_telemetry(timeout \\ 60_000) do + receive do + {:telemetry, m, meta} -> {m, meta} + after + timeout -> flunk("expected telemetry event") + end + end + + test "empty token list errors immediately on both call types" do + server = start_server([]) + + assert {:error, :empty_prompt} = Server.generate_tokens(server, []) + assert [{:error, :empty_prompt}] = Server.stream_tokens(server, []) |> Enum.to_list() + end + + test "per-request cache_prompt overrides: hit, no-reuse, no-retention" do + server = start_server(n_parallel: 1) + attach_collector([:llama_cpp_ex, :server, :request, :start]) + + prompt = "User: Name three colors.\nAssistant:" + {:ok, reply} = Server.generate(server, prompt, max_tokens: 12) + {%{prefix_cache_tokens: 0}, _} = next_telemetry() + + # Exact-prefix continuation hits the cache (works on :full models too). + turn2 = prompt <> reply <> "\nUser: One more.\nAssistant:" + {:ok, reply2} = Server.generate(server, turn2, max_tokens: 12) + {%{prefix_cache_tokens: hit}, _} = next_telemetry() + assert hit > 0 + + # cache_prompt: false → no reuse now, no retention after. + turn3 = turn2 <> reply2 <> "\nUser: Again.\nAssistant:" + {:ok, _} = Server.generate(server, turn3, max_tokens: 8, cache_prompt: false) + {%{prefix_cache_tokens: 0}, _} = next_telemetry() + + {:ok, _} = Server.generate(server, turn3, max_tokens: 4) + {%{prefix_cache_tokens: 0}, _} = next_telemetry() + end + + test "session affinity keeps interleaved conversations on their slots" do + server = start_server([]) + attach_collector([:llama_cpp_ex, :server, :request, :start]) + + a1 = "Chat A. User: Name three colors.\nAssistant:" + b1 = "Chat B, unrelated. User: Name three animals.\nAssistant:" + + {:ok, ra} = Server.generate(server, a1, max_tokens: 12, session: :a) + {_, %{seq_id: slot_a}} = next_telemetry() + {:ok, rb} = Server.generate(server, b1, max_tokens: 12, session: :b) + {_, %{seq_id: slot_b}} = next_telemetry() + assert slot_a != slot_b + + {:ok, _} = + Server.generate(server, a1 <> ra <> "\nUser: More.\nAssistant:", + max_tokens: 8, + session: :a + ) + + {%{prefix_cache_tokens: hit_a}, %{seq_id: slot_a2}} = next_telemetry() + + {:ok, _} = + Server.generate(server, b1 <> rb <> "\nUser: More.\nAssistant:", + max_tokens: 8, + session: :b + ) + + {%{prefix_cache_tokens: hit_b}, %{seq_id: slot_b2}} = next_telemetry() + + assert slot_a2 == slot_a + assert slot_b2 == slot_b + assert hit_a > 0 + assert hit_b > 0 + end + + test "max_queue rejects overflow immediately, queued work still completes" do + server = start_server(n_parallel: 1, max_queue: 1) + + long = + Task.async(fn -> + Server.generate(server, "Write a long story:", max_tokens: 200, timeout: 240_000) + end) + + Process.sleep(300) + + queued = + Task.async(fn -> Server.generate(server, "2+2=", max_tokens: 4, timeout: 240_000) end) + + Process.sleep(200) + + assert {:error, :queue_full} = Server.generate(server, "3+3=", max_tokens: 4) + + assert [{:error, :queue_full}] = + Server.stream(server, "4+4=", max_tokens: 4) |> Enum.to_list() + + assert {:ok, _} = Task.await(long, 240_000) + assert {:ok, _} = Task.await(queued, 240_000) + end + + test "halting a stream early cancels generation and frees the slot" do + server = start_server(n_parallel: 1) + attach_collector([:llama_cpp_ex, :server, :request, :done]) + + _ = Server.stream(server, "Write an endless story:", max_tokens: 400) |> Enum.take(3) + + {%{generated_tokens: n}, %{stop_reason: :cancelled}} = next_telemetry() + assert n < 50 + + # Slot is immediately usable. + assert {:ok, _} = Server.generate(server, "2+2=", max_tokens: 4) + end + + test "a request exceeding its context budget fails alone" do + server = start_server(n_parallel: 2, n_ctx: 256) + + long = + Task.async(fn -> + Server.generate(server, "Write a story:", max_tokens: 400, timeout: 240_000) + end) + + short = + Task.async(fn -> + Process.sleep(200) + Server.generate(server, "2+2=", max_tokens: 4, timeout: 240_000) + end) + + assert {:error, :context_full} = Task.await(long, 240_000) + assert {:ok, _} = Task.await(short, 240_000) + + # Server is still healthy afterwards. + assert {:ok, _} = Server.generate(server, "The sky is", max_tokens: 4) + end + else + @tag :skip + test "server smoke tests skipped — set LLAMA_SMOKE_GEN_MODEL" do + :ok + end + end +end diff --git a/test/server_test.exs b/test/server_test.exs new file mode 100644 index 0000000..143552e --- /dev/null +++ b/test/server_test.exs @@ -0,0 +1,125 @@ +defmodule LlamaCppEx.ServerTest do + use ExUnit.Case, async: true + + alias LlamaCppEx.Server + + # --- Pure slot/cache logic (no model files needed) --- + + defp idle_slot(seq_id, cached_tokens, t_last_used) do + {seq_id, + %{ + state: :idle, + cached_tokens: cached_tokens, + cached_pos: length(cached_tokens), + t_last_used: t_last_used + }} + end + + describe "pick_cached_slot/2 (similarity threshold + LRU fallback)" do + test "picks the slot with the best match when it clears the 10% threshold" do + slots = [ + idle_slot(0, [1, 2, 3, 4, 5, 6, 7, 8], 100), + idle_slot(1, [1, 2, 9, 9], 200) + ] + + # 8/10 of the prompt matches slot 0 + assert Server.pick_cached_slot(slots, [1, 2, 3, 4, 5, 6, 7, 8, 20, 21]) == 0 + end + + test "falls back to LRU when the best match is below the threshold" do + # Slot 1 holds a long valuable cache; the new prompt barely matches it. + slots = [ + idle_slot(0, [], 100), + idle_slot(1, Enum.to_list(1..200), 200) + ] + + prompt = [1] ++ Enum.to_list(900..950) + # 1/52 match < 0.1 → LRU (slot 0, oldest) — protects slot 1's cache. + assert Server.pick_cached_slot(slots, prompt) == 0 + end + + test "pick_lru_slot returns the least recently used" do + slots = [idle_slot(0, [], 300), idle_slot(1, [], 100), idle_slot(2, [], 200)] + assert Server.pick_lru_slot(slots) == 1 + end + end + + describe "session_slot_if_idle/3" do + test "returns the session's slot when idle" do + state = %{sessions: %{"conv-a" => 1}} + idle = [idle_slot(0, [], 0), idle_slot(1, [], 0)] + assert Server.session_slot_if_idle(state, "conv-a", idle) == 1 + end + + test "returns nil when the session's slot is busy" do + state = %{sessions: %{"conv-a" => 1}} + idle = [idle_slot(0, [], 0)] + assert Server.session_slot_if_idle(state, "conv-a", idle) == nil + end + + test "returns nil for unknown or nil sessions" do + state = %{sessions: %{}} + idle = [idle_slot(0, [], 0)] + assert Server.session_slot_if_idle(state, "ghost", idle) == nil + assert Server.session_slot_if_idle(state, nil, idle) == nil + end + end + + describe "donor_prefix_match/2 (only fed tokens count)" do + test "idle donor matches against its cached tokens" do + slot = %{state: :idle, cached_tokens: [1, 2, 3, 4]} + assert Server.donor_prefix_match(slot, [1, 2, 3, 9]) == 3 + end + + test "prefilling donor is capped at prefill_pos" do + slot = %{state: :prefilling, prompt_tokens: [1, 2, 3, 4, 5, 6], prefill_pos: 2} + # Prompt matches 6 tokens, but only 2 are in the KV so far. + assert Server.donor_prefix_match(slot, [1, 2, 3, 4, 5, 6]) == 2 + end + + test "generating donor is capped at fed position" do + slot = %{ + state: :generating, + prompt_tokens: [1, 2, 3], + generated_token_ids: [5, 4], + pos: 4 + } + + # Fed history is [1, 2, 3, 4, 5] but only pos=4 tokens are in KV. + assert Server.donor_prefix_match(slot, [1, 2, 3, 4, 5, 6]) == 4 + end + end + + describe "RAM cache bookkeeping" do + defp entry(tokens, bytes), do: %{tokens: tokens, len: length(tokens), bytes: bytes, bin: ""} + + test "ram_cache_covers?/3 detects a covering entry" do + entries = [entry([1, 2, 3, 4, 5], 100)] + assert Server.ram_cache_covers?(entries, [1, 2, 3], 3) + refute Server.ram_cache_covers?(entries, [1, 9, 3], 3) + refute Server.ram_cache_covers?(entries, [1, 2, 3, 4, 5, 6, 7], 7) + end + + test "evict_ram_cache_to_budget/2 drops oldest entries first" do + state = %{ + ram_cache: [entry([1], 400), entry([2], 300), entry([3], 200)], + ram_cache_bytes: 900 + } + + result = Server.evict_ram_cache_to_budget(state, 500) + + assert [%{tokens: [2]}, %{tokens: [3]}] = result.ram_cache + assert result.ram_cache_bytes == 500 + end + + test "evict_ram_cache_to_budget/2 is a no-op within budget" do + state = %{ram_cache: [entry([1], 100)], ram_cache_bytes: 100} + assert Server.evict_ram_cache_to_budget(state, 100) == state + end + + test "evict_ram_cache_to_budget/2 survives an inconsistent empty cache" do + state = %{ram_cache: [], ram_cache_bytes: 999} + assert %{ram_cache: [], ram_cache_bytes: 0} = Server.evict_ram_cache_to_budget(state, 10) + end + end +end diff --git a/test/utf8_stream_test.exs b/test/utf8_stream_test.exs new file mode 100644 index 0000000..f6d11c8 --- /dev/null +++ b/test/utf8_stream_test.exs @@ -0,0 +1,51 @@ +defmodule LlamaCppEx.UTF8StreamTest do + use ExUnit.Case, async: true + + alias LlamaCppEx.UTF8Stream + + test "plain ASCII passes through" do + assert {"hello", ""} = UTF8Stream.push("", "hello") + end + + test "complete multibyte codepoint passes through" do + assert {"héllo", ""} = UTF8Stream.push("", "héllo") + assert {"🎉", ""} = UTF8Stream.push("", "🎉") + end + + test "holds back a split 2-byte codepoint" do + <> = "é" + assert {"", <<^a>>} = UTF8Stream.push("", <>) + assert {"é", ""} = UTF8Stream.push(<>, <>) + end + + test "holds back a 4-byte emoji split across four pieces" do + <> = "🎉" + + {out1, p1} = UTF8Stream.push("", <>) + {out2, p2} = UTF8Stream.push(p1, <>) + {out3, p3} = UTF8Stream.push(p2, <>) + {out4, p4} = UTF8Stream.push(p3, <>) + + assert out1 == "" and out2 == "" and out3 == "" + assert out4 == "🎉" + assert p4 == "" + end + + test "emits text before a trailing partial codepoint" do + <> = "🎉" + assert {"yay ", <<^a, ^b>>} = UTF8Stream.push("", "yay " <> <>) + assert {"🎉!", ""} = UTF8Stream.push(<>, <> <> "!") + end + + test "genuinely invalid bytes pass through rather than stall" do + # 0xFF can never start a valid codepoint; 0x80 is a bare continuation. + assert {<<0xFF>>, ""} = UTF8Stream.push("", <<0xFF>>) + assert {<<0x80, 0x80, 0x80, 0x80>>, ""} = UTF8Stream.push("", <<0x80, 0x80, 0x80, 0x80>>) + end + + test "long valid text with trailing lead byte" do + text = String.duplicate("日本語", 10) + <> = "本" + assert {^text, <<^lead>>} = UTF8Stream.push("", text <> <>) + end +end