KV Cache 优化与内存管理深度解析：PagedAttention 原理与实践

1
┌─────────────────────────────────────────────────────────────┐
2
│               自回归生成的内存问题                           │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  Prefill 阶段（一次计算）：                                  │
6
│  ┌──────────────────────────────────────────────────────┐   │
7
│  │  Input: [T1, T2, T3, ..., T_n]                       │   │
8
│  │  Output: 所有位置的 KV hidden states                  │   │
9
│  │  特点：高度并行，适合 GPU                              │   │
10
│  └──────────────────────────────────────────────────────┘   │
11
│                                                              │
12
│  Decode 阶段（逐 token 生成）：                              │
13
│  ┌──────────────────────────────────────────────────────┐   │
14
│  │  Step 1: [T1] → T2                                    │   │
15
│  │  Step 2: [T1, T2] → T3                                │   │
16
│  │  Step 3: [T1, T2, T3] → T4                            │   │
17
│  │  ...                                                  │   │
18
│  │  问题：每步都需要访问所有历史 KV！                      │   │
19
│  └──────────────────────────────────────────────────────┘   │
20
│                                                              │
21
│  关键问题：                                                  │
22
│  → KV 序列长度随生成不断增长                                │
23
│  → 无法预知最终长度                                         │
24
│  → 需要高效存储和检索历史 KV                                │
25
│                                                              │
26
└─────────────────────────────────────────────────────────────┘

1
class MemoryWallAnalysis:
2
    """
3
    内存墙分析
4
    """
5

6
    def model_vs_memory_growth(self):
7
        """
8
        模型规模 vs 内存需求增长
9
        """
10
        return {
11
            "model_scaling": {
12
                "LLaMA-2 7B": "14 GB weights",
13
                "LLaMA-2 13B": "26 GB weights",
14
                "LLaMA-2 70B": "140 GB weights",
15
            },
16
            "kv_cache_scaling": {
17
                "per_token_kv_fp16": "2 × 2 × layers × hidden_size × 2 bytes",
18
                "7B_model_per_token": "512 KB",
19
                "8k_context": "4 GB KV cache",
20
                "32k_context": "16 GB KV cache",
21
            },
22
            "total_memory_7b": {
23
                "weights": "14 GB",
24
                "8k_kv_cache": "4 GB",
25
                "activations": "2 GB",
26
                "total": "20 GB (RTX 3090 barely fits)",
27
            },
28
        }
29

30
    def gpu_memory_timeline(self):
31
        """
32
        GPU 显存发展 vs 模型需求
33
        """
34
        return """
35
        时间线对比：
36

37
        2020: GPT-3 (175B) 需要 ~350 GB
38
              A100 80GB × 5 = 勉强
39

40
        2022: LLaMA 65B 需要 ~130 GB
41
              A100 80GB × 2 = 可以
42

43
        2023: LLaMA-2 70B 需要 ~160 GB
44
              A100 80GB × 2 + 优化 = 可以
45

46
        2024: Claude 200K context 需要 ~800 GB KV
47
              → 单 GPU 不可能
48

49
        趋势：
50
        → 上下文窗口越来越长
51
        → KV Cache 成为主要瓶颈
52
        → 内存优化刻不容缓
53
        """

1
class StandardAttentionMemory:
2
    """
3
    标准 Attention 的内存问题
4
    """
5

6
    def attention_computation(self):
7
        """
8
        标准 Attention 计算
9

10
        Scaled Dot-Product Attention:
11

12
        Attention(Q, K, V) = softmax(QK^T / √d) V
13
        """
14
        return """
15
        对于序列长度 N：
16

17
        Q: (batch, num_heads, N, head_dim)
18
        K: (batch, num_heads, N, head_dim)
19
        V: (batch, num_heads, N, head_dim)
20

21
        QK^T: (batch, num_heads, N, N)  ← N×N 矩阵！
22

23
        问题：
24
        → N = 8192 时，QK^T 需要 8192 × 8192 × 4 bytes = 256 MB per head
25
        → 32 heads = 8 GB per layer
26
        → 32 layers = 256 GB total
27
        → 完全不可接受
28
        """
29

30
    def kv_cache_storage(self):
31
        """
32
        KV Cache 存储问题
33
        """
34
        return {
35
            "store_all_kv": "需要保存所有历史 K 和 V",
36
            "access_pattern": "Decode 阶段需要 O(1) 随机访问",
37
            "fp16_storage": "每个值 2 bytes",
38

39
            "example_7b": {
40
                "layers": 32,
41
                "heads": 32,
42
                "head_dim": 128,
43
                "per_token_kv": "2 × 32 × 128 × 2 = 16 KB",
44
                "per_layer_per_token": "512 KB",
45
                "total_per_token": "32 × 512 KB = 16 MB",
46
            },
47

48
            "problem": "长序列 + 多请求 = 内存爆炸",
49
        }

