首页 / 文章

推测解码、提前退出与量化:更快的大语言模型解码方式

推测解码与提前退出如何降低大语言模型推理延迟,何时采用初步接受策略更有效,以及它们与量化及KV缓存分页的结合方式。

发布者:REACTAPP.TOP阅读约 8 分钟1877 词

来源: **Foundations of LLM Inference Optimization: Speculative Decoding and Early Exit | Part 3D**

大语言模型推理分为两个阶段。预填充阶段会并行处理整个提示语,从而使GPU持续处于工作状态。解码阶段则是一次生成一个标记,且每个标记都依赖于前一个标记,因此硬件大多处于等待状态。正是这种顺序依赖性,导致即便从理论上看浮点运算次数足够多,交互延迟依然很高。

Prefill phase:
  Input: [token_1, token_2, ..., token_512]  → all 512 tokens processed in parallel
  Matrix shape: [batch, 512, 4096]
  GPU utilization: high — large matrix, full Tensor Core throughput

Decode phase (one step):
  Input: [token_513]  → one token processed
  Matrix shape: [batch, 1, 4096]
  GPU utilization: low - tiny matrix, most CUDA cores idle

为何解码位于关键路径上

在每个解码步骤中,模型都会对不断增长的KV缓存执行查询。随着上下文长度的增加,每一步的处理量也会上升,而在聊天场景中无法通过批量处理来缓解这一问题,因为延迟服务水平协议限制了请求在公司内部等待的时间。相关计算数据显示,单个步骤实际上只占用芯片很少的资源:

One attention head, decode step:

Query Q: [8, 1, 128]   →  8 × 128 = 1,024 elements
Keys  K: [8, 2048, 128] →  8 × 2048 × 128 = 2M elements
Values V: [8, 2048, 128] →  same

Matrix multiply (QKᵀ) per head:
  Shape: [8, 1, 128] × [8, 128, 2048] → [8, 1, 2048]
  FLOPs: 8 × 1 × 128 × 2048 ≈ 2.1M FLOPs per head

For 32 heads: ≈ 67M FLOPs per layer
For 32 layers: ≈ 2.1B FLOPs total per decode step

A100 peak at BF16: ~312 TFLOPS = 312 × 10¹² FLOPs/sec

Time to execute 2.1B FLOPs at 100% utilization:
  2.1 × 10⁹ / 312 × 10¹² ≈ 0.0067 ms

Actual observed decode latency per step: ~10–30ms

Effective compute utilization: < 0.1%

有效利用率低于0.1%意味着解码受内存流量和步数限制,而非单纯的计算能力限制。有两种应对方法:在每一步中生成更多可被接受的标记,或者每步执行更少的操作。推测性解码属于前者,而提前终止则属于后者。

推测性解码:先草拟再验证

一个小型草拟模型会提出若干未来的标记。目标模型会在一次并行前向传播中检查所有这些标记,保留最长的匹配前缀,并在首次出现不匹配时重新采样。每一个被接受的草拟标记都意味着你无需为它支付完整的目标步骤成本。

Without speculative decoding:
  Generate 5 tokens: 5 × 30ms = 150ms

With speculative decoding (K=4):
  Draft 4 tokens: 4 × 1.5ms = 6ms
  1 target verification pass: ~35ms  (slightly longer than decode,
                                       processes K+1=5 positions)
  Expected accepted tokens per round:
    (1 - 0.80^5) / (1 - 0.80) ≈ 3.36 tokens
  Time per round: 6ms + 35ms = 41ms
  Time per token: 41ms / 3.36 ≈ 12.2ms
Speedup: 30ms → 12.2ms ≈ 2.5×

加速程度完全取决于草稿与目标结果的匹配频率。简单、熵值较低的文本可以容纳较长的草稿;而某个异常的标记则会提前终止本轮处理。以下是该流程的概念性描述(实际实现还会在两个模型之间管理KV缓存):

import torch
import torch.nn.functional as F

