LLM 推理工程极限：Continuous Batching 与 GPU 利用率优化

背景：为什么 LLM 推理这么难优化

大语言模型推理和传统深度学习推理有个本质区别：输入输出长度不固定。一个请求可能只输出 20 个 token，另一个请求要输出 2000 个 token。如果简单按请求排队，等短请求的人被长请求堵死，GPU 利用率惨不忍睹。

2023 年之前，主流做法是静态分桶（Static Batching）：把请求按长度分组，每批凑满固定长度再处理。效果差强人意。2024 年，Orca 论文引入了迭代级调度（Iteration-Level Scheduling），后来的 vLLM 把它实现为 Continuous Batching，GPU 利用率直接上一个台阶。

今天我们来深度拆解这套机制，配合 PyTorch 代码示例，搞清楚它为什么有效。

1. 经典 Static Batching 的瓶颈

先看静态分批的问题。假设我们有 4 个请求，长度各不相同：

`

Request A: 512 input + 64 output

Request B: 512 input + 128 output

Request C: 1024 input + 32 output

Request D: 1024 input + 256 output

`

静态分批要把它们凑成一批处理，必须 padding 到最大长度：

`

Batch: 1024 (max) input + 256 (max) output = 1280 tokens per iteration

`

每个请求实际有用的 token 比例很低，大部分计算是无效的 padding。更糟糕的是，Request A 在 64 步后就可以返回了，但它必须等 Batch 里最慢的 D 跑完 256 步。这是 死等问题（Head-of-Line Blocking）。

`

Timeline (iteration steps):

Step 1-64: [A B C D] 四路并行

Step 65-128: [ B C D ] A 结束但占着位置

Step 129-256:[ C D ] B 结束但占着位置

Step 257-320:[ D ] 只有 D 在跑

`

GPU 利用率曲线就是一个波峰然后漫长的低谷。

2. Continuous Batching：迭代级调度

Continuous Batching 的核心思想：不再等待整个批次完成才加入新请求，而是在每个 generation step（每个 token 生成）结束后，检查是否有请求已完成，如果有就立即移出，同时把新请求加进来。

这是迭代级调度，不是请求级调度。调度发生在每个生成步，而不是每个请求结束时。

2.1 调度循环

`python

import torch

from typing import List, Dict

class ContinuousBatchingScheduler:

"""

迭代级调度器：每个 forward pass 后决定谁能进谁要出

"""

def __init__(self, max_batch_size: int, device: str = "cuda"):

self.max_batch_size = max_batch_size

self.device = device

# running 表示正在生成的请求

self.running: List[GenerationRequest] = []

def step(self,

requests: List[GenerationRequest],

logits: torch.Tensor) -> List[int]:

"""

每个生成步的调度决策

返回: 已完成请求的 indices，要从批次中移除

"""

finished = []

for i, req in enumerate(self.running):

# 计算这个请求在此步的采样

next_token = self.sample(logits[i], req.temperature)

req.generated_tokens.append(next_token)

# 检查是否结束

if self.is_finished(req, next_token):

finished.append(i)

# 核心操作：移除已完成的，加入等待的

# 这行是 Continuous Batching 的关键所在

self._evict_and_fill(finished, requests)

return finished

def _evict_and_fill(self,

finished_indices: List[int],

pending_requests: List[GenerationRequest]):

"""移除完成的，填入新请求"""

# 从后往前删，避免 index 错位

for i in sorted(finished_indices, reverse=True):

self.running.pop(i)

# 填入新请求直到满批

while (len(self.running) < self.max_batch_size

and pending_requests):

self.running.append(pending_requests.pop(0))

`

2.2 GPU 时间线对比

Continuous Batching 让 GPU 保持高利用率：

`

Static Batching:

[======ABCD======] [==EFGH==]

busy idle gap busy

Continuous Batching:

[==A==] [==E==]

[==B==] [==AB==] [==ABC==] [==FG==] ...

[==C==] [==CD==] [==BCD==] [==GH==]

[==D==] [==D===] [==D====]

`

从四路并行逐渐变成一路，然后立刻有新请求填进来。没有长空闲。

3. PagedAttention：vLLM 的内存革命

Continuous Batching 解决了调度问题，但还有一个瓶颈：KV Cache 内存管理。

Attention 计算需要存储 Key 和 Value 向量。对于一个 4096 token 上下文、70B 参数的模型，单个请求的 KV Cache 就能达到：

