LLM 推理优化：KV Cache · Flash Attention · 量化

用 NumPy 实现单层单头的简化 Transformer，分别跑 无 KV Cache 和 有 KV Cache 两个版本，对比 K/V 投影次数。两者输出完全相同，KV Cache 是无损优化。

import numpy as np

np.random.seed(42)
d_model = 8   # embedding 维度（单头，d_h == d_model）

# 固定权重矩阵（推理时只读，不会更新）
W_Q = np.random.randn(d_model, d_model) * 0.3
W_K = np.random.randn(d_model, d_model) * 0.3
W_V = np.random.randn(d_model, d_model) * 0.3

# 模拟 5 个 token 的 embedding：「我爱北京天」
tokens = ["我", "爱", "北", "京", "天"]
embeddings = np.random.randn(len(tokens), d_model) * 0.5


def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()


def attention(q, K_mat, V_mat):
    """单个 query 对整个 K/V 做 scaled dot-product attention"""
    scores  = K_mat @ q / np.sqrt(d_model)  # (seq_len,)
    weights = softmax(scores)               # (seq_len,)
    return weights @ V_mat                  # (d_model,)


# ─────────────────────────────────────────────────────────
# 版本 A：无 KV Cache —— 每次生成都重新计算所有历史 token 的 K/V
# ─────────────────────────────────────────────────────────
print("=" * 60)
print("版本 A：无 KV Cache")
print("=" * 60)

total_proj_no_cache = 0

for i, token in enumerate(tokens):
    # 每次都对当前 token 及之前所有 token 重新做 K/V 投影
    context = embeddings[:i+1]              # (i+1, d_model)
    K_all = context @ W_K                   # 重新算！(i+1, d_model)
    V_all = context @ W_V                   # 重新算！(i+1, d_model)
    total_proj_no_cache += 2 * (i + 1)      # K+V 各 1 次，共 2*(i+1) 次

    q   = embeddings[i] @ W_Q
    out = attention(q, K_all, V_all)

    print(f"  生成「{token}」: 重新算了 {2*(i+1):2d} 次 K/V 投影  "
          f"out[0]={out[0]:.4f}")

print(f"\n  ✗ 总 K/V 投影次数：{total_proj_no_cache}\n")


# ─────────────────────────────────────────────────────────
# 版本 B：有 KV Cache —— 只算新 token 的 K/V，历史的直接复用
# ─────────────────────────────────────────────────────────
print("=" * 60)
print("版本 B：有 KV Cache")
print("=" * 60)

cache_K = []   # 每步 append 新 token 的 K 向量
cache_V = []   # 每步 append 新 token 的 V 向量
total_proj_cache = 0

for i, token in enumerate(tokens):
    # 只计算当前新 token 的 K、V（1 次投影，不管历史多长）
    k_new = embeddings[i] @ W_K             # 只算自己的！
    v_new = embeddings[i] @ W_V             # 只算自己的！
    total_proj_cache += 2                   # 永远只有 2 次

    cache_K.append(k_new)                   # 追加进缓存
    cache_V.append(v_new)

    K_all = np.stack(cache_K)               # 从缓存读，无需重算
    V_all = np.stack(cache_V)

    q   = embeddings[i] @ W_Q
    out = attention(q, K_all, V_all)

    print(f"  生成「{token}」: 只算了   2 次 K/V 投影  "
          f"out[0]={out[0]:.4f}  (cache 中 {len(cache_K)} 个 token)")

print(f"\n  ✓ 总 K/V 投影次数：{total_proj_cache}\n")


# ─────────────────────────────────────────────────────────
# 对比总结
# ─────────────────────────────────────────────────────────
print("=" * 60)
saving = total_proj_no_cache - total_proj_cache
pct    = saving / total_proj_no_cache * 100
print(f"  无 KV Cache：{total_proj_no_cache} 次投影")
print(f"  有 KV Cache：{total_proj_cache} 次投影")
print(f"  节省：{saving} 次（{pct:.0f}%）")
print(f"  结论：两个版本 out[0] 完全相同 → KV Cache 是无损优化")
print("=" * 60)

# 预期输出：
#   版本 A 总投影次数：2+4+6+8+10 = 30
#   版本 B 总投影次数：2×5 = 10
#   节省 20 次（67%）；序列 100 token 时节省 ~99%

复制到本地，pip install numpy 后直接 python kv_cache_demo.py 即可运行，亲眼验证节省次数和输出一致性。

用 Python 模拟 vLLM 的 Block Table 管理：3 条并发请求，Block Size=4，展示逻辑→物理映射、动态分配、请求完成后物理 Block 回收复用。

from collections import defaultdict

BLOCK_SIZE = 4       # 每个 Block 存 4 个 token 的 KV
TOTAL_PHYS = 9       # 模拟只有 9 个物理 Block（GPU 显存有限）


