LLM推理优化:vLLM PagedAttention深度解析与工程实践

一、排了两个月的队,我决定自己动手

2024年底,我给团队搭了一套推理服务,基于 Transformers + HuggingFace 的 naive 实现。QPS 大概在 0.8 左右——跑 LLaMA-13B,A100 单卡。用户一多,请求开始排队。最长的一次,一个用户等了 47 秒才看到第一个 token。

排队的根因不是模型慢。模型本身的前向计算差不多 120ms/token,瓶颈在显存。

传统推理框架里,每个请求来了先分配一块连续的 KV Cache 空间。假设一个请求生成长度 2048 的序列,KV Cache 要占用约 2.8GB 显存(FP16,13B 模型)。问题是:你不知道最终生多长,于是只能按最大长度预分配。用户只说了 200 个字,你给它留了 2048 个 token 的位置。碎片率和浪费率惨不忍睹。

直到我看到 vLLM 那篇 PagedAttention 论文。它解决的就是这个问题——把 KV Cache 切成固定大小的"页",像操作系统的虚拟内存一样管理。我当时的第一反应是:这不就是数据库里早就玩烂了的分页吗?但认真看完实现,发现把操作系统的内存管理思想搬到 GPU 显存上,工程落地的细节比想象中多得多。

本文从工程实现的角度,拆解 PagedAttention 的设计思路、核心数据结构和我在接入过程中踩过的坑。

二、KV Cache 为什么是瓶颈

先算一笔账。一个 Transformer decoder layer 的 self-attention 计算中,对于每个 token,我们要算:

Q = x * W_Q, K = x * W_K, V = x * W_V
attn = softmax(Q * K^T / sqrt(d)) * V

对于自回归生成,生成 token i 时,需要用到之前所有 token 的 K 和 V。如果每次重新算,复杂度是 O(n²) 的——生成到第 2048 个 token 时,前面的都要重算一遍,这是不可接受的。

于是有了 KV Cache:把每个 layer 的 K 和 V 矩阵存下来,每次追加新 token 的 K 和 V。

显存计算:

KV Cache per token = 2 (K 和 V) × num_layers × d_model × dtype_bytes

以 LLaMA-13B 为例:
num_layers = 40
d_model = 5120
dtype = FP16 = 2 bytes

每个 token 的 KV Cache = 2 × 40 × 5120 × 2 = 819,200 bytes ≈ 0.8MB

生成 2048 个 tokens → 2048 × 0.8MB ≈ 1.6GB
加上 batch 维度,batch_size = 4 时 → 6.4GB

A100 80GB 显存,模型权重占约 26GB(FP16),剩下的 54GB 用来做 KV Cache 和中间激活。你猜怎么着?大部分推理框架在 batch_size=8 时就把显存吃光了,不是因为模型算力不够,而是KV Cache 的分配策略太浪费

传统方案的问题:

  1. 预分配最大长度:每个请求按 max_seq_len 预留空间,实际用的可能只有 10%
  2. 外部碎片:请求长度不一,先来的释放了空间,但留下的空洞不连续,没法给新请求用
  3. 内部碎片:预留了 2048 slot 但只用了 300,那 1700 个 slot 就浪费了

三、PagedAttention 的核心思想

PagedAttention 的核心就一句话:把 KV Cache 切成固定大小的物理块(Block),通过逻辑到物理的映射表来管理,按需分配,用完即还

像极了操作系统的分页内存管理。但 GPU 上没有 MMU,所以 vLLM 自己做了一套 Block Manager。

3.1 Block Table

每个请求(vLLM 里叫 Sequence)维护一个逻辑 Block Table:

逻辑 Block ID | 物理 Block ID | 已占用的 slot 数
     0        |     47        |      16
     1        |     23        |      16
     2        |     89        |      8

物理 block 大小为 16 个 token 的 KV 数据。Block 满了(16/16)就分配下一个。最后一个 block 可能不满(如上图 block 2 只用了 8 个 slot)。

这种设计带来的好处:

  • 无外部碎片:任何大小的释放都能被复用,因为 block 是等长的
  • 按需分配:只分配实际使用的 block,不预分配
  • Copy-on-Write:同一个 block 可以被多个请求共享,在 beam search 场景下特别有用

3.2 Block Manager 的核心流程

