吞吐量批处理策略

优化单个令牌生成（通过KV缓存和专用注意力机制等方法）能够显著降低每一步的延迟和内存开销，但要最大化LLM推理服务器的总体吞吐量，就需要同时处理多个请求。GPU是深度学习的运算主力，在进行大型矩阵乘法时能达到最佳性能。逐个处理请求通常会导致GPU利用率不足，因为控制逻辑和步骤间的数据传输可能占据单个序列的大部分计算时间。批处理策略通过将多个推理请求组合在一起，让GPU并行处理它们，从而分摊开销并增加每秒生成的令牌数量。

静态批处理

静态批处理是最简单的方法。在这种方法中，推理服务器会等待，直到预设数量的请求（batch_size）到达或发生超时。然后这些请求被分组，填充到批处理中最长序列的长度，并一起处理。

填充是指在批处理中较短的序列后添加特殊令牌，使所有序列具有相同的长度。这会生成GPU上进行高效矩阵乘法所需的统一张量形状。注意力掩码在这里很实用，它能确保模型在自注意力计算期间不会关注这些填充令牌。

import torch

# 示例：准备一个静态批处理
requests = [
    "This is sequence one.",
    "A shorter sequence.",
    "This is the third sequence, and it's quite long."
]

# 假设分词器添加特殊令牌并转换为ID
# 分词后的序列（简化ID）
seq1_ids = [101, 2023, 2003, 5537, 2028, 1012, 102] # len=7
seq2_ids = [101, 1037, 9087, 5537, 1012, 102]      # len=6
seq3_ids = [
    101, 2023, 2003, 1996, 2353, 5537, 1010, 1998, 2009, 1005, 1055,
    3747, 2146, 1012, 102
] # len=15

sequences = [seq1_ids, seq2_ids, seq3_ids]
max_len = max(len(seq) for seq in sequences) # max_len = 15
padding_token_id = 0 # 假设PAD令牌ID为0

padded_sequences = []
attention_masks = []

for seq in sequences:
    pad_len = max_len - len(seq)
    padded_seq = seq + [padding_token_id] * pad_len
    # 屏蔽填充部分
    attention_mask = ([1] * len(seq) +
                      [0] * pad_len)

    padded_sequences.append(padded_seq)
    attention_masks.append(attention_mask)

# 转换为张量作为模型输入
input_ids = torch.tensor(padded_sequences)
attention_mask = torch.tensor(attention_masks)

print("Input IDs Shape:", input_ids.shape) # torch.Size([3, 15])
print("Attention Mask Shape:", attention_mask.shape) # torch.Size([3, 15])

# 现在 input_ids 和 attention_mask 可以作为模型输入
# model(input_ids=input_ids,
#       attention_mask=attention_mask, ...)

静态批处理的局限：

• 低效： 大量计算资源浪费在填充令牌上，尤其当批处理中序列长度差异很大时。批处理时间由最长序列决定。
• 队头阻塞： 较短的请求可能会经历更高的延迟，因为它们需要等待足够多的请求来形成一个完整批处理，或者等待同一批处理中的一个长请求完成。

动态批处理

动态批处理旨在通过在小时间窗内对到达的请求进行分组，而不是等待固定的批处理大小，从而提高GPU利用率，优于静态批处理。这通常会导致批处理中的序列长度差异更大。尽管它仍然需要填充，但服务器可以更频繁地开始处理批处理。

其主要思路是灵活性。服务器不使用固定的batch_size，而是在短时间内（例如10毫秒）积累请求，然后将所有已到达的请求作为一个批处理进行处理。

# 动态批处理的服务器逻辑
import time
import queue

request_queue = queue.Queue()
MAX_WAIT_TIME_MS = 10
MAX_BATCH_SIZE = 16 # 可选的上限

def process_batch(batch):
    # （与静态批处理类似的填充逻辑）
    # ... 分词、填充、创建掩码 ...
    # input_ids, attention_mask = prepare_batch(batch)
    # with torch.no_grad():
    #   outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    # ... 处理输出 ...
    print(f"Processed batch of size {len(batch)}")

