批处理和调度策略
批处理是提高LLM推理吞吐量的关键技术。不同的批处理策略适用于不同的场景。
2.7.1 Static Batching (静态批处理)
最简单的批处理方式,所有请求一起开始、一起结束。
工作流程
特点
| 特性 | 描述 |
|---|---|
| 实现复杂度 | 低 |
| 吞吐量 | 低(受最长请求限制) |
| GPU利用率 | 低(尾部空闲) |
| 适用场景 | 简单原型、请求长度相近 |
2.7.2 Dynamic Batching (动态批处理)
动态批处理允许新请求加入正在进行的批次,但实现复杂。
工作流程
latex
┌─────────────────────────────────────────────────────────────────┐
│ Dynamic Batching │
├─────────────────────────────────────────────────────────────────┤
│ │
│ Time 0: Batch [Req1, Req2, Req3] 开始 │
│ │
│ Time t1: Req1完成,新请求Req4加入 │
│ Batch [Req2, Req3, Req4] │
│ │
│ Time t2: Req2完成,新请求Req5加入 │
│ Batch [Req3, Req4, Req5] │
│ │
│ 挑战: │
- 不同请求的KV Cache长度不同 │
│ - 需要动态内存管理 │
│ - 实现复杂度高 │
│ │
└─────────────────────────────────────────────────────────────────┘实现挑战
python
# 动态批处理的核心挑战:变长序列处理
def dynamic_batch_forward(model, batch_requests):
"""
batch_requests: 不同长度的请求列表
"""
# 1. 找到最大长度
max_len = max(len(r['input_ids']) for r in batch_requests)
# 2. Padding 到最大长度
padded_inputs = []
attention_masks = []
for req in batch_requests:
padding_length = max_len - len(req['input_ids'])
padded = req['input_ids'] + [PAD_TOKEN] * padding_length
mask = [1] * len(req['input_ids']) + [0] * padding_length
padded_inputs.append(padded)
attention_masks.append(mask)
# 3. 前向传播(浪费计算在padding上)
outputs = model(torch.tensor(padded_inputs),
attention_mask=torch.tensor(attention_masks))
return outputs2.7.3 Continuous Batching (连续批处理)
Continuous Batching (也称为 Inflight Batching或 )是目前最先进的批处理技术。
核心思想
在每个iteration(生成一个token)的边界进行调度,而非请求级别:
vLLM的Continuous Batching实现
vLLM通过PagedAttention实现高效的Continuous Batching:
python
class ContinuousBatchingScheduler:
def __init__(self, model, max_batch_size, max_seq_len):
self.model = model
self.max_batch_size = max_batch_size
self.max_seq_len = max_seq_len
self.waiting_queue = [] # 等待队列
self.running_batch = [] # 运行中的批次
def schedule(self):
"""每个iteration调用一次"""
# 1. 检查完成的请求
completed = [req for req in self.running_batch if req.is_done()]
self.running_batch = [req for req in self.running_batch if not req
# 2. 尝试添加新请求
while (len(self.running_batch) < self.max_batch_size and
self.waiting_queue and
self.can_allocate(self.waiting_queue[0])):
new_req = self.waiting_queue.pop(0)
self.running_batch.append(new_req)
self.allocate_kv_cache(new_req)
# 3. 执行前向传播
if self.running_batch:
outputs = self.model.forward(self.running_batch)
self.update_requests(outputs)
return completed
def can_allocate(self, request):
"""检查是否有足够的KV Cache空间"""
# 使用分页内存管理检查
return self.kv_cache_manager.has_space(request)2.7.4 三种批处理策略对比
| 特性 | Static Batching | Dynamic Batching | Continuous Batching |
|---|---|---|---|
| 调度粒度 | 请求级别 | 请求级别 | Token级别 |
| GPU利用率 | 低 | 中 | 高 |
| 实现复杂度 | 低 | 中 | 高 |
| 吞吐量 | 基准 | 1.5-2× | 5-10× |
| 延迟稳定性 | 差 | 中 | 好 |
| 内存效率 | 低 | 中 | 高(配合分页) |
| 代表系统 | 简单实现 | TensorRT-LLM | vLLM, TGI |
2.7.5 调度策略
FCFS(First Come First Serve)
最简单的调度策略,按到达顺序处理:
python
def fcfs_schedule(waiting_queue):
""" 先来先服务"""
return waiting_queue.pop(0)优点:实现简单、公平
缺点:不考虑请求特性,可能导致长请求阻塞短请求
Shortest Job First (SJF)
优先处理预计完成时间短的请求:
python
def sjf_schedule(waiting_queue):
""" 短作业优先"""
# 按预估生成长度排序
waiting_queue.sort(key=lambda r: r.expected_output_length)
return waiting_queue.pop(0)优点:平均等待时间短
缺点:长请求可能饿死、需要预估输出长度
基于优先级的调度
为不同优先级请求分配不同资源:
python
def priority_schedule(waiting_queue, running_batch):
"""优先级调度"""
# 高优先级请求可以抢占低优先级请求
high_priority = [r for r in waiting_queue if r.priority == 'high']
if high_priority:
# 尝试抢占资源
for req in running_batch:
if req.priority == 'low' and req.can_preempt():
preempt_request(req)
return high_priority[0]
return waiting_queue.pop(0)2.7.6 性能数据对比
在A100上测试7B模型的推理性能:
| 批处理策略 | Batch Size | 吞吐量 (tokens/s) | 平均延迟 (ms/token) |
|---|---|---|---|
| 无批处理 | 1 | 15 | 67 |
| Static | 8 | 80 | 100 |
| Static | 32 | 200 | 160 |
| Dynamic | 8 | 120 | 67 |
| Dynamic | 32 | 350 | 91 |
| Continuous | 动态 | 800+ | 40 |