`

hidden_size = 8192 # 70B 模型

num_heads = 64

head_dim = 128

kv_cache_per_token = 2 * num_heads * head_dim * 2bytes(fp16)

= 2 * 64 * 128 * 2 = 32KB per token

4096 tokens * 32KB = 128MB per request

`

如果同时跑 100 个请求，仅 KV Cache 就要 12.8GB，加上模型权重 140GB... GPU 显存根本装不下。

vLLM 提出的 PagedAttention 灵感来自操作系统的分页管理：不再一次性分配一大块连续显存给 KV Cache，而是按 block 分页，按需分配。

3.1 Block 管理逻辑

`python

from dataclasses import dataclass

from typing import Dict, List

@dataclass

class KVCacheBlock:

"""KV Cache 物理块，类似 OS 的内存页"""

block_id: int

num_slots: int = 16 # 每个 block 16 个 token slot

num_free_slots: int = 16

# 物理显存指针

k_ptr: int = 0

v_ptr: int = 0

class KVCacheManager:

"""

类 OS 页表的 KV Cache 管理器

逻辑块 -> 物理块的映射，允许非连续存储

"""

def __init__(self, total_blocks: int, block_slots: int):

self.total_blocks = total_blocks

self.block_slots = block_slots

# 逻辑块到物理块的映射表

self.block_tables: Dict[int, List[int]] = {}

# 物理块分配状态

self.physical_blocks: List[KVCacheBlock] = [

KVCacheBlock(block_id=i, num_slots=block_slots)

for i in range(total_blocks)

]

self.free_blocks = set(range(total_blocks))

def allocate(self, num_tokens: int) -> List[int]:

"""为新请求分配物理块，返回块 ID 列表"""

num_blocks_needed = (num_tokens + self.block_slots - 1) // self.block_slots

allocated = []

for _ in range(num_blocks_needed):

if not self.free_blocks:

raise RuntimeError("KV Cache 内存耗尽，需要 evict 策略")

block_id = self.free_blocks.pop()

allocated.append(block_id)

return allocated

def free(self, block_list: List[int]):

"""释放物理块"""

for b in block_list:

self.free_blocks.add(b)

class GenerationRequest:

"""请求对应的逻辑块序列"""

def __init__(self, request_id: int, prompt_tokens: List[int]):

self.request_id = request_id

self.prompt_tokens = prompt_tokens

self.generated_tokens: List[int] = []

self.block_ids: List[int] = [] # 逻辑块序列 -> 物理块

self.num_generated: int = 0

`

3.2 PagedAttention Kernel

`python

# 简化的 PagedAttention 计算逻辑

# 实际实现需要用 CUDA C++ 或 Triton's custom kernel

def paged_attention(

query: torch.Tensor, # [batch, num_heads, seq_len, head_dim]

key_cache: torch.Tensor, # [num_blocks, num_heads, block_size, head_dim]

block_tables: List[List[int]], # 每请求的物理块列表

seq_lens: List[int]

) -> torch.Tensor:

"""

query: 当前 step 的查询向量

key_cache: 按 block 存储的 KV Cache

block_tables: 逻辑位置 -> 物理块的映射

"""

output = torch.zeros_like(query)

for batch_idx in range(query.shape[0]):

seq_len = seq_lens[batch_idx]

blocks = block_tables[batch_idx]

block_size = key_cache.shape[2]

# 把逻辑序列映射到物理块

num_blocks = len(blocks)

key_states = torch.zeros(

seq_len, query.shape[2], query.shape[3],

device=query.device, dtype=query.dtype

)

block_offset = 0

for phys_block_id in blocks:

# 从物理块读取数据（可以非连续）

phys_block = key_cache[phys_block_id]

copy_len = min(block_size, seq_len - block_offset)

key_states[block_offset:block_offset+copy_len] = phys_block[:copy_len]

block_offset += copy_len

# 标准 attention 计算

attn_weights = torch.matmul(query[batch_idx], key_states.transpose(-2, -1))

attn_weights = attn_weights / (query.shape[3] ** 0.5)

attn_weights = torch.softmax(attn_weights, dim=-1)

output[batch_idx] = torch.matmul(attn_weights, key_states)

return output

`

关键洞察：物理块可以非连续分配，但逻辑上是连续的。这解决了传统方案"一次性分配大连续显存"的浪费问题。