# 伪代码,表达核心逻辑
class BlockManager:
    def __init__(self, num_gpu_blocks, block_size=16):
        self.free_blocks = list(range(num_gpu_blocks))
        self.allocated = {}  # seq_id -> [physical_block_ids]
        self.block_size = block_size
    
    def allocate(self, seq_id, num_tokens):
        """为 seq 分配容纳 num_tokens 所需的物理块"""
        needed_blocks = ceil(num_tokens / self.block_size)
        already_used = len(self.allocated.get(seq_id, [])) * self.block_size
        
        if already_used >= num_tokens:
            return
        
        new_blocks_needed = ceil((num_tokens - already_used) / self.block_size)
        
        if len(self.free_blocks) < new_blocks_needed:
            raise OOM("显存不足,需要执行 swap 或 preemption")
        
        for _ in range(new_blocks_needed):
            block = self.free_blocks.pop(0)
            self.allocated.setdefault(seq_id, []).append(block)
    
    def free(self, seq_id):
        for block_id in self.allocated.get(seq_id, []):
            self.free_blocks.append(block_id)
        del self.allocated[seq_id]

四、工程实现细节

4.1 注意力计算的修改

PagedAttention 最 tricky 的部分在 CUDA kernel 层面。标准 multi-head attention 假设 K 和 V 是连续的——[num_tokens, num_heads, head_dim]。但有了分页之后,物理上 K 和 V 的存储是离散的:

# 标准 attention:K 是连续 tensor [total_tokens, num_heads, head_dim]
# PagedAttention:K 是 [num_blocks, block_size, num_heads, head_dim]
#                       其中 block 在物理上不连续

所以 vLLM 自己写了两个 CUDA kernel:

  1. paged_attention_v1:每个 block 单独触发一个 block-level GEMM,然后累加。适合 block 数量少的情况。
  2. paged_attention_v2:先 partial accumulate,再 merge。通过减少 kernel launch 次数来降低 overhead。

实际线上用的是 v2。从 A100 的 nsys profile 结果来看,v2 相比 v1 减少了约 30% 的 kernel launch 时间。

4.2 Prefix Caching(自动前缀缓存)

vLLM 0.4.0 之后引入了 automatic prefix caching。同一个 block 的 KV 如果和之前某个请求的前缀相同,可以直接复用。

请求1: "介绍一下强化学习的基本原理"
请求2: "介绍一下强化学习的应用场景"
              ^ 前缀 token 的 block 是相同的

开启方式:

export VLLM_ENABLE_PREFIX_CACHING=1

实测数据:在 multi-turn conversation 场景下(共享 system prompt),prefix cache hit rate 能达到 60-80%,prefill 阶段的延迟降低约 40%。

4.3 Block 大小的选择

Block size 是 vLLM 的关键超参数。vLLM 默认 16,但这个值的影响很微妙:

  • block_size 越大:Block Table 越小(内存开销低),但内部碎片更多,浪费率更高
  • block_size 越小:碎片率低,但 Block Table 变大,管理开销增加

在 A100 上做过几组 A/B 测试,结论是:

block_size 平均显存利用率 QPS (batch=8) TFLOPS
8 86.2% 1.41 38.2%
16 84.7% 1.48 39.1%
32 78.3% 1.44 37.8%
64 69.5% 1.35 35.1%

block_size=16 是 sweet spot——QPS 最高且显存利用率足够好。

五、接入实战:从 Transformers 迁移到 vLLM

5.1 最小接入代码

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-13b-chat-hf",
    tensor_parallel_size=2,
    gpu_memory_utilization=0.90,
    max_num_seqs=256,
    enable_prefix_caching=True,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=2048,
    stop=["</s>"],
)

outputs = llm.generate(prompts, sampling_params)

5.2 性能压测

在 2×A100-80GB, LLaMA-2-13B 环境下:

指标 Transformers vLLM (block=16) 提升倍数
单请求延迟 (50 token 输出) 1.2s 0.9s 1.33x
Batch=8 吞吐 (token/s) 128 384 3.0x
Batch=32 吞吐 (token/s) 224 1,024 4.57x
最大支持 batch size 12 256 21.3x
KV Cache 利用率 ~45% ~85% 1.89x

5.3 Serving 部署

python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-13b-chat-hf \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.90 \
    --max-num-seqs 128 \
    --port 8000
from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="sk-xxx",
)

response = client.chat.completions.create(
    model="meta-llama/Llama-2-13b-chat-hf",
    messages=[{"role": "user", "content": "解释一下 PagedAttention"}],
    max_tokens=1024,
)

六、生产环境中踩过的坑

坑 1:gpu_memory_utilization 调大不一定好

我把 gpu_memory_utilization 设到 0.95,结果跑了一周,频繁出现 CUDA OOM。

排查后发现:这个参数只控制了 KV Cache 分配的显存上限,但模型跑起来之后,中间激活(activation memory)也是动态的。如果某个请求有很长的 prompt(比如 8K+),中间激活 tensor 会撑爆剩下的那点空间。

