3、显存优化策略：KV Cache量化（FP8/INT8）、内存池复用、显存碎片整理

显存，说白了就是大模型推理的命根子。你想想看，一个70B的模型，光参数就要140GB显存，再加上KV Cache，稍微不注意就OOM了。我这些年调过的推理系统，十次有八次的问题都出在显存上。今天咱们就聊聊怎么把显存这块「硬骨头」啃下来。

3.1 KV Cache量化：FP8/INT8

先说说KV Cache。这东西是Transformer推理时的「临时记忆」，存的是每层计算出来的Key和Value。随着序列变长，它占的空间会越来越大。我见过一个场景，生成长文本时KV Cache占了总显存的60%以上，你说吓不吓人？

那怎么压缩它呢？量化。说白了就是把高精度的浮点数，用低精度的格式存起来。

3.1.1 FP8量化

FP8是NVIDIA H100开始支持的新格式。它有两种变体：E4M3和E5M2。E4M3精度高一些，适合前向计算；E5M2动态范围大，适合反向传播。在KV Cache场景下，我们主要用E4M3。

我个人习惯的做法是：

• 对每个head单独做量化，因为不同head的分布差异很大
• 使用per-token的缩放因子，而不是per-tensor
• 量化前先做一次校准，收集几百个样本的统计量

# 伪代码：FP8 KV Cache量化
def quantize_kv_cache(kv_tensor, scale):
    # kv_tensor: [batch, num_heads, seq_len, head_dim]
    # 对每个head计算缩放因子
    scales = []
    for h in range(num_heads):
        head_data = kv_tensor[:, h, :, :]
        max_val = head_data.abs().max()
        scale = max_val / 448.0  # FP8 E4M3最大值
        scales.append(scale)
    # 量化
    quantized = (kv_tensor / scales).to(torch.float8_e4m3fn)
    return quantized, scales

3.1.2 INT8量化

INT8量化更成熟，兼容性也更好。vLLM里默认就支持INT8的KV Cache。做法和FP8类似，但缩放因子计算方式不同。

我曾经踩过一个坑：直接用min-max量化，结果遇到异常值（outlier）时，整个量化都崩了。后来改用百分位法，取99.9%的分位数作为最大值，效果好了很多。

# INT8量化：百分位法
def int8_quantize(kv, percentile=99.9):
    flat = kv.flatten()
    k = int(len(flat) * percentile / 100)
    max_val = torch.kthvalue(flat, k).values
    scale = max_val / 127.0
    quantized = torch.clamp(torch.round(kv / scale), -128, 127).to(torch.int8)
    return quantized, scale

3.2 内存池复用

说到内存池，我得先吐槽一下。很多新手写推理代码，每次请求都重新分配显存，用完就释放。这就像你每次吃饭都重新买一套碗筷，吃完就扔——太浪费了！

内存池复用的核心思想就一句话：提前分配，重复使用。

vLLM里用的是PagedAttention，它把KV Cache分成固定大小的块（block）。每个block是16个token的KV值。这样做的好处是：

• 按需分配，不会浪费
• 物理上不连续，但逻辑上连续
• 支持共享，比如beam search时多个候选可以共用前缀

我建议你在实现内存池时，注意这几点：

1. 预分配策略：根据最大并发数和最大序列长度，提前算好需要多少block。比如最大并发32，最大长度4096，每个block 16 token，那就需要32 * 4096 / 16 = 8192个block。
2. 空闲链表：用一个链表维护空闲block。分配时从链表头部取，释放时插回链表。O(1)复杂度。
3. 引用计数：对于共享的block，用引用计数管理。计数为0时才真正释放。

# 内存池简化实现
class KVCachePool:
    def __init__(self, num_blocks, block_size):
        self.free_list = list(range(num_blocks))
        self.ref_counts = [0] * num_blocks
        self.block_size = block_size
    
    def alloc(self):
        if not self.free_list:
            raise OOMError("显存不足！")
        block_id = self.free_list.pop()
        self.ref_counts[block_id] = 1
        return block_id
    
    def free(self, block_id):
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            self.free_list.append(block_id)

3.3 显存碎片整理

显存碎片，这是个老生常谈的问题。你想想看，频繁地分配和释放不同大小的显存块，时间长了就会产生很多「小洞」。这些小洞单独看都不大，但加起来可能占了几十GB。

为什么会这样？因为GPU显存分配器（比如cudaMalloc）是「最佳适配」策略。它会在空闲块里找大小最合适的。如果找不到，就报OOM。哪怕总空闲显存还有10GB，但每个空闲块都只有1GB，你要分配一个2GB的块，照样OOM。

我遇到过最夸张的一次：显存总共80GB，空闲显示还有35GB，但就是分配不出一个4GB的tensor。查了半天，发现空闲块被分割成了200多个小碎片。

怎么解决？三个办法：

• 预分配大块：一次性分配一个大块，然后自己管理。vLLM就是这么干的。
• 定期整理：把碎片合并成大块。但注意，整理时需要暂停推理，所以频率不能太高。
• 对齐分配：所有分配都按256字节对齐，减少碎片产生。

我个人推荐的做法是「预分配+对齐」。具体来说：

1. 启动时分配一个大的显存池，比如总显存的90%
2. 内部按block大小（比如2MB）切分
3. 所有分配都按block的整数倍来

这样基本不会产生碎片。如果实在需要整理，可以用下面的策略：

# 显存碎片整理策略
def defragment(pool):
    # 1. 标记所有空闲块
    free_blocks = [b for b in pool.blocks if b.ref_count == 0]
    # 2. 按地址排序
    free_blocks.sort(key=lambda b: b.addr)
    # 3. 合并相邻块
    merged = []
    for b in free_blocks:
        if merged and merged[-1].addr + merged[-1].size == b.addr:
            merged[-1].size += b.size
        else:
            merged.append(b)
    # 4. 更新空闲链表
    pool.free_list = merged

3.4 三者如何配合

这三个策略不是孤立的，它们需要配合使用。我画了一张图，你看看就明白了：

你看这个流程：先量化压缩，再通过内存池复用，最后定期整理碎片。三步走下来，显存利用率能提升30%-50%。

最后说一句：这些策略不是银弹。如果你的模型本身就很大，序列又很长，该加显存还是得加。但在这之前，先把这些优化做足，能省不少钱。