FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍

FAST-Prefill 并非简单的算法硬件映射,而是一个硬件-算法深度协同设计的范例。它基于 Flex-Prefill 算法的计算模式与数据流特征,定制了专用的数据通路与访存策略。这项工作不仅提供了一个实用的加速方案,更确立了一种方法论启示:针对特定工作负载,精巧的架构设计往往比通用算力的堆砌更为高效。

关键词:FPGA 加速、稀疏注意力、长上下文 LLM、预填充阶段动态稀疏模式

当大语言模型需要处理长达 128K Token 的文本时,预填充阶段的计算量会爆炸式增长。FAST-Prefill 提出了一种基于 FPGA 的动态稀疏注意力加速方案,通过创新的硬件架构设计,在保持模型精度的同时,实现了 2.5 倍的首令牌生成时间(TTFT)速度提升4.5 倍的能效提升

让大语言模型阅读一本数百页的小说并回答问题,模型首先需要处理整本书的内容——这就是 LLM 推理中的“预填充阶段”。当输入文本变长时,预填充阶段的自注意力计算量呈平方级增长,成为整个推理过程的性能瓶颈。

近年来,稀疏注意力技术被提出来解决这个问题。其核心思想很直观:每个 Token 不需要与所有其他 Token 交互,只需关注其中的一小部分。但问题在于,这种稀疏模式是“动态”的——不同输入文本有不同的稀疏模式,同一文本中不同的注意力头也有各自独特的模式。

这种动态特性让 GPU 陷入了尴尬境地:
* 一方面,生成稀疏索引需要复杂的控制流和数据依赖;
* 另一方面,稀疏的 KV 缓存访问模式导致数据重用率低、内存带宽利用率不足。

GPU 虽然计算能力强,但在这种内存密集型任务中难以发挥优势。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍

这正是 FPGA 的机会所在。FPGA 具有可定制的存储层次和计算流水线,能够为特定工作负载设计专用的数据通路。来自南加州大学的研究团队在 FPGA‘24 会议上提出的 FAST-Prefill,正是首个针对动态稀疏注意力长上下文预填充阶段的 FPGA 加速器。

本文将深入解读 FAST-Prefill 的工作机制,探讨它如何通过创新的硬件架构设计,在保持模型精度的同时,实现比 NVIDIA A5000 GPU 快 2.5 倍的性能提升和 4.5 倍的能效提升。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 5 | FAST-Prefill 与基准 GPU 实现的首令牌生成时间(TTFT)对比。该图对比了 Nvidia A5000 GPU 与 FAST-Prefill 在 Llama3.2-1B、Qwen-2-1.5B、Llama3.2-3B 模型上,4K 至 128K 上下文长度的归一化 TTFT。结果显示 FAST-Prefill 在所有模型和上下文长度下均表现更优,且随上下文长度增加优势更显著,核心原因是其定制化的稀疏索引核与双级缓存大幅减少了片外访存,而 GPU 存在不规则访存与算力利用率低的问题。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 6 | FAST-Prefill 与基准 GPU 实现的能效对比。该图以每焦耳生成令牌数为指标,对比了 GPU 与 FAST-Prefill 在三款模型不同上下文长度下的能效表现。FAST-Prefill 凭借 FPGA 的低功耗特性与架构级的访存优化,实现了最高 4.5 倍的能效提升,且长上下文场景下能效优势更明显。GPU 因密集的片外访存和高功耗计算单元,在长上下文稀疏注意力推理中能效大幅下降,而 FPGA 的流处理与片上存储适配了该场景的低算力密度需求。

本文目录

  • 一、背景:理解 LLM 推理的两个阶段
    • 1.1 预填充阶段 vs. 解码阶段
    • 1.2 动态稀疏注意力:解决方案与挑战
  • 二、三大核心挑战
    • 2.1 挑战 1:稀疏索引生成产生大量中间张量
    • 2.2 挑战 2:稀疏模式依赖的 KV 缓存访问
    • 2.3 挑战 3:DSP 资源限制矩阵乘法吞吐量
  • 三、FAST-Prefill 的架构创新
    • 3.1 创新点 1:流式内存感知的稀疏索引生成
    • 3.2 创新点 2:活性驱动的双级 KV 缓存
    • 3.3 创新点 3:混合矩阵处理单元
  • 四、实验评估
    • 4.1 实验设置
    • 4.2 精度评估
    • 4.3 性能评估
    • 4.4 能效评估
    • 4.5 消融研究
  • 五、相关工作
  • 六、结论与展望

