跳至主要内容

KV Cache 模式

TinyLoop 支援五種 KV cache 儲存模式。它們在「每層每 token 究竟儲存什麼」、「快取在 VRAM 中有多大」，以及「每個 decode token 需要額外支付多少計算量」三件事上各不相同。這一頁是完整參考資料；如果你正在為某個部署選擇 cache mode，判斷所需的所有資訊都在這裡。

本頁假設你已經有一個可用的 TinyLoop 執行時（Model.generate、Model.score_last 等）。周邊 API 請見 Python API。

五種模式​

名稱啟用環境變數每個 token 每層實際儲存內容品質代價
FP16 KV（預設）—兩個 FP16 向量：K（[D]）、V（[D]）無（參考基線）
INT8 KVTINYLOOP_KV_INT8=1兩個 INT8 向量 + 每個 head 一個 FP16 scale大約到 40 token 前 PPL 幾乎無漂移，之後慢慢分歧
FP16 hTINYLOOP_KV_H_MODE=1一個 FP16 向量：h = LN(residual)（[D]）無（K/V 重建位元級一致）
INT8 hTINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT8=1一個 INT8 的 h 向量 + per-head FP16 scalesK/V L2 相對誤差約 0.6 %，沒有可量測的 attention argmax 漂移
INT4 hTINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1一個 nibble-packed h 向量 + per-head FP16 scalesK/V L2 相對誤差約 10 %（INT8 有 255 級，INT4 對稱量化只有 15 級）

三種 h 模式都是只有權重共享迴圈式 Transformer 才有的架構專屬能力。它們之所以能成立，是因為每次 loop iteration 都使用同一組 W_k / W_v，因此每個快取位置的 K 和 V 都能在需要時透過 W_k · h / W_v · h 即時重建。標準深層 Transformer 做不到這點，因為每一層都有自己的 W_k。

在 TINYLOOP_KV_H_MODE=1 下，TINYLOOP_KV_INT4 的優先級高於 TINYLOOP_KV_INT8。如果兩者都設了，會選 INT4-h。若沒開 TINYLOOP_KV_H_MODE=1 只設 TINYLOOP_KV_INT4=1，則會被靜默忽略（TinyLoop 沒有 INT4 K/V 路徑；INT4 只支援儲存 h）。

記憶體公式​

令 D 為 hidden dimension，n_heads 為 attention heads 數，max_seq_len 為 cache 容量（token 數），而 n_layers = n_pre_blocks + L 為總 cached-attention layer 數。以 D = 2048、n_heads = 16 為例，每個 token 每層的成本如下：

模式每個 token 每層位元組公式D=2048, n_heads=16 時的位元組數
FP16 KV2 × D × 28192
INT8 KV2 × D × 1 + 2 × n_heads × 24160
FP16 hD × 24096
INT8 hD × 1 + n_heads × 22080
INT4 hD / 2 + n_heads × 21056

將上面的數值再乘上 n_layers × max_seq_len，即可得到整份快取的總成本。在 max_seq_len=1024、L=32（因此 n_layers=34）時，2026-04-18 在 H100 上量測到的 cache footprint 為：

模式VRAM（seq=1024, L=32）相對於 FP16 KV 的節省
FP16 KV272 MB—
INT8 KV138 MB49 %
FP16 h136 MB + 20 MB scratch = 156 MB43 %
INT8 h68 MB + 24 MB scratch = 92 MB66 %
INT4 h34 MB + 24 MB scratch = 59 MB78 %

三種 h 模式中的 scratch buffer 是一筆固定約 20–24 MB 的配置，並由所有 layer 共用（用來在注意力時計算 K/V 重建）。這筆成本和 L 無關，所以 L 越大，它所占比重越低。例如在 L=64, n_layers=66 時，INT4-h 的節省可達 82 %，逐漸逼近理論上約 87 % 的極限。

延遲特性​

2026-04-18 量測條件：H100 / 407M GPTQ-INT4 / L=8 / prompt=32 / decode=64（完成 skip-Q + partial-fp16-attn_qkv 最佳化之後）：

儲存模式scalar recon+ partial-attn_qkv+ full fp16_body
FP16 KV（基線）9.48——
FP16 h16.48（+74 %）9.02（−5 %） ⚡—
INT8 h16.48（+74 %）9.06（−4 %） ⚡—
INT4 h16.54（+74 %）9.09（−4 %） ⚡2.92（−69 %）
407M 上權重 VRAM 增量0+75 MB+406 MB

