TriForce: 突破长序列生成瓶颈的分层投机解码技术

在人工智能领域，大型语言模型（LLMs）的长序列生成能力一直是研究的热点。然而，随着模型规模的增长，推理过程中的内存和计算瓶颈成为了限制其应用的主要障碍。为了解决这一问题，Carnegie Mellon University和Meta AI (FAIR)的研究团队联合推出了TriForce，一种创新的分层推测解码系统。TriForce通过利用注意力稀疏性和上下文局部性，有效地缓解了关键的键值（KV）缓存和模型权重瓶颈，实现了对长序列生成的无损加速。在A100 GPU上，TriForce为Llama2-7B-128K模型带来了高达2.31倍的推理速度提升，在RTX 4090 GPU上更是达到了惊人的7.78倍加速，显著减少了生成每个token所需的时间。而且TriForce在处理大型批量数据时也展现了出色的性能，在实际应用中有着高效性和可扩展性。

Figure 1展示了 TriForce 如何使用基于检索的草图和分层推测来有效解决双重瓶颈问题。整合了两个模型和三种缓存，包括草图模型、目标模型以及草图模型的 StreamingLLM 缓存，以及目标模型的检索缓存和完整缓存。过程首先通过重复草图 γ1 步，帮助目标模型使用检索的部分 KV 缓存生成超过 γ2 个 token，然后由目标模型使用完整的 KV 缓存进一步验证这些 token

TriForce的设计深受对大模型（LLMs）处理长上下文时注意力机制的观察所启发。研究揭示了一个关键的现象：在预训练的LLMs中，即使面对长达120K的上下文，通过仅使用4K的token，就能恢复超过96%的注意力得分。这一发现表明，LLMs在处理长序列时，其注意力分布呈现出显著的稀疏性。注意力稀疏性的发现对于TriForce来说至关重要，因为它意味着在进行推测解码时，可以只利用一小部分关键信息，而不必加载整个KV缓存。

a. 展示了随着序列长度增加，使用 StreamingLLM 或 Sliding-window 与 Llama2-7B-128K 之间的自然发散持续增加，表明推测解码的接受率随着更长上下文的增加而下降。温度敏感性表明缺乏鲁棒性。 b. 与模型权重相比，随着上下文长度的增加，KV 缓存逐渐成为另一个瓶颈。 c. 随着上下文长度的增加，KV 缓存占据了大部分内存

另一个关键的观察是上下文局部性。这一概念指出，相邻的token在长上下文中往往需要相似的信息。也就是说，对于任何给定的token，与其相关的信息在上下文中往往集中在一个局部区域内。这种局部性允许TriForce在多个解码步骤中重用特定的缓存段，而不是每次都重新构建整个缓存。通过这种方式，TriForce能够显著提高解码效率，同时保持对上下文的敏感度。

a. Llama2-7B-128K 模型在 120K 上下文中展示了显著的注意力稀疏性。除了前两层外，其余层表现出显著的稀疏性。 b. 可以使用部分 KV 缓存和完整模型权重进行自我推测。使用有限预算的现有方法可以实现高接受率。 c. 在大多数层中观察到明显的局部性，随着上下文的发展这种局部性逐渐减少

这两个观察结果为TriForce提供了坚实的理论基础，使其能够通过智能地选择和重用信息来优化长序列生成过程，从而在保持生成质量的同时，显著提高推理速度。

TriForce: 创新的分层推测解码系统

TriForce系统的核心在于其创新的检索式草稿方法。这种方法通过将KV缓存分割成小块，并计算给定查询与每个块中平均key缓存的注意力，有效地识别和选择最相关的信息。这一策略不仅直观，而且无损，与现有的基于逐出策略的方法相比，如StreamingLLM和H2O，它在保持低延迟的同时，能够更准确地检索详细的上下文信息。

基于检索的草图方法。在需要长期上下文依赖的情况下，StreamingLLM 和 H2O 等方法由于其缓存更新策略而表现不佳，因为它们不可避免地且不可恢复地丢弃 KV 对

为了应对KV缓存瓶颈之外的模型权重瓶颈，TriForce采用了分层推测架构。这一架构首先使用一个轻量级模型，配备StreamingLLM缓存，对目标模型进行初步推测。这种分层的方法显著降低了草稿阶段的延迟，从而加速了整个推理过程。

TriForce的算法伪代码进一步阐释了其工作流程。系统首先初始化目标模型和草稿模型，然后构建检索缓存，并根据上下文局部性的洞察进行更新。在每次迭代中，草稿模型首先预测一定步数的token，然后目标模型使用完整的KV缓存进行自我验证。这个过程不仅构建了一个效率层级，还通过两个瓶颈的有效管理，提高了整体的推理速度。

以下是一个简化的伪代码示例，概述TriForce算法的主要步骤和结构：

# 输入参数：
# Prefix: 输入前缀序列 [x1,...,xt]
# Mp: 目标模型与完整缓存 Cp
# Mq: 草稿模型与StreamingLLM缓存 Cq
# T: 目标序列长度
# γ1, γ2: 推测长度参数
# Draft: 自回归草稿阶段
# Verify: 验证阶段
# Correct: 校正阶段

# 初始化阶段
Initialize(Prefix, Mp, Mq, T)
Prefill(Mp, C_p, Prefix)
Prefill(Mq, C_q, Prefix)
Construct(Cr, Prefix[-1])

# 迭代推断，直到达到目标序列长度 T
while N < T:
    # 草稿模型推测阶段
    for i in 1 to γ1:
        qi = Draft(Mq, Cq, x≤N)
    
    # 采样与验证阶段
    for i in 1 to γ1:
        ~xi ~ qi
        p_hat = Mp(Cr, x≤N, ~x≤i)
        if Verify(~xi, qi, p_hat):
            n += 1
        else:
            xi = Correct(qi, p_hat)
            n += 1
            break
    
    # 如果所有推测的token都被接受，则采样下一个token
    if all tokens accepted:
        xn+1 ~ p_hat
        N += 1
    
    # 更新缓存
    Update(Cr, Cq, accepted tokens)

# 返回生成的序列
return Generated_Sequence

系统实现采用了Transformer架构，并利用PyTorch CUDA图形和FlashAttention技术来最小化推测解码期间的内核启动开销，并加速注意力操作。为了保持在固定的KV缓存预算内，系统实现了额外的推测缓存结构，包括检索缓存和StreamingLLM缓存。

实验

在实证评估部分，研究团队展示了TriForce在不同硬件配置下对长序列生成的加速能力。实验结果表明，在A100 GPU上，TriForce为Llama2-7B-128K模型实现了高达2.31倍的速度提升。而在两个RTX 4090 GPU上，通过卸载设置，TriForce实现了7.78倍的速度提升，将每个token的生成时间降至0.108秒，仅是A100 GPU上自回归基线的一半。

除了速度提升，TriForce在处理更长上下文和不同温度设置下也展现了出色的扩展性和稳健性。实验数据显示，随着上下文长度的增加，TriForce能够保持高接受率，理论上甚至可以达到13.1倍的速度提升。即使在温度参数设置为1.0的情况下，TriForce的接受率仍然能够保持在0.9以上，显示了其在不同条件下的稳定性和可靠性。TriForce在处理大批量数据时同样表现出色，实验结果显示，在批量大小为六的情况下，每个样本包含19K上下文，TriForce能够实现1.9倍的效率提升。