KV snapshots on the switch: save idle slots on the drain, restore them on the next start

CHANGELOG.md

@@ -29,6 +29,17 @@ has a matching GitHub release whose body comes from its section below.
   scopeless lease. Without a user manager, inside a `Type=notify` service (where
   systemd only accepts readiness from the service's own cgroup), and on Windows, the run
   starts as before and the window stands.
+- `kv_snapshots` and `kv_snapshots_mib` on an `[[engine]]` entry: the drain saves an idle slot's
+  prompt cache into the engine's `--slot-save-path` directory once the boundary is reached, keeps
+  the directory within its budget (16 GiB by default, above the library free floor), and the
+  gateway restores the newest snapshot of the engine's model into its empty slots before the first
+  request, so a sequence of agents skips the prefill after each switch. The snapshots are keyed by
+  the model the engine itself names on `/v1/models`, so an engine that lists none or many
+  (llama-swap) never snapshots, and nothing crosses models. The restore also warms the files into
+  the page cache during the start's wait: cached, never reserved. A save is bounded at 60 s and a
+  restore at 30 s, so whatever goes wrong along the way costs one line and the switch goes on as
+  before. A snapshot holds the conversation's tokens: the ledger is owner-only, and the README asks
+  for a private directory.
 
 ## [v0.3.3] - 2026-10-21
 

README.md

@@ -603,6 +603,8 @@ start    = ["/path/to/strata/.venv/bin/python", "serve/server.py", "--engine", "
 serves   = ["router-strata", "qwen3.8-flash-next-iq2_xs"]
 priority = "foreground"   # equal to the engine it replaces, so the last model you ask for wins
 prefetch = ["/models/strata"]   # read into the page cache while the holder drains
+# kv_snapshots    = "/home/you/.cache/berth/kv/strata"  # the --slot-save-path the saves and restores share
+# kv_snapshots_mib = 16384
 ```
 
 Run `berth gateway` (loopback `127.0.0.1:9799`; a user unit is in `contrib/berth-gateway.service`) and send
@@ -672,6 +674,28 @@ larger than the RAM the machine can spare gains little: its pages are pushed out
 its own files and not its subdirectories, and a file that cannot be asked for is one WARN line in the
 gateway's journal while the start goes on. `berth prefetch MODEL` makes the same ask by hand.
 
+**An agent context comes back without the prefill.** Give a llama-server engine `kv_snapshots`,
+the same directory its `--slot-save-path` points at. On a drain, once every slot is idle, Berth asks
+the engine to save each slot's prompt cache into that directory, records what each save wrote in a
+ledger there, skips slots whose cache is empty, and trims the directory to `kv_snapshots_mib`
+(16 GiB by default). When the gateway starts that engine again for the same model, it warms the
+snapshots into the page cache during the wait, and the engine's empty slots restore their newest
+snapshot once it answers, before the first request. A long agent conversation then switches and
+switches back without prefilling again.
+
+A warm cache is cached, never reserved: the pages sit in the kernel's page cache until the machine
+needs the RAM, and a restore afterwards reads its bytes from the SSD again. Snapshots never cross
+models: the restore takes only the files the ledger holds for the model the engine itself reports on
+`/v1/models`. An engine that lists no model, or more than one (llama-swap lists every model it can
+load), gets no snapshots, so a KV snapshot covers llama-swap's many-model entries only when
+llama-swap gains per-model slots. An engine started without `--slot-save-path` says so, the drain
+logs one line, and the switch is a plain one. The save step is bounded at 60 s and the restore at
+30 s, so a failure on either side leaves the switch exactly as it was before.
+
+A snapshot holds the slot's prompt tokens, so it is the conversation itself. Keep the directory
+private to the user the engine runs as (`mkdir -p DIR && chmod 700 DIR`); Berth writes its ledger
+there owner-only.
+
 Measured on the reference RTX 4090 with the [swap bench](contrib/bench/README.md) switching between
 NInfer (through llama-swap) and Strata (262K context), both on an NVMe SSD, 2026-10-20 (two runs,
 five switches each way):
@@ -1198,10 +1222,8 @@ engine's files during a drain (`prefetch`) and serves measured switch costs at `
 
 In the next v0.3.x releases. Each item has its building blocks merged or a spec written:
 
-- KV snapshots wired into the drain and start paths, so a sequence of agents skips the prefill
-  after each switch.
 - Measured switch costs on the live daemon.
-- Gateway fallback routing, and models started on demand for a standalone llama-server.
+- Models started on demand for a standalone llama-server.
 - Parked resume numbers in [Swap numbers](#swap-numbers), which take the vLLM and SGLang park
   adapters out of experimental, and a measurement of NVIDIA's `cuda-checkpoint` for unmodified
   CUDA engines on Linux.

internal/config/config.go

@@ -68,6 +68,20 @@ type Engine struct {
 	DrainComfyUI string   `toml:"drain_comfyui,omitempty" json:"drain_comfyui,omitempty" yaml:"drain_comfyui,omitempty"` // as berth run --drain-comfyui
 	// KeepComfyUI is berth run --keep-comfyui-up: unload and stay up.
 	KeepComfyUI bool `toml:"keep_comfyui_up,omitempty" json:"keep_comfyui_up,omitempty" yaml:"keep_comfyui_up,omitempty"`
+	// KVSnapshots is the snapshot directory the engine's --slot-save-path
+	// points at. On a drain, Berth saves every idle slot's prompt cache
+	// into it; when this gateway later starts the engine again, the empty
+	// slots restore the newest snapshot for that model before the first
+	// request, so a long agent context is not prefilled again after every
+	// switch. Snapshots never cross models, and an engine that does not
+	// name exactly one model on /v1/models (llama-swap lists many) never
+	// snapshots. Absolute path: the gateway starts engines from its own
+	// working directory, and the drained engine writes through the same
+	// absolute name.
+	KVSnapshots string `toml:"kv_snapshots,omitempty" json:"kv_snapshots,omitempty" yaml:"kv_snapshots,omitempty"`
+	// KVSnapshotsMiB is how many MiB of snapshots the directory may hold
+	// before the least recently saved are removed. 0 = the default.
+	KVSnapshotsMiB uint64 `toml:"kv_snapshots_mib,omitzero" json:"kv_snapshots_mib,omitempty" yaml:"kv_snapshots_mib,omitempty"`
 	// DrainMax is berth run --drain-max: a ceiling on how long a drain may
 	// wait at a boundary (a processing slot, a running ComfyUI job) before
 	// it stops the engine anyway. 0 = unbounded, the default.
@@ -241,6 +255,10 @@ type ModelsConfig struct {
 	Dirs []string `toml:"dirs,omitempty" json:"dirs,omitempty" yaml:"dirs,omitempty"`
 }
 
+// DefaultKVSnapshotsMiB is the snapshot budget an engine with kv_snapshots
+// and no kv_snapshots_mib gets.
+const DefaultKVSnapshotsMiB = 16384
+
 // DefaultFreeFloor is the share of a root's total space a pull must leave
 // free. A drive never fills: a pull that would go under it is refused.
 const DefaultFreeFloor = 0.10
@@ -574,6 +592,9 @@ func (c *Config) Validate() error {
 				errs = append(errs, fmt.Errorf("%s: drain_slots %q must be the engine's http(s)://host:port", where, e.DrainSlots))
 			}
 		}
+		if e.KVSnapshots != "" && !filepath.IsAbs(e.KVSnapshots) {
+			errs = append(errs, fmt.Errorf("%s: kv_snapshots %q must be an absolute path: the gateway starts engines from its own working directory", where, e.KVSnapshots))
+		}
 		if e.DrainMax.Duration < 0 {
 			errs = append(errs, fmt.Errorf("%s: drain_max %s must be >= 0", where, e.DrainMax))
 		}

internal/config/config_test.go

@@ -720,6 +720,47 @@ func TestLibraryFloorDefaults(t *testing.T) {
 	}
 }
 
+// kv_snapshots is an absolute directory the engine's --slot-save-path points
+// at, with an open budget that takes the default; a relative path is refused
+// and the size parses as given.
+func TestEngineKVSnapshotsParsedAndValidated(t *testing.T) {
+	abs := filepath.Join(t.TempDir(), "kv")
+	c, err := Load(write(t, "[[engine]]\nname = \"llama-server\"\ntype = \"llamacpp\"\nurl = \"http://127.0.0.1:8080\"\nkv_snapshots = '"+abs+"'\nkv_snapshots_mib = 4096\n"))
+	if err != nil {
+		t.Fatalf("a valid kv_snapshots was refused: %v", err)
+	}
+	if c.Engine[0].KVSnapshots != abs {
+		t.Errorf("kv_snapshots = %q, want %q", c.Engine[0].KVSnapshots, abs)
+	}
+	if c.Engine[0].KVSnapshotsMiB != 4096 {
+		t.Errorf("kv_snapshots_mib = %d, want 4096", c.Engine[0].KVSnapshotsMiB)
+	}
+	c, err = Load(write(t, "[[engine]]\nname = \"x\"\ntype = \"llamacpp\"\nurl = \"http://127.0.0.1:8080\"\n"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c.Engine[0].KVSnapshots != "" || c.Engine[0].KVSnapshotsMiB != 0 {
+		t.Errorf("an engine that names no snapshot directory carries %q and %d",
+			c.Engine[0].KVSnapshots, c.Engine[0].KVSnapshotsMiB)
+	}
+	if DefaultKVSnapshotsMiB != 16384 {
+		t.Errorf("the open budget default is %d, want 16384", DefaultKVSnapshotsMiB)
+	}
+
+	for _, tc := range []struct {
+		dir     string
+		wantErr string
+	}{
+		{"kv", `engine[0] "x": kv_snapshots "kv" must be an absolute path`},
+		{"./kv", `engine[0] "x": kv_snapshots "./kv" must be an absolute path`},
+	} {
+		_, err := Load(write(t, "[[engine]]\nname = \"x\"\ntype = \"llamacpp\"\nurl = \"http://127.0.0.1:8080\"\nkv_snapshots = \""+tc.dir+"\"\n"))
+		if err == nil || !strings.Contains(err.Error(), tc.wantErr) {
+			t.Errorf("kv_snapshots %q: error lacks %q: %v", tc.dir, tc.wantErr, err)
+		}
+	}
+}
+
 func TestEngineFallbackParsedAndValidated(t *testing.T) {
 	const engines = `
 [[engine]]

internal/gateway/gateway.go

@@ -779,13 +779,23 @@ func (g *Gateway) ensure(ctx context.Context, e *config.Engine) error {
 			// free while that happens, so the reads go first.
 			go g.prefetch(e)
 		}
+		if e.KVSnapshots != "" {
+			// The same wait warms the engine's saved prompt caches, so the
+			// restore below reads the page cache where it can.
+			go g.kvWarm(e, ready)
+		}
 		go g.run(e, s, ready)
 	}
 
 	deadline := time.Now().Add(ready)
 	for {
 		if probe(ctx, e.URL) {
 			g.markAnswered(e, s)
+			if starting && e.KVSnapshots != "" {
+				// Only the start this gateway ran restores: its slots are
+				// empty, and the restore lands before the first request.
+				g.kvRestore(e)
+			}
 			return nil
 		}
 		select {

internal/gateway/gateway_test.go

@@ -8,6 +8,7 @@ package gateway
 import (
 	"bufio"
 	"context"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -66,6 +67,20 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+// slotsSpec parses the env a fake engine reads for its slot surface:
+// BERTH_GW_SLOTS is the idle slot count to list (unset: no slot endpoints
+// at all), and BERTH_GW_KVLOG the file its slot calls and chat requests
+// are appended to, one line each in arrival order.
+func slotsSpec() (n int, kvlog string) {
+	if spec := os.Getenv("BERTH_GW_SLOTS"); spec != "" {
+		if _, err := fmt.Sscanf(spec, "%d", &n); err != nil {
+			fmt.Fprintf(os.Stderr, "bad BERTH_GW_SLOTS %q: %v\n", spec, err)
+			os.Exit(2)
+		}
+	}
+	return n, os.Getenv("BERTH_GW_KVLOG")
+}
+
 // testRun is the `berth run` a spawned gateway engine process stands in
 // for: the same flags the gateway passes, then the engine's command after
 // "--". It takes a real lease from the test daemon.
@@ -100,6 +115,20 @@ func testRun(argv []string) {
 }
 
 func fakeEngine(addr, model string, load time.Duration) {
+	nSlots, kvlog := slotsSpec()
+	record := func(line string) {
+		if kvlog == "" {
+			return
+		}
+		// One small Fprintln behind O_APPEND: each line lands whole, so
+		// order between handlers is arrival order.
+		f, err := os.OpenFile(kvlog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
+		if err != nil {
+			return
+		}
+		defer f.Close()
+		fmt.Fprintln(f, line)
+	}
 	mux := http.NewServeMux()
 	// The port is bound at once and answers 503 until load passes, as a real
 	// engine does while it loads: binding late would leave a window in which
@@ -113,6 +142,7 @@ func fakeEngine(addr, model string, load time.Duration) {
 		fmt.Fprintf(w, `{"object":"list","data":[{"id":%q}]}`, model)
 	})
 	mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
+		record("chat " + model)
 		w.Header().Set("Content-Type", "text/event-stream")
 		for _, tok := range []string{"lease", "hand off", "never collide"} {
 			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":%q}}]}\n\n", tok)
@@ -121,6 +151,40 @@ func fakeEngine(addr, model string, load time.Duration) {
 		}
 		fmt.Fprint(w, "data: [DONE]\n\n")
 	})
+	if nSlots > 0 || kvlog != "" {
+		// The slot surface, off unless the test asks: an idle slot list,
+		// and save/restore that answer the engine's own numbers and log
+		// the call in arrival order.
+		mux.HandleFunc("/slots", func(w http.ResponseWriter, r *http.Request) {
+			ss := []string{}
+			for i := 0; i < nSlots; i++ {
+				ss = append(ss, fmt.Sprintf(`{"id":%d,"is_processing":false}`, i))
+			}
+			fmt.Fprintf(w, "[%s]", strings.Join(ss, ","))
+		})
+		mux.HandleFunc("/slots/", func(w http.ResponseWriter, r *http.Request) {
+			slot, err := strconv.Atoi(strings.TrimPrefix(r.URL.Path, "/slots/"))
+			if err != nil {
+				http.NotFound(w, r)
+				return
+			}
+			action := r.URL.Query().Get("action")
+			var b struct {
+				Filename string `json:"filename"`
+			}
+			_ = json.NewDecoder(r.Body).Decode(&b)
+			switch action {
+			case "restore":
+				record(fmt.Sprintf("restore %d %s", slot, b.Filename))
+				fmt.Fprintf(w, `{"id_slot":%d,"filename":%q,"n_restored":12,"n_read":4096,"timings":{"restore_ms":1.0}}`, slot, b.Filename)
+			case "save":
+				record(fmt.Sprintf("save %d %s", slot, b.Filename))
+				fmt.Fprintf(w, `{"id_slot":%d,"filename":%q,"n_saved":1000,"n_written":8192,"timings":{"save_ms":1.0}}`, slot, b.Filename)
+			default:
+				http.NotFound(w, r)
+			}
+		})
+	}
 	if every := os.Getenv("BERTH_GW_LOG_EVERY"); every != "" {
 		d, err := time.ParseDuration(every)
 		if err != nil {