def speculative_decode(target_model, draft_model, input_ids,
                       max_new_tokens, K=4, temperature=1.0):
    """
    Conceptual speculative decoding loop.
    Real implementations handle KV cache management across both models.
    """
    generated = input_ids.clone()
    while generated.shape[1] - input_ids.shape[1] < max_new_tokens:
        # --- Draft phase ---
        draft_tokens = []
        draft_probs = []
        draft_input = generated.clone()
        for _ in range(K):
            with torch.no_grad():
                draft_logits = draft_model(draft_input).logits[:, -1, :]
            q = F.softmax(draft_logits / temperature, dim=-1)
            token = torch.multinomial(q, num_samples=1)
            draft_tokens.append(token)
            draft_probs.append(q)
            draft_input = torch.cat([draft_input, token], dim=1)
        # --- Verify phase: one target forward pass over all K+1 positions ---
        verify_input = torch.cat([generated] + draft_tokens, dim=1)
        with torch.no_grad():
            target_logits = target_model(verify_input).logits
        # target_logits[:, -K-1:, :] covers all K draft positions + bonus
        # --- Accept/reject ---
        accepted = 0
        for i in range(K):
            p = F.softmax(target_logits[:, -(K+1)+i, :] / temperature, dim=-1)
            q = draft_probs[i]
            token = draft_tokens[i]
            # Acceptance probability
            accept_prob = torch.min(
                torch.ones_like(p.gather(1, token)),
                p.gather(1, token) / (q.gather(1, token) + 1e-9)
            )
            if torch.rand(1) < accept_prob:
                generated = torch.cat([generated, token], dim=1)
                accepted += 1
            else:
                # Sample corrected token and stop this round
                corrected_dist = F.relu(p - q)
                corrected_dist = corrected_dist / corrected_dist.sum(dim=-1, keepdim=True)
                corrected_token = torch.multinomial(corrected_dist, num_samples=1)
                generated = torch.cat([generated, corrected_token], dim=1)
                break
        else:
            # All K accepted - take bonus token
            bonus_logits = target_logits[:, -1, :]
            p_bonus = F.softmax(bonus_logits / temperature, dim=-1)
            bonus_token = torch.multinomial(p_bonus, num_samples=1)
            generated = torch.cat([generated, bonus_token], dim=1)
    return generated

尽可能选择同一系列中的候选版本,比如经过蒸馏或量化的简化版本,然后在自己的流量上测试其接受率,而非依赖公共博客的示例。

当候选版本出现偏差时

如果候选版本的分布与目标版本偏离,接受率会骤降,而你几乎要为它支付完整的成本。需持续监控接受率,一旦下降就立即恢复到普通解码方式。

可以用整数进行快速验证:假设目标步骤的成本为10毫秒,候选版本提出5个标记,平均约有3个被接受(60%)。只要接受率保持良好,即使扣除候选版本的额外开销,每个被接受标记的成本仍会低于使用单个标记的步骤。但如果接受率降至约1个标记,该方案就会失效。正是这种敏感性使得数据面板比个案描述更具参考价值。

提前终止:对简单标记跳过深层处理

某些架构允许在置信度足够高时让标记在中间层离开网络,从而节省了对标点符号、冠词以及其他可预测标记的处理计算资源。

Layer distribution of exits:Exit at layers 1–8  (very easy tokens like punctuation, articles):  15%
Exit at layers 9–16 (medium tokens, common continuations):           35%
Exit at layers 17–24 (harder tokens, named entities, numbers):       30%
Exit at layers 25–32 (full computation required):                    20%Weighted average layers executed:
  0.15 × 6 + 0.35 × 12 + 0.30 × 20 + 0.20 × 32
  = 0.90 + 4.20 + 6.00 + 6.40
  = 17.5 layers averageSpeedup vs always running 32 layers:
  32 / 17.5 ≈ 1.83×

与推测性解码不同,提前退出机制按设计会在跳过某些层时改变输出结果。它以速度为代价换取质量,且高度依赖于置信度的校准程度,因此需要预估精度差异并对其进行测量。

量化、批量处理与KV缓存结合使用

推测性解码是通过添加第二个模型来改变令牌提议循环;而提前退出则是改变单个模型内的深度。两者以不同方式面临类似的瓶颈,有时都可以与量化、连续批量处理以及KV缓存分页技术结合使用:

Combined stack example:

Target: 7B model, BF16, FlashAttention, PagedAttention
  → Model: ~14 GB, memory-efficient attention, paged KV cache

Draft: 70M model, INT4, FlashAttention
  → Model: ~35 MB, near-zero memory overhead

Speculative decode:
  → with K=4, α=0.80
  → ~2.5× token generation speedup on long outputs

Full stack speedup vs FP32 no-optimization baseline:
  - Quantization: 2–2.3× throughput
  - FlashAttention: 20–40% attention latency reduction
  - Speculative decoding: 2–2.5× decode speedup (on eligible requests)

Combined: 5–8× improvement in end-to-end tokens/second

