wllm

wllm turns batches of prompts into reproducible activation datasets. It gives researchers an HTTP, Python, and JSONL workflow for extracting scientifically identified intermediate block states, retaining row order and metadata, and loading checksum-verified NumPy matrices later without a model or GPU.

The primary 0.3 workflow is: prompts in, a probe-ready [N,D] or [N,L,D] matrix out. wllm also retains a small non-streaming OpenAI-shaped generation surface and the 0.2 single-trace APIs. It is not a drop-in vLLM server, does not claim full OpenAI parity, and does not provide continuous HTTP batching.

0.3 status: this repository is on an unreleased 0.3 development path. The source version remains 0.2.0 until a clean commit binds installed-wheel evidence. Live activation gates and TruthGen 300-row matrices for Qwen3-0.6B and Qwen3-8B are recorded under reports/. Numerical cosine/L2 and token-alignment targets pass. Remaining release blockers: a clean-tree commit for identity binding, probe paired-CI tightness on some seeds, GPU-matched no-persistence wall ratio at larger HF batches on 0.6B, and a commit-bound release-candidate package. On Blackwell GPUs set VLLM_USE_FLASHINFER_SAMPLER=0 so generation warm-up uses the portable sampler.

Choose the right tool

Tool	Prefer it when	Important distinction
wllm	You need a served, repeatable batch-to-matrix job with row identity, exact selectors, provenance, transactional shards, digests, and CPU-only reload	Narrowly targets identified intermediate block outputs and activation-dataset lifecycle
Hugging Face Transformers	You are working directly in one notebook/process and will manage batching, storage, and provenance yourself	Usually the simplest reference path for output_hidden_states=True
raw vLLM	You need high-throughput inference or its supported pooling/token-embedding outputs	Current vLLM can expose some final token embeddings, but its OpenAI server is not a stable arbitrary-intermediate-block API
NNsight	Interactive model tracing, interventions, or its vLLM-backed tracing abstraction is the main job	It supports activation access and batching; compare semantics and overhead rather than claiming wllm is the only option
TransformerLens	Circuit analysis and interventions on one of its supported architectures are the main job	It is a different model-analysis abstraction, not a vLLM activation-dataset serving contract

wllm's distinction is not merely “vLLM with hidden states.” It binds ordered feature rows to model/tokenizer identity, layer and token-position semantics, request hashes, runtime versions, and verified storage. If those guarantees are unnecessary, Transformers may be the better tool.

Ten-minute activation-dataset workflow

The package has no mandatory Torch or vLLM dependency. From a checkout, install the lightweight client/core or one exact serving profile:

python -m pip install .
python -m pip install '.[serve]'          # current alias: vLLM 0.26.0
python -m pip install '.[serve-current]'  # explicit current profile
python -m pip install '.[serve-legacy]'   # legacy/reproduction profile: vLLM 0.10.2

The current and legacy extras conflict and belong in separate environments. Contributors using the lockfile can run make bootstrap; that installs the development environment and current serving profile.

Set a model, public model name, and durable artifact root:

export WLLM_MODEL=Qwen/Qwen3-0.6B
export WLLM_SERVED_MODEL=research-model
export WLLM_ARTIFACT_ROOT="$PWD/wllm-artifacts"

Preflight the intended single-model activation configuration. Omit --local-files-only unless the model is already cached or WLLM_MODEL is a local directory.

wllm doctor \
  --json \
  --model "$WLLM_MODEL" \
  --runtime-mode activation \
  --capture-mode replay \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.9 \
  --artifact-dir "$WLLM_ARTIFACT_ROOT"

Start a prompt-only activation server. Activation mode owns one vLLM pooling model and does not load a second full-weight generation model. Completion routes are intentionally unavailable in this mode.

wllm serve "$WLLM_MODEL" \
  --served-model-name "$WLLM_SERVED_MODEL" \
  --runtime-mode activation \
  --dtype bfloat16 \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.9 \
  --artifact-dir "$WLLM_ARTIFACT_ROOT"

