KV Cache Poisoning & Exploitation

# Without KV cache: O(n^2) total compute for n tokens
for t in range(seq_len):
    # Recompute K, V for ALL tokens 0..t at every step
    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)
 
# With KV cache: O(n) total compute
kv_cache = {}
for t in range(seq_len):
    # Only compute K, V for the NEW token
    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'])

# Conceptual: how prefix caching works in vLLM-style systems
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:
            # Compute and cache KV states for this prefix
            self.cache[key] = model.compute_kv(prefix_tokens)
        return self.cache[key]
 
    # VULNERABILITY: if cache entries are mutable references,
    # a request that modifies kv_states in-place will poison
    # all future requests sharing that prefix

import time
import httpx
 
async def probe_prefix_cache(api_url: str, test_prefix: str) -> float:
    """Measure response latency to detect cached prefixes.
    A cache hit (shared prefix) produces noticeably lower latency."""
    start = time.perf_counter()
    response = await httpx.AsyncClient().post(api_url, json={
        "prompt": test_prefix + " Continue.",
        "max_tokens": 1,  # minimize generation time
    })
    elapsed = time.perf_counter() - start
    return elapsed
 
# Compare latencies for candidate system prompts
# Lower latency = likely a cached prefix = another tenant's system prompt
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]}...")

# Detect potential KV cache residue from previous requests
# by analyzing attention patterns for anomalous historical context
def detect_cache_residue(model, clean_prompt, suspicious_prompt):
    """Compare attention patterns between clean and potentially
    contaminated cache states."""
    clean_attn = get_attention_weights(model, clean_prompt, use_cache=False)
    cached_attn = get_attention_weights(model, suspicious_prompt, use_cache=True)
 
    # Residue manifests as attention to positions beyond the current prompt
    divergence = compute_kl_divergence(clean_attn, cached_attn)
    return divergence