一、背景:理解 LLM 推理的两个阶段

在深入 FAST-Prefill 的设计之前,我们需要先理解 LLM 推理的基本流程。

1.1 预填充阶段 vs. 解码阶段

LLM 推理包含两个截然不同的计算阶段:

预填充阶段:模型处理完整的输入序列,生成第一个输出 Token。在这个阶段,输入是一个完整的 Token 序列,模型通过矩阵乘法生成 Query(Q)、Key(K)和 Value(V)矩阵,然后计算因果多头自注意力,最后通过前馈网络(FFN)处理。

解码阶段:模型自回归地逐个生成后续 Token,每次迭代只处理最新生成的 Token。这意味着预填充阶段的矩阵-矩阵乘法变成了解码阶段的矩阵-向量乘法。

对于长上下文应用如文档摘要、代码生成、多轮问答等,预填充阶段的计算复杂度随上下文长度呈平方级增长,导致生成第一个 Token 的延迟(TTFT)急剧增加。

1.2 动态稀疏注意力:解决方案与挑战

为了缓解这个问题,研究者提出了稀疏注意力。其工作原理是将注意力矩阵划分为固定大小的块,然后为每个注意力头选择一部分块进行计算。

二、三大核心挑战

研究团队识别出在 FPGA 上加速动态稀疏注意力的三个关键挑战:

2.1 挑战一:稀疏索引生成产生大量中间张量

生成稀疏索引需要对所有 Key 块与最后一个查询块向量进行评分,以识别最相关的块。以 128K Token 的上下文长度为例,每个注意力头会产生 128×128K 的中间结果,存储需求超过 2GB。若加上 Softmax、指数运算和块池化等操作,总中间张量可达 4GB。

关键洞察:这些庞大的中间张量无法存储在 FPGA 片上存储器中,而频繁的片外访存会急剧增加延迟。传统的算子融合技术在此场景下效果有限,因其主要针对逐元素操作,难以融合跨不同维度的复杂操作。

2.2 挑战二:稀疏模式依赖的 KV 缓存访问

KV 缓存的大小通常在 3-4GB 量级,只能存储在片外高带宽内存中。然而,稀疏注意力机制导致每个查询仅访问 KV 缓存的一个子集,这与密集注意力中所有查询访问完整 KV 缓存的行为模式截然不同,从而引发以下问题:

  1. 预取时机难以把握:过早预取会浪费宝贵的片上存储资源,而过晚预取则会导致计算流水线停顿。
  2. 访存带宽利用率低:按需、小块的非连续内存读取,难以充分利用 HBM 的高带宽特性。
  3. 数据重用性差:在组查询注意力等机制下,KV 块在不同注意力头之间缺乏有效的数据重用机会。

2.3 挑战三:DSP 资源限制矩阵乘法吞吐量

长上下文推理涉及大量的块级矩阵乘法运算。以 Xilinx Alveo U280 FPGA 为例,若仅依赖 DSP 资源构建脉动阵列,最多只能实现约 6 个 32×32 的阵列规模。这远远不足以并行处理稀疏注意力计算中所需的大量矩阵乘法操作,成为性能提升的主要瓶颈。

三、FAST-Prefill 的架构创新

面对上述挑战,FAST-Prefill 提出了一套完整的硬件架构解决方案。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 1 | FAST-Prefill 整体架构。该架构是专为长上下文大模型预填充阶段设计的 FPGA 加速器,整合了稀疏索引生成、稀疏注意力计算等核心单元,由全局有限状态机协调全流程调度。各单元通过 AXI 互连网络与 DDR、HBM 内存控制器通信。CPU 完成分词等预处理后,数据传入 FPGA,由混合矩阵处理单元等核心模块完成计算,并通过优化的片上存储层级管理数据流,实现对动态稀疏注意力的端到端高效加速。

架构主要组件包括:
* 全局有限状态机:协调控制各计算阶段的执行顺序。
* 特殊功能单元:处理 Softmax、层归一化及激活函数等非线性操作。
* 稀疏索引生成单元:动态计算并生成注意力稀疏模式。
* 稀疏注意力单元:执行基于稀疏索引的注意力计算。
* 混合矩阵处理单元:提供高吞吐量的矩阵乘法计算能力。