目前建議的 h-mode 生產設定是：TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 TINYLOOP_H_MODE_FP16_ATTN_QKV=1。 這個組合在量測誤差內幾乎與 FP16 KV 基線同速，同時只需要在 407M 上增加約 75 MB 權重 VRAM（相比 full fp16_body 要加 406 MB），並保留 INT4-h 的完整 78 % KV cache 節省。只要 seq_len ≥ ~512，淨 VRAM 就會轉為正收益，因為 KV cache 節省大於額外權重成本。相同旗標也適用於 FP16-h（43 % cache 節省、K/V 位元級一致）與 INT8-h（66 % cache 節省、K/V 誤差約 0.6 %）。

運作原理。 當 h-mode 啟用時，每次 attention call 都會透過 W_k · h / W_v · h 從快取的 h 重建 K 與 V，這本質上是一個對 attn_qkv 權重做的 GEMM。在 scalar INT4 路徑上，這一步是 compute-bound，因此會主導 decode 時間（+74 % 回歸）。若開啟 TINYLOOP_H_MODE_FP16_ATTN_QKV=1，則只有 attn_qkv 會在模型載入時預先 dequant 成 FP16，重建 GEMM 也就能改走 cuBLAS tensor cores，也就是 full fp16_body 同樣的高速路徑，但 VRAM 成本小得多（attn_qkv 約只佔一個 block 權重的五分之一，而且權重共享 loop 架構通常只有 3 個獨特 block：n_pre_blocks=2 + 1 個 loop block，因此 407M 上大約是 3 × 25 MB = 75 MB）。

更早版本的 +108 % 延遲回歸，已在 2026-04-18 另一項修改中降到 +74 %：重建 GEMM 現在會跳過 attn_qkv 中屬於 Q 的那一段（我們只需要 K/V 來重建 cache），而且一個 kernel 就能在共享記憶體中同時完成 K/V split。

怎麼選​

  • 品質優先 / 研究基準： 使用 FP16 KV（預設）。
  • 想用更小 cache、但仍保留相當穩定的 greedy 輸出： 使用 INT8 KV。
  • 想拿到架構專屬的最大 VRAM 節省，而且願意支付額外重建計算：
    • FP16 h：43–50 % 節省，位元級一致。
    • INT8 h：66–75 % 節省，K/V 相對誤差很小。
    • INT4 h：78 % 節省，對 eval 友善，但生成容易提早分歧。
  • 想同時兼顧低延遲與低 KV VRAM： h_mode + fp16 attn_qkv reconstruction 是目前最好的折衷。

Python 啟用範例​

import os

# 選擇其中一組設定，務必在 import tinyloop_py 前設置。

# 1) INT8 KV
os.environ["TINYLOOP_KV_INT8"] = "1"

# 2) FP16 h
# os.environ["TINYLOOP_KV_H_MODE"] = "1"

# 3) INT8 h
# os.environ["TINYLOOP_KV_H_MODE"] = "1"
# os.environ["TINYLOOP_KV_INT8"] = "1"

# 4) INT4 h
# os.environ["TINYLOOP_KV_H_MODE"] = "1"
# os.environ["TINYLOOP_KV_INT4"] = "1"

# 5) h mode + tensor-core reconstruction
# os.environ["TINYLOOP_KV_H_MODE"] = "1"
# os.environ["TINYLOOP_KV_INT4"] = "1"
# os.environ["TINYLOOP_H_MODE_FP16_ATTN_QKV"] = "1"

import tinyloop_py as tl
model = tl.Model("model.tinyloop", max_seq_len=2048)

不變量：所有模式都應給出同一個 argmax stream​

在 greedy decoding（temperature=0.0）下，四種壓縮模式與基線的關係應如下：

  • FP16 KV 與 FP16 h 完全位元級一致：每個 token id 都相同。
  • INT8 KV 通常在前約 40 個 token 內與 FP16 KV 的 argmax 保持一致，之後才可能慢慢分歧；兩條 continuation 都是合法樣本。
  • INT8 h 在 prompt 位置與前幾個 decode step 上通常能保持一致，但 0.6 % 的 K/V 相對誤差最終仍可能在接近分數的 logit 上翻轉 argmax。

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

