大模型推理引擎高频优化：连续批处理、PagedAttention 与 KV Cache 实战拆解

大模型推理引擎高频优化：连续批处理、PagedAttention 与 KV Cache 实战拆解

当模型参数量以每年 10 倍的速度增长，而 GPU 显存带宽仅提升 1.5 倍时，大模型工程师的核心使命不再是“调参炼丹”，而是在算力与存力的夹缝中榨取硬件的每一丝潜力。本文不聊基础 API 调用，直击 LLM 推理服务的三大命门——显存管理、计算吞吐与延迟毛刺，结合 vLLM、TensorRT-LLM 等主流引擎的底层原理，给出可落地的性能压榨方案。

1. 引言：大模型推理的“内存墙”困境

一个 70B 参数的模型仅权重加载就需要 140GB（FP16）。在 A100（80GB）上单卡无法装载，必须依赖张量并行（TP）。然而，推理中最隐蔽的“显存杀手”并非权重，而是 KV Cache。生成 2048 个 tokens 时，KV Cache 占用可达 ~3GB / 请求。若采用静态 Batching，Padding 导致的浪费常超过 40%。

大模型工程师面临的数学不等式：

总显存≥模型权重+KV Cache+激活值总显存≥模型权重+KV Cache+激活值

我们的优化目标，就是在满足 SLA（首 Token 延迟 < 500ms，Decode 吞吐 > 2000 tokens/s）的前提下，最大化 吞吐量（Throughput） 与 GPU 利用率。

2. 批处理策略进化：从 Static 到 Continuous Batching

2.1 Static Batching 的“木桶效应”

传统静态批处理（如 HuggingFace 默认 Pipeline）必须等 Batch 中所有序列完成 Decode 才能统一返回。由于请求长度分布极不均匀（幂律分布），短序列被迫等待长序列完成，GPU 算力空转严重。更致命的是，为对齐张量形状，大量 <PAD> Token 参与矩阵运算，浪费了 30%~50% 的算力。

2.2 Continuous Batching（又名 In-flight Batching）

主流引擎（vLLM、TensorRT-LLM、TGI）已全面转向迭代级调度。核心逻辑是：每个 Decode 迭代结束后，动态检查是否有请求生成完毕（遇到 EOS）或新请求到达，立即插入空位继续计算。

调度伪代码：

class Scheduler:
    def schedule(self, waiting_queue, running_requests):
        # 1. 移除已完成的请求，释放 KV Cache 块
        self.running = [r for r in self.running if not r.finished]
        
        # 2. 计算剩余可用显存块 (Free Slots)
        free_blocks = self.kv_cache_manager.get_free_blocks()
        
        # 3. 从等待队列中尽可能多地调度新请求（受限于 free_blocks）
        while waiting_queue and free_blocks > 0:
            req = waiting_queue.pop(0)
            if req.required_blocks <= free_blocks:
                self.running.append(req)
                free_blocks -= req.required_blocks
            else:
                break  # 显存不足，暂停调度
        return self.running

收益：在某金融问答场景中，我们将 Max Batch Size 从静态的 8 提升至动态平均 24，吞吐量从 120 tokens/s 跃升至 450 tokens/s，且 P99 延迟降低了 40%（避免了尾部等待）。

3. 显存管理的“操作系统级”革命：PagedAttention

3.1 KV Cache 的碎片化灾难

PyTorch 的传统显存分配器（Cache Allocator）在处理变长 KV Cache 时会产生大量内存碎片。即便总显存有剩余，因无法找到连续大块内存，OOM（Out of Memory）依然频发。

3.2 虚拟内存式分页管理

PagedAttention（vLLM 的核心技术）借鉴操作系统虚拟内存思想，将 KV Cache 切分为固定大小的 Page Blocks（默认 16 tokens/block）。

• 逻辑块：每个请求维护一个逻辑块列表。
• 物理块：GPU 显存被划分为等大的物理块池。

物理块无需连续存储，通过块表（Block Table） 映射逻辑顺序。当新生成 Token 时，仅动态申请一个物理块（若当前块写满）。这彻底解决了外部碎片问题，且支持同一物理块在多个请求间共享（如系统 Prompt 前缀缓存）。

实测内存节省：在 ShareGPT 数据集上，PagedAttention 相比 vLLM 之前的 PyTorch 实现，显存占用减少了 54%，吞吐量提升 22 倍（在 2.4K 并发请求下）。

3.3 代码层窥探（Block 管理）

# 简化自 vLLM core/block_manager.py
class BlockSpaceManager:
    def allocate(self, request_id, num_tokens):
        num_required_blocks = ceil_div(num_tokens, self.block_size)
        # 从空闲列表取出物理块
        physical_blocks = [self.free_blocks.pop() for _ in range(num_required_blocks)]
        self.block_tables[request_id] = physical_blocks
        return physical_blocks
    
    def append_token(self, request_id, token_id):
        last_block = self.block_tables[request_id][-1]
        if self._is_block_full(last_block):
            new_block = self.free_blocks.pop()
            self.block_tables[request_id].append(new_block)
            # 写入新块
        else:
            # 写入当前块末尾