安全值是 0.85-0.90,留出 10-15% 给中间激活和 CUDA context。

坑 2:max_num_seqs 不是设得越大越好

把 max_num_seqs 设到 256,结果 QPS 反而下降了。

原因在于 vLLM 的调度策略是 iterate-batch-level scheduling——每个 decode step 都把 batch 里所有 sequence 拿出来一起算。256 个 sequence 虽然不 OOM,但算力分摊开之后,每个 sequence 的延迟从 50ms 涨到了 300ms。

从实测来看,LLaMA-13B 在 A100 的 sweet spot 是 batch_size=64~128 之间。

坑 3:量化的坑

vLLM 支持 AWQ 和 GPTQ 量化模型。我用 AWQ 4bit 量化 LLaMA-13B,模型文件从 26GB 降到了 7.2GB。但精度下降在某些任务上很明显——GSM8K 准确率从 82% 降到 71%,HumanEval pass@1 从 34% 降到 27%。

vLLM 的 AWQ kernel 对 group size 有要求:必须能被 128 整除,且 group size 不能超过 256。如果量化时用了 group_size=32,vLLM 直接报错加载不了。

坑 4:Prefix Caching 的内存开销

开启 prefix caching 后,hash table 本身也吃显存。如果 prompt 几乎都不一样,cache hit 率不到 5%,hash table 反而浪费了空间。这个功能只有在共享前缀比例高的时候才有价值

七、性能调优实践

7.1 调度器参数

llm = LLM(
    ...,
    max_num_batched_tokens=4096,
    max_num_seqs=256,
    scheduler_delay_factor=0.1,
)
  • max_num_batched_tokens 控制 prefill 阶段的 batch 大小。经验值:4096-8192。
  • scheduler_delay_factor 控制调度器"等一等"的意愿。0.1 表示等待时间占 decode iteration 时间的 10%。

7.2 实测调优流程

python -m vllm.entrypoints.openai.run_batch \
    --model meta-llama/Llama-2-13b-chat-hf \
    --input-file requests.jsonl \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.90 \
    --max-num-seqs 128

requests.jsonl 格式:

{"prompt": "Hello, how are you?", "max_tokens": 256, "temperature": 0.7}
{"prompt": "Write a poem about AI", "max_tokens": 512, "temperature": 0.8}

7.3 最终部署配置

模型:Llama-2-13b-chat-hf
硬件:2×A100-80GB (NVLink)
TP:2
gpu_memory_utilization:0.88
max_num_seqs:128
enable_prefix_caching:true
block_size:16
max_num_batched_tokens:6144

实测:
- P50 TTFT:380ms
- P95 TTFT:1.2s
- TPOT:52ms per token
- QPS:约 3.3
- 单卡显存峰值:74.2GB (92.75%)

八、与其它推理框架的对比

特性 vLLM TensorRT-LLM TGI
PagedAttention ✅ 原生
量化支持 AWQ/GPTQ/FP8 AWQ/FP8/INT4 AWQ/GPTQ
调度策略 基于分页调度 静态 batch 动态 batch
OOM 恢复 Preemption
Prefix Caching 有限
易用性 ⭐⭐⭐⭐⭐ ⭐⭐⭐ ⭐⭐⭐⭐

TensorRT-LLM 的优势在于推理速度——kernel 手工优化得更彻底,同样的模型和硬件,通过能达到 vLLM 的 1.1-1.2x。但上手成本高,没有 preemption 机制,显存不够直接崩。

vLLM 胜在工程友好:一键启动、自动调度、自动 prefix cache、graceful OOM 处理。

九、总结与建议

部署 LLM 推理服务的建议顺序:

  1. 先用 vLLM 上线——15 分钟跑起来,稳定性够用
  2. 再加 Prefix Caching——prompt 有共享前缀时提效最高
  3. 再考虑量化——延迟不是瓶颈就不要量化;选 AWQ 4bit
  4. 最后再考虑 TensorRT-LLM——只有需要极致吞吐、愿意花两周调优时才有价值

PagedAttention 给我的最大启发是:AI 系统的瓶颈往往不在算法本身,而在资源管理的粒度上。把 KV Cache 从"连续大块"切到"小页管理",不改变任何数学计算,就带来了几十倍的吞吐提升。这种"计算不变,存储重构"的思路,在 AI 工程化中值得反复使用。

最后留一条建议:不要在生产环境用最新版 vLLM。vLLM 迭代极快,每个 release 都可能引入 regression。我们的做法是锁定一个大版本(比如 0.6.x),小版本只打 patch 不追新,等社区跑稳了再跳版本。

Logo

免费领 100 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