Skip to content

pion-vllm-mlx

The client for the prompt-cache path on Apple Silicon: PionPromptCache is the four-line drop-in for mlx_lm.make_prompt_cache, HybridRetrievalCache does the same for retrieved RAG chunks, and the Stage-2 attention patch lets Pion run the attention over the cached prefix itself. Apache-2.0; the signatures below are read from the source at build time.

pip install 'pion-vllm-mlx[mlx]'

PionPromptCache

PionPromptCache 

PionPromptCache(model=None, vquant: str = 'fp16', host: str = '127.0.0.1', port: int = 1974, stage2: bool = False, admission_threshold: int = 1, boundary_protect: int = 0, prefill_chunk_size: int | None = None, softmax_bitexact: bool = False)

Drop-in replacement for make_prompt_cache that backs onto Pion's V-store.

The first time a namespace is seen, the prefix is prefilled locally and pushed to Pion via KV.PREFIX.REGISTER + V.STOREBATCH per layer. Subsequent sessions with the same namespace fetch K/V via V.FETCH RANGE and skip prefill entirely.

Multi-worker deployments

Pion's V-store is per-worker (no shared mmap), so a multi-worker pion-server fragments the cache. The canonical pattern is to run two pion-server instances:

./pion-server -w N                      # primary KV/vector, e.g. :1974
./pion-server --kvcache -w 1 -p 1984    # shared KV cache only

Then construct PionPromptCache pointing at the KV cache instance:

pc = PionPromptCache(model, vquant="fp16", port=1984)

pion-server with --kvcache automatically caps itself at -w 1 since 0.501.

Pion (Q-only on the wire, no K/V re-fetch). Stage 1 path is unchanged — mlx-lm still consumes the returned MLX cache. Stage 2 is reached through .attend_query() / .attend_drop().

IMPORTANT: stage2=True is bridge plumbing, not an mlx-lm speedup. mlx-lm's model(suffix, cache=cache) computes attention itself on the K/V tensors in the MLX cache — it does not consume external attention output. For a standalone Mac + mlx-lm user calling only get_or_prefill(): cold path: slightly slower (extra ATTEND.PREFIX.STORE per layer) warm path: unchanged net latency change: zero, slightly negative on first request

Stage 2 produces a measurable end-to-end win only when something downstream of get_or_prefill() actually calls attend_query(): - pion-exo's gpu_attention mode (Mac cluster inference) - a custom vLLM CacheEngine (Linux GPU) - monkey-patched mlx-lm attention layer (single-Mac, measurable)

See doc/shared_kv_cache.md "Honest caveat" for the full accounting.

Requires pion-server --kvcache --metal-attention -w 1.

admission_threshold

Number of times a namespace must be observed before its K/V is actually pushed to Pion. Default 1 means "cache on first miss" (current behaviour, matches §22 measurements). Set to 2 to skip one-off prefixes — prevents single-shot prompts from evicting long-lived hot prefixes under LRU pressure. Recommended for multi-tenant workloads with a long tail of unique prompts.

boundary_protect

When > 0, the K-side and V-side V-store sessions are created via V.CREATE … SCHEMA (per-layer formats) instead of the legacy uniform KV.PREFIX.REGISTER. K stays fp16 across all layers (K drives softmax routing — keep it conservative). V uses fp16 for the first boundary_protect + last boundary_protect layers and the user's vquant for the middle. Reduces V-side storage when vquant is fp8/turbo4 while keeping the boundary layers' precision intact.

Trade-off: skips KV.PREFIX.REGISTER's cross-worker directory publish — the SCHEMA path is single-worker visibility only. With --kvcache forcing -w 1 today, that's a no-op. A future KV.PREFIX.REGISTER.SCHEMA server command would lift this.

Default 0 = uniform vquant across all layers (current behavior, backward compatible).

get_or_prefill 

get_or_prefill(prefix_token_ids, namespace: str)

Return an MLX prompt cache populated with this prefix's K/V.

Hits Pion if the namespace was registered before; otherwise prefills the prefix locally, registers + stores, returns the populated cache.

Admission policy (admission_threshold > 1): observed-miss count must reach the threshold before Pion is touched. Below threshold the prefill runs locally and the cache is returned without registering — keeps one-off prompts out of the V-store under LRU pressure.

lookup 

lookup(namespace: str) -> bool

Check whether the namespace's K/V is hot on the server.

Two states need to agree for the cache to be HIT-usable
  1. V-store registration (KV.PREFIX.LOOKUP) — used by Stage-1 fetch path.
  2. ATTEND.PREFIX.* Metal session cache (ATTEND.PREFIX.LOOKUP) — used by Stage-2 wire-mode consumers (wire-lane sparse auto), which read from the Metal cache, NOT V-store.

Pre-gh-#65-follow-on this checked only (1); a stale V-store hit while the Metal cache was cold caused _stage2_push_cold to skip on re-instantiated PionPromptCache objects (test runs, server restarts). Now: when stage2=True, BOTH must hit. The probe on (2) checks layer 0 of the attend session; full-layer audit would over-spend wire bandwidth and is unnecessary since _stage2_push_cold writes all layers atomically.

