论文

CATS:用于内存有限的 LLM 推理加速的级联自适应树推测

CATS: Cascaded Adaptive Tree Speculation for Memory-Limited LLM Inference Acceleration

模型推理投机采样

摘要

大语言模型 (LLM) 中的自回归解码本质上是受内存限制的:每个生成步骤都需要从内存(例如,GPU 服务器的高带宽内存 (HBM))加载模型权重和中间结果,从而使吞吐量的瓶颈受到内存带宽而不是计算的限制。 推测解码 通过启用多个草稿词元的并行验证来解决这个问题,从而有效地分摊每个目标模型调用的成本。然而,现有的 推测解码 方法是在 HBM 足够大以同时容纳目标模型和辅助草图模型的假设下设计的 - 这一假设在内存受限的设备(例如具有有限 DRAM 的边缘平台)上会失败。我们分析了这种内存受限机制中的推理瓶颈,并提出了 CATS,这是一种自 推测解码 框架,可根据内存有限设备上的内存预算和参数卸载模式进行级联验证和校正。此设计最大限度地提高了词元接受率和端到端加速,同时保持设备上的峰值内存占用量等于目标模型的峰值内存占用量。我们在真实边缘设备上通过五个基准测试不同模型来评估 CATS。 CATS 可以实现高达 5.08 倍的挂钟加速,而不会降低生成质量,在边缘内存限制下,性能比 SOTA 方法高出高达 1.45 倍。