1
class KVCacheEssence:
2
    """
3
    KV Cache 本质
4
    """
5

6
    def definition(self):
7
        """
8
        KV Cache 定义
9

10
        在 Prefill 阶段预计算并缓存所有历史 Key 和 Value
11
        Decode 阶段直接复用，避免重复计算
12
        """
13
        return """
14
        无 Cache:
15
        ┌─────────────────────────────────────────────────────┐
16
        │ Step 1: Attention([T1])                             │
17
        │ Step 2: Attention([T1, T2]) = Attention([T1]) +  │
18
        │              Attention(Q2, K1, V1)                  │
19
        │ Step 3: Attention([T1, T2, T3]) = ...             │
20
        │         ↑ 每次都重新计算 K1, V1!                    │
21
        └─────────────────────────────────────────────────────┘
22

23
        有 Cache:
24
        ┌─────────────────────────────────────────────────────┐
25
        │ Prefill: 计算并缓存 K=[K1,K2,K3], V=[V1,V2,V3]      │
26
        │                                                          │
27
        │ Step 4: Attention([T1..T4])                           │
28
        │         = Attention(Q4, [K1..K3], [V1..V3]) +       │
29
        │           Attention(Q4, K4, V4)                      │
30
        │         ↑ 只计算 K4, V4!                              │
31
        └─────────────────────────────────────────────────────┘
32
        """
33

34
    def memory_breakdown(self):
35
        """
36
        内存分解（LLaMA-2 7B）
37
        """
38
        return {
39
            "weights_fp16": "14 GB",
40

41
            "kv_cache_8k": {
42
                "per_token_kb": 512,
43
                "total_gb": 4,
44
                "note": "8k tokens × 512 KB",
45
            },
46

47
            "kv_cache_32k": {
48
                "per_token_kb": 512,
49
                "total_gb": 16,
50
                "note": "32k tokens × 512 KB",
51
            },
52

53
            "activations": "1-3 GB (取决于 batch size)",
54

55
            "total_8k": "14 + 4 + 2 = 20 GB",
56
            "total_32k": "14 + 16 + 3 = 33 GB",
57

58
            "problem_70b": {
59
                "weights": "140 GB",
60
                "kv_cache_8k": "32 GB",
61
                "total": "172 GB",
62
                "a100_80gb_needed": "3 cards minimum",
63
            },
64
        }

1
class KVCacheImplementation:
2
    """
3
    KV Cache 实现
4
    """
5

6
    def naive_implementation(self):
7
        """
8
        朴素实现
9
        """
10
        return '''
11
        class NaiveKVCache:
12
            """朴素实现 - 每个序列独立缓存"""
13

14
            def __init__(self, max_seq_len):
15
                self.k_cache = {}  # {seq_id: tensor}
16
                self.v_cache = {}  # {seq_id: tensor}
17
                self.max_seq_len = max_seq_len
18

19
            def update(self, seq_id, new_k, new_v):
20
                """追加新的 KV"""
21
                if seq_id not in self.k_cache:
22
                    self.k_cache[seq_id] = []
23
                    self.v_cache[seq_id] = []
24

25
                self.k_cache[seq_id].append(new_k)
26
                self.v_cache[seq_id].append(new_v)
27

28
            def get(self, seq_id):
29
                """获取完整 KV"""
30
                return (
31
                    torch.cat(self.k_cache[seq_id], dim=1),
32
                    torch.cat(self.v_cache[seq_id], dim=1),
33
                )
34

35
        问题：
36
        → 列表追加效率低
37
        → torch.cat 会复制大量数据
38
        → 内存碎片
39
        → 不支持随机访问
40
        '''
41

42
    def preallocated_implementation(self):
43
        """
44
        预分配实现
45
        """
46
        return '''
47
        class PreallocatedKVCache:
48
            """预分配连续内存"""
49

50
            def __init__(self, max_seq_len, num_heads, head_dim):
51
                self.max_seq_len = max_seq_len
52

53
                # 预分配最大长度
54
                self.k_cache = torch.zeros(
55
                    max_seq_len, num_heads, head_dim
56
                )
57
                self.v_cache = torch.zeros(
58
                    max_seq_len, num_heads, head_dim
59
                )
60

61
                self.seq_lengths = {}  # 跟踪每个序列的长度
62

63
            def update(self, seq_id, position, new_k, new_v):
64
                """在指定位置更新"""
65
                self.k_cache[position] = new_k
66
                self.v_cache[position] = new_v
67
                self.seq_lengths[seq_id] = position + 1
68

69
            def get(self, seq_id, length):
70
                """获取指定长度的 KV"""
71
                return (
72
                    self.k_cache[:length],
73
                    self.v_cache[:length],
74
                )
75

76
        问题：
77
        → 必须预先知道 max_seq_len
78
        → 不同长度的序列都会占用最大长度
79
        → 内存浪费严重
80
        '''

