Batched Speculative Decoding：在Batch推理场景下应用投机采样的复杂调度

Batched Speculative Decoding：在Batch推理场景下应用投机采样的复杂调度

大家好，今天我们来深入探讨一个前沿的LLM推理加速技术——Batched Speculative Decoding。投机采样 (Speculative Decoding) 已经成为加速LLM推理的热门方法，它通过引入一个小的“草稿模型 (Draft Model)”来预测多个后续token，然后用一个大的“目标模型 (Target Model)”来并行验证这些预测。如果预测正确，则可以显著减少Target Model的调用次数，从而加速推理。

然而，在实际应用中，尤其是在高吞吐量的Batch推理场景下，如何高效地调度和管理这些投机采样过程，以最大化加速效果，是一个具有挑战性的问题。这就是我们要讨论的重点：Batched Speculative Decoding中的复杂调度。

1. 投机采样 (Speculative Decoding) 基础回顾

为了更好地理解Batched Speculative Decoding，我们首先回顾一下其核心思想。传统的自回归解码过程是串行的，每次只能生成一个token，这限制了推理速度。投机采样通过以下步骤来加速这个过程：

1. 草稿生成 (Drafting)： 使用一个较小的、速度较快的Draft Model，基于当前已生成的token序列，生成一个候选token序列 (例如，生成 k 个 token)。
2. 验证 (Verification)： 将 Draft Model 生成的候选序列与当前序列拼接，然后使用 Target Model 一次性计算所有候选 token 的概率。
3. 接受与拒绝 (Acceptance & Rejection)： 根据 Target Model 的概率分布，决定接受或拒绝 Draft Model 生成的 token。如果 Draft Model 的预测与 Target Model 的预测足够接近，则接受该 token。
4. 更新状态 (State Update)： 将接受的 token 添加到已生成序列中，并重复上述过程。

以下是一个简单的投机采样伪代码：

def speculative_decode(target_model, draft_model, prompt, k):
  """
  投机解码算法.

  Args:
    target_model: 目标模型.
    draft_model: 草稿模型.
    prompt: 输入提示.
    k: 草稿模型预测的token数量.

  Returns:
    生成的文本序列.
  """

  generated_tokens = list(prompt)  # 将 prompt 转换为 token 列表
  while not stop_condition(generated_tokens):
    # 1. 草稿生成
    draft_tokens = draft_model.generate(generated_tokens, max_length=k)

    # 2. 验证
    combined_tokens = generated_tokens + draft_tokens
    target_logprobs = target_model.get_logprobs(combined_tokens)

    # 3. 接受与拒绝
    accepted_tokens = []
    for i in range(len(draft_tokens)):
      # 使用 Target Model 和 Draft Model 的概率分布来决定是否接受 token
      if accept_token(target_logprobs[len(generated_tokens) + i],
                       draft_model.get_logprobs(combined_tokens[:len(generated_tokens) + i+1])[-1]):
        accepted_tokens.append(draft_tokens[i])
      else:
        break

    # 4. 更新状态
    generated_tokens.extend(accepted_tokens)

    # 如果没有接受任何 token, 则必须从 Target Model 生成一个 token
    if not accepted_tokens:
      next_token = target_model.generate(generated_tokens, max_length=1)[0]
      generated_tokens.append(next_token)

  return detokenize(generated_tokens)  # 将 token 列表转换为文本字符串

def accept_token(target_logprob, draft_logprob):
    """
    判断是否接受草稿模型生成的token. 这里只是一个简化的例子，
    实际应用中会更复杂，需要考虑temperature等参数.
    """
    import random
    acceptance_probability = min(1.0, math.exp(target_logprob - draft_logprob))
    return random.random() < acceptance_probability

def stop_condition(generated_tokens):
  # 定义停止生成的条件，例如达到最大长度或生成了结束符
  return len(generated_tokens) > MAX_LENGTH or generated_tokens[-1] == END_TOKEN

2. Batch 推理的挑战与机遇