In another shell, wait for the model to be genuinely usable instead of backgrounding and guessing at startup time:

wllm wait-ready --base-url http://127.0.0.1:8000 --timeout 300

Extract a persisted one-layer matrix in one Python call:

from wllm import WLLMClient

statements = [
    "The review was thoughtful and useful.",
    "The experience was frustrating and slow.",
]

with WLLMClient(
    "http://127.0.0.1:8000",
    artifact_root="./wllm-artifacts",
) as client:
    client.wait_until_ready(timeout=300)
    dataset = client.extract_activations(
        model="research-model",
        inputs=statements,
        row_ids=["statement-001", "statement-002"],
        metadata=[{"label": 1, "split": "train"}, {"label": 0, "split": "test"}],
        layers=13,
        position="last_prompt",
        batch_size="auto",
        persist="dataset",
    )

X = dataset.features
print(X.shape, dataset.row_ids, dataset.manifest.dataset_id)

For one selected layer, dataset.features is [N,D]. For several layers it is [N,L,D]; dataset.features_3d is always the canonical [N,L,D] view. batch_size="auto" currently means the deterministic bounded value 16; it is not a hardware autotuner.

Persisted Python/HTTP results are local same-filesystem artifacts. The client's artifact_root must resolve to the same storage mounted as the server's --artifact-dir. There is no arbitrary artifact-download endpoint. Use persist="none" for a bounded inline matrix when a shared mount is unavailable.

JSONL to a local matrix

The CLI accepts strict UTF-8 JSONL. Every row needs a unique id, exactly one of prompt or messages, and optional finite JSON metadata:

{"id":"statement-001","prompt":"The review was thoughtful and useful.","metadata":{"label":1,"split":"train"}}
{"id":"statement-002","messages":[{"role":"user","content":"Classify this statement: the experience was frustrating."}],"metadata":{"label":0,"split":"test"}}

Run batched extraction and write a local manifest plus NumPy shards:

wllm activations statements.jsonl \
  --base-url http://127.0.0.1:8000 \
  --model "$WLLM_SERVED_MODEL" \
  --layer 13 \
  --position last-prompt \
  --batch-size 32 \
  --rows-per-shard 64 \
  --dataset-id sentiment-v1 \
  --output ./features

The CLI waits for readiness, requests bounded inline chunks, and writes the final dataset on the CLI machine. Therefore its --output does not need to be the server's artifact mount. It emits a machine-readable summary to stdout and writes:

features/
├── datasets/sentiment-v1/manifest.json
├── datasets/sentiment-v1/shard-00000.npy
├── failures.json
└── summary.json

Additional shards appear according to --rows-per-shard. With the default 64 rows per shard, a successful 300-row extraction is five .npy shards plus one manifest—not hundreds of per-row trace directories.

Resume is enabled by default. Completed rows are skipped only when row ID, prompt/messages, metadata, model, layers, position, site, and capture mode retain the same request identity. Changed rows, stale ordering, incompatible manifests, and corruption fail closed. Failed rows remain stage-labeled in failures.json and are retried on the next run. The command exits 1 while row failures remain and 2 for input, output, transport, or integrity errors. --no-resume does not overwrite an existing dataset; select a new output or dataset ID instead.

Load the result in a CPU-only process. This path imports NumPy and wllm core, not a model runtime:

from wllm.datasets import load_dataset

dataset = load_dataset("./features", "sentiment-v1")
X = dataset.features
labels = [row["label"] for row in dataset.metadata]
print(X.shape, labels)

The maintained trained_activation_probe.py example aligns labels by row ID, performs a group-aware train/test split, fits logistic regression using training-only preprocessing, saves the probe and manifest, and verifies identical predictions after reload. It can perform extraction and training as one maintained workflow:

python scripts/prepare_truthgen.py \
  --workload /tmp/truthgen-workload.jsonl \
  --labels /tmp/truthgen-labels.jsonl \
  --provenance /tmp/truthgen-provenance.json

