KV 快取投毒與利用

# 無 KV 快取：n 個 token 總計算 O(n^2)
for t in range(seq_len):
    # 每步對所有 token 0..t 重算 K、V
    K = W_k @ X[:, :t+1, :]    # [batch, t+1, d_k]
    V = W_v @ X[:, :t+1, :]    # [batch, t+1, d_v]
    Q = W_q @ X[:, t:t+1, :]   # [batch, 1, d_k]
    output = attention(Q, K, V)
 
# 有 KV 快取:總計算 O(n)
kv_cache = {}
for t in range(seq_len):
    # 只計算「新」token 的 K、V
    k_new = W_k @ X[:, t:t+1, :]
    v_new = W_v @ X[:, t:t+1, :]
    kv_cache['K'] = torch.cat([kv_cache.get('K', empty), k_new], dim=1)
    kv_cache['V'] = torch.cat([kv_cache.get('V', empty), v_new], dim=1)
    Q = W_q @ X[:, t:t+1, :]
    output = attention(Q, kv_cache['K'], kv_cache['V'])

# 概念：vLLM 式系統中前綴快取如何運作
class PrefixCache:
    def __init__(self):
        self.cache = {}  # hash(prefix_tokens) -> kv_states
 
    def get_or_compute(self, prefix_tokens, model):
        key = hash(tuple(prefix_tokens))
        if key not in self.cache:
            # 計算並快取此前綴的 KV 狀態
            self.cache[key] = model.compute_kv(prefix_tokens)
        return self.cache[key]
 
    # 漏洞：若快取條目是可變參照，
    # 某請求就地修改 kv_states 將投毒所有未來
    # 共享該前綴的請求

import time
import httpx
 
async def probe_prefix_cache(api_url: str, test_prefix: str) -> float:
    """量測回應延遲以偵測被快取的前綴。
    快取命中（共享前綴）會產生明顯較低延遲。"""
    start = time.perf_counter()
    response = await httpx.AsyncClient().post(api_url, json={
        "prompt": test_prefix + " Continue.",
        "max_tokens": 1,  # 最小化生成時間
    })
    elapsed = time.perf_counter() - start
    return elapsed
 
# 比較各候選系統提示的延遲
# 延遲較低 = 可能為被快取前綴 = 另一租戶的系統提示
candidates = [
    "You are a helpful assistant for AcmeCorp...",
    "You are a financial advisor. Never disclose...",
    "You are a medical chatbot. Always recommend..."
]
for prefix in candidates:
    latency = await probe_prefix_cache(api_url, prefix)
    print(f"Latency: {latency:.4f}s - {prefix[:50]}...")

# 透過分析歷史脈絡上的異常注意力模式，
# 偵測可能自前一請求留下的 KV 快取殘留
def detect_cache_residue(model, clean_prompt, suspicious_prompt):
    """比較乾淨與疑遭汙染快取狀態之注意力模式。"""
    clean_attn = get_attention_weights(model, clean_prompt, use_cache=False)
    cached_attn = get_attention_weights(model, suspicious_prompt, use_cache=True)
 
    # 殘留會以「注意到目前提示之外位置」呈現
    divergence = compute_kl_divergence(clean_attn, cached_attn)
    return divergence