wkvm: state-slot inference for RWKV, GDN and Mamba hybrids

The allocation unit wkvm changes

Serving a transformer usually means a paged KV cache: memory is sliced into blocks, blocks are chained per sequence, and the scheduler juggles a growing chain. wkvm argues that for RWKV-7, GDN, Mamba2 and hybrid-linear models this is the wrong primitive. Those architectures carry a fixed-size recurrent state, so the README frames the primary allocation object as a fixed-size per-request state slot rather than a paged KV block chain. The project calls itself a hypervisor for model state, and the name is literal: WKV plus KVM, where states are the VMs and the engine is the hypervisor.

The verbs it exposes are create, snapshot, fork, hibernate, resume and live-migrate. That vocabulary only makes sense because a state slot has a bounded size. A paged KV chain grows with context, so forking it means duplicating an unbounded structure; a fixed-size slot can be copied, parked and restored cheaply. The intended audience is therefore narrow and specific: people serving linear-attention or hybrid-linear checkpoints who want session-level state control rather than raw throughput on a standard transformer. If your workload is a plain dense transformer, the central abstraction buys you nothing.

Two engines under one name, and the contract between them

The README is explicit that the Open WebUI path and the RWKV-7 path are separate. The UI demo exercises approximate Gemma routed-span mode; the native RWKV-7 durable-state engine and API are described as separate paths. A reader who skims the quick start and assumes the demo describes the whole engine will be misled.

The interesting mechanism is the stateful contract the helper configures, called parent-token-v1. WKVM binds each parked state to the model, the user, the chat, the current assistant message, the previous assistant parent, the exact visible history, and an internal raw-token digest. On a follow-up request it compares those bindings. A hit reuses the parked state; an edit, a branch, a stale parent or missing metadata triggers a safe restart rather than a wrong answer. The README states the contract preserves the exact provider-generated token history, including hidden or noncanonical tokens, instead of pretending a decode-to-text-to-token round trip is identical. That is a defensible design position, and it is also why content-changing filters or edits fail the visible-history check. It is a correctness-first choice that trades away silent cache reuse.

Installing wkvm and running a first chat

The README gives one tested target: Linux, an NVIDIA GPU with 24 GB VRAM, Python 3.12, and roughly 40 GB of free disk for the checkpoint, environments and caches. It points at docs/OPEN_WEBUI_DEMO.md as the complete installation and troubleshooting guide before you adapt it to another machine. The first command builds isolated WKVM and Open WebUI 0.10.2 environments.

./scripts/open_webui_demo.sh install

Next you accept google/gemma-4-E4B-it on Hugging Face, authenticate, and download the checkpoint to a local directory.

hf auth login
hf download google/gemma-4-E4B-it \
  --local-dir "$HOME/models/gemma-4-E4B-it"

The doctor step checks the machine before anything starts, then start brings up both loopback services, and smoke verifies a real chat. Note that WKVM_MODEL_DIR is the environment variable the helper reads.

WKVM_MODEL_DIR="$HOME/models/gemma-4-E4B-it" \
  ./scripts/open_webui_demo.sh doctor
WKVM_MODEL_DIR="$HOME/models/gemma-4-E4B-it" \
  ./scripts/open_webui_demo.sh start
./scripts/open_webui_demo.sh smoke

Open http://127.0.0.1:3000, create the first local account, and select wkvm-gemma-4-e4b-it. Stop both services with ./scripts/open_webui_demo.sh stop. A pinned Docker Compose alternative lives in deploy/open-webui/compose.yaml and uses Linux host networking because the WKVM CLI binds to 127.0.0.1. The helper starts WKVM with --enable-openai-chat, four conservative state slots, normal EOS handling and the checkpoint-native production profile. It does not reuse the high-memory B32 benchmark recipe or --ignore-eos.

What the chat endpoint will and will not accept

The current chat contract is a text-only greedy subset, and the README says so plainly. Blocking and streaming /v1/chat/completions, /v1/models, /health and /metrics are available. Requests must use text messages, temperature=0, top_p=1 and n=1. Tools, images, logprobs and custom stop sequences are not supported.

That is a hard boundary, not a configuration detail. If your application depends on tool calling, multimodal input, or sampling diversity, wkvm's chat API is the wrong surface today. The README also notes that WKVM mirrors Open WebUI 0.10.2's persisted outer-whitespace normalization and that the helper disables reasoning-tag extraction. Those two behaviours exist to keep the visible-history check honest, and they mean a prompt pipeline that rewrites text between turns will invalidate parked states more often than you might expect. The upside is that invalidation is safe rather than silent: the README describes restarts from the full rendered prompt rather than reuse of a mismatched state.

The benchmark numbers and the limits the README puts on them

This is the part of the repository worth reading before any other. The README states there is no honest workload-independent claim that WKVM is 10x faster, that the evidence supports scoped workload-specific speedups, and that some workloads still favour vLLM. It also records that a later exact-trace vLLM optimization audit superseded an earlier RTX 4090 10x-vLLM result. A project that retracts its own headline number in its own README is telling you how to read the rest.

