CATS：面向内存受限 LLM 推理加速的级联自适应树猜测

# Cats：用于内存受限 LLM 推理加速的级联自适应树推测
来源：https://arxiv.org/html/2605.11186
Yuning Han 佛罗里达大学 盖恩斯维尔, 佛罗里达州 32611 yuninghan@ufl\.edu &Yangchenchen Jin 佛罗里达大学 盖恩斯维尔, 佛罗里达州 32611 yangchenchen\.jin@ufl\.edu &Dylan Zhao 佛罗里达大学 盖恩斯维尔, 佛罗里达州 32611 dylan\.zhao@ufl\.edu &Jingwei Sun 佛罗里达大学 盖恩斯维尔, 佛罗里达州 32611 sun\.jingwei@ufl\.edu

###### 摘要

大型语言模型（LLM）中的自回归解码在本质上受内存限制：每个生成步骤都需要从内存（例如，GPU 服务器的高带宽内存 HBM）中加载模型权重和中间结果，使得吞吐量受限于内存带宽而非计算能力。推测解码通过并行验证多个草稿 token 来解决这一问题，有效分摊了每次目标模型调用的成本。然而，现有的推测解码方法基于这样一个假设：HBM 足够大，可以同时容纳目标模型和辅助草稿模型——这一假设在内存受限的设备（如具有有限 DRAM 的边缘平台）上不再成立。我们分析了这种内存受限 regime 下的推理瓶颈，并提出了 **Cats**，一种自推测解码框架，它根据内存受限设备上的内存预算和参数卸载模式进行级联验证和纠正。该设计在保持设备上峰值内存占用与仅目标模型相同的情况下，最大化了 token 接受率和端到端加速比。我们在真实边缘设备上的五个基准测试中针对不同模型评估了 Cats。Cats 可实现高达 5.08 倍的时钟加速，且生成质量没有下降，在边缘内存约束下比最先进的方法快高达 1.45 倍。代码可用 https://github.com/ElizaFuLan/CATS.git。

## 1 引言

部署大型语言模型（LLM）以进行高效推理仍然是现代机器学习系统面临的核心挑战之一。LLM 推理必须在自回归循环中逐个生成 token——每一步都依赖于所有之前的输出，并需要完整的模型前向传播。这种顺序结构在本质上受内存限制：每个生成步骤都将模型权重和中间结果从内存（例如，GPU 服务器的高带宽内存 HBM）加载到片上计算单元，吞吐量由内存带宽而非算术能力决定 Vellaisamy et al.(2026); Alizadeh et al.(2024)。推测解码旨在通过*并行解码*突破这一瓶颈 Leviathan et al.(2023); Chen et al.(2023); Stern et al.(2018)：一个轻量级的草稿模型推测性地生成一系列候选 token，然后较大的目标模型在单个批量前向传播中对其进行验证。因为目标模型同时处理多个 token，每个 token 的内存带宽成本被分摊到几个已接受的 token 上，在保持精确输出分布的同时显著提高了吞吐量。

然而，这种吞吐量增益带来了结构性成本：经典的推测解码引入了*第二*组权重——辅助草稿模型——该模型必须与目标模型共存于设备内存中，从而增加了静态内存占用和每步带宽流量。自推测解码方法 Zhang et al.(2024); Elhoushi et al.(2024); Xia et al.(2025); Liu et al.(2024a) 作为对这种额外内存成本的实用回应而出现：通过从目标模型自身的浅层子网络生成草稿，它们消除了对单独草稿模型的需求。然而，它们仍然需要引入额外的适配器参数进行草稿生成，并且这些方法的性能往往受到子网络有限的草稿生成能力的限制。