while True:
    batch = []
    start_time = time.time()
    while True:
        try:
            # 非阻塞检查新请求
            request = request_queue.get_nowait()
            batch.append(request)
            if len(batch) >= MAX_BATCH_SIZE: # 如果达到最大大小
                break
        except queue.Empty:
            # 检查是否超出等待时间
            if (time.time() - start_time) * 1000 > MAX_WAIT_TIME_MS:
                break
            # 可选的短时间睡眠以避免忙等
            time.sleep(0.001)

    if batch:
        process_batch(batch)
    else:
        # 没有请求，可能睡更长时间
        time.sleep(0.01)

尽管动态批处理优于静态批处理，但它仍然存在填充效率低下的问题，并且在自回归生成过程中，整个批处理的进度由最慢（最长）的序列决定。如果一个序列需要500个令牌而其他序列只需要50个，那么批处理槽位将一直被占用，直到500个令牌的生成完成，这会导致批处理中已完成序列的GPU空闲（通常称为“气泡”）。

连续批处理

连续批处理（也称为迭代级调度或动态拆分融合）是一种更精巧的方法，旨在通过解决简单批处理方法的不足来最大化吞吐量。它将服务器级别的批处理与生成过程的迭代步骤分离。

连续批处理不是处理一个固定批处理直到所有序列都完成，而是按迭代进行操作：

1. 请求池： 维护一个当前正在生成中的活动请求池。
2. 迭代批处理： 在每个生成步骤（即为所有活动序列生成下一个令牌时），从池中选择准备好生成下一个令牌的序列子集。然后从这些序列形成一个批处理。
3. 前向传播： 对本次迭代的批处理执行一次前向传播（一个令牌生成步骤）。
4. 更新池：
   • 将新生成的令牌添加到已处理批处理中的每个序列。
   • 检查是否有任何序列达到其停止条件（例如，EOS令牌、最大长度）。如果达到，则将它们从活动池中移除并返回完成的结果。
   • 检查是否有新的传入请求，如果容量允许，则将其添加到活动池中。
5. 重复： 返回步骤2进行下一次生成迭代。

连续批处理的简化流程。从活动请求中形成迭代批处理，由GPU处理，然后用新令牌、已完成序列和传入请求更新请求池。

主要的优点是，只要有任何活动序列准备好生成，GPU就会保持忙碌。当一个序列完成时，它在下一次迭代批处理中的位置可以立即被池中另一个序列或新到达的请求占用。这避免了静态/动态批处理中GPU等待批处理中最长序列完成的“气泡”问题。Orca、vLLM、TensorRT-LLM和Text Generation Inference (TGI) 等系统都实现了连续批处理的不同形式。

实现考量与权衡

• KV缓存管理： 批处理，特别是连续批处理，需要细致地管理KV缓存。批处理中的每个序列都需要在生成步骤中维护自己独立的KV缓存状态。这些可能数量众多且动态变化的缓存的内存分配和访问效率至关重要。PagedAttention是vLLM等系统中用于更高效管理KV缓存内存，减少碎片化的一种方法。
• 吞吐量与延迟： 批处理在根本上是用增加延迟来换取更高的吞吐量。将请求分组意味着与立即处理相比，某些请求在开始处理前会等待更长时间。然而，整个系统的容量（令牌/秒）会显著增加。批处理策略和参数（例如，最大批处理大小，动态批处理的等待时间）的选择取决于特定应用程序对延迟和吞吐量的要求。更大的批处理大小通常会增加吞吐量，但也会增加平均延迟。
• 框架支持： 深度学习框架和专用服务工具包通常提供批处理实用程序。例如，当你向模型或分词器提供输入列表时，Hugging Face的transformers库会自动处理填充和注意力掩码。专用的推理服务器（Triton、TorchServe、TGI、vLLM）提供更先进的批处理功能，包括动态和连续批处理，并针对生产工作负载进行了优化。

总而言之，批处理对于高效的LLM推理不可或缺。尽管静态批处理简单，但动态批处理，尤其是连续批处理，通过并行处理多个请求并动态管理跨生成步骤的工作负载，在GPU利用率和总体吞吐量方面提供了显著的改进。选择和调整正确的批处理策略是在大规模部署LLM时一个重要考量。