vLLM 的实测数据：比 HuggingFace TF 提升2-3 倍吞吐量，比 Text Generation Inference (TGI) 提升 1.5-2 倍。

4. Prefill 阶段与 Decode 阶段的差异优化

LLM 推理分两个阶段：

• **Prefill**：处理输入 prompt，一次性计算 attention（是并行计算，适合大 batch）
• **Decode**：逐 token 生成（是自回归的，每步只处理一个 token，attention 计算量小但调度开销大）

两个阶段的特性差异巨大，混合在一起处理会互相影响。生产系统通常分离 Prefill 和 Decode 节点，或者使用 disaggregated prefill/decode 架构。

`python

class DisaggScheduler:

"""

分离式预填充/解码调度器

Prefill 节点处理新请求的 prompt

Decode 节点处理自回归生成

"""

def __init__(self, prefill_nodes: int, decode_nodes: int):

self.prefill_servers = prefill_nodes

self.decode_servers = decode_nodes

self.prefill_queue: List[GenerationRequest] = []

self.decode_queue: List[GenerationRequest] = []

def add_request(self, req: GenerationRequest):

# 新请求先进 Prefill 处理 prompt

self.prefill_queue.append(req)

def step(self):

# Prefill 完成后转 Decode 队列

finished_prefill = self._drain_prefill()

self.decode_queue.extend(finished_prefill)

# Decode 节点处理生成

self._process_decode()

def _drain_prefill(self) -> List[GenerationRequest]:

# Prefill 是一次性计算，可以大 batch

batch = self.prefill_queue[:self.prefill_batch_size]

# ... 调用 Prefill 节点

return [r for r in batch if r.is_prefill_done()]

def _process_decode(self):

# Decode 逐 token，需要 Continuous Batching

# ... Continuous Batching 调度逻辑

`

5. 实战调优：让 GPU 利用率从 40% 到 85%

光有算法不够，还要知道怎么调参。以下是我们在 8x H100 集群上的调优经验：

5.1 Batch Size 设置

`

公式：max_batch_size ≈ GPU显存 / (model_params * 2 + kv_cache_per_token * max_seq_len)

`

对于 70B FP16 模型（140GB），8x H100 80GB：

• 权重：140GB
• 激活值 + Attention 输出：约 20GB
• KV Cache 可用：640 - 160 = 480GB
• 每请求 4096 上下文 KV Cache：4096 * 32KB = 128MB
• max_batch_size ≈ 480GB / 128MB ≈ 3600 并发请求（理论值）

实际因为调度开销和碎片化，建议设置 256-512。

5.2 CUDA Graph 优化

每步执行 kernel 调用有开销。可以用 CUDA Graph 捕获一系列操作，一次性提交：

`python

# 使用 torch.cuda.CUDAGraph 减少 kernel 启动开销

graph = torch.cuda.CUDAGraph()

# 第一次运行 - 捕获

with torch.cuda.graph(graph):

static_input = torch.empty_like(dynamic_input)

static_output = model.forward(static_input)

# 后续每次只需 replay，避免 kernel 调度开销

dynamic_input.copy_(real_input)

graph.replay()

output.copy_(static_output)

`

实测可减少 10-15% 的调度 overhead。

5.3 量化配合

INT8 量化可以在精度损失可接受的情况下大幅减少显存占用：

`python

# 使用 BitsAndBytes 进行 INT8 量化

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(

load_in_8bit=True,

llm_int8_threshold=6.0, # outlier 阈值

llm_int8_has_fp16_weight=False

)

model = AutoModelForCausalLM.from_pretrained(

"meta-llama/Llama-3-70b",

quantization_config=bnb_config,

device_map="auto"

)

`

70B FP16 → 70B INT8，显存从 140GB 降到约 70GB，可以跑更大的 batch size。

6. 总结：推理优化的几个层次

GPU 利用率从 40% 到 85%，不是靠某一个 trick，而是这一套组合拳打出来的。每一层都有独立的paper和开源实现可以深入研究。

核心教训：LLM 推理不是"把模型跑起来"这么简单，它是一个系统工程问题。调度、内存、计算、架构四个层次缺一不可，你的瓶颈在哪一层，决定了你该优先优化什么。

---

*调优环境：8x NVIDIA H100 80GB, PyTorch 2.4, vLLM 0.6, Llama-3-70B*