1
class CacheAccessPatterns:
2
    """
3
    Cache 访问模式分析
4
    """
5

6
    def prefill_access(self):
7
        """
8
        Prefill 阶段的访问模式
9

10
        计算所有 token 的 KV
11
        """
12
        return """
13
        Prefill: [T1, T2, T3, ..., T_n]
14

15
        需要计算：
16
        - 每个 token i 的 K_i, V_i
17
        - 每个 token i 的 attention 输出
18

19
        访问模式：
20
        - Sequential: 顺序遍历所有 tokens
21
        - Heavy compute: 计算密集型
22
        - One-shot: 一次性完成
23
        """
24

25
    def decode_access(self):
26
        """
27
        Decode 阶段的访问模式
28

29
        只计算新 token 的 KV
30
        """
31
        return """
32
        Decode step t:
33

34
        Input: Q_t (单个 token)
35
        Cache: K_[1..t-1], V_[1..t-1]
36

37
        需要：
38
        1. 计算 K_t, V_t (新 token)
39
        2. 读取 K_[1..t-1], V_[1..t-1] (全部历史)
40
        3. 计算 attention
41

42
        访问模式：
43
        - Random read: 读取整个 KV cache
44
        - Sequential write: 顺序写入新 KV
45
        - Memory bound: 内存访问是瓶颈
46

47
        关键洞察：
48
        Decode 是 memory-bound 操作！
49
        """
50

51
    def attention_access(self):
52
        """
53
        Attention 计算的内存访问
54
        """
55
        return """
56
        Decode step t 时的 Attention 计算:
57

58
        Q_t: (1, num_heads, head_dim)
59
        K_cache: (t-1, num_heads, head_dim)
60
        V_cache: (t-1, num_heads, head_dim)
61

62
        内存访问量:
63
        - 读取 Q_t: ~32 KB
64
        - 读取 K_cache: (t-1) × 16 KB
65
        - 读取 V_cache: (t-1) × 16 KB
66
        - 写入 output: ~32 KB
67

68
        总计: ~32 × (t-1) KB
69

70
        当 t=4096:
71
        内存访问 = 128 KB
72
        计算量 = num_heads × t × head_dim × 2 = 2M FLOPs
73

74
        计算/内存比很低 → memory bound
75
        """

1
class PagingInspiration:
2
    """
3
    分页管理的启发
4
    """
5

6
    def os_paging_concept(self):
7
        """
8
        操作系统分页概念
9
        """
10
        return """
11
        OS 虚拟内存管理的核心思想：
12

13
        1. 虚拟地址空间
14
           - 进程看到的是连续的虚拟地址
15
           - 可以远大于物理内存
16

17
        2. 物理内存管理
18
           - 物理内存被分成固定大小的 pages
19
           - 4KB (x86) 或 2MB/1GB (huge pages)
20

21
        3. Page Table
22
           - 映射虚拟页到物理页
23
           - 允许虚拟地址连续但物理地址离散
24

25
        4. 按需加载
26
           - 不需要预先分配整个空间
27
           - 页面按需分配和换入/换出
28

29
        优势：
30
        → 消除外部碎片
31
        → 支持超内存运行
32
        → 灵活共享内存
33
        → 高效内存利用
34
        """
35

36
    def llm_analogy(self):
37
        """
38
        LLM 推理中的类比
39
        """
40
        return """
41
        PagedAttention 的类比：
42

43
        OS 虚拟内存          LLM 推理
44
        ─────────────────────────────────────────────
45
        虚拟地址空间     →    Token 序列位置
46
        物理页           →    KV Cache Blocks
47
        Page Table       →    Block Table
48
        虚拟页号         →    Token 索引
49
        物理页号         →    Block 索引
50

51
        关键对应：
52

53
        OS:  virtual_page_id → physical_page_id
54

55
        vLLM: token_position → kv_block_id
56

57
        Block 0: [token 0-15]
58
        Block 1: [token 16-31]
59
        Block 2: [token 32-47]
60
        ...
61
        """

1
class BlockTableDesign:
2
    """
3
    Block Table 设计
4
    """
5

6
    def table_structure(self):
7
        """
8
        Block Table 结构
9
        """
10
        return """
11
        概念：
12

13
        Token 序列: [T0, T1, T2, ..., T67] (68 tokens)
14
        Block Size: 16 tokens/block
15

16
        需要 Blocks: ceil(68/16) = 5 blocks
17

18
        Block Table (逻辑 → 物理映射):
19
        ┌───────────┬─────────────┬────────────────────┐
20
        │ Block ID  │ Physical ID │ Content            │
21
        ├───────────┼─────────────┼────────────────────┤
22
        │    0      │      7      │ [T0, T1, ..., T15] │
23
        │    1      │      3      │ [T16,...,T31]      │
24
        │    2      │      9      │ [T32,...,T47]      │
25
        │    3      │      2      │ [T48,...,T63]      │
26
        │    4      │     12      │ [T64,...,T67]      │ ← 部分使用
27
        └───────────┴─────────────┴────────────────────┘
28

29
        特点：
30
        → Block 物理位置可以不连续
31
        → 最后一个 block 可以部分填充
32
        → 动态扩展无需重新分配整个序列
33
        """
