Prefix pool

PrefixCache handles one shared prefix at a time. Real workloads often carry multiple reusable prefixes — a different system prompt per tenant, a bank of few-shot templates, a collection of function-calling preambles — and the hot path wants to auto-pick the best match for each incoming request without the caller doing the bookkeeping.

tinyloop_py.PrefixPool is that bookkeeping. You register tokenized prefixes into a pool; on every Model.generate_with_pool call, the runtime walks the pool for the longest registered prefix that the incoming prompt starts with, clones its cache, decodes the unmatched suffix tokens, and proceeds to sampling — all in one call. When nothing matches, it transparently falls through to Model.generate.

Architecturally this is §16.7 stage 4 in the project's internal tracker: the serving-side primitive that lets a multi-tenant service keep every tenant's system prompt cached concurrently without re-implementing LRU/longest-prefix logic above the runtime.

Quickstart

import numpy as np
import tinyloop_py as tl

with tl.Model("model.tinyloop", max_seq_len=2048) as model:
    pool = tl.PrefixPool(capacity=16)

    # Register a few tokenized system prompts once at startup.
    for system_prompt_tokens in system_prompts:
        model.register_prefix(pool, system_prompt_tokens, loops=8)

    # Every request: hand the full conversation to the pool; the runtime
    # picks the longest registered prefix the prompt starts with, reuses
    # its cache, and decodes from there.
    tokens = model.generate_with_pool(
        pool, conversation_tokens,
        max_tokens=128, loops=8, temperature=0.0, top_k=50,
    )

No explicit cache handle leaves the pool. You never touch a PrefixCache object directly — the pool owns and frees them.

tinyloop_py.PrefixPool

class PrefixPool:
    def __init__(self, capacity: int = 16): ...
    def size(self) -> int: ...
    def total_cached_tokens(self) -> int: ...

Opaque object holding up to capacity prebuilt PrefixCache entries, keyed on the tuple (exact token prefix, loops, cache_window).

PrefixPool(capacity=16)

Construct a new, empty pool.

Name	Type	Default	Description
capacity	int	16	Maximum number of prefix entries the pool holds. When an insertion would exceed this, the shortest registered prefix is evicted (a cheap approximation of LRU — longer prefixes tend to be more valuable and are preserved). 0 is coerced to 1 (the pool always holds at least one entry).

Constructing a pool does not allocate any GPU memory — memory is allocated only when you call Model.register_prefix.

pool.size() -> int

Number of registered prefixes currently in the pool. Increases with each successful register_prefix; decreases when eviction fires on a capacity-exceeding insert.

pool.total_cached_tokens() -> int

Sum of token lengths across all registered prefixes. Useful as a rough "how much prefix coverage does the pool hold?" metric for dashboards or admission control.

Lifetime and thread safety

• The pool is owned by Python reference count. When it goes out of scope, every cached entry (and its GPU memory) is freed.
• Pool operations are not thread-safe. The underlying PrefixCache objects hold CUDA state and should be accessed from one thread at a time. Serialise via a lock if you share a pool across threads — the typical server pattern is one pool per Engine, behind an asyncio.Lock (see the HTTP server code path for an example).
• The pool must outlive any in-flight generate_with_pool call — entries are looked up at call time and decoded against.

Model.register_prefix(pool, tokens, loops, cache_window)

Register a tokenized prefix in pool. Builds a fresh PrefixCache for tokens and inserts it into the pool; if the exact (tokens, loops, cache_window) key is already present, it is a no-op (LRU-touch placeholder; the current implementation does not reorder the pool on touch).

matched = model.register_prefix(
    pool, tokens,
    loops=8,
    cache_window=0,
)

Parameters

Name	Type	Default	Description
pool	tl.PrefixPool	—	Pool to insert into.
tokens	np.ndarray[int32]	—	Prefix token ids, shape [seq_len]. Copied into the pool — caller is free to mutate or drop the buffer afterwards. Must satisfy 1 <= seq_len <= max_seq_len.
loops	int	8	Loop depth the prefix is built at. The entry is specialised to this depth: a subsequent generate_with_pool call with a different loops value will not hit this entry (it may hit a different one at that depth, or fall through to plain generate).
cache_window	int	0	Sliding-window cap used when building the prefix. 0 = full cache. Part of the lookup key.