4. 权重量化与计算加速：FP8 与 AWQ 的工程取舍

为了放下 70B+ 的大模型并加速 Decode（内存带宽瓶颈），量化是必修课。

实战坑点：AWQ 在 Text Generation 任务上表现优异，但当 Batch Size 增大时，反量化（Dequantization）开销会抵消权值节省的带宽收益。建议：若 Batch Size > 64，使用 INT8 往往比 INT4 更快（因为 INT8 可用 Tensor Core 直接计算，INT4 需额外解包指令）。

量化启动示例（vLLM）：

python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-hf \
    --quantization awq \
    --tensor-parallel-size 2 \
    --max-model-len 4096 \
    --block-size 16

5. 投机采样（Speculative Decoding）：打破自回归的“奈奎斯特极限”

自回归生成（Autoregressive）的硬伤在于内存带宽主导（Memory Bound），算力利用率极低。投机采样（又称辅助生成）利用小模型（Draft Model）快速生成 5~6 个候选 Token，再由大模型（Target Model）并行验证。

5.1 数学加速原理

假设大模型验证一次的时间等于小模型生成 K 个 Token 的时间，且接受率为 αα，则加速比近似为：

S≈1+α+α2+...+αK−11+K⋅TdraftTtargetS≈1+K⋅Ttarget​Tdraft​​1+α+α2+...+αK−1​

当 Tdraft≪TtargetTdraft​≪Ttarget​（例如用 68M 的草稿模型辅助 7B 大模型），且 αα 约 0.7 时，Decode 速度可提升 2.5~3 倍。

5.2 工程集成（HuggingFace Assisted Generation）

from transformers import AutoModelForCausalLM, AutoTokenizer

assistant_model = AutoModelForCausalLM.from_pretrained("distilbert/gpt2") # 小模型
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b")

outputs = model.generate(
    input_ids,
    assistant_model=assistant_model,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
    num_assistant_tokens=5  # 每次草稿生成 5 个候选
)

注意：投机采样在 贪婪解码（Greedy） 下加速最明显，且必须保证草稿模型与目标模型共用 Tokenizer，否则验证阶段会出现 Token 对齐错误。

6. 分布式推理：张量并行（TP）与流水线并行（PP）的抉择

对于 70B+ 模型，单卡绝对装不下。我们在 Megatron-LM 和 vLLM 中常见的并行策略：

• 张量并行（TP）：将 Attention Head 和 FFN 的权重切分到多卡。优点：无跨机通信延迟，单次 Forward 极快。缺点：All-Reduce 通信量大（每层都要同步），跨机（> 8卡）时通信带宽（NVLink vs PCIe）成为瓶颈。
• 流水线并行（PP）：按层切分，不同 GPU 负责不同 Layer。优点：通信量极小（仅传递中间激活值）。缺点：存在“气泡”（Bubble）空闲时间，且需要 Micro-batch 调度。

生产建议：

• 单机（8×A100）：优先使用 TP=8，利用 NVLink 全互联。
• 跨机（多节点）：采用 TP=4 + PP=2 混合并行。vLLM 0.6.0+ 已支持 --pipeline-parallel-size 参数。
• 通信优化：设置 NCCL_ALGO=Tree 并启用 NVSwitch 亲和性绑定。

7. 可观测性与优雅降级：SLA 的最后一道防线

大模型推理的“长尾延迟”（Long-tail Latency）非常讨厌。我们通过 FastAPI + Prometheus + Grafana 监控以下核心指标：

1. TTFT (Time To First Token)：用户感知卡顿的关键，必须 < 500ms。
2. TPOT (Time Per Output Token)：生成每个 Token 的间隔，稳定在 < 30ms 为佳。
3. KV Cache 使用率：超过 85% 时触发 HPA（Horizontal Pod Autoscaler）扩容。

降级策略（当 GPU 排队请求 > 阈值时）：

• 自动将 max_tokens 限制减半（从 2048 降到 1024）。
• 禁用投机采样（降低计算负载）。
• 将采样温度强制设为 0（贪婪解码，减少计算变数）。

8. 实战 Benchmark：我们在 Llama-3-70B 上的压测数据

环境：2 台 AWS p4d.24xlarge（8×A100 80GB each），网络 400 Gbps EFA。

结论：组合优化下，吞吐量提升近 10 倍，且显存利用率稳定在 95% 以上。

9. 总结与未来趋势

大模型工程师的职责边界正在模糊——我们既要懂 CUDA Kernel 的 Launch 参数，又要懂 Kubernetes 的弹性调度。本文覆盖的 Continuous Batching、PagedAttention、量化与投机采样，已是生产环境的标准标配。

接下来的技术浪潮：

• SSM（如 Mamba） 架构对 Attention 的替代，将彻底改变 KV Cache 的管理范式。
• 异构计算（CPU + GPU）卸载部分 Layer，支持无限长上下文。
• 端到端编译（如 Mojo + MLIR）将取代手动算子融合，让大模型自动寻找最优计算图。

AI 大模型工程师，本质上是在不确定性（模型黑盒）中构建确定性（SLA 保障）的工匠。希望这份实战笔记能为你跳过一些深坑。