KV Cache 与推理服务内存管理

大家好，我是芯缘，是 Datawhale 社区发起的 “llm-algo-leetcode 推理优化方向”组队学习活动的助教。

2026年8月运营助教

2026年9月专业助教

PagedAttention#

问题：

在线请求不会同时开始或同时结束，静态预留造成的空洞。
最大序列长度来分配浪费，KV Cache 的显存碎片化

解法：PagedAttention (vLLM):

借鉴操作系统的虚拟内存管理。把显存切分成固定大小的 Block（比如 1 个 Block 存 16 个 Token）。在生成时，按需分配物理 Block，并通过 Block Table 记录逻辑 Token 序列到物理块的映射。

\text{Block Table}[logical\_block\_id] = physical\_block\_id

KV Cache 变成“可按块寻址、可按需增长”的结构：prefill 申请所需 Block，decode 只在跨越边界时补 Block，计算阶段再依据块表访问离散缓存。

vLLM 的层次如下：

1
vLLM Serving
2
│
3
├── Scheduler
4
│      └── 动态 batching / 调度
5
│
6
├── KV Cache Manager
7
│      └── Block 分配 / 释放 / Block Table
8
│
9
└── Attention
10
       └── 根据 Block Table
11
           对非连续 KV Cache 进行 Attention
12
           ↑
13
       PagedAttention 核心计算思想

Scheduler部分如下：

1
                    Scheduler
2
                        │
3
              schedule 一个 batch
4
                        │
5
            ┌───────────┴───────────┐
6
            ↓                       ↓
7
        WAITING                   RUNNING
8
            │                       │
9
            ↓                       ↓
10
       prefill(req)         decode_one_token(req)
11
            │                       │
12
            │                       ↓
13
            │             paged_attention_forward()
14
            │                       │
15
            └───────────┬───────────┘
16
                        ↓
17
                 GPU KV Cache
18
                        │
19
                  Block Table
20
                        │
21
                        ↓
22
                Physical Blocks
23

24

25
如果生成结束：
26
        ↓
27
    finish(req)
28
        ↓
29
Block → FREE_BLOCKS

阶段	输入/触发	状态变化	需要检查的边界
Prefill	prompt tokens	按 prompt 所需容量申请多个物理 Block，并建立 Block Table	空闲物理 Block 是否足够；最后一个 Block 可不满
Decode	每次新增 token	写入当前 Block；当前 Block 满时再申请新物理 Block，并更新 Block Table	新 Block 分配失败时如何处理/调度
Release	请求完成/终止	归还该请求占用的物理 Block，并删除/清空对应 Block Table	防止重复释放；共享 Block 时需考虑引用计数
Reuse	后续请求需要空间	从空闲 Block 池重新分配之前释放的物理 Block	同一请求获得的物理 Block无需连续

KV Cache Manager如下：

1
                GPU KV Cache
2
        ┌─────────────────────────┐
3
        │ Physical Block Pool     │
4
        │                         │
5
        │ 0  1  2  3  4 ... 17    │
6
        └─────────────────────────┘
7
                    ↑
8
                    │ allocate / free
9
                    │
10
              KV Cache Manager
11
                    ↑
12
                    │
13
         ┌──────────┴──────────┐
14
         │                     │
15
      Request A             Request B
16

17
Block Table              Block Table
18
[17, 3, 25]              [8, 11]
19
    │                        │
20
    ↓                        ↓
21

22
逻辑 KV：                 逻辑 KV：
23

24
Block 0 → 17             Block 0 → 8
25
Block 1 → 3              Block 1 → 11
26
Block 2 → 25
27

28

29
Decode：
30
当前 Block 有空间
31
        ↓
32
      继续写
33

34
当前 Block 满
35
        ↓
36
从 Pool 申请新 Block
37
        ↓
38
追加到 Block Table
39

40

41
Request Finish：
42
Block Table 中所有 Block
43
        ↓
44
归还 Pool

伪代码如下：

1
# ==========================================
2
# 0. 初始化：GPU 上预留 KV Cache Block Pool
3
# ==========================================
4

5
GPU_KV_CACHE = PhysicalBlockPool(
6
    num_blocks = N,
7
    block_size = 16 tokens
8
)
9

10
FREE_BLOCKS = [0, 1, 2, ..., N-1]
11

12
# 每个 Request 维护自己的 Block Table
13
Request:
14
    tokens
15
    block_table = []
16
# Logical Block 0 -> Physical Block 17
17
# Logical Block 1 -> Physical Block 3
18
# Logical Block 2 -> Physical Block 25
19
# block_table = [17, 3, 25]
20

21
# ==========================================
22
# 1. Prefill
23
# ==========================================
24

25
prefill(request):
26

27
    prompt_len = len(request.tokens)