Returns

int — the number of tokens that will be matchable from this entry. On success this is exactly seq_len; on failure (e.g. invalid args, build_prefix_cache returned null) returns 0.

Raises

• RuntimeError — model has been closed; out-of-memory during the underlying build_prefix_cache.

Eviction policy

When a new entry would push the pool over capacity, the shortest existing prefix is dropped. This is a rough approximation of LRU — long prefixes dominate the prefill savings that the pool exists to preserve, so the policy keeps them preferentially. Full timestamp-ordered LRU is tracked as a stage-4b follow-up.

Matching semantics

On lookup (inside generate_with_pool), an entry matches an incoming prompt when:

1. The entry's loops and cache_window equal the call's loops and cache_window. The pool never mixes depths — changing loops misses even if the tokens are identical.
2. The incoming prompt's first len(entry.tokens) token ids are byte-identical to the entry. No BPE normalization, no whitespace handling — this is a raw token-id prefix match.

The pool holds its entries in longest-first order, so when multiple entries match, the longest one wins. That is the one whose cache skips the most prefill work.

Model.generate_with_pool(pool, prompt, ...)

Generate from a prompt, automatically using the longest matching cached prefix in pool to skip prefill on the matched tokens. When no match is found, behaves like Model.generate.

tokens = model.generate_with_pool(
    pool, prompt,
    max_tokens=128,
    loops=8,
    temperature=0.0,
    top_k=50,
    cache_window=0,
)

Parameters

Name	Type	Default	Description
pool	tl.PrefixPool	—	Pool to consult. Passing an empty pool falls through to generate on every call.
prompt	np.ndarray[int32]	—	Full prompt token ids, shape [seq_len]. seq_len + max_tokens <= model.max_seq_len.
max_tokens	int	128	Upper bound on generated tokens (excludes the prompt). Same semantics as generate.
loops	int	8	Loop depth. Must match the loops of the pool entry you expect to hit; mismatches cause a miss and fall through.
temperature	float	1.0	Sampling temperature. 0.0 = greedy.
top_k	int	50	Top-k cutoff.
cache_window	int	0	Sliding-window cap. Same semantics as generate; must match the entry to hit.

Returns

List[int] — the full token stream (prompt + decoded). Length is seq_len + emitted.

Raises

• RuntimeError — model has been closed; out-of-memory during cache clone or decode; seq_len + max_tokens > max_seq_len.

Execution flow

1. Look up the longest matching entry for (prompt prefix, loops, cache_window) in pool.
2. If no match, call generate(model, prompt, seq_len, config) and return.
3. If the entire prompt was cached (match length == seq_len), delegate to generate_from_prefix_cache and return — zero prefill work beyond what was already amortised at register_prefix time.
4. If the match is a strict prefix of the prompt (suffix non-empty): clone the cached RuntimeKVCache, then for each suffix token call decode_with_kv_cache_capture to extend the cache by one position (cheap single-token decode — no re-prefill of the matched prefix). Finally, seed a synthetic PrefixCache from the extended cache + final hidden state and delegate to generate_from_prefix_cache for the main decode loop.
5. Any allocation / clone failure during step 4 is caught and the call falls through to plain generate — the pool never aborts user-visible generation on an internal failure.

The partial-match path is the common case when you register a short system prompt and then receive longer user conversations that start with it.

Composability

• With sliding-window cache. Register and query both at the same cache_window. Mixing values is treated as a miss.
• With TINYLOOP_KV_H_MODE / TINYLOOP_KV_INT8 / TINYLOOP_KV_INT4. Cache-mode env vars are read when the model is constructed; every pool entry shares the model's mode. The pool itself is storage-mode-agnostic — see KV cache modes for how each mode changes the per-entry memory cost.
• With generate_stream, beam search, warm-start. Not yet composed — those paths take their own prompt directly. Use generate_with_pool when you want the pool savings; use the other APIs for the features they expose. Combined paths (e.g. streaming tokens from a pool-matched cache) are tracked as follow-ups.