参见图 1 说明：Vicuna-7B 在五个基准测试上的端到端加速比。更根本的限制贯穿于*所有*现有的推测解码方法：它们假设模型权重在整个推理过程中驻留在设备内存中（例如，服务器上的 HBM 和边缘设备上的 DRAM）。在大规模服务器上，这是成立的，但在内存受限的平台（如边缘设备）上，DRAM 容量远低于单个 7B 参数模型，并且必须与操作系统共享——因此权重必须在每次前向传播时从闪存流式传输 Alizadeh et al.(2024)，使得 flash↔DRAM 成为速率限制传输，而不是 DRAM↔片上内存。现有方法中的每一个设计选择——草稿深度、验证策略、接受策略——都针对内存驻留 regime，并与这种内存受限设置不匹配。辅助模型方法受到双重惩罚：维护单独的草稿模型增加了内存容量需求和每个推测步骤的额外传输流量，加剧了它们本应缓解的瓶颈。

我们研究了这种内存受限的推理瓶颈，并提出了 **Cats**（**C**ascaded **A**daptive **T**ree **S**peculation），一种自推测解码框架，它执行级联验证以在内存预算内最小化推理时间。Cats 将推理组织为三个阶段，其层边界由可用的 DRAM 预算决定：（1）*草稿*阶段，使用浅层子网络通过重复草稿迭代来生成候选 token；（2）*浅层验证*阶段，使用 DRAM 预算下的中间层（每个周期从闪存流式传输一次），并行检查草稿 token 并生成纠正候选；（3）*目标验证*阶段，剩余层从闪存中分块卸载，以处理草稿跟踪和所有纠正分支上的树状输入，选择最长的接受前缀。我们还提出了*减少 KL 损失*，将蒸馏监督集中在高概率 token 上，以最大化子网络的草稿和验证能力。

我们的主要贡献如下：

- •我们确定了 flash↔DRAM 数据移动是边缘设备上推测解码的主要瓶颈——这是一个被现有方法完全忽视的约束——并提出了 Cats，一种用于这种内存受限 regime 的级联自推测解码框架。
- •Cats 是模型无关的：级联框架适用于任何基于 Transformer 的 LLM，无需进行架构修改。
- •我们在边缘设备上实现了我们的框架。在 Vicuna-7B/13B 和 LLaMA2-7B/13B 上的五个基准测试中的广泛实验表明，在内存受限设置下，Cats 实现了高达 5.08 倍的加速，优于所有比较的自推测解码基线，包括 Kangaroo Liu et al.(2024a)、Medusa Cai et al.(2024) 和 Eagle Li et al.(2024b)。

## 2 相关工作

##### 带辅助草稿模型的推测解码。

规范的推测解码框架 Leviathan et al.(2023); Chen et al.(2023); Xia et al.(2023) 依赖于单独训练的小模型来提出候选项，目标模型并行验证。后续工作将其扩展到树状验证 Miao et al.(2024); Chen et al.(2024a); Wang et al.(2024); Xiong et al.(2024)、知识蒸馏草稿模型 Zhou et al.(2024); Liu et al.(2024b); Du et al.(2024)、基于检索的草稿 He et al.(2024); Li et al.(2024a); Oliaro et al.(2025) 和前瞻生成 Fu et al.(2024)；参见 Xia et al.(2024) 进行全面综述。依赖独立训练的小模型的方法 Kim et al.(2023); Chen et al.(2024b); Zhao et al.(2024); Wang et al.(2025b); Zhong et al.(2025); Sun et al.(2023), 2025; Bachmann et al.(2025); Wang et al.(2025a); Sun et al.(2024); Chen et al.(2025b) 产生显著的内存开销，这是在边缘场景中的一项挑战，即使目标模型也勉强适合设备内存。

##### 自推测解码。

