session: snapshot and restore a session's decoding state (#548)

[geisten]

Requested by Plasma · project thread

Summary

Before: a Qwen3.5 (DeltaNet hybrid) session had no reusable prefix. pin_prefix returns GEIST_E_UNSUPPORTED there, so every request re-prefilled the whole constant context.

After: geist_session_snapshot saves a session's complete decoding state into a caller buffer. geist_session_restore puts it back into the same session or into another session of the same model (a fork). Prefill the constant part once, snapshot, then restore + prefill only the new tokens per request. Restored sessions continue bit-identically (tokens and logits), with greedy decoding and with sampling.

How: the image holds the used KV rows [0, kv_len) of each attention layer, the conv and recurrent state of each DeltaNet layer, the pending logits, a decode step whose forward is still owed, the pinned prefix and the sampler RNG. Every KV representation is position-major from offset 0, so only the used prefix of each buffer is copied. Device-local buffers (Vulkan) go through buffer_download/upload. Restore validates everything before it writes: magic/version, a per-model snapshot_id, the KV mode, max_seq_len, the pending token and the exact size. A bad image returns GEIST_E_FORMAT and leaves the session untouched.

Closes #548.

Changes

• include/geist_util.h: geist_session_snapshot_size, geist_session_snapshot, geist_session_restore (EXPERIMENTAL). Out-size-first, handle last, like peek_logits.
• include/geist_arch.h: three optional decoder ops, appended so existing offsets stay put.
• src/archs/transformer/snapshot.c: the implementation. snapshot_id on the arch state binds an image to its model.
• src/engine/session.c: public wrappers. They refuse while an audio stream is open.
• Not covered, GEIST_E_UNSUPPORTED: KIVI KV cache, an enabled MTP drafter (GEIST_MTP=1), embedding models.
• CHANGELOG, API_CONTRACT (listed as not in the contract), pin_prefix docs point to it.

Testing

• make test-unit: 92 passed, 31 skipped (x86-64 AVX-512, gcc 14)
• make MODE=asan test-unit: same
• make format-check clean (clang-format 22.1.5); check-headers and check-api-contract pass
• make BACKENDS="vulkan cpu_x86 cpu_scalar" test-unit FILTER=snapshot on lavapipe: the Vulkan path passes too
• The new tests/test_session_snapshot_unit builds a llama and a qwen35 hybrid in memory and runs on every backend in the build. It covers FP32/F16/INT8/INT4 KV, greedy and sampling. Cases: restore after a prefill into the same and another session, a second turn, restoring twice, mid-decode with an owed forward, prefix reuse across interleaved sessions, an empty image, and a pinned prefix carried by the image. It also checks refusals: short buffer, truncated or corrupted image, another model, another KV mode, too small a max_seq_len, and KIVI. Mutation-checked: dropping the DeltaNet state, the RNG or the owed-forward flag on restore fails it.

No perf impact: no existing path changed. Not measured on Metal (no Mac here). The test lists metal, so the macOS leg runs it.

API impact

• No change to include/geist.h
• New EXPERIMENTAL surface only
• Breaks a STABLE symbol

🤖 Generated with Claude Code

https://claude.ai/code/session_01Nb4xVBy5t3fTdSStHPPKPg

---

Generated by Claude Code