Batch 推理是指同时处理多个独立的推理请求。与单次推理相比，Batch 推理可以显著提高硬件利用率，从而提高吞吐量。然而，将投机采样应用到 Batch 推理中，会带来新的挑战：

• 长度不一致： 不同的输入 prompt 可能会生成不同长度的 token 序列。这使得在一个 Batch 中有效地管理和调度投机采样过程变得复杂。
• 计算资源分配： 如何在不同的推理请求之间合理地分配计算资源，以最大化整体吞吐量，是一个关键问题。
• 同步问题： 由于每个请求的接受/拒绝 token 数量不同，需要仔细处理Batch内请求之间的同步问题，避免资源浪费和死锁。

尽管存在这些挑战，Batched Speculative Decoding 也带来了巨大的机遇：

• 更高的硬件利用率： 通过将多个投机采样过程合并到一个 Batch 中，可以更充分地利用 GPU 等硬件资源。
• 更高的吞吐量： 优化的调度策略可以显著提高整体吞吐量，满足大规模应用的需求。

3. Batched Speculative Decoding 的调度策略

Batched Speculative Decoding 的核心在于如何有效地调度和管理多个投机采样过程。以下是一些常用的调度策略：

• 静态 Batching： 在开始推理之前，将固定数量的请求组成一个 Batch。这种方法简单易行，但可能无法充分利用硬件资源，因为不同请求的计算复杂度可能差异很大。
• 动态 Batching： 根据当前系统的负载和请求的特征，动态地调整 Batch 的大小。这种方法可以更好地适应不同的 workload，但需要更复杂的调度算法。
• 优先队列调度： 维护一个请求队列，并根据优先级（例如，请求的延迟要求）来调度请求。这种方法可以保证高优先级请求的及时处理。

更进一步，考虑到投机采样自身的特性，可以设计更精细的调度策略：

• 基于接受率的调度： 监控每个请求的接受率（即，Draft Model 预测的 token 被 Target Model 接受的比例）。如果一个请求的接受率较低，则可以减少其投机采样的次数，从而将更多资源分配给接受率较高的请求。
• 基于 token 长度的调度： 根据当前已生成的 token 序列的长度来调整投机采样的次数。对于较短的序列，可以增加投机采样的次数，以加速生成过程。

4. 复杂调度算法示例

下面我们给出一个基于动态 Batching 和接受率的 Batched Speculative Decoding 调度算法的示例。为了简化起见，我们假设有一个固定的 Batch 大小，并且所有请求的优先级相同。

import torch
import math
import random