When to use which

Scenario	Recommended API
Single fixed prefix; every request starts with it	build_prefix_cache + generate_from_prefix_cache (no pool overhead)
Multiple fixed prefixes, longest-match dispatch per request	PrefixPool + generate_with_pool (this page)
Same prompt, different loops depths	build_resume_handle + resume_generate
One-off prompts with no reuse	Model.generate

The pool has a constant-factor overhead over a single PrefixCache (the longest-first scan is linear in pool.size(), plus the clone-and-extend step when the match is a strict prefix). For typical pool sizes (up to the default 16 entries) this is negligible compared to the prefill work the pool avoids. Expect the same 2–3× throughput asymptote as a bare PrefixCache on shared-prefix RAG workloads — see the PrefixCache measured numbers.

Example patterns

Multi-tenant system prompts

import numpy as np
import tinyloop_py as tl

# Load one-model-many-tenants: each tenant has a distinct system prompt.
with tl.Model("model.tinyloop", max_seq_len=4096) as model:
    pool = tl.PrefixPool(capacity=64)

    # Register every tenant's tokenised system prompt at startup.
    for tenant_id, sys_tokens in tenants.items():
        n = model.register_prefix(pool, sys_tokens, loops=8)
        print(f"tenant {tenant_id}: {n} prefix tokens cached")

    def answer(request_tokens: np.ndarray) -> list[int]:
        # `request_tokens` already starts with the tenant's system prompt.
        # The pool picks the right one; no tenant bookkeeping here.
        return model.generate_with_pool(
            pool, request_tokens, max_tokens=256, loops=8, temperature=0.0,
        )

Few-shot template library

# Register several few-shot templates, each at a different length.
pool = tl.PrefixPool(capacity=8)
for name, tmpl_tokens in [
    ("zero_shot",    zero_shot_tokens),
    ("one_shot",     one_shot_tokens),
    ("three_shot",   three_shot_tokens),
    ("chain_of_thought", cot_tokens),
]:
    model.register_prefix(pool, tmpl_tokens, loops=8)

# Callers just hand in the whole prompt. The pool picks the deepest
# matching template automatically — no manual `name` → cache lookup.
out = model.generate_with_pool(pool, full_prompt_tokens, max_tokens=128, loops=8)

Inspecting pool state

print(f"entries: {pool.size()}  "
      f"total cached tokens: {pool.total_cached_tokens()}")

Use these to emit ops metrics or to gate admission when total_cached_tokens approaches a VRAM budget.

Error conditions

Error	Cause	Fix
RuntimeError: tinyloop.Model is closed; cannot call register_prefix	Registered against a closed Model	Construct a fresh Model; move the pool's lifetime inside the with block
RuntimeError: out of memory	Cumulative prefix cache VRAM exceeds available	Reduce capacity, shorten prefixes, or enable KV-compression env vars before import tinyloop_py
register_prefix returns 0	Invalid args (seq_len <= 0, null pool) or underlying build_prefix_cache failed	Check the tokens array shape and that the model loaded successfully
generate_with_pool runs at full prefill cost despite a registered prefix	loops or cache_window mismatched between register and generate calls	Align both kwargs. Pool uses an exact (tokens, loops, cache_window) key
Pool entries disappear unexpectedly	Capacity exceeded; shortest entry was evicted	Increase capacity, or register entries in longest-first order so early evictions don't drop the ones you care about

See also

• Prefix cache — single-prefix primitive underneath the pool
• Generation — standard autoregressive decode (what the pool falls through to)
• KV cache modes — per-entry memory cost under FP16 / INT8 / store-h variants
• HTTP server — the pool is the intended building block for a multi-tenant serving frontend