堆叠处理需要谨慎考虑。草稿长度会影响连续批处理调度器的假设,而糟糕的调度器会导致批次碎片化,降低资源利用率。KV缓存淘汰策略也需要调整,以避免长时间对话时系统性能下降。预填充阶段的提示词缓存也很重要,否则即便解决了解码问题,预填充功能仍会造成资源浪费。此外,将提前退出与激进量化结合使用可能会迅速增加幻觉现象。合理的服务设计远比直接将某篇论文中的方法应用到核心流程中更有效。

每种技术何时能发挥作用

当草稿一致性较高且目标步数对延迟影响较大时,推测性解码能发挥作用;但若草稿匹配度低或草稿处理开销超过了其带来的收益,该方法就会失效。在许多标记易于处理且精度预算允许的情况下,提前退出能起到作用;而对于难度较高的标记或置信度校准不准确的情况,则此方法无效。

具体来说:一个70亿参数的模型用于聊天服务,每次回复的token数在50到150之间。预先填充2000个token的系统提示语虽然工作量大,但每轮只会发生一次,而解码步骤才是导致用户感知延迟的主要原因。通过推测性接受来减少解码步骤,或通过提前退出来降低每步的处理量,才是用户真正能感受到的改进。

对于小型设备端模型来说这重要吗?通常不如大型模型那么重要。草稿版本能在CPU上运行吗?有时可以,需要实际测试。按照优先级,先解决批量处理和缓存问题,再加入推测性解码,如果技术栈支持的话,最后再实现提前退出功能。

将其应用于实际的大语言模型推理

  • 在做出任何更改之前,先记录基准的p95延迟值和响应质量。
  • 将草稿模型与目标模型部署在同一节点上。跨主机部署草稿模型会产生网络抖动,从而抵消原有的优势;此外,两个模型加上KV缓存可能会耗尽本可轻松容纳一个模型的内存。
  • 记录接受率直方图、平均接受长度、质量评估差异值、GPU利用率以及p95延迟时间。
  • 当草稿模型在不同设备上运行时,需关注CPU与GPU的负载平衡情况。
  • 针对事实类任务测试A/B版本的质量。在全面推广之前,先在事实问答、编程及拒绝处理类任务中使用固定提示词进行测试,并在配置推测性解码以实现精确分布匹配的情况下,比较令牌完全相同的响应率。为提前终止测试,需跟踪分级任务中的胜率以及现有条件下的人类偏好数据。
  • 将每项优化措施置于特征标志之后,并设置可恢复到原始解码方式的关闭开关,同时做好回滚准备。
  • 在更改分词器或微调参数后,应重新检查所有内容。
  • 正确实现的推测解码器应与目标分布相符,但应通过配对测试来验证,而非仅凭直觉判断。使用无关的草稿版本、忽略温度参数对接受率的影响,或仅通过吞吐量基准就宣布成功而不设置质量检测,这些都是常见的错误做法,而且每一种错误都是可以测量的。

    平台团队应负责服务层中的这些优化工作,避免各应用团队自行设计草稿选择方式。应与服务层维护人员协调,而非在应用代码中修改相关设置;同时需暴露表明是否使用了推测解码或提前终止的头部信息或追踪属性,并按优化类型对令牌费用进行标记,以便财务部门了解优化带来的收益。

    依然存在的瓶颈

    即便解码速度加快,用户仍可能需等待工具调用、数据检索或客户端端的 Markdown 渲染。此时的限制因素可能是内存带宽、网络或客户端本身,而非 FLOPs,因此在应用论文中的各种技巧之前,应先进行端到端的跟踪分析与性能评估。推测性解码也无法替代高效的检索机制或优质的提示词;它仅能在模型已明确知晓要表达的内容时,降低生成后续标记的成本。

    顺序解码是自回归生成模式固有的结构性开销。在可量化的条件下,推测性解码与提前终止机制能够降低这一开销,因此应像处理其他生产级功能一样对其加以实现:设定相关指标、控制开关、提供回滚机制,并明确服务平台团队中的责任人员。

    相关阅读

  • LLM端到端推理:预填充、解码、KV缓存与采样:从分词器到流式处理的完整讲解,阐述为何预填充受计算能力限制、解码受带宽限制,以及GQA、分页机制和推测算法如何影响推理效率。
  • KV缓存是LLM解码过程中消耗GPU内存的隐形数据库:预填充与解码的对比、KV缓存相关数学原理、GQA/MQA、PagedAttention、量化技术以及用于生产环境推理的缓存淘汰策略。