28
    num_blocks_needed = ceil(prompt_len / BLOCK_SIZE)
29

30
    if FREE_BLOCKS.size < num_blocks_needed:
31
        无法调度该请求
32
        return
33

34
    for i in range(num_blocks_needed):
35

36
        physical_block = FREE_BLOCKS.pop()
37
        request.block_table.append(physical_block)
38

39
    model_forward(request)
40

41

42
# ==========================================
43
# 2. Decode：每生成一个 token
44
# ==========================================
45

46
decode_one_token(request):
47

48
    next_position = request.seq_len
49
    logical_block_id = next_position // BLOCK_SIZE
50
    offset_in_block  = next_position % BLOCK_SIZE
51

52
    # 情况 A：已有 Block 还有位置
53
    if logical_block_id < len(request.block_table):
54
        physical_block =
55
            request.block_table[logical_block_id]
56

57
    # 情况 B：跨越 Block 边界->按需申请新的物理 Block
58
    else:
59
        if FREE_BLOCKS.empty():
60
            无法继续调度 / 需要抢占其他请求
61
            return
62
        physical_block = FREE_BLOCKS.pop()
63
        request.block_table.append(physical_block)
64

65

66
    # Attention 根据 Block Table
67
    # 找到这个请求之前所有 KV 的物理位置
68
    attention(
69
        query = current_query,
70
        block_table = request.block_table,
71
        kv_cache = GPU_KV_CACHE
72
    )
73
    # 新 token 产生的 K/V 写入对应物理 Block
74
    GPU_KV_CACHE[
75
        physical_block,
76
        offset_in_block
77
    ] = new_KV
78
    request.seq_len += 1
79

80

81
# ==========================================
82
# 3. Attention：根据 Block Table 访问 KV
83
# ==========================================
84

85
paged_attention(query, request):
86
    for logical_block_id in request需要访问的所有逻辑Block:
87
        # 查表
88
        physical_block_id = request.block_table[logical_block_id]
89
        # 直接读取对应物理 Block
90
        K, V = GPU_KV_CACHE[physical_block_id]
91
        # 对该 Block 做 Attention
92
        accumulate_attention(query, K, V)
93

94

95
# ==========================================
96
# 4. 请求结束
97
# ==========================================
98
finish(request):
99
    for physical_block in request.block_table:
100
        FREE_BLOCKS.push(physical_block)
101
    request.block_table.clear()

RadixAttention#

问题：

多个请求可能包含相同的系统提示词、对话历史、Few-shot 示例或搜索路径，但普通推理会为每个请求重复执行相同前缀的 Prefill。
请求结束后若立即释放全部 KV Cache，后续请求即使拥有相同前缀，也无法复用已经计算过的 KV。

解法：RadixAttention (SGLang)：

系统维护一棵以 Token 序列为 Key、以 KV Cache 物理位置为 Value 的 Radix Tree（压缩前缀树）。树中的一条边可以保存一段 Token，而不是像普通 Trie 那样每个 Token 都建立一个节点。公共 Token 前缀因此只对应一条共享路径，其 KV Cache 也可以被多个请求复用。

新请求到达时，系统先用完整 Prompt 在树中做最长前缀匹配（Longest Prefix Match）：

H = \max_j \operatorname{lcp}(prompt\_tokens, cached\_path_j)

其中 H 是命中的 Token 数量，即 hit_len。Prompt 随后被划分为：

\text{Prompt} = \text{Cached Prefix} + \text{Uncached Suffix}

RadixAttention 先用 Radix Tree 找到最长缓存前缀，复用 Prefix 的 KV，只计算 Suffix；计算结果继续插入树中，请求结束后也不立即释放，显存不足时再淘汰当前没有请求使用的缓存节点。

1
SGLang Serving
2
│
3
├── Scheduler
4
│      └── 请求排序 / continuous batching
5
│
6
├── Radix Cache
7
│      └── Match / Insert / Split / Evict
8
│          Token Prefix → KV Cache indices
9
│
10
├── KV Cache Allocator
11
│      └── KV slot / Page 的分配与释放
12
│
13
└── Attention
14
       └── 根据命中的 KV indices
15
           读取分页存放的 GPU KV Cache

Scheduler 部分如下：

1
                    Scheduler
2
                        │
3
                  接收新请求 req
4
                        │
5
                        ↓
6
              match_prefix(req.prompt)
7
                        │
8
            ┌───────────┴───────────┐
9
            ↓                       ↓
10
      Cached Prefix           Uncached Suffix
11
       直接复用 KV             prefill_suffix()
12
            │                       │
13
            └───────────┬───────────┘
14
                        ↓
15
                 insert() / Split
16
                        │
17
                        ↓
18
                    Decode
19
                        │
20
                请求结束 finish()
21
                        │
22
            KV 保留在 Radix Tree 中
23
                        │
