Prefix pool

PrefixCache 一次只管一個共享前綴。實際工作負載常有多個可重用的前綴 — 每個租戶不同的 system prompt、一批 few-shot 模板、一組 function-calling 前言 — 而 hot path 需要在每個 request 進來時自動挑最佳匹配，不要呼叫端做這些帳。

tinyloop_py.PrefixPool 就是這層帳務。你把已 tokenize 的前綴註冊進 pool；每次呼叫 Model.generate_with_pool，runtime 會在 pool 中找出最長的已註冊前綴（進來的 prompt 必須以它開頭），clone 它的 cache，decode 掉未匹配的後段，然後開始 sampling — 全部一次呼叫完成。完全沒匹配時，自動退回 Model.generate。

在專案內部追蹤系統裡這是 §16.7 stage 4：讓多租戶服務可以同時 cache 住每個租戶的 system prompt，而不需要在 runtime 之上自己再實作 LRU / longest-prefix 邏輯。

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)

    # 啟動時把每個已 tokenize 的 system prompt 註冊進去。
    for system_prompt_tokens in system_prompts:
        model.register_prefix(pool, system_prompt_tokens, loops=8)

    # 每個 request：把整個對話交給 pool；runtime 自動選出最長的
    # 已註冊前綴、重用它的 cache，從那裡開始 decode。
    tokens = model.generate_with_pool(
        pool, conversation_tokens,
        max_tokens=128, loops=8, temperature=0.0, top_k=50,
    )

沒有任何 explicit cache handle 會離開 pool。你從頭到尾不會直接碰 PrefixCache 物件 — pool 擁有並負責釋放它們。

tinyloop_py.PrefixPool

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

Opaque 物件，持有最多 capacity 個預先建好的 PrefixCache，以 (完整 token 前綴, loops, cache_window) tuple 為 key。

PrefixPool(capacity=16)

建立一個空 pool。

名稱	型別	預設	說明
capacity	int	16	Pool 最多持有的前綴項數。當插入會超過上限時，最短的已註冊前綴會被 evict（這是 LRU 的近似 — 較長的前綴通常價值更高，所以會被保留）。0 會被 coerce 成 1（pool 至少持有一項）。

建立 pool 不會配置 GPU 記憶體 — 記憶體只在呼叫 Model.register_prefix 時才會配置。

pool.size() -> int

目前 pool 中註冊的前綴項數。每次成功 register_prefix 會遞增；超過容量觸發 eviction 時會遞減。

pool.total_cached_tokens() -> int

所有已註冊前綴的 token 長度總和。用來粗略回答「這個 pool 目前涵蓋了多少前綴長度？」— 適合 dashboard 或 admission control。

生命週期與 thread safety

• Pool 由 Python 參考計數管理。離開 scope 時每個 cache entry（與其 GPU 記憶體）會被釋放。
• Pool 操作不是 thread-safe。底層的 PrefixCache 持有 CUDA 狀態，同一時間應該只有一條 thread 在用。如果跨 thread 共用 pool，請用 lock 串列化 — 典型的 server 模式是一個 Engine 一個 pool，背後以 asyncio.Lock 保護（見 HTTP server 範例）。
• Pool 必須 outlive 所有進行中的 generate_with_pool 呼叫 — entries 是在呼叫時查找並用於 decode。

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

把 tokenized 前綴註冊進 pool。為 tokens 建立新的 PrefixCache 並插入 pool；若 (tokens, loops, cache_window) 已存在，則是 no-op（LRU-touch 佔位；目前實作不會在 touch 時重新排序）。

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

參數

名稱	型別	預設	說明
pool	tl.PrefixPool	—	插入目標 pool。
tokens	np.ndarray[int32]	—	前綴 token id，形狀 [seq_len]。會被 copy 進 pool — 呼叫端之後可自由修改或釋放 buffer。1 <= seq_len <= max_seq_len。
loops	int	8	前綴建構時的 loop depth。entry 專屬於此深度：generate_with_pool 若 loops 不同將 miss（可能匹配到另一個同深度 entry，或退回 plain generate）。
cache_window	int	0	Sliding-window 上限。0 = full cache。屬於查找 key 的一部分。

回傳

int — 此 entry 可匹配的 token 數。成功時等於 seq_len；失敗（無效參數、build_prefix_cache 回 null）時回 0。

Raises

• RuntimeError — model 已關閉；底層 build_prefix_cache 記憶體不足。

Eviction 政策

新 entry 會把 pool 推過 capacity 時，最短的現有前綴會被丟掉。這是 LRU 的粗略近似 — 長前綴主導 pool 的 prefill 節省價值，所以政策傾向保留它們。完整時間戳排序的 LRU 追蹤為 stage-4b follow-up。

匹配語意

generate_with_pool 內部查找 entry 時，一個 entry 匹配進來的 prompt 條件是：

1. Entry 的 loops 與 cache_window 與呼叫端的 loops / cache_window 完全一致。pool 不會跨深度混用 — 改 loops 即使 tokens 相同也 miss。
2. 進來 prompt 前 len(entry.tokens) 個 token id 與 entry byte-identical。沒有 BPE 正規化、沒有空白處理 — 就是原始 token-id 前綴比對。

Pool 以 最長優先 順序持有 entry，所以多個匹配時，最長的贏。這也是 cache 能跳過最多 prefill 工作的那個。

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

從 prompt 生成；自動用 pool 中最長匹配的已 cache 前綴跳過匹配段的 prefill。沒匹配時，行為與 Model.generate 相同。

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

參數

