pion-vllm-mlx

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:

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

• Main project: ../README.md
• Shared KV cache design: ../doc/shared_kv_cache.md
• 64K NIAH reproducer (Gemma 4): ../examples/sparse_mask_64k_niah.py

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.