The current Open WebUI checkpoint is a strict 2026-07-23 B32 x 8 run: all 256 requests passed, all 224 eligible continuations were reused, 32 sessions opened, zero closed, zero parent-history rejections. R5 completed in 94.953s at 345.097 output tok/s, measuring 2.153x the tested vLLM mode-0 profile and 5.305x the tested SGLang profile. The README immediately qualifies this as one controlled cross-run comparison on an active RTX 4090 desktop, with repeated rotated runs still required before treating the ratios as a publication-grade envelope. The A800 scout is described as a single exploratory B32, 98,304-token, 12-turn run measuring 12.107x versus vLLM, explicitly not a publication-grade 10x claim, and it ran on a dirty worktree.

The four-chat live demo carries similar framing: on one RTX 4090, four concurrent Open WebUI chats completed and validated first turns and follow-ups, with whole-GPU memory moving from 18,117 to 18,360 MiB and zero provider, capture, probe or validation errors. The README calls it a functional four-slot browser demo, not a controlled load test and not a 10x Open WebUI measurement, and says it does not establish quality equivalence or a universal engine ranking. Treat the ratios as hypotheses with reproduction instructions attached, not as settled throughput.

Where wkvm is the wrong tool

The README's own framing of the Open WebUI path is the clearest limitation: it is a compatibility demo, not a same-semantics replacement for full-KV Gemma serving, because routed-span keeps a bounded approximate memory. Approximate memory means the recall behaviour is not the same as a full-KV implementation, and the long-context evidence is a single needle test over a 12,000-token excerpt of Alice's Adventures in Wonderland from common-pile/project_gutenberg document 11, returning BLUE-742, Samarkand and lantern at a browser-observed TTFT of 2.162 s and E2E of 3.705 s. One needle in one document is a demonstration, not a quality envelope. If your application needs exact long-context retrieval, verify it on your own corpus first.

The security posture is the second hard limit. Both services bind to loopback, and the README states WKVM currently has no API-key enforcement, so port 8000 must not be exposed to a LAN or the Internet. That rules out any multi-tenant or shared-network deployment as shipped. Third, the chat API subset excludes tools, images, logprobs and custom stop sequences, so agent frameworks and multimodal front ends cannot use it without changes to the client. Fourth, the tested hardware is a 24 GB NVIDIA card on Linux with Python 3.12; the README directs you to the full guide before adapting to another machine, which is an acknowledgement that other setups are untested rather than merely undocumented.

How wkvm differs from vLLM and SGLang

vLLM and SGLang both organise serving around a paged KV cache and continuous batching for transformer attention. wkvm's difference is not a faster kernel; it is a different allocation object. When the model is RWKV-7, GDN, Mamba2 or another hybrid-linear architecture, there is no growing KV chain to page, so wkvm allocates a fixed-size state slot per request and makes that slot the thing you can snapshot, fork, hibernate, resume and live-migrate. The README's own comparisons against vLLM and SGLang are therefore workload-scoped, and it records that vLLM remains faster on some workloads. That is the honest reading: the state-slot model wins where state reuse across turns dominates, and loses where a mature paged-attention scheduler is already well matched to the work.

The second difference is the parent-token-v1 contract. A paged-KV server typically recomputes or reuses prefixes by token identity. wkvm additionally binds a parked state to the user, chat, assistant message, previous assistant parent, visible history and a raw-token digest, and restarts when any of those disagree. That is stricter than prefix matching and it is the mechanism behind the README's claim of zero parent-history rejections in the B32 x 8 run, alongside 224 reused continuations.

Maintenance, licence and upgrade cost

The last push to the default branch was on 2026-09-10, and the repository is not archived. The version in pyproject.toml is 0.0.1, and no releases were retrieved, so there is no tagged upgrade path to follow. Installation pins Open WebUI 0.10.2 and the helper is written against that version; the README also notes that WKVM mirrors Open WebUI 0.10.2's persisted outer-whitespace normalization. Upgrading Open WebUI past that pin is therefore a compatibility question you would have to answer yourself, because the visible-history check depends on matching normalisation behaviour.

On licensing, pyproject.toml declares license = { text = "Apache-2.0" } under [project], and the repository's top-level entries do not include a LICENSE file. A declared SPDX string in packaging metadata and a shipped licence text are not the same artefact, so confirm which one applies before you redistribute anything. The Python requirement is >=3.11 while the tested target is 3.12. The dependency split is deliberate: the M0 core is dependency-free so the scheduler and arena stay unit-testable without a GPU or torch, while the gpu and gemma-server extras pull numpy, torch, transformers and safetensors. That means the upgrade surface for the core is small, and the upgrade surface for the Gemma server tracks transformers and torch releases. The transformers constraint is >=5.7,<6, so a major bump is fenced off.

Editorial conclusion

Adopt wkvm if you are working with RWKV-7, GDN, Mamba2 or hybrid-linear models and want state handles you can snapshot, fork and resume, and if you can run its tested target: Linux, an NVIDIA GPU with 24 GB VRAM, Python 3.12 and roughly 40 GB of free disk. Do not adopt it if you need a same-semantics drop-in replacement for full-KV Gemma serving, tools, images or logprobs, or any exposure beyond loopback, because the README states there is no API-key enforcement. Before committing, read docs/OPEN_WEBUI_DEMO.md end to end, confirm the Apache-2.0 declaration in pyproject.toml against the repository's own LICENSE file, and reproduce the B32 x 8 run yourself rather than trusting the single 2.153x ratio.

Frequently asked questions

Official sources

1. Issues
2. README
3. xiaol/wkvm on GitHub