名稱	型別	預設	說明
pool	tl.PrefixPool	—	要查找的 pool。空 pool 代表每次都 fall through 到 generate。
prompt	np.ndarray[int32]	—	完整 prompt token id，形狀 [seq_len]。seq_len + max_tokens <= model.max_seq_len。
max_tokens	int	128	生成 token 數上限（不含 prompt）。語意同 generate。
loops	int	8	Loop depth。必須與你期望命中的 pool entry 的 loops 一致；不匹配就 miss 並退回。
temperature	float	1.0	Sampling temperature。0.0 = greedy。
top_k	int	50	Top-k 截斷。
cache_window	int	0	Sliding-window 上限。語意同 generate；必須匹配才命中。

回傳

List[int] — 完整 token stream（prompt + decoded）。長度為 seq_len + emitted。

Raises

• RuntimeError — model 已關閉；cache clone 或 decode 記憶體不足；seq_len + max_tokens > max_seq_len。

執行流程

1. 在 pool 中為 (prompt prefix, loops, cache_window) 查最長匹配 entry。
2. 無匹配：呼叫 generate(model, prompt, seq_len, config) 回傳。
3. 整個 prompt 都被 cache 住（匹配長度 == seq_len）：delegate 到 generate_from_prefix_cache 回傳 — 除了 register_prefix 時已分攤的工作外，prefill 零成本。
4. 匹配是 prompt 的嚴格前綴（後段非空）：clone 已 cache 的 RuntimeKVCache，為每個後段 token 呼叫 decode_with_kv_cache_capture 把 cache 延長一位（便宜的單 token decode — 不重跑匹配段的 prefill）。最後用延長後的 cache + 最終 hidden state 建一個 synthetic PrefixCache，delegate 到 generate_from_prefix_cache 走主要 decode loop。
5. 步驟 4 中任何 allocation / clone 失敗都會被 catch 住並退回 plain generate — pool 永遠不會因內部失敗而讓使用者可見的 generation abort。

Partial-match 路徑是常見情況：註冊短 system prompt 後，收到以它為前綴的較長使用者對話。

Composability

• 與 sliding-window cache。 註冊與查詢時用相同的 cache_window。混用視為 miss。
• 與 TINYLOOP_KV_H_MODE / TINYLOOP_KV_INT8 / TINYLOOP_KV_INT4。 Cache mode env var 在 model 建構時讀取；每個 pool entry 共享 model 的 mode。pool 本身對儲存 mode 無知 — 見 KV cache modes 了解每個 mode 對 per-entry 記憶體成本的影響。
• 與 generate_stream、beam search、warm-start。 尚未組合 — 這些路徑自己吃 prompt。要 pool 的節省就用 generate_with_pool；要那些功能就用該 API。組合路徑（例如從 pool-matched cache 串流 token）追蹤為 follow-up。

何時用哪個

情境	建議 API
單一固定前綴；每次 request 都以它開頭	build_prefix_cache + generate_from_prefix_cache（沒有 pool overhead）
多個固定前綴，每個 request 做 longest-match dispatch	PrefixPool + generate_with_pool（本頁）
相同 prompt，不同 loops 深度	build_resume_handle + resume_generate
一次性、不重用的 prompt	Model.generate

Pool 相對單個 PrefixCache 有常數因子 overhead（longest-first scan 對 pool.size() 線性，加上 strict-prefix 命中時的 clone-and-extend）。對典型 pool size（預設 16 以內）相較於 pool 節省的 prefill 工作可忽略。在共享前綴的 RAG 工作負載上，可預期與裸 PrefixCache 相同的 2–3× 吞吐 asymptote — 見 PrefixCache 的量測章節。

範例樣式

多租戶 system prompt

import numpy as np
import tinyloop_py as tl

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

    # 啟動時把每個租戶的 tokenized system prompt 註冊進去。
    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 已經以租戶的 system prompt 開頭。
        # pool 會自動挑對的那個；這裡不需做租戶帳務。
        return model.generate_with_pool(
            pool, request_tokens, max_tokens=256, loops=8, temperature=0.0,
        )

Few-shot 模板庫

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)

# 呼叫端直接交出整段 prompt。pool 自動挑匹配最深的模板
# — 不需手動做 `name` → cache 的查找。
out = model.generate_with_pool(pool, full_prompt_tokens, max_tokens=128, loops=8)

檢視 pool 狀態

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

用來在 dashboard 上輸出維運指標，或當 total_cached_tokens 接近 VRAM 預算時做 admission control。

錯誤條件

錯誤	原因	修法
RuntimeError: tinyloop.Model is closed; cannot call register_prefix	對已關閉的 Model 註冊	建一個新的 Model；把 pool 的生命週期移進 with block
RuntimeError: out of memory	所有前綴 cache VRAM 累積超過可用量	減少 capacity、縮短前綴，或在 import tinyloop_py 前啟用 KV 壓縮 env var
register_prefix 回 0	無效參數（seq_len <= 0、null pool）或底層 build_prefix_cache 失敗	檢查 tokens array 形狀與 model 是否成功 load
generate_with_pool 明明有註冊前綴卻跑完整 prefill	register 與 generate 的 loops 或 cache_window 不一致	對齊雙方 kwargs。pool 用的是精確 (tokens, loops, cache_window) key
Pool entries 莫名消失	超過容量；最短的 entry 被 evict	加大 capacity，或以最長優先順序註冊 entry，讓早期 eviction 不會丟掉重要的

See also

• Prefix cache — pool 底下的單前綴原語
• Generation — 標準自回歸 decode（pool 會退回的目標）
• KV cache modes — FP16 / INT8 / store-h 變體下的 per-entry 記憶體成本
• HTTP server — 多租戶 serving frontend 預期會包裝這個原語