34

35
    def implementation(self):
36
        """
37
        实现
38
        """
39
        return '''
40
        class BlockTable:
41
            """Block Table 实现"""
42

43
            def __init__(self, block_size=16):
44
                self.block_size = block_size
45
                self.tables = {}  # {seq_id: {"physical_blocks": [], "num_full_slots": int}}
46

47
            def add_block(self, seq_id, physical_block_id):
48
                """添加一个新的 block"""
49
                if seq_id not in self.tables:
50
                    self.tables[seq_id] = {
51
                        "physical_blocks": [],
52
                        "num_full_slots": 0,
53
                    }
54

55
                self.tables[seq_id]["physical_blocks"].append(physical_block_id)
56

57
            def get_physical_block(self, seq_id, block_idx):
58
                """获取指定 block 的物理位置"""
59
                return self.tables[seq_id]["physical_blocks"][block_idx]
60

61
            def get_token_block(self, seq_id, token_idx):
62
                """获取 token 所在的 block"""
63
                block_idx = token_idx // self.block_size
64
                return self.get_physical_block(seq_id, block_idx)
65

66
        # 使用示例
67
        table = BlockTable(block_size=16)
68
        table.add_block("seq_1", 7)  # tokens 0-15 → block 7
69
        table.add_block("seq_1", 3)  # tokens 16-31 → block 3
70
        table.add_block("seq_1", 9)  # tokens 32-47 → block 9
71

72
        # 读取 token 25
73
        block_idx = 25 // 16 = 1
74
        physical = table.get_physical_block("seq_1", 1)  # → 3
75
        offset = 25 % 16 = 9
76
        # → 读取 block 3, offset 9
77
        '''
78

79
    def dynamic_growth(self):
80
        """
81
        动态增长
82
        """
83
        return """
84
        动态扩展示例：
85

86
        初始状态（生成 50 tokens）:
87
        Block Table: [0, 1, 2, 3]  (4 blocks)
88

89
        继续生成（需要 60 tokens）:
90
        Block Table: [0, 1, 2, 3] + [4]  (追加)
91

92
        继续生成（需要 100 tokens）:
93
        Block Table: [0, 1, 2, 3] + [4, 5, 6]  (继续追加)
94

95
        优势：
96
        → 无需预先知道最终长度
97
        → 按需分配 blocks
98
        → 已完成的序列释放 blocks
99
        → 新请求复用释放的 blocks
100

101
        对比预分配：
102
        preallocate 100 tokens worth → 浪费 50 tokens 的空间
103
        paged 动态增长 → 只用多少分配多少
104
        """

1
class PagedAttentionKernel:
2
    """
3
    PagedAttention CUDA Kernel
4
    """
5

6
    def kernel_design(self):
7
        """
8
        Kernel 设计
9
        """
10
        return '''
11
        __global__ void paged_attention_kernel(
12
            const half* __restrict__ query,      // (batch, heads, dim)
13
            const half* __restrict__ key_cache,  // 存储在 blocks 中
14
            const half* __restrict__ value_cache,
15
            const int* block_tables,             // (batch, max_blocks)
16
            const int* seq_lens,                  // 每个序列的实际长度
17
            half* output,
18
            int block_size,
19
            int num_heads,
20
            int head_dim,
21
        ) {
22
            // 获取当前线程负责的 query
23
            int batch_idx = blockIdx.x;
24
            int head_idx = threadIdx.x;
25

26
            // 获取该序列的 block table
27
            int* table = &block_tables[batch_idx * max_blocks];
28
            int num_blocks = (seq_lens[batch_idx] + block_size - 1) / block_size;
29

30
            // 分块处理 keys
31
            float acc = 0.0f;
32
            float max_val = -INFINITY;
33

34
            for (int block = 0; block < num_blocks; block++) {
35
                int physical_block = table[block];
36

37
                // 从物理 block 读取 keys
38
                half4 k_block = *((half4*)&key_cache[
39
                    physical_block * block_size * num_heads * head_dim
40
                    + head_idx * block_size * head_dim
41
                ]);
42

43
                // 计算 partial attention
44
                // ...
45
            }
46

47
            // 保存结果
48
            // ...
49
        }
50
        '''
51

52
    def memory_access_optimization(self):
53
        """
54
        内存访问优化
55
        """
