# Phase 9: KV Cache + Autoregressive Generation — Design Document

## Goal

实现 KV Cache，将 decode 从每步 full forward (O(S²)) 降为增量计算 (O(S))。这是最大的单点性能提升。

## 核心变化

### Before (no cache)
```
每生成一个 token:
  forward(all_tokens)  → 重新计算所有层的 Q/K/V/attention
  开销: O(S²) attention per step, S 递增
```

### After (with cache)
```
Prefill:
  forward(prompt_tokens)  → 计算并缓存所有层的 K/V
  
Decode (per token):
  forward(last_token_only) → 只计算新 token 的 Q/K/V
  Q: [1, H, 1, D]          → 新 token 的 query
  K: append to cache        → cache 变为 [1, H, S+1, D]
  V: append to cache
  attention: Q @ K_cache^T  → [1, H, 1, S+1], O(S) not O(S²)
```

## KVCache 数据结构

```rust
pub struct KVCache {
    k: Vec<Tensor>,  // per layer, shape [1, num_heads, current_len, head_dim]
    v: Vec<Tensor>,
    len: usize,      // current sequence length
}
```

## Forward Pass 变化

模型需要两种 forward 模式:
1. **prefill(tokens)**: 处理完整 prompt，填充 KV cache
2. **decode(token, cache)**: 处理单个 token，读写 KV cache

## 实现策略

为了最小化改动，在 GPT-2 forward 中加入可选的 `&mut KVCache` 参数：
- cache=None → 现有行为（full forward）
- cache=Some → prefill 或 decode 模式

CPU round-trip 问题暂不修复（Phase 15），先让 KV cache 逻辑正确。

## Test Plan

- [x] KV cache vs no-cache: 50/50 bit-identical output
- [x] Benchmark: 18x decode speedup (407ms → 22ms TBT)
- [x] 50 prompt validation: 40/50 vs HF (10 are FP divergence, gap 0.04-0.56)

## Takeaways

1. **KV cache 数据布局是核心难点**：初始实现直接 append flat bytes 导致 head 维度交错错误。正确做法：per-head 独立存储，reconstruct 时按 `[1, H, S, D]` layout 组装。这是一个非常容易犯的 layout bug，调试时输出看起来"几乎对"但不完全对。

2. **18x 提速 > 理论预期**：理论上 KV cache 将 decode 从 O(S²) 降到 O(S)，对 S=20-25 的序列预期 ~20x 提速。实测 18x 符合预期。TTFT 也从 400ms 降到 24ms，因为 prefill 只跑一次而不是每步重跑。

3. **xserv vs HF 的 10 个 mismatch 不是 bug**：logit gap 仅 0.04-0.56（在 -80 到 -140 的 logit 值上），是不同 CUDA kernel 实现间的浮点累积误差导致 argmax 翻转。重要验证：**xserv KV-cache vs xserv no-cache 是 50/50 完全一致的**——证明 KV cache 实现本身无误。

4. **CPU round-trip 仍是主要瓶颈**：KV cache 的 per-head 数据存在 CPU Vec 中，每步 decode 都要重新组装成 GPU tensor。这意味着每步仍有 24 次 GPU→CPU→GPU 传输（12 层 × 2 KV）。Phase 15 需要将 KV cache 直接放在 GPU 上。