class BatchedSpeculativeDecoder:
    def __init__(self, target_model, draft_model, batch_size, k):
        """
        Batched Speculative Decoding 的实现.

        Args:
            target_model: 目标模型.
            draft_model: 草稿模型.
            batch_size: Batch 大小.
            k: 草稿模型预测的token数量.
        """
        self.target_model = target_model
        self.draft_model = draft_model
        self.batch_size = batch_size
        self.k = k
        self.active_requests = [] # 存储正在处理的请求
        self.completed_requests = [] # 存储已经完成的请求

    def add_request(self, prompt):
        """
        添加一个新的推理请求.

        Args:
            prompt: 输入提示.
        """
        self.active_requests.append({
            'prompt': prompt,
            'generated_tokens': list(prompt), # 将 prompt 转换为 token 列表
            'acceptance_rate': 1.0, # 初始接受率
            'num_speculative_steps': 0 # 投机采样的次数
        })

    def step(self):
        """
        执行一个推理步骤.
        """

        # 1. 创建 Batch
        batch = self.active_requests[:self.batch_size] # 取前 batch_size 个请求
        if not batch:
            return # 如果没有请求，则返回

        # 2. 草稿生成
        draft_tokens_list = []
        for request in batch:
            draft_tokens = self.draft_model.generate(request['generated_tokens'], max_length=self.k)
            draft_tokens_list.append(draft_tokens)

        # 3. 验证 (Batch 方式)
        combined_tokens_list = [request['generated_tokens'] + draft_tokens_list[i] for i, request in enumerate(batch)]
        max_len = max(len(tokens) for tokens in combined_tokens_list)

        # padding 使长度一致
        padded_combined_tokens_list = [tokens + [0] * (max_len - len(tokens)) for tokens in combined_tokens_list]

        # 将 list of list 转换为 tensor
        combined_tokens_tensor = torch.tensor(padded_combined_tokens_list)

        target_logprobs = self.target_model.get_logprobs(combined_tokens_tensor) #形状为 [batch_size, max_len, vocab_size]

        # 4. 接受与拒绝
        for i, request in enumerate(batch):
            accepted_tokens = []
            num_accepted = 0
            for j in range(len(draft_tokens_list[i])):
                target_logprob = target_logprobs[i, len(request['generated_tokens']) + j] #形状为[vocab_size]
                draft_logprob = self.draft_model.get_logprobs(combined_tokens_list[i][:len(request['generated_tokens']) + j+1])[-1]

                if accept_token(target_logprob, draft_logprob):
                    accepted_tokens.append(draft_tokens_list[i][j])
                    num_accepted += 1
                else:
                    break

            # 5. 更新状态
            request['generated_tokens'].extend(accepted_tokens)
            request['num_speculative_steps'] += 1

            # 更新接受率 (使用滑动平均)
            request['acceptance_rate'] = 0.9 * request['acceptance_rate'] + 0.1 * (num_accepted / self.k if self.k > 0 else 0)

            # 如果没有接受任何 token, 则必须从 Target Model 生成一个 token
            if not accepted_tokens:
                next_token = self.target_model.generate(request['generated_tokens'], max_length=1)[0]
                request['generated_tokens'].append(next_token)

            # 检查是否完成
            if stop_condition(request['generated_tokens']):
                self.completed_requests.append(request)
                self.active_requests.remove(request)

        # 6. 动态调整 Batch (简单示例)  --  这里可以根据 acceptance_rate 调整 batch 成员
        # 例如，将 acceptance_rate 低的请求移到后面，优先处理 acceptance_rate 高的请求
        self.active_requests.sort(key=lambda x: x['acceptance_rate'], reverse=True)

    def get_results(self):
        """
        获取已完成请求的结果.
        """
        return [detokenize(request['generated_tokens']) for request in self.completed_requests]

def accept_token(target_logprob, draft_logprob):
    """
    判断是否接受草稿模型生成的token. 这里只是一个简化的例子，
    实际应用中会更复杂，需要考虑temperature等参数.
    """
    #  target_logprob 是一个向量，我们需要提取对应 token 的 logprob
    #  假设 draft_logprob 对应的是 token 在 target_logprob 中的索引
    #  这个假设在实际中是不成立的，需要根据你的模型输出结构进行调整
    import random
    predicted_token_index = torch.argmax(torch.softmax(torch.tensor(draft_logprob), dim=0)).item() #找到 draft_logprob 概率最高的 token 的索引
    acceptance_probability = min(1.0, math.exp(target_logprob[predicted_token_index] - draft_logprob[predicted_token_index]))
    return random.random() < acceptance_probability

def stop_condition(generated_tokens):
  # 定义停止生成的条件，例如达到最大长度或生成了结束符
  MAX_LENGTH = 100
  END_TOKEN = 2 # 假设结束符的 token id 是 2
  return len(generated_tokens) > MAX_LENGTH or generated_tokens[-1] == END_TOKEN

def detokenize(tokens):
    """
    将 token 列表转换为文本字符串.  这只是一个占位符，需要根据你的 tokenizer 实现
    """
    return " ".join(map(str, tokens))

代码解释:

• BatchedSpeculativeDecoder 类: 封装了 Batched Speculative Decoding 的核心逻辑。
• add_request(self, prompt): 将新的推理请求添加到活跃请求列表中。
• step(self): 执行一个推理步骤，包括草稿生成、验证、接受/拒绝和状态更新。
• accept_token(target_logprob, draft_logprob): 判断是否接受 Draft Model 生成的 token。这里需要仔细设计接受策略，以平衡速度和准确性。
• stop_condition(generated_tokens): 定义停止生成的条件，例如达到最大长度或生成了结束符。
• detokenize(tokens): 将 token 列表转换为文本字符串。