3.1 创新点一:流式内存感知的稀疏索引生成

FAST-Prefill 最核心的创新之一,是将稀疏索引生成从传统的“生成-聚合”两步操作,转变为“流式-累积”的单步操作。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 3 | 稀疏索引生成单元工作流。该单元通过并行处理键块与查询块的提取和池化,结合点积运算生成块级相关性分数。其摒弃了存储完整大张量的传统方式,采用增量聚合方式计算分数,并基于 Jensen-Shannon 散度动态选择稀疏模式,最终仅输出高相关 KV 块的索引集合,极大降低了后续计算的访存与运算开销。

传统方法需要显式计算并存储完整的中间注意力张量(例如大小为 [查询块数,键块数]),再进行池化。对于 128K 上下文,这意味每头数十 MB 的存储需求。

FAST-Prefill 的设计截然不同:
1. 顺序流式获取:Key 块严格按照递增的块索引顺序从 HBM 中读取,确保了长连续突发传输,避免了随机访问。
2. 增量分数更新:对于每个流入的 Key 块,混合矩阵处理单元计算其与当前查询块的点积分数。系统并不存储完整的 [查询块,键块] 张量,而是增量更新一个紧凑的、存储在片上的逐块统计信息缓冲区。
3. 片上累积与选择:分数在片上的 Score Buffer 中累积。索引选择阶段采用流式 Top-k 算法,顺序处理分数并维护一个小型候选列表,当累积分数超过预设阈值 γ 时即完成选择。

关键成果:稀疏索引生成单元不向片外写入任何大型中间张量,每个 Key 块仅需读取一次。通过将分数计算、池化、散度评估和 Top-k 选择融合到一个流式流水线中,FAST-Prefill 将稀疏索引生成的中间存储需求从 GB 级降至 KB 级,实现了两个数量级的优化,并将其从内存受限的不规则操作转变为确定性、高带宽效率的数据通路。

3.2 创新点二:活性驱动的双级 KV 缓存

第二个核心创新是针对稀疏访问模式设计的智能 KV 缓存管理系统。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 4 | 稀疏注意力单元工作流。该单元依据稀疏索引集精准调度 KV 块,通过预取单元从 HBM 获取数据并存入冷热双级缓存。缓存控制逻辑根据块的剩余复用次数管理其驻留状态。查询块与缓存中的 KV 块完成注意力计算后,结果直接输出。采用块优先的调度策略,并通过键控累积解决输出乱序问题,从而最大化缓存命中率,最小化片外访存。

块优先调度:稀疏注意力单元并非按查询块或注意力头顺序遍历,而是按照 KV 块的索引升序进行调度。稀疏索引集被转换为紧凑的作业列表,其中每个 KV 块标识符关联一个消费者列表,明确指出哪些注意力头和查询块需要该 KV 块。

活性驱动的缓存替换:在作业构建阶段,为每个 KV 块计算一个“剩余使用计数器”,精确记录该块还将被后续多少轮注意力计算所使用。每次该块被消费,计数器递减;当计数器归零时,该块即可被立即标记为可驱逐,实现了“用完即弃”的高效管理。

双级缓存结构:为进一步提升效率,片上 KV 缓存被划分为“热”和“冷”两个区域。剩余使用次数超过静态阈值(例如总查询块数的 50%)的高重用块被存入热区;重用次数有限的块则存入冷区或直接旁路缓存。这种设计有效防止了中等重用块驱逐高重用块,避免了缓存颠簸,显著提升了缓存利用效率。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
(注:此处 IMAGE_8 的图注与 IMAGE_7 完全重复,疑似原文错误。在实际文档中,此处应替换为展示双级缓存结构或块优先调度细节的示意图。)

3.3 创新点三:混合矩阵处理单元

第三个创新点旨在解决 FPGA 中 DSP 资源受限的问题。其核心洞察在于:INT8 乘法运算可以分解为位平面操作,从而能够利用查找表(LUT)实现高效运算。

数学基础:给定两个 8 位有符号整数 ab,其乘积可分解为位级运算的和。每个部分积本质上可简化为按位与(AND)操作后进行移位,这种特性非常适合用 LUT 实现。然而,直接实例化所有 8×8 位平面的交互既不面积高效,延迟也非最优。

