4、显存池与大模型推理：KV Cache的显存管理、动态Batch与显存调度
聊到大模型推理，显存永远是个绕不开的话题。我刚开始做推理优化那会儿，最头疼的就是显存碎片化问题。明明总显存还有不少，但就是分配不出一个连续的KV Cache块。后来我意识到，没有显存池，大模型推理就像在堆满杂物的房间里跳舞——迟早要绊倒。
4.1 KV Cache：推理中的显存大户
先说说KV Cache是什么。说白了，它就是Transformer解码过程中，每一层自注意力计算产生的Key和Value矩阵的缓存。你想想看，模型每生成一个token，都要重新计算前面所有token的注意力。如果没有缓存，每次都得从头算，那计算量就爆炸了。
KV Cache的大小怎么算？我直接给公式：
KV Cache大小 = 2 × 层数 × 头数 × 序列长度 × 头维度 × 精度字节数举个例子，一个70B参数的模型，层数80，头数64，序列长度4096，头维度128，用FP16存储。算下来：
2 × 80 × 64 × 4096 × 128 × 2 = 10.7 GB嗯，一个请求就要吃掉10GB+显存。如果并发32个请求，光KV Cache就要340GB。这还没算模型权重和中间激活呢。
核心痛点：KV Cache是动态增长的，每个请求的序列长度不一样，而且请求会随时到达、随时结束。这就导致显存分配和释放非常频繁，碎片化严重。
4.2 显存池：从“随用随分”到“预分配复用”
传统的做法是每个请求来了，动态malloc一块显存。我在项目中遇到过，跑着跑着突然OOM，但一看显存使用率才70%。为什么？碎片化！
显存池的思路很简单：提前申请一大块连续显存，切成固定大小的块（比如每个块存一个token的KV）。请求来了，从池子里拿块；请求走了，把块还回去。这样避免了频繁的系统级显存分配，也消除了碎片。
我的经验：块大小怎么定？我一般按最大序列长度来切。比如最大支持4096个token，那就预分配4096个块。但实际请求可能只用512个token，剩下的块就浪费了。所以后来我改用“按需增长”策略——先给每个请求分配少量块，不够了再从池子里追加。
显存池的核心数据结构，我画了个图：
4.3 动态Batch：让显存利用率翻倍
静态Batch有个问题：你得提前定好batch size。设小了，吞吐上不去；设大了，显存不够用。我见过有人把batch size设为64，结果大部分请求序列很短，显存浪费了一半。
动态Batch的思路是：不固定batch size，而是根据当前显存池的空闲块数量，动态决定能塞进多少个请求。说白了，就是“有多少米煮多少饭”。
核心算法：
- 每个请求到达时，先估算它需要的最大块数（根据输入长度 + 最大生成长度）
- 检查显存池空闲块是否足够
- 如果够，分配块并加入当前batch
- 如果不够，把请求放入等待队列
- 每完成一个请求，释放块，唤醒等待队列中的请求
代码实现大概长这样：
class DynamicBatchScheduler:
def __init__(self, memory_pool, max_batch_tokens):
self.pool = memory_pool
self.max_batch_tokens = max_batch_tokens
self.active_requests = []
self.waiting_queue = []
def add_request(self, req):
needed_blocks = req.input_len + req.max_output_len
if self.pool.available_blocks() >= needed_blocks:
blocks = self.pool.allocate(needed_blocks)
req.assign_blocks(blocks)
self.active_requests.append(req)
else:
self.waiting_queue.append(req)
def step(self):
# 每步推理后，检查完成的请求
for req in self.active_requests:
if req.is_finished():
self.pool.free(req.blocks)
self.active_requests.remove(req)
# 尝试从等待队列拉入新请求
while self.waiting_queue:
req = self.waiting_queue[0]
needed_blocks = req.input_len + req.max_output_len
if self.pool.available_blocks() >= needed_blocks:
self.waiting_queue.pop(0)
self.add_request(req)
else:
break注意：动态Batch不是万能的。如果请求的序列长度差异太大，短请求会被长请求“拖累”——因为整个batch要等最长的那个请求做完才能一起释放。我建议配合“请求分桶”策略，把序列长度相近的请求放在同一个batch里。
4.4 显存调度：从“被动等待”到“主动管理”
有了显存池和动态Batch，是不是就万事大吉了？还差一步——调度策略。
我见过最朴素的调度是FIFO：先来先服务。但问题来了，如果先来的请求要生成1024个token，后来的请求只生成10个token，那后来的请求得等死。用户体验极差。
更好的做法是引入优先级调度。我常用的策略：
| 调度策略 | 核心思想 | 适用场景 |
|---|---|---|
| 最短剩余时间优先（SRTF） | 优先处理剩余生成token最少的请求 | 在线服务，追求低延迟 |
| 最大吞吐优先（MTF） | 优先处理能填满显存池的请求组合 | 离线批处理，追求高吞吐 |
| 公平调度（Fair Share） | 按用户/租户分配显存配额 | 多租户场景 |
我的实战经验：在vLLM项目中，我们用的是“先来先服务 + 显存池水位线”的混合策略。当显存池空闲块低于30%时，启动“紧急回收模式”——暂停新请求的调度，优先完成短请求。这招在流量突增时特别管用。
还有一个容易被忽略的点：显存池的碎片整理。即使有预分配，长时间运行后还是会产生碎片。我建议每隔一段时间（比如每1000次推理）做一次“显存压缩”——把分散的空闲块合并成大块。代价是暂停服务几十毫秒，但能避免显存利用率持续下降。
4.5 总结一下
KV Cache的显存管理，说白了就是三个层次：
- 分配层：显存池预分配，避免碎片
- 调度层：动态Batch，按需分配显存
- 策略层：优先级调度，平衡延迟和吞吐
我刚开始做的时候，总想着用更复杂的算法。后来发现，很多时候一个简单的显存池 + 动态Batch就能解决80%的问题。剩下的20%，靠监控和调参慢慢磨。嗯，工程嘛，就是在理想和现实之间找平衡。
想要更多资料，联系可可