为了消除对单独草稿模型的依赖，一系列*自推测*方法直接从目标模型自身的层派生草稿模型。早期退出方法，如 Draft & Verify Zhang et al.(2024) 和 LayerSkip Elhoushi et al.(2024)，将 token 路由通过前几个 Transformer 层进行草稿生成，并使用完整模型进行验证。Swift Xia et al.(2025) 在没有微调的情况下为每个输入动态选择跳过层，而 KnapSpec Cha et al.(2026) 将层选择表述为背包优化问题。Kangaroo Liu et al.(2024a) 在浅层子网络之上引入了一个轻量级适配器，训练其模仿目标模型的输出分布，以最小的开销实现强大的接受率。多头方法，如 Medusa Cai et al.(2024) 和 Hydra Ankner et al.(2024)，将并行预测头附加到目标模型的最后一层，使得在单次前向传播中提出几个未来 token 成为可能。Eagle 系列 Li et al.(2024b), 2025a, 2025b 在特征级别进行草稿生成并构建动态 token 树，在服务器级硬件上实现了最先进的加速。基于一致性的训练 Kou et al.(2024); Guo and Ermon (2025) 和多 token 预测目标 Gloeckl et al.(2024); Qin et al.(2024); Monea et al.(2023) 提供了并行生成的互补方法。所有这些方法的一个共同局限是它们仍然需要额外的适配器参数进行草稿生成，这在内存受限设置下可能会削弱现有推测解码方法获得的性能增益。并且这些方法的性能往往受到子网络有限的草稿生成能力的损害。

##### 内存受限硬件上的 LLM 推理。

LLM in a Flash Alizadeh et al.(2024) 表征了边缘设备上的闪存↔DRAM 带宽瓶颈，其中模型权重不能驻留在 DRAM 中，并提出了窗口加载和行列捆绑以减少每次前向传播的传输量。PowerInfer Song et al.(2024) 利用激活稀疏性选择性地仅加载热神经元。这些工作 Ren et al.(2025); Chen et al.(2025a) 认识到内存瓶颈是边缘推理的关键约束——我们试图在相同设置为现有推测解码方法解决这一约束。

## 3 预备知识与动机

参见图 2 说明：服务器与边缘的内存层次结构，以及自回归 Vicuna-7B 推理的每 token 延迟分解测量值。*服务器（B200）：*模型权重驻留在 HBM 中；绑定瓶颈是 HBM↔SRAM 带宽，每 token 延迟由计算主导。*边缘（Jetson AGX Orin）：*DRAM 无法容纳完整模型，因此权重必须在每次前向传播时从闪存暂存；flash↔DRAM 传输主导每 token 延迟。第 1 节中确定的结构不匹配——所有推测解码方法都假设内存驻留权重，这一条件在边缘设备上失败——源于两种部署 regime 之间内存层次结构的根本差异，如图 2 所示。在传统 GPU 平台（例如，NVIDIA B200）上，HBM 提供约 180 GB 和约 8 TB/s，可以轻松容纳目标模型、辅助草稿模型和所有中间状态；解码期间的绑定成本是快速的 HBM↔SRAM 路径。在具有统一 DRAM 架构的边缘平台上，GPU/NPU 和 CPU 共享一个小于 16 GB 的 DRAM 池——难以容纳单个 7B 参数模型（约 14 GB），更不用说额外的草稿模型。DRAM 和闪存之间的数据移动速度仅为约 2 GB/s。在推理期间，模型权重必须在每次前向传播时从闪存分块加载到 DRAM 中 Alizadeh et al.(2024)，绑定瓶颈

参见图 3 说明：推测解码方法在 B200 与 Jetson AGX Orin 上的端到端加速比。转移到这种 flash↔DRAM 移动，如图 2 中的延迟分解所示。这种瓶颈转移使得现有的推测解码方法在内存受限设置下表现不佳。我们对具有有限内存预算的 AGX Orin 设备进行了性能分析，并在图 3 中可视化了这一效果：当从服务器移动到边缘时，与相应平台上的普通推理相比，Kangaroo 和 EAGLE 的端到端加速比分别下降了 19.8% 和 24.5%。对于引入额外草稿模型的方法（如 EAGLE），这种影响更为严重，因为这些组件也必须通过受限的边缘内存层次结构进行暂存。

## 4 方法论

在本节中，我们形式化了 Cats 核心的内存自适应三阶段推理框架，并描述其三阶段推理流水线。然后我们详细说明验证树的构建以及用于训练