核心调度逻辑：

1. 动态 Batching: 每次取活跃请求列表的前 batch_size 个请求组成一个 Batch。
2. 接受率调整: 使用滑动平均来更新每个请求的接受率。
3. 请求排序: 根据接受率对活跃请求列表进行排序，优先处理接受率较高的请求。

5. 高级优化技巧

除了上述基本的调度策略之外，还可以采用一些高级优化技巧来进一步提高 Batched Speculative Decoding 的性能：

• 共享 Key-Value Cache： 在 Target Model 和 Draft Model 之间共享 Key-Value Cache，可以减少内存占用和计算量。
• Kernel Fusion： 将多个操作合并到一个 Kernel 中，可以减少 Kernel Launch 的开销。
• 量化 (Quantization)： 使用更低精度的数据类型（例如，FP16 或 INT8）来表示模型参数和激活值，可以减少内存占用和计算量。
• 模型蒸馏 (Model Distillation)： 使用 Target Model 来训练 Draft Model，可以提高 Draft Model 的预测准确性，从而提高接受率。

表格：不同调度策略的比较

调度策略	优点	缺点	适用场景
静态 Batching	简单易实现	无法充分利用硬件资源，对不同长度的请求处理效率低	请求长度基本一致，对延迟不敏感的场景
动态 Batching	可以更好地适应不同的 workload，提高硬件利用率	需要更复杂的调度算法，实现难度较高	请求长度差异大，需要平衡吞吐量和延迟的场景
优先队列调度	可以保证高优先级请求的及时处理	需要维护优先级队列，增加管理开销，可能导致低优先级请求长时间得不到处理	需要区分请求优先级，保证重要请求的及时处理的场景
基于接受率调度	可以将更多资源分配给接受率较高的请求，提高整体吞吐量	需要实时监控接受率，可能引入额外的计算开销，对接受率波动敏感	投机采样接受率差异较大，希望最大化吞吐量的场景
基于 token 长度调度	可以根据序列长度动态调整投机采样次数，加速生成过程	需要评估不同长度序列的最佳投机采样次数，可能引入额外的调优成本，对序列长度变化敏感	序列长度差异较大，希望针对不同长度的序列进行优化的场景

6. 实际应用中的考虑因素

在实际应用 Batched Speculative Decoding 时，还需要考虑以下因素：

• 模型选择： 选择合适的 Target Model 和 Draft Model 是至关重要的。Target Model 的准确性直接影响生成质量，而 Draft Model 的速度和准确性则影响加速效果。
• 硬件平台： 不同的硬件平台对 Batched Speculative Decoding 的支持程度不同。需要根据具体的硬件平台进行优化。
• 延迟要求： 不同的应用场景对延迟的要求不同。需要在吞吐量和延迟之间进行权衡。
• Tokenizer: Tokenization 的效率对整体推理性能有重要影响。选择一个高效的 Tokenizer 至关重要。

7. 未来发展趋势

Batched Speculative Decoding 仍然是一个快速发展的领域。未来的发展趋势可能包括：

• 自适应投机采样： 根据当前系统的状态和请求的特征，自动调整投机采样的参数。
• 混合精度计算： 结合使用不同的精度来表示模型参数和激活值，以进一步提高性能。
• 分布式投机采样： 将投机采样过程分布到多个设备上，以提高并行度和可扩展性。
• 更强的 Draft Model： 通过知识蒸馏等技术训练更强大的 Draft Model， 从而降低 target model 的计算负担。

总结：灵活调度，高效推理

Batched Speculative Decoding 通过复杂的调度策略，充分利用硬件资源，能够显著提高 LLM 推理的吞吐量。通过动态调整 Batch 大小、基于接受率进行调度等方法，可以更好地适应不同的 workload。

结语：持续优化，迎接挑战

尽管 Batched Speculative Decoding 取得了显著的进展，但仍然存在许多挑战。希望通过不断的研究和优化，能够进一步提高其性能，并将其应用到更广泛的场景中。 谢谢大家！