56
        return {
57
            "coalesced_access": "合并相邻线程的内存访问",
58
            "shared_memory": "使用 shared memory 缓存频繁访问的数据",
59
            "vectorized_load": "使用 float4/half4 向量化加载",
60
            "bank_conflict_free": "避免 shared memory bank 冲突",
61

62
            "optimal_blocking": "每个 block 处理多个 token",
63
            "thread_per_head": "每个 head 一个 thread block",
64
            "warp_level_reduction": "Warp 级别的 softmax 归约",
65
        }

1
class MemoryFragmentation:
2
    """
3
    内存碎片类型
4
    """
5

6
    def internal_fragmentation(self):
7
        """
8
        内部碎片
9

10
        预分配但不使用的空间
11
        """
12
        return """
13
        场景：预分配 max_seq_len = 8192
14

15
        Request A: 生成 100 tokens
16
        实际使用: 100 / 8192 = 1.2%
17
        浪费: 99%
18

19
        Request B: 生成 5000 tokens
20
        实际使用: 5000 / 8192 = 61%
21
        浪费: 39%
22

23
        问题根源：
24
        → 必须按最大可能长度分配
25
        → 大多数请求远小于最大值
26
        → 大量内存永久浪费
27
        """
28

29
    def external_fragmentation(self):
30
        """
31
        外部碎片
32

33
        可用但不连续的内存
34
        """
35
        return """
36
        场景：多请求混合长度
37

38
        内存池状态随时间变化：
39

40
        T1: [Req1:100][Req2:500][Req3:2000][    free    ]
41

42
        T2: [Req1:done ][Req2:500][Req3:2500][    free    ]
43
               ↑释放      ↑增长      ↑增长
44
               但不能      不能       不能
45
               给Req4      扩展       扩展
46

47
        问题：
48
        → 虽然有 free 空间
49
        → 但分散在不连续位置
50
        → 无法满足大请求的连续需求
51
        → 只能 OOM 或等待
52
        """
53

54
    def paged_solution(self):
55
        """
56
        Paged 方案如何解决
57
        """
58
        return """
59
        PagedAttention 的解决方案：
60

61
        1. 内部碎片 → 消除
62
           - Block Size 固定但较小 (16 tokens)
63
           - 只分配实际需要的 blocks
64
           - 最后一个 block 只用必要的空间
65

66
        2. 外部碎片 → 消除
67
           - Blocks 可以放在任何位置
68
           - 物理位置可以不连续
69
           - Block Table 提供逻辑连续视图
70

71
        结果：
72
        → 内存利用率接近 100%
73
        → 任意长度的请求都能高效处理
74
        → 碎片成为历史
75
        """

1
┌─────────────────────────────────────────────────────────────┐
2
│              内存利用率对比                                  │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  预分配策略（max_seq_len = 8192）：                          │
6
│  ┌──────────────────────────────────────────────────────┐   │
7
│  │ Request 1 (100 tokens):   ██░░░░░░░░░░░░░░░░░░░░░░░ │   │
8
│  │                           1.2% used, 98.8% waste     │   │
9
│  │                                                          │   │
10
│  │ Request 2 (5000 tokens):  ██████████████████░░░░░░ │   │
11
│  │                           61% used, 39% waste        │   │
12
│  │                                                          │   │
13
│  │ Request 3 (200 tokens):    ██░░░░░░░░░░░░░░░░░░░░░░░ │   │
14
│  │                           2.4% used, 97.6% waste     │   │
15
│  └──────────────────────────────────────────────────────┘   │
16
│                                                              │
17
│  PagedAttention 策略（block_size = 16）：                   │
18
│  ┌──────────────────────────────────────────────────────┐   │
19
│  │ Request 1 (100 tokens):   ██████████                 │   │
20
│  │                           100% used (7 blocks)       │   │
21
│  │                                                          │   │
22
│  │ Request 2 (5000 tokens):  ██████████████████████████ │   │
23
│  │                           100% used (313 blocks)      │   │
24
│  │                                                          │   │
25
│  │ Request 3 (200 tokens):   ██████████████████           │   │
26
│  │                           100% used (13 blocks)      │   │
27
│  └──────────────────────────────────────────────────────┘   │
28
│                                                              │
29
│  总结：                                                      │
30
│  预分配：平均利用率 ~20-40%                                  │
31
│  Paged：平均利用率 ~90-95%                                   │
32
│  提升：2-4x                                                   │
33
│                                                              │
34
└─────────────────────────────────────────────────────────────┘

1
class BlockSizeSelection:
2
    """
3
    Block Size 选择
4
    """
5

6
    def tradeoffs(self):
7
        """
8
        Block Size 的权衡
9
        """