python examples/trained_activation_probe.py \
  --labels /tmp/truthgen-labels.jsonl \
  --dataset-root "$WLLM_ARTIFACT_ROOT" \
  --extract-model "$WLLM_SERVED_MODEL" \
  --layer 13 \
  --batch-size 32 \
  --output-dir ./truthgen-probe

The example waits for readiness, persists the feature matrix, discards the response-owned array, reloads the dataset through its checksums, trains the probe, saves it, reloads it, and requires the held-out probabilities and predictions to match exactly. The saved dataset and probe can then be validated end to end in a fresh CPU-only process after the server stops:

CUDA_VISIBLE_DEVICES="" python examples/trained_activation_probe.py \
  --labels /tmp/truthgen-labels.jsonl \
  --dataset-root "$WLLM_ARTIFACT_ROOT" \
  --dataset DATASET_ID_FROM_THE_TRAINING_SUMMARY \
  --output-dir ./truthgen-probe \
  --verify-reload-only

That command checksum-loads the real persisted matrix and saved probe and asserts byte-digest- bound exact held-out probabilities and predictions without model inference. Probe training remains example code rather than a server abstraction.

HTTP contract

POST /v1/activation-datasets is the dedicated batch operation. This request selects the last prompt token after complete transformer block 13 and asks the server to persist a sharded matrix:

{
  "model": "research-model",
  "rows": [
    {"id": "row-001", "prompt": "Calibration matters because", "metadata": {"split": "train"}},
    {"id": "row-002", "prompt": "Confidence is useful when", "metadata": {"split": "test"}}
  ],
  "layers": [13],
  "position": "last_prompt",
  "site": "block",
  "capture_mode": "replay",
  "persistence": "dataset",
  "batch_size": 16
}

The response preserves every input ID in request order, records successful feature offsets and structured row failures, declares [row,layer,hidden] axes, and returns either inline features or a digest-bound manifest reference. Model work is one batched LLM.encode operation per chunk, not one encode per row. Requests to one app are serialized around unsafe runtime state, while /livez and /readyz remain event-loop state snapshots and should remain responsive during a long extraction.

The current dataset endpoint accepts capture_mode="replay". Batched online capture is rejected before inference. In activation runtime mode, “replay” describes prompt/block semantics; the forward pass uses the one resident pooling model. In generation mode, replay initializes an isolated pooling runner and therefore has a different memory profile.

Tensor identity and BF16 semantics

site="block" means the output of the complete selected transformer block after residual updates. Profile-guarded architectures that return a fused (branch_output, residual) pair are reconstructed by summing the pair. position="last_prompt" means the last token ID of the exact rendered prompt passed to capture. It is not the final normalized token embedding and not the logits-predicting state for that same token. Chat rows use the served tokenizer's chat template with a generation prompt.

The activation-dataset wire and storage dtype is currently float32. On a BF16 server, values originate from a BF16 model/capture path and are widened for NumPy transport; widening does not restore FP32 precision. That is why an approximately 1% relative-L2 difference from an HF FP32 reference can coexist with cosine similarity above 0.999. Compare identical token IDs, resolved layers, capture site, chat rendering, model revision, and dtype before interpreting either metric.

Every manifest binds:

• ordered row IDs, metadata, offsets, and request hashes;
• selected resolved layers, position, site, and capture mode;
• prompt token IDs and their digests;
• model/tokenizer names and resolved immutable revisions when the backend exposes them; mutable requested branches or tags are retained separately and never presented as resolved commits;
• wllm, vLLM, Transformers, Torch, NumPy, FastAPI, Pydantic, CUDA, and device provenance;
• feature axes, shape, dtype, byte counts, per-shard SHA-256 digests, and a manifest content digest;
• completed pre-commit phases (including measured shard-write time). The HTTP response separately reports end-to-end persistence and total request time because those values are knowable only after the immutable manifest has been committed.