24
                 显存不足时 evict()

操作	输入	核心行为	结果
Match	新请求 prompt	在 Radix Tree 中做最长前缀匹配	得到可直接复用的 KV 和 `hit_len`
Prefill	未命中的 suffix	只计算 `prompt[hit_len:]`	避免重复计算已命中 prefix
Insert	新计算出的 token + KV	把新的 KV 路径加入 Radix Tree；必要时内部 Split	新 KV 可以被未来请求复用
Finish	请求完成/终止	把可复用路径留在树中，并解除该请求对路径的锁定	KV 不因请求结束而立即释放
Evict	显存压力	从可淘汰叶子中选择 Victim，删除节点并释放 KV	给新请求腾出 KV 空间；正在被请求使用的路径不能淘汰

这里的 Split 不是一个独立的 Serving 阶段。它是 Match 或 Insert 在压缩边内部遇到匹配终点或分叉点时，用来维护树结构的内部操作。

例如树中原来只有一条路径：

1
root
2
  │
3
[A B C D E]

现在插入 [A B C X Y]，两条路径在 C 之后分叉，原节点会被拆成：

1
             root
2
               │
3
            [A B C]
4
            /     \
5
         [D E]   [X Y]

[A B C] 及其 KV 只保留一份，两个请求分别沿不同的 Suffix 继续向下。

Radix Cache 的结构如下：

1
RadixNode:
2
    tokens          # 从 parent 到当前节点的压缩 Token 段
3
    kv_indices      # 这段 Token 对应的 GPU KV 位置
4
    children        # 下一 Token → 子节点
5
    parent
6
    lock_ref        # 是否有正在运行的请求使用该路径
7
    last_access     # 供 LRU 等淘汰策略使用
8

9

10
                         ROOT
11
                           │
12
                [system prompt tokens]
13
                 KV: [7, 12, 31, ...]
14
                    /             \
15
          [user question A]   [user question B]
16
          KV: [18, 22, ...]   KV: [5, 16, ...]
17
                 │
18
          [model answer A]
19
          KV: [9, 27, ...]

伪代码如下。为突出主线，省略了 Page 对齐、Chunked Prefill、分布式同步和 Host Cache 等工程细节。

1
# ==========================================
2
# 0. 初始化：Radix Tree + GPU KV Allocator
3
# ==========================================
4

5
KV_ALLOCATOR = KVCacheAllocator(total_slots = N)
6

7
ROOT = RadixNode(
8
    tokens = [],
9
    kv_indices = [],
10
    children = {},
11
    parent = None,
12
    lock_ref = 1
13
)
14

15

16
# ==========================================
17
# 1. Split：在压缩边内部暴露分叉点
18
#    这是 Match / Insert 的内部操作
19
# ==========================================
20

21
split_node(node, split_pos):
22

23
    # 原节点： [A B C D E]
24
    # split_pos = 3
25
    # 拆分后： [A B C] -> [D E]
26

27
    suffix_node = RadixNode(
28
        tokens = node.tokens[split_pos:],
29
        kv_indices = node.kv_indices[split_pos:],
30
        children = node.children,
31
        parent = node
32
    )
33

34
    for child in suffix_node.children.values():
35
        child.parent = suffix_node
36

37
    node.tokens = node.tokens[:split_pos]
38
    node.kv_indices = node.kv_indices[:split_pos]
39
    node.children = {
40
        suffix_node.tokens[0]: suffix_node
41
    }
42

43
    return node
44

45

46
# ==========================================
47
# 2. Match：寻找最长缓存前缀
48
# ==========================================
49

50
match_prefix(prompt):
51

52
    node = ROOT
53
    position = 0
54
    matched_kv = []
55

56
    while position < len(prompt):
57

58
        child = node.children.get(prompt[position])
59

60
        if child 不存在:
61
            break
62

63
        common_len = longest_common_prefix(
64
            prompt[position:],
65
            child.tokens
66
        )
67

68
        if common_len == 0:
69
            break
70

71
        # 整条压缩边都匹配，继续向下
72
        if common_len == len(child.tokens):
73
            matched_kv.extend(child.kv_indices)
74
            position += common_len
75
            node = child
76
            continue
77

78
        # 命中终点落在压缩边内部
79
        # Split 后，命中的 Prefix 形成完整节点
80
        prefix_node = split_node(child, common_len)
81
        matched_kv.extend(prefix_node.kv_indices)
82
        position += common_len
83
        node = prefix_node
84
        break
85

86
    return MatchResult(
87
        hit_len = position,
88
        kv_indices = matched_kv,
89
        last_node = node
90
    )
91

92

93
# ==========================================
94
# 3. Prefill：只计算没有命中的 Suffix
95
# ==========================================
96

97
prefill_suffix(request, match_result):
98