半字节划分:为了优化实现,将每个 8 位操作数划分为高 4 位和低 4 位的“半字节”:
a = a_h * 2^4 + a_l
b = b_h * 2^4 + b_l

则乘法可展开为:
a * b = (a_h * b_h) * 2^8 + (a_h * b_l + a_l * b_h) * 2^4 + (a_l * b_l)
展开后的每一项都对应一个 INT4 乘法,可以使用小型 LUT 高效实现。这种半字节级的分解在保持精确计算语义的同时,显著降低了运算延迟并优化了硬件资源使用。

混合实现:在 Xilinx Alveo U280 FPGA 上,混合矩阵处理单元(MPU)被设计为包含 6 个基于 DSP 的 32×32 计算阵列和 6 个基于位平面的 32×32 计算阵列。该处理单元执行 INT8 精度的乘法运算,并在 INT32 精度下进行累加。

四、实验评估

4.1 实验设置

研究团队在 Xilinx Alveo U280 FPGA 上实现了 FAST-Prefill 系统,并与在 NVIDIA RTX A5000 GPU 上运行的 Flex-Prefill 实现进行了性能对比。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
表 1:GPU 与 FPGA 平台的硬件参数对比。

上表对比了两种平台的核心硬件参数。可以看出:
* GPU 拥有更多的计算单元和更高的峰值算力,但其主频高,功耗也更大。
* FPGA 的 DSP 单元数量与 GPU 的 CUDA 核心数接近,并具备 HBM 与 DDR 的混合内存架构。虽然其峰值算力较低,但通过定制化架构能够显著提升在稀疏场景下的实际计算利用率,其差异化的内存带宽也更好地适配了不同的访存需求。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
表 2:FAST-Prefill 在 FPGA 上的资源利用率。

上表展示了设计在目标 FPGA 上的资源使用情况。其中,URAM(超 RAM)利用率达到 95%,充分用于片上大容量 KV 缓存优化;DSP 利用率为 71.6%,LUT 利用率为 64.3%。资源分配与架构设计高度匹配,既满足了稀疏索引生成和矩阵运算的算力需求,又未出现资源过度占用的情况,实现了资源的高效利用。

4.2 精度评估

研究团队使用 RULER 基准测试评估模型精度,比较了以下三种配置:
* Flex-Prefill (BF16):原始 BrainFloat 16 精度实现。
* Flex-Prefill (INT8):权重被量化至 INT8,但矩阵乘法时需要反量化为 16 位进行计算。
* FAST-Prefill:采用 W8A8(权重和激活值均为 INT8)量化,所有计算在 INT8 精度下完成。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
表 3:RULER 基准测试的精度对比(数值越高越好)。

实验结果显示,FAST-Prefill 采用的 W8A8 量化方案达到了与 INT8 版本的 Flex-Prefill 相似的精度水平,验证了其硬件设计在数值上的准确性。虽然量化不可避免地导致精度低于原始的浮点(BF16)版本,但其精度完全满足实际推理任务的需求,实现了精度与性能的有效平衡。

4.3 性能评估

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 5:FAST-Prefill 与基线 GPU 实现的首令牌生成时间(TTFT)对比。

在 4K 到 128K 的不同上下文长度下,FAST-Prefill 相比 GPU 实现了 1.5 倍到 2.5 倍的 TTFT 加速。这一性能提升主要归功于:
1. 定制化的稀疏索引生成内核:通过操作融合,显著减少了片外内存访问流量。
2. 活性驱动的双级缓存:利用片上缓冲区大幅增加了数据重用机会。
3. 混合矩阵处理单元设计:提高了矩阵乘法的计算吞吐量。

相比之下,GPU 实现由于缺乏高效的、面向数据重用的 KV 缓存块预取机制,导致了频繁且不规则的片外内存访问。此外,GPU 方案将大部分稀疏索引生成逻辑卸载至 CPU 执行,也引入了额外的延迟。

4.4 能效评估

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 6:FAST-Prefill 与基线 GPU 实现的能效对比(Token/Joule)。

