从源码剖析 vLLM 的显存管理与连续批处理：大模型推理吞吐量提升 20 倍的底层密码

从源码剖析 vLLM 的显存管理与连续批处理：大模型推理吞吐量提升 20 倍的底层密码

引言：为什么 vLLM 成为生产环境的事实标准？

当我们将大模型（LLM）从实验环境推向生产时，最先遭遇的“三座大山”是：

1. 显存墙：一块 A100（80GB）甚至装不下 70B 模型的单个副本。
2. 吞吐低：Naive 的 PyTorch 实现，GPU 利用率常年在 30% 以下。
3. 调度僵：静态批处理（Static Batching）受限于最短序列长度，导致计算资源极大浪费。

vLLM 通过 PagedAttention 和 Continuous Batching 两大杀手锏，将吞吐量提升 20 倍以上。本文将深入 vLLM 源码目录（vllm/），带大家看透这套“虚拟内存式”的显存管理系统究竟是如何运作的。

第一部分：传统推理的“显存黑洞”——KV Cache 管理困境

1.1 什么是 KV Cache，为什么它如此耗显存？

在 Decode-only 模型中，每个 Token 生成时都需要计算 Attention。为了避免重复计算历史 Token 的 Key 和 Value，业界采用 KV Cache 策略。

• 对于 Llama3-70B，每生成 1 个 Token，KV Cache 占用约 2 * layers * hidden_size * precision。
• 在动态生成场景下，每个请求的长度不一，显存分配呈现严重的碎片化（External Fragmentation）。

1.2 原生 PyTorch 的痛点

原生实现通常预先申请最大长度（如 2048）的连续显存张量（Contiguous Tensor）。

后果：哪怕请求只生成了 10 个 Token，也要占满 2048 个 Token 的显存空间，浪费率高达 99%。且频繁的 torch.cuda.malloc 极易触发 CUDA OOM。

第二部分：PagedAttention 源码级解构（操作系统思想移植）

vLLM 的核心创新在于将操作系统的分页内存（Paging） 机制引入 GPU 显存管理。

2.1 逻辑块与物理块的映射（BlockTable）

在源码 vllm/block.py 和 vllm/core/block_manager.py 中，vLLM 将 KV Cache 切分为固定大小的 物理块（Physical Blocks）（默认 16 个 Token/块）。

• 逻辑块（Logical Blocks）：属于请求自身的连续 Token 序列。
• 物理块（Physical Blocks）：GPU 显存上的离散内存块。

核心数据结构源码映射（简化伪代码）：