Loading uses allow_pickle=False and verifies strict JSON, canonical paths, root containment, manifest identity, byte counts, SHA-256, shard ranges, row order, shape, dtype, and finite values. It never chooses a tensor through filename fuzzing. Dataset commits use a staging directory and atomic replacement; failed commits restore an earlier valid dataset or remove the incomplete new one.

Persistence policies

Policy	Result	Intended use
none	Bounded inline float32 features; no durable tensor files	Interactive checks or a client that will persist its own chunks
dataset	Manifest plus uncompressed .npy shards	Default research dataset, probe training, integrity-checked reload
full_trace	The same usable dataset plus one exact audit-trace reference per successful row	Maximum auditability when extra model passes and per-row artifacts are acceptable

full_trace deliberately performs the normal batched matrix capture first, then an additional prompt-only audit replay per successful row. It is slower and creates O(rows) trace material. New trace material is rolled back if full-trace construction or dataset persistence fails. Do not use it as the throughput baseline.

Activation mode versus generation mode

Runtime mode	Full-weight runners	Completion routes	Activation-dataset behavior
activation	One pooling runner	Unavailable	Prompt-only batched extraction on the resident model; recommended for 7B/8B activation jobs
generation	One generation runner; replay may add one pooling runner	Non-streaming chat/text completions	Replay needs --gpu-memory-utilization <= 0.5; optional online trace capture uses the eager generation runner but is not available to the batch dataset endpoint

Normal generation requests do not install extraction hooks, copy activation tensors to CPU, initialize replay resources, or write artifacts. Generation mode retains /v1/models, /v1/chat/completions, /v1/completions, /v1/extract, and /v1/traces. The OpenAI-shaped routes are a tested subset, not full OpenAI compatibility.

Compatibility and current evidence

The machine-readable source of support claims is wllm.support_matrix, exposed through /v1/extraction-schema. Its schema is wllm.support.v2; extraction and trace schemas remain wllm.extraction.v2 and wllm.trace.v2.

Profile	Exact vLLM	Extra	0.3 evidence status
current	0.26.0	serve / serve-current	Dirty-tree installed-wheel live gates recorded for Qwen3-0.6B (full) and Qwen3-8B (activation-only, one model on 24 GB); cosine/L2 and token alignment pass. Clean-commit release binding still required
legacy-0.10	0.10.2	serve-legacy	Dirty-tree installed-wheel activation-only live gate recorded for Qwen3-0.6B. The 0.3 release claim remains activation-mode-only; keep a pinned 0.2 environment for legacy generation/trace reproduction

Unknown vLLM versions fail closed. Hidden-state capture is tensor-parallel-size 1 only. Public 0.3 release claims must come from exact-commit machine-readable evidence after a clean identity bind, not from this table alone; see docs/support.md for the complete boundary.

Historical correctness baseline—not a 0.3 benchmark

An external 0.2 test used Qwen3-0.6B revision c1899de289a04d12100db370d81485cdf75e47ca, layer 13, 300 statements, and an RTX 5090 Laptop. All 664 compared vectors had cosine similarity above 0.999 against HF FP32 and BF16. A sentiment positive control produced accuracy/AUC 1.0 through both extractors; an arithmetic negative control remained chance-like. wllm took 23–27 ms per statement versus roughly 9 ms for HF and created about 600 per-row objects. Those results establish a useful correctness baseline and the reason for this batch/shard pivot; they do not establish 0.3 throughput or broad probe accuracy.

The executable activation-dataset benchmark writes machine-readable JSON and renders Markdown from it. HF is the reference adapter. Raw vLLM and NNsight have executable capability probes: they report only the APIs and semantics they can validate and emit no fabricated comparison vectors when the required intermediate-layer contract is unavailable. Dirty-tree harness reports under reports/ already exercise HF agreement, trained probe, storage count, and peak-VRAM sampling; publishable release numbers still require a clean commit identity and a rebind of those gates to the exact 0.3.0 wheel.

Security and privacy