class BlockAllocator:
    """物理 Block 分配器（空闲链表）"""
    def __init__(self, n):
        self.free = list(range(n))   # 物理 Block 编号 0..n-1
        self.ref_count = defaultdict(int)

    def alloc(self):
        if not self.free:
            raise RuntimeError("OOM: 无空闲物理 Block")
        blk = self.free.pop(0)
        self.ref_count[blk] = 1
        return blk

    def free_block(self, blk):
        self.ref_count[blk] -= 1
        if self.ref_count[blk] == 0:
            self.free.append(blk)
            print(f"  [Allocator] 物理 Block {blk} 已回收，当前空闲: {sorted(self.free)}")

    def share(self, blk):
        """共享一个 Block（ref count +1）"""
        self.ref_count[blk] += 1


class Sequence:
    """一条请求的 KV Cache 状态"""
    def __init__(self, name, allocator):
        self.name = name
        self.allocator = allocator
        self.block_table = []   # 逻辑 Block → 物理 Block 编号
        self.filled = []        # 每个逻辑 Block 已填充的 slot 数
        self.total_tokens = 0

    def append_tokens(self, n_tokens):
        """追加 n_tokens 个 token 的 KV，自动管理 Block"""
        for _ in range(n_tokens):
            # 如果当前最后一个 Block 已满，分配新 Block
            if not self.block_table or self.filled[-1] == BLOCK_SIZE:
                phys = self.allocator.alloc()
                self.block_table.append(phys)
                self.filled.append(0)
                print(f"  [{self.name}] 分配新物理 Block {phys}  "
                      f"(逻辑 Block {len(self.block_table)-1})")
            self.filled[-1] += 1
            self.total_tokens += 1

    def release(self):
        """请求完成，释放所有物理 Block"""
        for phys in self.block_table:
            self.allocator.free_block(phys)
        self.block_table.clear()
        self.filled.clear()
        print(f"  [{self.name}] 请求完成，全部 Block 已释放")

    def status(self):
        mapping = [f"L{i}→P{p}({f}/{BLOCK_SIZE})"
                   for i, (p, f) in enumerate(zip(self.block_table, self.filled))]
        return f"{self.name}: tokens={self.total_tokens}  [{', '.join(mapping)}]"


# ────────────────────────────────────────────────────
# 模拟 3 条并发请求
# ────────────────────────────────────────────────────
allocator = BlockAllocator(TOTAL_PHYS)
print("=" * 62)
print(f"初始空闲物理 Block: {sorted(allocator.free)}")
print("=" * 62)

# 请求 A：prompt 7 tokens（需 2 个 Block：4 + 3）
print("\n--- 请求 A Prefill（7 tokens）---")
req_a = Sequence("ReqA", allocator)
req_a.append_tokens(7)
print(" ", req_a.status())

# 请求 B：prompt 5 tokens
print("\n--- 请求 B Prefill（5 tokens）---")
req_b = Sequence("ReqB", allocator)
req_b.append_tokens(5)
print(" ", req_b.status())

# 请求 A Decode：再生成 3 个 token
print("\n--- 请求 A Decode（再生成 3 tokens）---")
req_a.append_tokens(3)
print(" ", req_a.status())

# 请求 C：prompt 4 tokens
print("\n--- 请求 C Prefill（4 tokens）---")
req_c = Sequence("ReqC", allocator)
req_c.append_tokens(4)
print(" ", req_c.status())

print(f"\n当前空闲物理 Block: {sorted(allocator.free)}")
print("=" * 62)
print("请求 B 完成，释放内存：")
req_b.release()
print(f"\n请求 B 释放后，空闲物理 Block: {sorted(allocator.free)}")

# 请求 D 进来，复用请求 B 释放的 Block
print("\n--- 请求 D Prefill（6 tokens，复用 ReqB 释放的 Block）---")
req_d = Sequence("ReqD", allocator)
req_d.append_tokens(6)
print(" ", req_d.status())
print("=" * 62)
print("最终所有请求状态：")
for seq in [req_a, req_c, req_d]:
    print(" ", seq.status())
print(f"\n剩余空闲物理 Block: {sorted(allocator.free)}")

# 结论：
# - 物理 Block 可被不同请求动态分配/回收，无需连续
# - 请求 D 复用了 ReqB 释放的 Block，GPU 显存零浪费
# - 每条请求最多浪费 1 个 Block 的未填充 slot

运行此代码无需任何依赖，python paged_attn_demo.py 即可。可以看到请求 B 释放后，其物理 Block 立刻被请求 D 复用——这正是 vLLM 显存利用率达 96% 的核心机制。

量化（Quantization）就是把连续的浮点数，映射到有限个离散的「格点」上存储——本质是有损压缩，bit 越少，格点越稀，误差越大，但存储越省。