從 Python 使用 KV cache 模式

TinyLoop 有五種 KV cache 儲存模式（預設 FP16、INT8、store-h FP16、store-h INT8、store-h INT4），透過環境變數選擇。它們是在 decode 延遲與 VRAM 之間做取捨：INT8 KV 可節省約 49 % KV 記憶體，只付出小約 8 % 的延遲代價；三種 h 模式則在純 scalar reconstruction 路徑下，能節省 43–78 %、但 decode 成本大約變成 3 倍，或在 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 時反而比 FP16 KV 基線更快。

這一頁講的是 Python 端表面。關於實際儲存版面、完整記憶體公式、可組合性與選型建議，請看框架層級的 KV Cache Modes。

快速模式表

模式	啟用方式	VRAM 節省（L=32, seq=1024）	K/V 雜訊	相對基線 decode（+`fp16_body`）
FP16 KV	預設	—	0	—（開 `fp16_body` 後 −22 %）
INT8 KV	`TINYLOOP_KV_INT8=1`	49 %	約前 40 tok 幾乎為 0，之後才開始漂移	+13 %
FP16 h	`TINYLOOP_KV_H_MODE=1`	43 %	0（位元級重建）	−38 %
INT8 h	`TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1`	66 %	L2 相對誤差約 0.6 %	−37 %
INT4 h	`TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1`	78 %	L2 相對誤差約 10 %	−37 %

如何設定模式

務必在 import tinyloop_py 之前 設定 env var：

import os
os.environ["TINYLOOP_KV_H_MODE"] = "1"      # 啟用 store-h 模式
os.environ["TINYLOOP_KV_INT4"]   = "1"      # ... 把 h 打包成 signed INT4 nibbles
import tinyloop_py as tl
model = tl.Model("model.tinyloop", max_seq_len=2048)

這些 env var 會在首次配置 cache 時被探測，並快取在模組層級的 static 變數中。若你在第一次 generate / build_prefix_cache / build_resume_handle 之後再修改它們，不會有任何效果。因此 TinyLoop 不支援在同一個 process 內熱切換模式。

優先順序

如果在 TINYLOOP_KV_H_MODE=1 下同時設了 TINYLOOP_KV_INT8 與 TINYLOOP_KV_INT4，INT4 優先，也就是會選更激進的壓縮。若沒有 TINYLOOP_KV_H_MODE=1 卻只設 TINYLOOP_KV_INT4=1，它會被靜默忽略（TinyLoop 沒有 INT4 K/V 路徑，INT4 只接到儲存 h 的模式）。

而 TINYLOOP_KV_INT8=1 在未開 TINYLOOP_KV_H_MODE=1 時，則單獨選擇 INT8 K/V 路徑，與 h mode 完全無關。

建議組合

品質優先、小模型（≤ 1B）、VRAM 很寬鬆： 不設任何 env var；若想再加快 decode，可另開 TINYLOOP_EXPERIMENTAL_FP16_BODY=1。
延遲與品質兼顧，而且有一些權重 VRAM 餘裕： TINYLOOP_KV_H_MODE=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1。FP16-h + tensor-core reconstruction 可在保留 43 % KV 節省的同時，達成 −38 % decode，且 K/V 完全無雜訊。
長上下文，可接受小量誤差： TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1。可節省 66 % KV，K/V 僅有 0.6 % L2 相對誤差，而且同樣大約 −37 % decode。
長上下文，VRAM 極度緊張： TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 TINYLOOP_EXPERIMENTAL_FP16_BODY=1。可節省 78 % KV（比 INT8-h 再小 37 %），但 K/V L2 相對誤差約 10 %。

在 Python 端可觀察到的行為

對 Python 呼叫端來說，儲存模式是透明的：本系列頁面中所有 API（score_*、generate、prefix cache、warm-start）在任一模式下都能使用。會改變的是：

VRAM 用量。 model.vram_usage_mb() 會隨著 cache 增長，在 INT8 / h 模式下回報更低數字。
Decode 延遲。 在 generate_stream 的 decode_steady_ms 欄位中可直接看出。
數值容忍度。 有損模式的 K/V 只會在各自的 L2 相對誤差範圍內逼近 FP16 參考值。位元級 parity test 只會在 FP16 模式下通過；有損模式有各自對應的容差測試。
建立 cache 的速度。 INT8 / INT4 打包會在 prefill append 路徑上增加些微 per-token 開銷（每個 head 多一次 reduction + 量化），但相較於 GEMM 幾乎可以忽略。

與其他 API 的組合

API	h mode	INT8 h	INT4 h	說明
`score`、`score_last`	✓	✓	✓	輸出只會受到各模式重建誤差包絡的影響。
`score_logit_lens`	✓	✓	✓	每次迭代的 logit lens 都會使用與 cache mode 對應的 K/V。
`score_trajectory`	✓	✓	✓	trajectory 擷取的是 FP32 residual，本身與 cache mode 無關；`hidden` 陣列是 cache-independent。
`score_with_uncertainty` / `score_with_consistency`	✓	✓	✓	KL / argmax 比較對小量 K/V 雜訊具韌性。
`generate`	✓	✓	✓	採樣是基於 cache logits；INT4-h 下的小差異有機會翻轉貼近的 argmax。
`generate_stream`	✓	✓	✓	與 `generate` 相同；時間欄位會把重建成本算進去。
`generate_speculative`	✓	✓	✓	draft 與 verifier 都跑在同一 cache mode 上。
`build_prefix_cache` / `generate_from_prefix_cache`	✓	✓	✓	prefix cache 會繼承該模式；`h` 模式能讓 prefix cache 大約再小一倍。
`build_resume_handle` / `resume_generate`	✓	✓	✓	resume handle 也會繼承此模式。FP16 模式保證位元級 parity；有損模式則只能保證落在各自誤差邊界內。