10
        return {
11
            "small_block_8": {
12
                "pros": "内存利用率高，最后一个 block 浪费少",
13
                "cons": "Block table 更大，管理开销高",
14
                "fragmentation": "内部碎片 < 8 tokens",
15
            },
16
            "medium_block_16": {
17
                "pros": "平衡内存利用率和管理开销",
18
                "cons": "略多于 8 的管理开销",
19
                "fragmentation": "内部碎片 < 16 tokens",
20
                "recommendation": "vLLM 默认推荐",
21
            },
22
            "large_block_32": {
23
                "pros": "管理开销低，缓存命中率高",
24
                "cons": "内部碎片最多 32 tokens",
25
                "fragmentation": "对于短序列浪费明显",
26
            },
27
        }
28

29
    def benchmark_recommendation(self):
30
        """
31
        基准测试建议
32
        """
33
        return """
34
        vLLM 官方推荐：
35

36
        block_size = 16
37

38
        原因：
39
        1. 大多数请求的 KV cache 远大于 16 tokens
40
        2. 内部碎片 < 16 tokens，可接受
41
        3. Block table 管理开销可忽略
42
        4. CUDA kernel 优化针对 block_size=16
43

44
        特殊情况：
45
        → 极短序列（< 32 tokens）：考虑 block_size=8
46
        → 超长序列 + 稀疏访问：考虑 block_size=32
47
        """

1
class MemoryPoolArchitecture:
2
    """
3
    内存池架构
4
    """
5

6
    def pool_structure(self):
7
        """
8
        内存池结构
9
        """
10
        return """
11
        vLLM GPU 内存布局：
12

13
        ┌──────────────────────────────────────────────────────────┐
14
        │                        GPU Memory                         │
15
        ├──────────────────────────────────────────────────────────┤
16
        │                                                              │
17
        │  ┌────────────────────────────────────────────────────┐   │
18
        │  │            Model Weights (固定)                     │   │
19
        │  │  - 7B FP16: 14 GB                                  │   │
20
        │  │  - 70B FP16: 140 GB                                │   │
21
        │  └────────────────────────────────────────────────────┘   │
22
        │                                                              │
23
        │  ┌────────────────────────────────────────────────────┐   │
24
        │  │            KV Cache Pool (动态)                      │   │
25
        │  │                                                       │   │
26
        │  │  Block 0  Block 1  Block 2  Block 3  Block 4  ...   │   │
27
        │  │  ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐           │   │
28
        │  │  │ KV  │ │ KV  │ │ KV  │ │ KV  │ │ KV  │           │   │
29
        │  │  │16Tk │ │16Tk │ │16Tk │ │16Tk │ │16Tk │           │   │
30
        │  │  └─────┘ └─────┘ └─────┘ └─────┘ └─────┘           │   │
31
        │  │                                                       │   │
32
        │  └────────────────────────────────────────────────────┘   │
33
        │                                                              │
34
        │  ┌────────────────────────────────────────────────────┐   │
35
        │  │            Activations (临时)                        │   │
36
        │  │  - Prefill 中间结果                                 │   │
37
        │  │  - Decode 临时 buffer                              │   │
38
        │  └────────────────────────────────────────────────────┘   │
39
        │                                                              │
40
        └──────────────────────────────────────────────────────────┘
41

42
        gpu_memory_utilization 参数控制 KV Cache Pool 大小：
43
        → 0.9 表示 90% 的剩余内存用于 KV Cache
44
        → 10% 预留给 Activations
45
        """
46

47
    def allocation_strategy(self):
48
        """
49
        分配策略
50
        """
51
        return """
52
        KV Cache 分配流程：
53

54
        1. 初始化时计算可用 blocks 数
55
           available_blocks = gpu_memory × utilization / block_size
56

57
        2. 维护空闲 blocks 列表
58
           free_blocks = {0, 1, 2, 3, ..., N-1}
59

60
        3. 请求到来时分配
61
           - 计算需要的 blocks 数
62
           - 从 free_blocks 取出
63
           - 分配给请求
64

65
        4. 请求完成时释放
66
           - 归还 blocks 到 free_blocks
67
           - 可立即被其他请求复用
68

69
        关键：无需磁盘交换，纯 GPU 内管理
70
        """

1
class EvictionStrategies:
2
    """
3
    Eviction 策略
4
    """
5

6
    def lru_strategy(self):
7
        """
8
        LRU (Least Recently Used)
9

10
        驱逐最久未使用的序列
11
        """
12
        return """
13
        LRU 策略：
14

15
        场景：内存池满，需要为新请求腾出空间
16

17
        当前状态：
18
        ┌──────────────────────────────────────────┐
19
        │ Req1 (3 blocks)  last_used: 10 iterations ago │
20
        │ Req2 (5 blocks)  last_used: 2 iterations ago  │
21
        │ Req3 (2 blocks)  last_used: 50 iterations ago │
22
        └──────────────────────────────────────────┘
23

24
        LRU 决策：
25
        → Req3 最久未使用
26
        → 驱逐 Req3 的所有 blocks
27
        → 释放 2 blocks
28

29
        原因：
30
        → 活跃请求可能继续生成
31
        → 不活跃请求可能已放弃或完成
32
        → 优先保护活跃请求
33
        """