class BlockTable:
    def __init__(self):
        # 核心映射表：逻辑块索引 -> 物理块地址
        self._blocks: Dict[int, PhysicalBlock] = {}
        self._num_full_slots = 0

    def append_token(self, token_id, kv_cache):
        # 当逻辑块写满 16 个 Token 时，申请新的物理块
        if self._num_full_slots % BLOCK_SIZE == 0:
            new_block = self._allocate_physical_block()  # 关键：从显存池取空闲块
            self._blocks[self._num_full_slots // BLOCK_SIZE] = new_block
        # 计算物理地址并写入 KV Cache
        physical_addr = self._blocks[block_idx].address + offset
        write_kv_to_gpu(physical_addr, kv_cache)

2.2 显存池（KV Cache Pool）的初始化

在 vllm/worker/cache_engine.py 中，vLLM 启动时会一次性向 CUDA 申请一大块连续显存（由 gpu_memory_utilization 控制，默认 0.9），然后自行切分为无数个等大的物理块。

优势：由 vLLM 自己管理内存池，彻底杜绝了向 CUDA Driver 频繁申请释放显存的开销，且物理上不连续，完美解决了外部碎片问题。

2.3 Attention 前向计算的魔改（FlashAttention 的变体）

在 vllm/attention/backends/flash_attn.py 中，计算 Attention 时不再直接传入连续的 (batch, seq_len, dim) 张量，而是传入物理块编号列表（Block Tables）。计算单元通过索引重映射（Index Remapping），在物理显存上直接 Gather 数据并完成矩阵乘。

第三部分：Continuous Batching（连续批处理）调度源码分析

如果说 PagedAttention 解决了显存“空间”问题，Continuous Batching 则解决了“时间”效率问题。

3.1 静态批处理的死穴

传统静态批处理必须等待批次中最慢的请求全部完成（Finish），才能返回结果并接收新请求。这导致 GPU 在迭代后期大量空闲。

3.2 迭代级调度（Iteration-level Scheduling）

在 vllm/core/scheduler.py 的 schedule() 函数中，vLLM 实现了在每个 Decode 迭代开始时动态决定本轮要执行的请求序列。

def schedule(self) -> SchedulerOutputs:
    # 1. 检查正在运行的请求，踢出已完成（Finished）或超时的请求
    self._update_running()
    
    # 2. 从等待队列（Waiting Queue）中尽可能多地取出请求
    #    关键约束：显存中是否有足够的物理块（Physical Blocks）来容纳新请求的前缀 Token
    while self.waiting:
        request = self.waiting[0]
        if self.block_manager.can_allocate(request):  # 检查显存余量
            self.waiting.pop()
            self.running.append(request)
        else:
            break
    
    # 3. 构造本轮执行的序列组（Sequence Groups）
    #    注意：不同请求的序列长度不同，但 PagedAttention 允许它们混杂计算
    return SchedulerOutputs(running=self.running)

核心价值：某个请求生成了 “EOS” 结束符，下一微秒立刻释放其显存，新请求即可抢占资源。GPU 永远处于满负荷运转状态。

第四部分：生产环境部署的“显存与吞吐”调参实战（基于腾讯云 TKE）

理论源码分析完毕，下面给出我们在腾讯云 A10-24GB 显卡上部署 Qwen-72B（AWQ 4bit 量化）时的实际调参记录。

4.1 必须调整的 3 个核心环境变量

vLLM 的默认参数偏向通用性，针对高并发场景必须修改：

# 启动 vLLM 服务
python -m vllm.entrypoints.openai.api_server \
    --model /data/models/Qwen-72B-AWQ \
    --tensor-parallel-size 2 \          # 双卡 A10
    --max-num-seqs 256 \                # 【关键】最大并发序列数，默认 256，但显存小需降低
    --max-num-batched-tokens 8192 \     # 【关键】每批次最大 Token 总数，影响吞吐
    --gpu-memory-utilization 0.85 \     # 不要拉满，留 15% 给系统开销
    --block-size 16 \                   # 默认 16，不建议改
    --enable-prefix-caching             # 【灵魂】开启前缀缓存，如果 Prompt 包含固定 System Prompt，命中率极高

4.2 压测数据对比（Offline vs Online）

第五部分：监控与排障——别等 OOM 了才看日志

5.1 vLLM 内置 Prometheus 指标

加上 --disable-log-requests 和 --enable-metrics 后，可通过 /metrics 接口接入 Grafana。 重点关注指标：

• vllm:num_requests_running（当前运行请求数）
• vllm:gpu_cache_usage_perc（显存池占用百分比，超过 95% 时需扩容或限流）
• vllm:time_to_first_token_seconds（TTFT 分位数）

5.2 常见 Crash 根源：Block Manager 死锁

若并发设置过大，block_manager.can_allocate 会频繁返回 False，导致调度器死循环。解决方案：在启动命令中加入 --max-num-seqs 严格限制并发上限，并在网关层（如 Nginx）配置 worker_connections 排队，而非将所有请求涌入 vLLM。

结语：理解底层，方能驾驭上层

vLLM 告诉我们一个深刻的道理：在大模型时代，系统优化（Systems Optimization）的重要性绝不亚于模型架构（Model Architecture）创新。

PagedAttention 并非复杂的数学公式，而是一个极其优雅的工程落地思路。当你的业务需要承载日活百万的 AIGC 应用时，千万别只盯着 transformers 库里的 generate() 函数。深入源码，调优 Scheduler 和 BlockManager，你也能将昂贵的 GPU 算力榨干到极致。

讨论点：各位在生产环境是选择 vLLM、TGI 还是 TensorRT-LLM？在处理超长上下文（>100k）时，大家遇到的最大显存瓶颈是在 Attention 计算部分，还是在 KV Cache 存储部分？欢迎在评论区深聊。

附录：源码调试环境搭建

git clone https://github.com/vllm-project/vllm.git
cd vllm
pip install -e .   # 开发模式安装，方便打断点
# 在 PyCharm/VSCode 中 attach 到 main.py，观察 Scheduler 的迭代调度过程