如何確認目前啟用的是哪個模式

Model.config() 不會直接回報 cache mode，但你可以透過 env 與時間側面確認：

import os
import time
import numpy as np

prompt = np.arange(64, dtype=np.int32)
t0 = time.perf_counter()
_ = model.generate(prompt, max_tokens=32, loops=8, temperature=0.0)
dt = time.perf_counter() - t0

print(f"{dt * 1000:.1f} ms for 32 tokens (config: "
      f"h_mode={os.environ.get('TINYLOOP_KV_H_MODE', '0')}, "
      f"int8={os.environ.get('TINYLOOP_KV_INT8', '0')}, "
      f"int4={os.environ.get('TINYLOOP_KV_INT4', '0')}, "
      f"fp16_body={os.environ.get('TINYLOOP_EXPERIMENTAL_FP16_BODY', '0')})")

如果你開了 TINYLOOP_KV_H_MODE=1 但沒開 TINYLOOP_EXPERIMENTAL_FP16_BODY=1，而壁鐘卻仍接近 FP16 KV 基線，那通常代表 env var 沒有真正生效（多半是因為它們在第一次建構 Model 之後才設定）。請重啟 process。

如何並排 benchmark 不同模式

在同一個 process 內切換模式是不安全的（env 會被快取）。正確做法是「每個模式一個子進程」：

import os, subprocess, sys

def bench(mode_env, model_path):
    code = f'''
import os
for k, v in {mode_env!r}.items(): os.environ[k] = v
import numpy as np, time, tinyloop_py as tl
model = tl.Model({model_path!r}, max_seq_len=2048)
prompt = np.arange(512, dtype=np.int32)
_ = model.generate(prompt, max_tokens=8, loops=8, temperature=0.0)  # warmup
t0 = time.perf_counter()
r = model.generate_stream(prompt, lambda t: True,
    max_tokens=64, loops=8, temperature=0.0)
print(f"decode_ms={{r['decode_steady_ms']:.2f}}")
'''
    return subprocess.check_output([sys.executable, "-c", code]).decode()

for label, env in [
    ("FP16 KV", {}),
    ("FP16 h + fp16_body", {"TINYLOOP_KV_H_MODE": "1",
                              "TINYLOOP_EXPERIMENTAL_FP16_BODY": "1"}),
    ("INT4 h + fp16_body", {"TINYLOOP_KV_H_MODE": "1",
                              "TINYLOOP_KV_INT4": "1",
                              "TINYLOOP_EXPERIMENTAL_FP16_BODY": "1"}),
]:
    print(f"{label}: {bench(env, '/workspace/model_1b_effective.tinyloop').strip()}")

現成版本可直接用 tests/bench_kv_h_mode.py，它會跑九種配置（三種 cache mode × 兩種 reconstruction path，再加控制組），並印出彙總表。

不變量與正確性

在 greedy decoding（temperature=0.0）下，五種模式之間的關係應該如下：

FP16 KV 與 FP16 h 位元級一致：每個 token id 都應完全相同；若看到分歧，就是 bug。
INT8 KV 通常在前大約 40 個 token 內與 FP16 KV argmax 一致，之後才可能分歧；兩條 continuation 都可視為合法樣本。
INT8 h 在 prompt 位置和前幾個 decode step 通常與 FP16 KV 一致，但 0.6 % K/V L2 相對誤差終究可能翻轉某些近距離 logit 的 argmax。
INT4 h 通常在 prompt 位置也能保持一致，但 10 % K/V L2 相對誤差更容易翻轉接近分數的 argmax。尚未完成大規模語意測試；下一個驗證項目是 1B WikiText-103 val 的 PPL sweep。

若你在第一個 decode token 就看到 FP16 KV 與其他任一模式分歧，那就是 bug，應該回報。

常見陷阱

匯入順序。 env var 一定要在 import tinyloop_py 之前設定，否則沒效果。
模式是 process-wide。 你不能在同一個 process 裡讓兩個 Model 實例使用不同 cache mode；要靠子進程分離。
INT4-h 要求 head_dim 為偶數。 head_dim = dim / n_heads。若不滿足，alloc path 會記警告並退回 FP16-h。現有 407M / 1B 參考模型都符合 head_dim = 128。
Scratch 是跨層共用。 h 模式只會為整個 RuntimeKVCache 配一組重建 scratch buffer，而不是每層一組。所以在低 n_layers 下，scratch 會占較大比例，估算記憶體時別忽略。

快速模式表​

如何設定模式​

優先順序​

建議組合​

在 Python 端可觀察到的行為​

與其他 API 的組合​

如何確認目前啟用的是哪個模式​

如何並排 benchmark 不同模式​

不變量與正確性​

常見陷阱​