34

35
    def fifo_strategy(self):
36
        """
37
        FIFO (First In First Out)
38

39
        驱逐最早的序列
40
        """
41
        return """
42
        FIFO 策略：
43

44
        场景：内存池满
45

46
        当前状态：
47
        ┌──────────────────────────────────────────┐
48
        │ Req1 created: 100s ago  (3 blocks)       │
49
        │ Req2 created: 50s ago   (5 blocks)       │
50
        │ Req3 created: 10s ago   (2 blocks)       │
51
        └──────────────────────────────────────────┘
52

53
        FIFO 决策：
54
        → Req1 创建最早
55
        → 驱逐 Req1
56

57
        问题：
58
        → 可能驱逐仍在活跃使用的序列
59
        → 不如 LRU 智能
60
        """
61

62
    def hybrid_strategy(self):
63
        """
64
        混合策略
65
        """
66
        return """
67
        vLLM 实际使用的策略：
68

69
        1. 主要使用 LRU
70
           - 跟踪每个序列的最后访问时间
71
           - 驱逐最久未使用的
72

73
        2. 优先级队列
74
           - 区分 prefill 和 decode 阶段
75
           - decode 中的序列优先级更高
76

77
        3. 部分驱逐
78
           - 可以只驱逐部分 blocks
79
           - 保留最近的 blocks
80
           - 用于处理超长序列
81

82
        4. 拒绝而非驱逐（极端情况）
83
           - 如果所有序列都很活跃
84
           - 拒绝新请求而非驱逐
85
           - 返回队列等待
86
        """

1
class CPUSwap:
2
    """
3
    CPU-GPU Swap
4
    """
5

6
    def when_needed(self):
7
        """
8
        何时需要 Swap
9
        """
10
        return """
11
        需要 Swap 的场景：
12

13
        1. 超长序列
14
           - 序列长度超过 GPU KV Cache 容量
15
           - 部分 KV 需要换出到 CPU
16

17
        2. 内存压力
18
           - 高并发时内存池耗尽
19
           - 驱逐可能导致结果不正确
20

21
        3. 优先级场景
22
           - 某些高优先级请求需要更多内存
23
           - 临时换出低优先级请求
24

25
        权衡：
26
        → Swap 开销：PCIe 带宽 (~64 GB/s)
27
        → 但比 OOM 崩溃好
28
        → 可以处理超长上下文
29
        """
30

31
    def implementation(self):
32
        """
33
        Swap 实现
34
        """
35
        return '''
36
        class KVCacheSwap:
37
            """KV Cache CPU-GPU Swap"""
38

39
            def __init__(self, gpu_pool, cpu_pool_size=10000):
40
                self.gpu_pool = gpu_pool
41
                self.cpu_pool = {}  # {seq_id: {blocks: [], in_gpu: bool}}
42
                self.cpu_pool_size = cpu_pool_size
43

44
            def swap_out(self, seq_id, blocks_to_swap):
45
                """将 GPU blocks 换出到 CPU"""
46
                for block_id in blocks_to_swap:
47
                    # GPU → CPU
48
                    cpu_data = self.gpu_pool.get_block_data(block_id)
49
                    self.cpu_pool[seq_id]["data"].append(cpu_data)
50
                    self.gpu_pool.free_block(block_id)
51

52
            def swap_in(self, seq_id, blocks_to_swap):
53
                """将 CPU blocks 换回 GPU"""
54
                for block_id in blocks_to_swap:
55
                    # CPU → GPU
56
                    gpu_block = self.gpu_pool.allocate_block()
57
                    gpu_block.load_data(self.cpu_pool[seq_id]["data"].pop())
58
        '''
59

60
    def performance_impact(self):
61
        """
62
        性能影响
63
        """
64
        return {
65
            "swap_out_latency": "~1-5 ms per block (PCIe)",
66
            "swap_in_latency": "~1-5 ms per block (PCIe)",
67
            "full_swap_1k_tokens": "~40 ms (8 blocks)",
68

69
            "without_swap": "OOM，任务失败",
70
            "with_swap": "延迟增加但可完成",
71

72
            "recommendation": "尽量避免 swap，预留足够 GPU 内存",
73
        }

1
    """
2
    共享前缀场景
3
    """
4

5
    def chat_completion(self):
6
        """
7
        Chat Completion 场景
8

9
        系统提示 + 对话历史 + 当前查询
10
        """