與其他模式疊加​

組合是否支援說明
只開 TINYLOOP_KV_INT8是走 INT8 K/V 路徑。
只開 TINYLOOP_KV_H_MODE是走 FP16 h 路徑。
TINYLOOP_KV_H_MODE + TINYLOOP_KV_INT8是走 INT8 h 路徑。
TINYLOOP_KV_H_MODE + TINYLOOP_KV_INT4是走 INT4 h 路徑；若同時也設 INT8，INT4 優先。
上述任一組合 + TINYLOOP_EXPERIMENTAL_FP16_BODY=1是讓 body GEMM 走 cuBLAS FP16 tensor cores；權重 VRAM 會明顯上升，但 decode 可大幅加速。
上述任一組合 + TINYLOOP_CROSS_ITER_ROUND=1僅 FP16 KV / INT8 KV 完整支援h 模式目前會退回 unfused reconstruction 路徑。
TINYLOOP_KV_H_FUSED=1是，但目前更慢僅在 h_mode + INT8 h + plain INT4 attn_qkv 下生效。

端到端選型範例​

import os, time, numpy as np

def bench(mode_env):
    for k, v in mode_env.items():
        os.environ[k] = v
    import importlib, sys
    if "tinyloop_py" in sys.modules:
        raise RuntimeError("Restart the process to change cache mode")
    import tinyloop_py as tl
    model = tl.Model("model.tinyloop", max_seq_len=2048)
    prompt = np.arange(512, dtype=np.int32)
    t0 = time.perf_counter()
    r = model.generate_stream(prompt, lambda t: True,
                               max_tokens=64, loops=8, temperature=0.0)
    return time.perf_counter() - t0, r["decode_steady_ms"]

# 在不同子進程裡對每個模式各自呼叫一次 `bench({...})`
# 再根據你的 (wall time, VRAM) 約束做選擇。

由於 TinyLoop 會在首次匯入後快取 env var，因此你不能在同一個 process 內切換模式。上述 helper 也因此被設計成「每個模式一個單獨 process」的 benchmark。

Hidden-state trajectory API​

當 h_mode 啟用時，每個位置的 hidden state 本身就成了第一級快取量。若你想直接觀察它（例如研究 iteration-axis 壓縮變體），可以使用 Model.score_trajectory，在一次 forward pass 中回傳完整的逐階段最後位置 hidden trajectory，而不需要顯式建立 cache。

何時不該用這些模式​

以上四種模式是執行時層級的 KV 壓縮選項。對應的權重壓縮（INT2 / INT4 / FP16-body cache / Marlin tensor-core GEMM）則屬於另一條正交軸，請見 runtime modes 與 CLI reference 中的 TINYLOOP_EXPERIMENTAL_FP16_BODY 一節。權重壓縮與 KV 壓縮可自由組合，兩者是獨立旋鈕。

已知問題與後續工作​

  1. Fused INT8-h kernel（TINYLOOP_KV_H_FUSED=1）目前比 unfused 路徑還慢（在 1B-effective 基準上大約是 4 s 對 2.1 s）。正確性已由 tests/test_kv_h_mode 驗證，因此它仍保留為 opt-in，後續會再做 vectorised-load + shared-memory-scale 最佳化。延遲敏感工作負載目前不要開。
  2. h_mode 的 tensor-core reconstruction 是把 decode 延遲降到競爭水位的自然下一步。當前可走的路徑是：開啟 TINYLOOP_EXPERIMENTAL_FP16_BODY=1，讓額外的 attn_qkv projection 透過 cuBLAS FP16 matmul；或者等 SM_80 硬體到位後，把 gptq-marlin 也接到 h_mode reconstruction 上。詳見 CHECKLIST 的 §16.14.b Phase 1F。
  3. INT8-h sliding-window parity test 目前只有 FP16-h 變體的端到端測試（tests/test_kv_h_mode 第 7 項）。wrap path 雖已接好，但尚未對 INT8-h 做完整驗證；理論上應該可用，但在測試落地前不建議用於正式生產。
  4. INT4-h 變體 會在 INT8-h 基礎上再省一半（記憶體大約只剩基線的 13 %）。最初追蹤為 Phase 1G，尚未實作。 已於 2026-04-18 出貨：開啟 TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 即可。詳見上方 modes table 的 INT4 h 列。

另見​