• Configure --api-key for every non-health route and use TLS at a trusted reverse proxy when crossing a network. /livez and /readyz intentionally remain unauthenticated.
• The server's structured request log records request/dataset/trace IDs, route, status, and duration—not prompt bodies or authorization headers. Unexpected responses expose exception type rather than arbitrary backend text.
• Artifact roots are local, unencrypted files. Use restrictive filesystem permissions. Token IDs, activations, metadata, and especially full_trace bundles may reveal sensitive input content.
• Treat manifests and probe files as untrusted input. Dataset loading rejects path traversal and integrity mismatches. The trained-probe example uses joblib/pickle semantics and must only load trusted probe files after digest verification.
• The project adds no telemetry and no remote object-store transport.

Limitations

• Linux, single node, NVIDIA CUDA for model serving; CPU-only dataset loading is supported.
• Exact profiled vLLM versions only; no promise for an arbitrary newer patch.
• Tensor parallel size 1 for activation capture; no multi-node activation serving.
• Block outputs at last_prompt for the batch API; no arbitrary token matrix, attention weights, post-attention/post-MLP sites, or interventions.
• Attention weights | unsupported. Attention-weight extraction returns an actionable error before inference.
• No streaming, tool calling, multimodality, LoRA contract, speculative decoding contract, raw logits, exact full-vocabulary entropy, or full OpenAI API parity.
• One app serializes model operations. Dataset requests batch within that operation, but wllm does not claim vLLM-style continuous batching across concurrent HTTP requests.
• Persisted Python/HTTP datasets are local same-filesystem artifacts; the CLI is the supported way to materialize bounded chunks on another machine.
• The loader currently concatenates shards into memory; it is not a streaming or memory-mapped analytics engine.
• Live installed-wheel activation evidence and TruthGen 300-row matrices for 0.6B/8B are recorded under reports/. A full release-candidate package still requires a clean commit identity bind.
• On Blackwell (sm_120) hosts, set VLLM_USE_FLASHINFER_SAMPLER=0 so generation warm-up uses the portable sampler (live validators set this automatically).
• Matched-batch GPU wall-time ratio vs HF can exceed 1.5× when HF benefits from large local batches while wllm pays HTTP/JSON costs; batch-1 prompt-only paths often meet the target.

See docs/migration-0.3.md when moving from 0.2.

Development and release gates

Run the local source/package gate:

make check

Emit GPU-independent evidence:

make release-cpu

That report must remain release_ready=false until exact-commit GPU evidence is present. On a CUDA runner with local models, the release validator is:

.venv/bin/python scripts/validate_release.py \
  --current-model /absolute/path/to/Qwen3-0.6B \
  --legacy-model /absolute/path/to/Qwen3-0.6B \
  --useful-model /absolute/path/to/Qwen3-8B \
  --smoke-benchmark-report /absolute/path/to/qwen3-0.6b/matrix.json \
  --useful-benchmark-report /absolute/path/to/qwen3-8b/matrix.json \
  --artifact-dir dist/gpu-tested-release \
  --report reports/release-candidate-current.json

The 0.2 single-trace examples remain available as scripts/prefill_actmap_extraction.py and scripts/dataset_workflow.py, but the latter's per-row JSONL workflow is historical; use wllm activations for batch-to-matrix extraction. Historical v0.1 and 0.2 reports do not override executable current-head evidence.

Before public release, the repository owner should set—after approval:

• GitHub description: Reproducible, checksum-verified activation datasets from vLLM-backed models.
• Homepage/docs URL: the repository README or future documentation site.
• Topics: vllm, hidden-states, interpretability, llm-evaluation, mechanistic-interpretability, uncertainty-quantification, activations, and reproducible-research.

Do not tag, publish to PyPI, or use GitHub metadata APIs until the changelog, support matrix, wheel/sdist, clean installs, current and legacy profiles, numerical checks, benchmark, 8B memory gate, trained probe, and exact-commit release evidence all agree.