W10 · 长上下文与 KV Cache

1. Prefill 与 Decode 是两种工作负载

Prefill 一次处理整个 prompt 并构造每层 K/V，矩阵大、并行度高，用户指标是 TTFT。Decode 每步输入一个新 token，复用历史 K/V，只追加当前状态，常更受内存带宽、调度和 batch 影响，指标是 TPOT。只报“每秒 token”会混合两个阶段。

2. 没有缓存会重复计算什么

生成第 t 个 token 时，历史位置的 K/V 与之前相同。若每步对完整序列前向，就重复生成历史 K/V：

K_cache[layer]: [batch, kv_heads, tokens, head_dim]
V_cache[layer]: [batch, kv_heads, tokens, head_dim]

新 query 仍要与历史 key 做注意力，所以 decode 不是总工作量 O(1)；cache 避免重复投影历史 token，并用内存换计算。

3. KV Cache 显存公式

KV bytes =
2 × num_layers × tokens × num_kv_heads
× head_dim × bytes_per_element × batch

32 层、32 KV heads、head_dim 128、4096 tokens、FP16：

2 × 32 × 4096 × 32 × 128 × 2 bytes
= 2,147,483,648 bytes
≈ 2 GiB / request

这还没包含权重、运行时 buffer、临时工作区和碎片。容量规划不能只看模型文件大小。

4. MHA、GQA 与 MQA

MHA：query heads 与 KV heads 相同。
GQA：多个 query heads 共享一组 K/V。
MQA：所有 query heads 共享单组 K/V。

若 32 query heads 改为 8 KV heads，理论缓存缩为四分之一。代价可能包含质量与 kernel 支持差异。

def kv_cache_gib(config, tokens, batch=1, bytes_per_element=2):
    layers = config.num_hidden_layers
    kv_heads = getattr(
        config, "num_key_value_heads", config.num_attention_heads
    )
    head_dim = config.hidden_size // config.num_attention_heads
    total = (
        2 * layers * tokens * kv_heads
        * head_dim * bytes_per_element * batch
    )
    return total / 2**30

5. 动态、静态与量化缓存

Dynamic：随 token 增长，灵活，但动态 shape 可能限制编译优化。
Static：预分配最大长度，可能浪费空间，但固定 shape 利于编译。
Quantized：低精度存 K/V，降低容量与带宽，需验证质量和量化开销。
Sliding window：只保留最近窗口，容量封顶，但丢失窗口外细节。

6. PagedAttention 解决碎片与调度

变长请求若连续预分配会产生碎片。PagedAttention 把 KV 切成固定 block，用逻辑页映射物理 block：

非连续物理内存表示连续逻辑序列。
请求结束后快速回收 block。
共享前缀可复用 block。
调度器能组合更多变长请求。

它主要优化内存管理和服务吞吐，不改变模型数学语义。

7. 测量 TTFT、TPOT 与吞吐

start = time.perf_counter()
first_token_at = None
count = 0

for token in stream_generate(prompt):
    now = time.perf_counter()
    if first_token_at is None:
        first_token_at = now
    count += 1

ttft = first_token_at - start
total = now - start
tpot = (now - first_token_at) / max(1, count - 1)
throughput = count / total

prompt length: 128 / 512 / 2048 / 8192
output length: 32 / 128 / 512
concurrency: 1 / 4 / 16 / 32
cache: on / off / quantized

报告 P50/P95/P99，不只平均数。长尾常来自长 prompt、排队或内存回收。

8. 理论值为何与实测不同

偏差来自 allocator 对齐与预留、cache block 元数据、特殊层、张量并行分片、dtype、临时 logits/attention buffer 和 CUDA graph 预留。报告要解释差值，不能为了吻合公式忽略系统开销。

9. 本课实验

读取模型 config，计算 128–8192 token 理论 KV。
实测 torch.cuda.max_memory_allocated()。
测 cache on/off 的 TTFT、TPOT、显存。
改变 batch/concurrency，找到容量拐点。
若支持，比较 dynamic/static/quantized。

交付容量表、偏差解释和指定显存下安全并发的计算过程。

10. 练习与答案提示

上下文翻倍，KV cache？**答：**近似线性翻倍。
标准 prefill attention？**答：**分数部分近似平方增长。
GQA 为何省缓存？**答：**减少 KV heads。
Static cache 一定省内存？**答：**不一定，预分配可能浪费。
PagedAttention 改变答案吗？**答：**理想情况下不改变数学结果。

长上下文与 KV Cache

学完这一课，你应该能：