Continuous batching perf: fused tick NIF, cross-slot prefix caching, server-routed chat completions, cancellation

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

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
 

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()

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

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()