99
    hit_len = match_result.hit_len
100
    suffix = request.prompt[hit_len:]
101

102
    if KV_ALLOCATOR.free_slots < len(suffix):
103
        evict(len(suffix) - KV_ALLOCATOR.free_slots)
104

105
    if KV_ALLOCATOR.free_slots < len(suffix):
106
        无法调度该请求
107
        return OOM
108

109
    new_kv_slots = KV_ALLOCATOR.allocate(len(suffix))
110

111
    # Prefix 的 KV 已经存在，只计算 Suffix
112
    model_forward_prefill(
113
        tokens = suffix,
114
        prefix_kv = match_result.kv_indices,
115
        output_kv = new_kv_slots
116
    )
117

118
    return new_kv_slots
119

120

121
# ==========================================
122
# 4. Insert：将 Token 路径与 KV 引用插入树中
123
# ==========================================
124

125
insert(tokens, kv_indices):
126

127
    node = ROOT
128
    position = 0
129

130
    while position < len(tokens):
131

132
        child = node.children.get(tokens[position])
133

134
        # 没有公共前缀：剩余部分直接成为新压缩边
135
        if child 不存在:
136
            new_node = RadixNode(
137
                tokens = tokens[position:],
138
                kv_indices = kv_indices[position:],
139
                parent = node
140
            )
141
            node.children[tokens[position]] = new_node
142
            return new_node
143

144
        common_len = longest_common_prefix(
145
            tokens[position:],
146
            child.tokens
147
        )
148

149
        # 整条已有边都匹配，继续向下
150
        if common_len == len(child.tokens):
151
            position += common_len
152
            node = child
153
            continue
154

155
        # 在已有边内部发生分叉
156
        branch_node = split_node(child, common_len)
157
        position += common_len
158

159
        # 新路径恰好结束在分叉点
160
        if position == len(tokens):
161
            return branch_node
162

163
        new_node = RadixNode(
164
            tokens = tokens[position:],
165
            kv_indices = kv_indices[position:],
166
            parent = branch_node
167
        )
168
        branch_node.children[tokens[position]] = new_node
169
        return new_node
170

171
    return node
172

173

174
# ==========================================
175
# 5. Evict：显存不足时淘汰未被使用的叶子
176
# ==========================================
177

178
evict(required_slots):
179

180
    released = 0
181

182
    while released < required_slots:
183

184
        victim = choose_evictable_leaf(
185
            policy = LRU,
186
            condition = node.lock_ref == 0
187
        )
188

189
        if victim 不存在:
190
            break
191

192
        KV_ALLOCATOR.free(victim.kv_indices)
193
        released += len(victim.kv_indices)
194

195
        parent = victim.parent
196
        del parent.children[victim.tokens[0]]
197

198
        # parent 失去最后一个 child 后，也可能成为可淘汰叶子
199
        update_evictable_state(parent)
200

201
    return released
202

203

204
# ==========================================
205
# 6. 请求结束：缓存 KV，而不是立即释放
206
# ==========================================
207

208
finish(request):
209

210
    insert(
211
        tokens = request.all_tokens,
212
        kv_indices = request.kv_indices
213
    )
214

215
    # 请求不再使用这条路径，允许未来在显存压力下淘汰
216
    unlock_path(request.last_node)
217
    request.last_node = None
218
    request.state = FINISHED
219

220
    # 注意：这里不调用 KV_ALLOCATOR.free()
221
    # KV 仍由 Radix Tree 持有，供未来请求复用

完整生命周期可以概括为：

1
新请求到达
2
    ↓
3
Match：找到最长缓存 Prefix
4
    ↓
5
Lock：防止正在使用的路径被淘汰
6
    ↓
7
Prefill：只计算未命中的 Suffix
8
    ↓
9
Insert：把新 KV 路径加入树，必要时内部 Split
10
    ↓
11
Decode
12
    ↓
13
Finish：解除锁定，但 KV 继续留在树中
14
    ↓
15
显存不足
16
    ↓
17
Evict：从未锁定的叶子开始回收 KV

PagedAttention 的 finish() 与 RadixAttention 的 finish() 有一个关键区别：

1
PagedAttention（不考虑 Prefix Cache）：
2
Request Finish
3
    ↓
4
物理 Block 立即归还 Free Pool
5

6

7
RadixAttention：
8
Request Finish
9
    ↓
10
KV 继续由 Radix Tree 持有
11
    ↓
12
未来相同 Prefix 可以复用
13
    ↓
14
显存有压力时才 Evict

因此，RadixAttention 不是简单地“多建一棵树”，而是改变了 KV Cache 的生命周期：从“请求私有、结束即释放”变为“按前缀共享、按需淘汰”。

音乐

音乐

PagedAttention#

RadixAttention#

参考资料#

文章分享

评论区

音乐

文章目录