FAST-Prefill 实现了高达 4.5 倍的能效提升。这凸显了 FPGA 在特定计算负载上的能效优势——通过定制化的数据通路和最大限度地减少不必要的片外数据访问,可以大幅降低系统功耗。在长上下文稀疏注意力推理场景中,GPU 因密集的片外访存和高功耗计算单元导致能效下降,而 FPGA 的流式处理架构与丰富的片上存储资源则更好地适配了该场景的低算力密度需求。

4.5 消融研究

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 7:消融实验——缓存对首令牌生成时间(TTFT)的影响(基于 Llama-3.2-3B 模型)。

该实验评估了活性驱动双级缓存的有效性。结果显示:
* 引入缓存的设计相比无缓存设计,TTFT 降低了 2.5 倍
* 高达 65% 的缓存命中率显著减少了耗时的片外内存访问次数,避免了因按需取块而导致的带宽利用率低下和计算流水线停滞问题。在长上下文场景下,缓存的性能收益更为突出。

FAST-Prefill:FPGA动态稀疏注意力加速器,突破长上下文LLM预填充瓶颈,性能提升2.5倍
图 8 | 消融实验:混合矩阵处理单元(MPU)对首令牌生成时间(TTFT)的影响(基于 Llama3.2-3B 模型)。 该实验对比了纯 DSP 矩阵处理单元与混合 MPU 在不同上下文长度下的归一化 TTFT。结果显示,混合 MPU 实现了 1.8 倍的延迟加速。混合 MPU 结合了 DSP 脉动阵列与基于 LUT 的位平面运算,充分利用了 FPGA 的 LUT 资源,避免了纯 DSP 设计的算力瓶颈。其 8 位定点运算设计适配了稀疏注意力的精度需求,并将约 85% 的 LUT 资源从闲置状态转变为有效利用

五、相关工作

近年来,FPGA 上的 LLM 推理加速成为研究热点。

| 加速方向 | 代表工作 | 主要技术特点 |
| :— | :— | :— |
| 解码阶段加速 | FlightLLM、CD-LLM | 定制解码数据流架构,优化吞吐与延迟,适用于 <1024 Token 的上下文。 |
| 量化技术 | Tereffic、LUT-LLM、Tellme | 支持三值等极端量化;定制内核适配量化权重;Tellme 统一优化预填充与解码,适配边缘 FPGA。 |
| 异构计算 | EdgeLLM、TeraFly | EdgeLLM 采用 CPU-FPGA 协同;TeraFly 实现多节点 FPGA 联合调度,提升解码并行性与扩展性。 |
| 长上下文加速 | AccLLM | 基于单一注意力稀疏模式优化,融合权重稀疏性与混合精度(FP16/INT4),支持 ≤8K Token 的长序列推理。 |

FAST-Prefill 的创新之处:与上述工作不同,FAST-Prefill 是首个面向动态稀疏注意力的长上下文预填充阶段 FPGA 加速器。它同时处理多种动态稀疏模式,支持高达 128K Token 的上下文长度,并通过创新的硬件架构解决了稀疏索引生成、KV 缓存访问和矩阵乘法吞吐量三大核心挑战

六、结论与展望

FAST-Prefill 展示了 FPGA 在动态稀疏注意力长上下文 LLM 预填充阶段的显著加速能力:相比 GPU,实现了 2.5 倍的 TTFT 加速4.5 倍的能效提升,同时保持模型精度。其核心贡献在于三项协同设计创新:

  • 流式内存感知的稀疏索引生成,降低索引开销;
  • 活性驱动的双级 KV 缓存,将不规则随机访问转化为可预测的流式访问;
  • 混合矩阵处理单元(MPU),以 LUT 补足 DSP 算力,提升资源均衡性。

该工作的本质是硬件-算法深度协同的范例:并非简单映射算法,而是基于 Flex-Prefill 的计算模式与数据流特征,定制数据通路与访存策略。

未来的拓展方向明确:集成 N:M 稀疏性、块剪枝等正交压缩技术,有望进一步降低延迟。随着长上下文应用(如全书理解、长视频分析、多轮深度对话)的兴起,预填充效率正成为关键瓶颈。FAST-Prefill 不仅提供了一个实用的加速方案,更确立了一种方法论启示——针对特定工作负载,精巧的架构设计往往比通用算力堆砌更高效。


关注“鲸栖”小程序,掌握最新AI资讯

本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/archives/23998

(0)
上一篇 2天前
下一篇 2天前

相关推荐