stats 

stats() -> dict

HybridRetrievalCache

HybridRetrievalCache 

HybridRetrievalCache(model, backend: Literal['inproc', 'pion'] = 'inproc', host: str = '127.0.0.1', port: int = 1974, vquant: str = 'fp16', manifest_path: str | PathLike | None = None)

Chunk-id-keyed K/V cache for RAG pipelines.

See module docstring for the two-backend design.

ingest 

ingest(chunk_id: str, chunk_token_ids: list[int]) -> dict

Encode the chunk through the model, store K/V. Idempotent.

For backend="inproc": held as MLX arrays in this process. For backend="pion": shipped via KV.PREFIX.REGISTER + V.STOREBATCH.

If chunk_id is already present with the same token count, this is a no-op (avoids re-encoding). Pass a fresh chunk_id to rotate content.

prepare 

prepare(chunk_id: str, suffix_token_ids: list[int], chunk_token_count: int | None = None)

Hydrate the chunk's K/V into a fresh mlx-lm cache and return (cache, suffix_token_ids). The cache is positioned at offset = chunk_token_count, so the next model forward will RoPE-rotate the suffix tokens at positions chunk_token_count..chunk_token_count+S-1 — identical to encoding (chunk + suffix) as one continuous text.

chunk_token_count is read from the local manifest if not provided; pass explicitly for cross-process consumers that maintain their own mapping alongside the embedding index.

Returns:

Type Description
tuple

(cache, suffix_token_ids) — pass to mlx-lm's generate loop. The cache already contains the chunk K/V; the consumer feeds suffix via model(suffix, cache=cache) directly, no extra prefill.

Stage 2 — the mlx-lm attention patch

install_pion_attention_patch 

install_pion_attention_patch() -> None

Replace mlx_lm.models.base.scaled_dot_product_attention with the Pion-aware version. Idempotent. Affects every model module that from .base import scaled_dot_product_attention AT IMPORT TIME — the monkey-patch reaches into each one's namespace.

The seam is checked BEFORE anything is patched, so an incompatible mlx-lm raises PionMlxCompatError naming the observed signature and leaves mlx-lm untouched — rather than surfacing as a TypeError inside generation, after the patch is already installed.

make_pion_prompt_cache 

make_pion_prompt_cache(model, namespace: str, prompt_cache, prefix_len: int, sparse_full_layers: Optional[Dict[str, int]] = None) -> List[PionPrefixCache]

Build a per-layer cache list pointing each layer at Pion.

Mirrors mlx_lm.models.cache.make_prompt_cache(model) but returns PionPrefixCache instances. Pass the result as cache= to model() for forward passes after a successful PionPromptCache.get_or_prefill().

Hybrid architectures (Gemma 4 / Mistral SWA / Qwen3.5) build a cache list that can be SHORTER than n_layers because KV-shared layers don't carry their own cache entry — the model routes them via shared_kv from earlier layers and pads the cache list with None internally. Probe model.model.make_cache() first so the length matches what mlx-lm actually expects; fall back to n_layers for flat-uniform models that don't expose make_cache().

Pass sparse_full_layers={"K_block": 64, "K_blocks": 8} to enable block-mean top-K sparse selection on FULL-attention layers only. Sliding-attention layers are already capped at sliding_window (typically 512) tokens by design and sparsifying below that crosses out of the training distribution, so they stay dense. Flat models without layer_types: every layer is treated as full and gets sparse.

PionPrefixCache 

PionPrefixCache(layer_idx: int, namespace: str, prompt_cache, prefix_len: int, fa_window: Optional[int] = None, sparse_mode: Optional[Dict[str, int]] = None)

Drop-in replacement for mlx-lm's KVCache that delegates the prefix portion of attention to Pion's MLX sidecar via ATTEND.PREFIX.QUERY.

The model's Attention layer still calls cache.update_and_fetch(K, V) and receives K/V back. We return ONLY the locally accumulated suffix K/V; the patched scaled_dot_product_attention then handles the merge with the Pion-resident prefix.

Identification: the patched SDPA detects this class via isinstance().

state property writable 

state

The seam with mlx-lm

PionMlxCompatError

Bases: RuntimeError

mlx-lm is present but its internals are not the shape the patch needs.

Deliberately NOT an ImportError: pion_vllm_mlx/__init__.py catches ImportError to make the patch optional, and an incompatibility swallowed there would resurface as a confusing missing-name error instead of this message.

Package README

The Pion prompt cache for mlx-lm on Apple Silicon: a prefix's K/V that a different process, a different model object, or a restarted server can reuse (Stage 1), plus an attention patch that lets Pion compute the attention over that prefix itself (Stage 2).

Stage 2 runs on three lanes, auto-selected by PionPromptCache:

Lane Where it runs TTFT p50 (Llama-3.2-1B-4bit, 5×20 prefill)
1. In-process (default, same-process consumer) MLX in the calling process; zero wire roundtrips 28.2 ms (6.51× vs vanilla cold)
2. Binary fast lane (port+1, 0xCA5E frames) Cross-process via sendmsg scatter-gather + single RTT 92.1 ms
3. RESP fallback Plain RESP, for older Pion servers without the binary listener 108.6 ms

Plus HybridRetrievalCache for RAG: chunk-id-keyed K/V hydration (4.5× p50 / 99.3 % token agreement on SQuAD v2).

Install

pip install -e pion-vllm-mlx/
# optional: dev tools + MLX
pip install -e 'pion-vllm-mlx/[dev,mlx]'

Not on PyPI yet — the distribution is built and verified, the upload is pending. When it lands this becomes pip install 'pion-vllm-mlx[mlx]'.

The mlx extra pins mlx-lm>=0.20.1,<0.32. That ceiling is not decoration: install_pion_attention_patch() replaces a private mlx-lm function and rebinds the snapshot-bound name inside every imported mlx_lm.models.* module. The signature is re-checked at patch time, so an incompatible mlx-lm raises PionMlxCompatError naming the observed signature and patches nothing, rather than failing as a TypeError inside your generation loop. To see the seam:

python -c "import pion_vllm_mlx as p; print(p.mlx_lm_seam_report())"
python pion-vllm-mlx/tests/test_mlx_lm_seam.py          # one version
pion-vllm-mlx/tests/run_mlx_version_matrix.sh           # the whole matrix

The MLX dependency is optional because the package can act as a wire-compatibility shim on machines that don't have MLX (e.g. a Linux test runner). Lanes 2 and 3 work without MLX in the consumer; lane 1 requires it.

Lanes 2 / 3 require Pion built with --kvcache --metal-attention:

./pion-server --kvcache --metal-attention -w 1

Usage

The four lines (Stage 1 — cross-process, native decode)
from mlx_lm import load, generate
from pion_vllm_mlx import PionPromptCache

model, tok = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
prefix_ids = tok.encode(SYSTEM_PROMPT, add_special_tokens=False)   # the part every request shares
suffix_ids = tok.encode(USER_TURN, add_special_tokens=False)       # the part that changes

pc = PionPromptCache(model, host="127.0.0.1", port=1974)                       # once per loaded model
cache = pc.get_or_prefill(prefix_ids, namespace="app|v1|llama-1b|system_v1")  # MISS: prefill once + store · HIT: fetch
text = generate(model, tok, prompt=suffix_ids, prompt_cache=cache)             # decode as usual, at native speed

get_or_prefill returns what mlx_lm.make_prompt_cache(model) would, with the prefix's K/V already in it, so mlx-lm decodes at native speed. The namespace names one exact token sequence for one model and quantization — key it on the tokens, never the text, and change it when either changes. The first process to ask pays the prefill once; every later call, from any process, fetches it. Runnable with timings: examples/prompt_cache_demo.py.

Stage 2 — Pion computes the attention over the prefix
from mlx_lm import load, generate
from pion_vllm_mlx import PionPromptCache, install_pion_attention_patch, make_pion_prompt_cache

install_pion_attention_patch()                                          # once per process
model, tok = load("mlx-community/gemma-4-e2b-it-4bit")
pc = PionPromptCache(model, stage2=True, host="127.0.0.1", port=1974)
ns = "rag/corpus_v1"

if not pc.lookup(ns):                       # cold: prefill once, push per-layer K/V to Pion
    pc.get_or_prefill(prefix_ids, namespace=ns)

# warm: no prefill anywhere; attention over the prefix runs in Pion
cache = make_pion_prompt_cache(model, namespace=ns, prompt_cache=pc, prefix_len=len(prefix_ids))
text = generate(model, tok, prompt=suffix_ids, prompt_cache=cache)

When the cold prefill ran in this same process the prefix stays resident as MLX arrays (lane 1, zero wire round trips — the 30.2 ms / 50.6× row in the main README). From any other process the patch uses lane 2 or 3, and then every decode step pays one round trip per layer: on a 1B model that roughly triples the per-token decode cost. So for a consumer that only generates text from another process, the four Stage-1 lines above are the faster end-to-end path; Stage 2 is for consumers that want the attention itself in Pion — the sparse long-context selectors (sparse_full_layers=), pion-exo, a custom vLLM CacheEngine. stage2=True without make_pion_prompt_cache changes nothing except a slightly slower cold path.

Hybrid models supported (per-layer routing): Llama-3.2-1B-Instruct-4bit, Qwen3.5-4B, Gemma-4-E2B-it-4bit, Gemma-4-12B-it-4bit. Add a new hybrid arch by exposing model.args.layer_types + sliding_window — the patch reads both directly.

Docs

License

Apache-2.0 — see LICENSE in this directory. Pion's client packages are deliberately permissive so they can be vendored into any stack; the Pion server this package talks to is Apache-2.0 too, with one closed binary library for its tuned vector kernels. The full map of what is under which licence is doc/licensing.md.