11
        return """
12
        典型 Chat 请求：
13

14
        Request 1:
15
        [System: "你是AI助手..." (500 tokens)]
16
        [History: "用户问题1" (100 tokens)]
17
        [Query: "当前问题" (50 tokens)]
18

19
        Request 2:
20
        [System: "你是AI助手..." (500 tokens)] ← 相同！
21
        [History: "用户问题2" (120 tokens)]
22
        [Query: "当前问题2" (50 tokens)]
23

24
        Request 3:
25
        [System: "你是AI助手..." (500 tokens)] ← 相同！
26
        [History: "用户问题3" (80 tokens)]
27
        [Query: "当前问题3" (50 tokens)]
28

29
        共享情况：
30
        → System Prompt 完全相同 → KV Cache 可共享
31
        → History 不同 → 不可共享
32

33
        内存节省：
34
        → 无共享：3 × 500 = 1500 tokens 的 KV
35
        → 有共享：500 + 300 = 800 tokens 的 KV
36
        → 节省：46%
37
        """
38

39
    def few_shot_learning(self):
40
        """
41
        Few-shot Learning 场景
42
        """
43
        return """
44
        Few-shot 示例：
45

46
        Request:
47
        [Example 1: Q + A]  (200 tokens)
48
        [Example 2: Q + A]  (200 tokens) ← 相同格式
49
        [Example 3: Q + A]  (200 tokens) ← 相同格式
50
        [Example N: Q]      (100 tokens)
51
        [Query]             (50 tokens)
52

53
        优化：
54
        → Examples 1..N-1 的 KV 可以共享
55
        → 只计算 Example N 和 Query 的 KV
56
        """
57

58
    def document_summarization(self):
59
        """
60
        文档摘要场景
61
        """
62
        return """
63
        批量摘要任务：
64

65
        Request 1:
66
        [Document A summary] + [User instruction] + [Document A]
67

68
        Request 2:
69
        [Document B summary] + [User instruction] + [Document B]
70

71
        Request 3:
72
        [Document C summary] + [User instruction] + [Document C]
73

74
        共享：
75
        → "User instruction" 完全相同
76
        → Summary 格式相同
77

78
        潜在优化：
79
        → 识别结构化前缀
80
        → 只计算文档内容的 KV
81
        """

1
class HashBasedCaching:
2
    """
3
    基于哈希的缓存
4
    """
5

6
    def block_hashing(self):
7
        """
8
        Block 哈希计算
9
        """
10
        return '''
11
        class BlockHashCache:
12
            """基于哈希的 Block 缓存"""
13

14
            def __init__(self):
15
                self.hash_to_block = {}  # hash → block_id
16
                self.block_to_hash = {}  # block_id → hash
17

18
            def compute_block_hash(self, tokens):
19
                """计算 block 内容哈希"""
20
                import hashlib
21
                content = bytes(tokens.cpu().numpy().tobytes())
22
                return hashlib.sha256(content).hexdigest()[:16]
23

24
            def add_block(self, tokens, block_id):
25
                """添加 block 并记录哈希"""
26
                h = self.compute_block_hash(tokens)
27

28
                if h in self.hash_to_block:
29
                    # 内容相同，共享已有的 block
30
                    return self.hash_to_block[h]
31
                else:
32
                    # 新内容
33
                    self.hash_to_block[h] = block_id
34
                    self.block_to_hash[block_id] = h
35
                    return block_id
36

37
            def lookup(self, tokens):
38
                """查找是否存在相同内容"""
39
                h = self.compute_block_hash(tokens)
40
                return self.hash_to_block.get(h)
41
        '''
42

43
    def cache_key_calculation(self):
44
        """
45
        Cache Key 计算
46
        """
47
        return """
48
        Cache Key 策略：
49

50
        1. 内容哈希（精确匹配）
51
           hash = SHA256(block_tokens)
52
           - 完全相同的内容才能命中
53
           - 精确但可能错过相似前缀
54

55
        2. 前缀哈希（前缀匹配）
56
           hash = SHA256(prefix_tokens)
57
           - 相同前缀的所有内容可复用
58
           - 需要额外的数据结构
59

60
        3. 多级哈希
61
           - 第一级：block 内容哈希
62
           - 第二级：token 级别哈希
63
           - 支持部分前缀共享
64

65
        vLLM 使用：
66
        → 第一级内容哈希
67
        → Block 为粒度的精确匹配
68
        → Prefix 共享通过 hash 碰撞检测
69
        """
70

71
    def prefix_lookup(self):
72
        """
73
        前缀查找
74
        """
75
        return """
76
        前缀查找流程：
77

78
        新请求：tokens = [A, B, C, D, E, F]
79

80
        Step 1: 计算每个 block 的哈希
81
        hash([A,B]) = abc
82
        hash([C,D]) = def
83
        hash([E,F]) = ghi
84

85
        Step 2: 查找缓存
86
        if abc in cache → 命中 block 7
87
        if def in cache → 命中 block 3
88
        if ghi not in cache → 需要计算
89

90
        Step 3: 构建 Block Table
91
        Block Table: [7, 3, new_block]
92

93
        结果：
94
        → 前缀 [A,B,C,D] 的 KV 来自缓存
95
        → 只需计算 [E,F] 的 KV
96
        """