在人工智能领域,被称为大语言模型的现代计算机程序在理解和生成人类语言方面已变得异常出色。这些系统通过逐个预测句子中的下一个词(即一个标记/token)来工作,从而逐步构建出连贯的响应。然而,随着这些模型能力的增强,当被要求处理非常长的文档或对话时,它们面临着一个显著的物理障碍。模型需要记忆的上下文越多,它就必须在快速内部存储器和主存储器之间不断地传输更多数据。这种数据的持续移动造成了瓶颈,就像试图用花园水管去填满一个游泳池,而排水口却开得很大一样。计算机大部分时间都在等待信息的到达,而不是在进行真正的思考,这减慢了整个过程,并限制了模型一次能够处理的文本量。
为了解决这个问题,来自复旦大学和上海人工智能实验室的研究人员开发了一种名为“快速闪存解码”(Faster Flash Decoding)的新方法。他们的方法通过改变模型决定保留哪些信息以及忽略哪些信息的方式来解决这一问题。该系统不再试图阅读大规模文档中的每一个单词来寻找相关的部分,而是使用了一个聪明的捷径。它首先为整个对话历史创建一个微小的、压缩后的“草图”。这个草图非常小,以至于计算机几乎可以瞬间扫描完成。通过查看这个草图,系统可以快速识别出对话历史中哪些部分可能是重要的,以及哪些可以被安全地忽略。只有在完成这次快速扫描后,模型才会检索所选部分的完整详细版本,以进行最终的计算。这个两步走的过程使模型能够在不丢失理解文本核心含义能力的前提下,跳过大量的无关数据。
研究人员在用于高端游戏和科学计算的强力显卡上测试了这种方法,发现它比目前的标准技术要快得多。在处理 256,000 个标记的上下文时,新系统将生成单个标记所需的时间从超过一毫秒降低到了仅为其中的一小部分。在整体速度方面,该系统生成文本的速度比以往的方法快了高达 2.37 倍,同时保持了相同的准确度水平。团队在包括复杂推理和从长文档中检索特定事实在内的广泛任务中验证了这一性能,确认了速度的提升并非以牺牲智能为代价。该系统无需重新训练模型即可运行,这意味着它可以立即插入现有的人工智能系统中以提高其效率。
这项工作的一个关键创新在于系统过滤信息的特定方式。传统方法通常依赖于固定规则,例如仅保留最重要的十个单词,或者需要整个系统在继续执行前暂停并同步的复杂计算。这种新方法使用了一种能够适应对话自然流动的动态阈值。它会寻找与当前上下文中最重要的单词相比具有显著重要性的单词,从而允许它根据注意力的集中程度来调整保留内容的多少。这种灵活性,结合在初始扫描中使用极低精度的数据,使得计算机能够绕过长期阻碍长上下文处理的内存瓶颈。其结果是,该系统能够以此前被认为不可能实现的、且不牺牲回答质量的方式,处理海量的文本。
技术摘要:Faster Than Flash (FFD)
问题陈述
长上下文大语言模型(LLM)的部署目前受到自回归解码阶段严重“内存墙”问题的阻碍。虽然预填充(prefill)阶段是计算密集型的,但解码阶段是严格受限于内存带宽的。随着序列长度的增加,标准的注意力机制需要为每个生成的 token 重复从高带宽内存(HBM)中加载整个键值(KV)缓存,从而导致线性延迟增长,限制了吞吐量。
现有的稀疏注意力方法试图缓解这一问题,但面临两个主要的困境:
- 度量困境(The Metric Dilemma): 依赖元数据(如聚类质心、几何边界)来估计 token 重要性的方法会引入额外的内存开销,或导致信息失真。
- 选择困境(The Selection Dilemma): 固定预算策略(Top-k)缺乏对变化注意熵的适应性;而分布自适应策略(Top-p)则需要全局同步(softmax)和排序,这破坏了流式流水线并引入了显著的开销。
方法论:Faster Flash Decoding (FFD)
FFD 是一个硬件-算法协同设计框架,它将选择器(selector)与计算器(computer)集成到一个完全融合的算子(kernel)中,消除了过滤与计算之间的分离。它通过以下三项核心创新解决了上述困境:
1. 通过低比特量化实现的感知内容扫描
FFD 不再存储额外的元数据索引,而是将键(Key)缓存划分为:
- 2-bit 量化缩略图(Thumbnails): 用于高吞吐量、感知内容的扫描,以估计注意力分数。
- 8-bit 残差(Residuals): 用于将键重构为接近 FP16 的精度,以进行最终计算。
这种方法在消除元数据内存开销的同时,保持了极高的信息保真度。扫描阶段以极低的算术强度运行,使用 2-bit 数据作为估计重要性的高速代理。
2. Top-δ 选择策略
FFD 引入了一种全新的选择准则,避免了 Top-k 的僵化和 Top-p 的全局同步问题。
- 机制: 如果 token j 的注意力分数 sij 满足 sij≥m~i−δ,则保留该 token,其中 m~i 是伪最大值估计,δ 是相对阈值。
- 伪最大值近似(Pseudo-Max Approximation): 为了避免计算真实全局最大值 (mi) 带来的全局规约(reduction)瓶颈,FFD 仅利用“汇点 token”(sink tokens,即初始 token)和局部上下文 token 来估计 m~i。实证分析表明,注意力分布主要由这些子集主导,使得该近似非常鲁棒。
- 自适应性: 阈值 δ(例如 5 或 7)控制了允许的注意力概率质量下降程度(例如 δ=5 保留了贡献超过峰值 e−5 的 token)。这使得稀疏率能够根据注意力分布的熵进行动态调整,而无需全局同步。
3. 融合算子与系统优化
FFD 实现为一个完全融合的 Triton 算子,包含三个阶段:
- 伪最大值估计: 使用汇点 token 和局部 token 计算阈值。
- Top-δ 选择: 流式传输 2-bit 键,计算初步分数,并根据阈值过滤块(blocks)。
- 细粒度精炼: 对于选定的块,加载 8-bit 残差和数值(values)以计算精确的注意力分数。
为了进一步降低延迟,作者采用了 基于块的 JIT 捕获(Block-wise JIT capture) 策略用于 CUDA Graphs。这使得整个解码步骤(包括 MLP 和归一化)可以被动态捕获,从而摊销编译成本,并消除小批量(small batch size)下的 CPU 启动开销。
关键结果
效率与吞吐量
- 算子加速: 在 NVIDIA RTX 4090 上,在 256K 上下文长度下,FFD 相比 FlashAttention-2 实现了高达 11.6倍 的算子级加速。
- 端到端吞吐量:
- RTX 4090: 在 16K 上下文时,FFD 的吞吐量比 FlashAttention-2 高出 2.37倍(51.8 tokens/s 对比 21.9 tokens/s)。
- H100: 在 16K 上下文时,FFD 比 FlashAttention-2 快 1.96倍(87.0 tokens/s 对比 44.5 tokens/s)。
- 可扩展性: 该方法能有效扩展至 256K 上下文长度,虽然延迟呈线性增长,但其斜率显著低于稠密注意力(dense attention)。
有效性与准确性
- RULER 基准测试(32K 上下文): FFD 在“大海捞针”(Needle In A Haystack)任务中保持了近乎完美的表现(单键检索为 100.0),并取得了 89.4 的综合得分(当 δ=7 时),优于 Quest (73.9) 和 KIVI (82.3) 等稀疏方法,并接近稠密基准线 (90.6)。
- LongBench: FFD 在所有稀疏方法中取得了最高的综合得分 (26.35),超过了 Quest、Twilight 和 KIVI,证明了高比例稀术并不会牺牲语义保真度。
- 泛化能力: 该方法在 Qwen2.5-7B-Instruct 上得到了验证,显示出一致的改进,表明其具有架构无关的适用性。
意义与主张
论文声称 FFD 代表了长上下文解码的一种范式转变,即将稀疏注意力重新定义为**几何过滤(geometric filtering)**而非元数据索引。其主要意义在于:
- 打破内存墙: 通过将选择与计算集成,FFD 实现了对扫描结果的复用,将解码瓶颈从内存受限问题转变为计算受限问题,其中低比特扫描是非常廉价的。
- 无需训练且即插即用: 该方案不需要模型重训,可作为标准注意力算子的掉落式替换(drop-in replacement)。
- 硬件-算法协同设计: 该工作证明,通过算法改变(2-bit 量化、伪最大值近似)针对特定硬件约束(内存带宽)进行优化,比纯算法或纯硬件优化能产生更优的结果。
作者总结道,未来的长上下文推理应优先考虑“以计算换 IO”(compute-for-IO)的权衡,利用低比特扫描中廉价的 FLOPs 来节省昂贵的 HBM 带宽。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。