理解并加速大语言模型推理中的内存处理流水线:一种 GPU-FPGA 异构方案 Understand and Accelerate Memory Processing Pipeline for Disaggregated LLM Inference
把稀疏注意力、RAG 等LLM优化统一为四阶段内存流水线,并用GPU-FPGA异构系统加速。
前置知识
KV Cache(键值缓存)
Transformer 在自回归解码时为每个历史 token 缓存 Key 和 Value 向量,避免重复计算注意力。缓存大小与上下文长度线性相关,1M token 的 GPT-OSS-120B 模型 KV 缓存可达 69 GB。
本文围绕 KV 缓存与衍生内存结构的处理开销展开,长上下文推理的内存成本是文章的核心动机。
稀疏注意力(Sparse Attention)
在标准注意力的二次复杂度基础上,选择性关注历史 token 的子集来降低计算与访存开销。代表方法包括 DeepSeek Attention(DSA)、SeerAttention-R(SA-R)、LServe 等,它们引入轻量索引器或分页 KV 缓存来加速检索。
稀疏注意力是本文重点评估的第一类 LLM 优化,其 Compute Relevancy 与 Retrieval 步骤正是异构加速的目标。
检索增强生成(RAG)
把外部知识库(文档、向量库)按查询相关性检索后拼接到 prompt 中,从而扩展 LLM 的有效上下文。代表系统包括 DRAGIN、FLARE、Fixed-Sentence RAG 和带重排序的两阶段 RAG。
RAG 检索的 BM25 评分与 top-k 选择具有不规则访存,是 FPGA 加速的主要收益来源。
FPGA(现场可编程门阵列)
可定制微架构的硬件,提供 BRAM/URAM 大容量片上存储(带宽可达 21.8 TB/s)、灵活的访存控制和低功耗优势,擅长处理稀疏、不规则、memory-bound 的工作负载,但峰值算力低于 GPU。
本文核心创新是把 LLM 内存处理流水线按计算特性切分到 FPGA 与 GPU,因此必须理解 FPGA 的存储层级与适用场景。
算术强度(Arithmetic Intensity)
定义为 FLOPs/byte,衡量一个计算步骤是 compute-bound 还是 memory-bound。值越大越偏计算密集,越小越偏访存密集;本文用该指标刻画四阶段流水线的异构性。
本文用它作为 GPU/FPGA 分工的客观依据,是判断哪个设备更适合执行某阶段的量化指标。
Prefill-Decode Disaggregation(预填-解码分离)
把 LLM 推理中计算密集的 prefill 阶段与访存密集的 decode 阶段分配到不同设备上执行,以避免互相干扰并提高资源利用率。
本文在 MemAgent 的合成内存场景中采用了这一思想,让 FPGA 负责解码生成文本记忆、GPU 负责 prefill 消费记忆。
研究动机
现代大语言模型(LLM)的上下文长度已扩展到 128k 到 1M token,对内存子系统造成巨大压力。以 GPT-OSS-120B 为例,存储 1M token 的 KV 缓存需要高达 69 GB 的显存;自回归解码时反复访问 KV 进一步放大了访存开销。为缓解这一问题,业界提出了稀疏注意力(如 DeepSeek Attention、LServe、SeerAttention-R)、RAG(DRAGIN、FLARE、两阶段 RAG)、压缩上下文记忆(Titans/HMT、MemAgent)、测试时训练(TTT/LaCT)等多种优化。然而既有研究多把这些方法当成彼此独立的技巧,缺乏对它们共性计算特性和硬件效率的系统性理解。通过在 AMD MI210 GPU 上的实际 profiling,本文发现:在 1M token 序列下,稀疏注意力的内存处理开销占解码时间的 22%–81%(在 4K 时仅 1%–11%);RAG 在 20M 文档场景下占 40%–61%;MemAgent 的 Prepare Memory 阶段最高可占 97%。也就是说,内存处理已成为 LLM 长文档推理的关键瓶颈,单点优化收益有限。
本文的目标是本文的目标是:(1)建立一个统一的分析框架,把现有 LLM 推理优化方法纳入一个可对比的内存处理流水线;(2)通过系统 profiling 量化各阶段的延迟占比与计算特性;(3)基于计算异构性设计一个 GPU-FPGA 异构系统,把不规则、memory-bound 的步骤下沉到 FPGA,把规则、compute-bound 的步骤留在 GPU;(4)端到端地证明该设计能在多个真实工作负载上同时带来显著的时延降低(1.04–2.2×)和能耗降低(1.07–4.66×)。
与已有工作不同的是,已有工作要么把 LLM 推理优化当作黑盒算法,缺乏统一的硬件视角;要么聚焦在 GPU 内部或专用 ASIC 加速器(如已有工作评估 LLM decode 的 GPU-FPGA 异构,但只覆盖单一工作负载)。本文的独特切入角度是:先从算法层面抽象出 Prepare→Compute Relevancy→Retrieval→Apply 的四阶段流水线作为通用骨架;再按算术强度、访存规律、数据依赖三个维度刻画每阶段的异构性;最后给出一套可复用的 GPU-FPGA kernel library 和三种部署范式(通用、合成内存、Memory as Context),能在同一个硬件节点上同时加速稀疏注意力、RAG、MemAgent、Titans/HMT 等异构算法。
核心方法
本文方法分两步:先建立分析骨架,再做硬件映射。骨架层面,作者把任意长上下文 LLM 推理优化抽象为生成模型 $L(g(\cdot), f(\cdot,\cdot), \{x_i\}_{i<t}, x_t) = y_t$,其中 $g$ 是记忆生成器、$f$ 是记忆处理器,并把它分解为四步流水线 $M_{<t} \xrightarrow{\text{prep}} I_{<t} \xrightarrow{\text{comp}} S \xrightarrow{\text{ret}} M'_{<t} \xrightarrow{\text{apply}} O_{<t}$,其中 Prepare Memory 把原始记忆压成可索引格式,Compute Relevancy 计算相关性分数,Retrieval 按分数与启发式选择子集,Apply to Inference 把检索结果融入解码。映射层面,作者把算术强度高、数据规则、依赖局部的步骤(Prepare/Apply、dense 线性层)放到 GPU;把算术强度低、不规则、数据依赖跨 memory 的步骤(BM25 评分、top-k 选择、最大归约、跨 memory 加权求和)下沉到 FPGA,并利用 FPGA 的 BRAM+URAM+HBM 三层存储与 streaming dataflow 实现 kernel fusion。具体给出三种部署范式:General Setup(sparse attention/RAG)、Synthesized Memory(MemAgent)、Memory as Context(Titans/HMT)。
本文的核心创新是把“算法分类“升级为“流水线分类“:此前评估 LLM 推理优化往往按 sparse attention/RAG/Memory-as-Context 等方法族分开做加速;本文证明它们都可以拆解为同一个四阶段流水线,从而可以用同一套硬件映射准则(compute-bound 在 GPU、irregular memory-bound 在 FPGA)来覆盖所有方法。这与既有 FPGA 加速 LLM 的工作(例如纯 decode 加速或单一方法加速)形成本质区别——后者无法同时处理 BM25 这种 token 级别不规则访存和 MemAgent 这种 LLM 解码负载。此外,本文提出的“以 kernel 为单位复用 + PCIe 通信最小化“的部署策略(即把跨设备通信开销远小于端到端延迟作为硬约束,约三个数量级的差距),也是该异构方案能稳定获益的关键。
方法步骤详情
方法实施分为五个层次。**第一层 Pipeline 定义**:把任意 LLM 推理方法统一表述为 $M_{ 10 的步骤以及规则顺序访存留在 GPU;算术强度在 1 附近或不规则、依赖跨 memory 的步骤下沉到 FPGA;PCIe 传输的数据量与延迟必须比端到端延迟小约三个数量级,否则保留在 GPU。**第四层 Kernel 设计**:在 FPGA 上实现 Compute Relevancy + Retrieval 的融合 kernel(Figure 7),包括 Inner Product Engine、Reduction Unit、Top-K Retriver 三级 streaming dataflow;query 与压缩 key 缓存在 BRAM/URAM(21.8 TB/s),溢出部分落到 HBM(460 GB/s);token id 越小的 key 越靠近 BRAM 以维持高吞吐。在 GPU 侧复用 vLLM、TileLang、HIPIFY 等现有优化库,仅在确有需要时新增 kernel。**第五层 三种部署范式**:General Setup(GPU 跑 Prepare/Apply,FPGA 跑 Compute Relevancy+Retrieval,FPGA 仅回传 top-k 索引)、Synthesized Memory(FPGA 跑 LLM decoding 生成文本记忆,GPU 跑 LLM prefilling 消费记忆)、Memory as Context(与 General Setup 类似但 memory 完全驻留 FPGA,且对每段输入做循环通信)。最终所有 kernel 形成可复用 library,用户可通过组合和跨设备通信 API 构建新算法。
技术新颖性
技术上,本文新颖性体现在三方面。第一,**统一的四阶段 Pipeline 抽象**——这是首次把 sparse attention、RAG、compressed memory、TTT 这四大类方法放进同一个数学骨架(Table 1 给出每种方法在每个阶段的具体算子映射)。第二,**基于计算异构性的 GPU-FPGA 协同设计**——提出按算术强度、访存规律、数据依赖三维度的映射准则,并通过 kernel fusion 把 BM25+top-k、Inner Product+Reduction+Top-K 等典型不规则流水线在 FPGA 上做成 streaming dataflow,比 GPU 实现获得 1.8–6.6× 的内存处理加速。第三,**针对不同方法族给出三套可复用的部署范式**:General Setup 适用于需要 GPU 端 dense 模型参与的稀疏类方法;Synthesized Memory 利用 prefill-decode 分离把 MemAgent 这类“用 LLM 生成记忆“的工作负载高效拆解;Memory as Context 通过把 memory 完全驻留 FPGA 并循环传输输入段,最小化 PCIe 开销。相比既有 FPGA 加速 LLM 的工作(多局限于 decode 阶段或单一方法),本设计在同一节点上同时覆盖六大类方法并提供可组合的 kernel library。
实验结果
本文在 AMD MI210 GPU + Alveo U55C FPGA 节点上系统评估了 sparse attention(DSA、SA-R、LServe)、RAG(DRAGIN、FLARE、FS-RAG、两阶段 RAG)、Memory as Context(HMT/Titans)、Synthesized Memory(MemAgent)四大类方法,主要实验结果可归纳为四点。**Profiling 揭示的瓶颈**:sparse attention 在 1M token 序列下内存处理占解码时间 22%–81%,RAG 在 20M 文档下占 40%–61%,MemAgent 的 Prepare Memory 最高占 97%,Titans/HMT 即使短上下文也主要由线性投影主导。**端到端加速**:稀疏注意力获得 1.04–1.49×(Figure 8),其中 DSA 1.1–1.2×、LServe 最显著(在 1M token 内持续提升);RAG 单阶段方法 5.1–6.6× 内存处理加速对应 1.10–1.45× 端到端加速,两阶段 RAG 受限于 reranker 仅 1.1–2.1×(Figure 10);Memory as Context 通过 query generation 与 cross attention 融合获得 3.1–4.0× 内存处理加速与 1.3–1.6× 端到端(Figure 11);MemAgent 在 prefill-decode 分离下稳定 1.8×(Figure 12)。**能耗效率**:Table 3 显示整体能耗下降 1.07–4.66×,其中 DSA 1.61×、LServe 1.43×、MemAgent 4.66×、HMT/Titans 1.65×,RAG 在 1.10–1.21× 之间。能耗收益不完全来自加速——FPGA kernel 的运行功耗本身就低于 GPU(见 Appendix G)。**Batch 扩展性(Table 4)**:稀疏注意力与 RAG 的加速比随 batch 增大而提升(DSA 从 BS=1 的 1.02× 增至 BS=32 的 1.15×;LServe 从 1.19× 增至 1.83×),原因是 dense 算子在 GPU 上复用收益更大,相对放大了 FPGA 在 retrieval 上的优势;Memory as Context 相反(1.48×→1.15×),因为它的 cross attention 中 dense 投影在 GPU 上的复用收益高;MemAgent 在 BS>2 时反而劣化(1.85×→0.13×),此时应动态切换到 GPU-only 部署。**降速原因**:DSA 与 LServe 在 >1M token 后会因 HBM 访问增多而掉速,系统需动态 fallback 到 GPU-only。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| DSA 内存处理加速(DeepSeek Attention) | memory-processing kernel speedup / end-to-end speedup | 1.3–2.2× / 1.1–1.2× | AMD MI210 GPU only | GPU-FPGA 异构相对 GPU 基线在内存处理步骤获得 1.3–2.2× 加速,端到端 1.1–1.2×;能耗 geomean 1.61×、最高 1.69× |
| SeerAttention-R Top-k 内存处理加速 | kernel speedup | 1.8–2.2×(Top-k)/ 2.6–4.9×(Threshold) | AMD MI210 GPU only(TileLang 优化 kernel) | Top-k 模式 1.8–2.2×、Threshold 模式 2.6–4.9×;端到端 LServe 1.04–1.49×;能耗 SA-R Top-k 1.11×、Threshold 1.14× |
| LServe 长序列稀疏注意力 | memory-processing kernel speedup / end-to-end speedup | 1.2–5.6× / 1.04–1.49× | AMD MI210 GPU only(HIPIFY 移植 CUDA kernel) | 依赖 paged KV cache 的 max-reduction 与 top-k 优势,在 BRAM+URAM 上获得最高 5.6× 内存处理加速;能耗 geomean 1.43×、max 1.62× |
| 单阶段 RAG 加速(DRAGIN/FLARE/FS-RAG) | memory-processing speedup / end-to-end speedup | 5.1–6.6× / 1.10–1.45× | AMD MI210 GPU only(Llama 2 7B 生成器) | BM25 评分的不规则 token 频率查找在 FPGA 上以 streaming dataflow 加速;能耗 geomean DRAGIN 1.10×、FLARE 1.14×、FS-RAG 1.21× |
| 两阶段 RAG 加速(带 Reranker) | memory-processing speedup / end-to-end speedup | 1.1–2.1× / 1.47–1.84× | AMD MI210 GPU only(Llama 3.1 8B 生成器) | reranker 留在 GPU(密度高),其余检索在 FPGA;能耗 geomean 1.07× |
| Memory as Context(HMT/Titans) | memory-processing speedup / end-to-end speedup | 3.1–4.0× / 1.3–1.6× | AMD MI210 GPU only | query generation 与 cross attention 融合带来 3.1–4.0× 内存处理加速;能耗 geomean 1.65×、max 1.74× |
| MemAgent 合成记忆 | end-to-end speedup | 1.8×(prefill-decode 分离) | AMD MI210 GPU only(Qwen 2.5 7B) | FPGA 跑 LLM decoding 生成文本记忆、GPU 跑 prefilling 消费记忆,能耗 geomean 4.66×、max 4.66×(最强) |
| Batch 扩展性(BS=1→32) | geomean speedup 变化 | DSA 1.02×→1.15×、LServe 1.19×→1.83×、DRAGIN 1.14×→1.92×、FLARE 1.19×→2.11×、FS-RAG 1.26×→2.10× | GPU-only 在各 batch 下的实现 | 稀疏注意力与 RAG 的相对优势随 batch 增大而放大;Memory as Context 1.48×→1.15× 略降;MemAgent 在 BS>2 劣化(1.85×→0.13×),需动态切换部署 |
局限与改进
作者明确指出的限制以及独立观察可归纳为以下几点。第一,**HBM 容量瓶颈**:当序列超过 1M token 时,压缩 key 无法全部驻留 BRAM+URAM(U55C 共约 40 MB),被迫访问 460 GB/s 的 HBM,导致 LServe 和 DSA 在 >1M 后掉速(geomean 能耗改善分别仅 1.43× 和 1.07×)。第二,**跨设备通信仍存在开销**:尽管系统刻意只传输索引或小段嵌入,仍需保留 PCIe DMA 与 P2P 调度开销,作者在 Appendix C.1 报告约三个数量级小于端到端延迟,但缺少对节点级联与 NVLink 替代方案的探索。第三,**方法覆盖不全**:作者在 TTT/LaCT 上明确放弃了 FPGA 异构部署,原因是其 Compute Relevancy(loss 计算)与 Apply(forward/backward pass)多为 compute-bound 且不规则程度不足,无法从 FPGA 获益;这意味着该异构方案并非对所有 LLM 优化普适。第四,**生态与可迁移性**:当前所有 FPGA kernel 用 Xilinx Vitis HLS 2024.2 + Vivado 2024.2 编写,迁移到 AMD Versal V80、Intel Agilex 等新一代 FPGA 或 ASIC 需要重写;同时与现有 GPU-only 推理栈(vLLM、SGLang、TensorRT-LLM)的兼容性有限。第五,**没有对比 NVIDIA A100 的实测**:作者在物理节点上只用了 AMD MI210,A100 数据是基于 profiling 的估算(Appendix H),这削弱了对 NVIDIA 生态下收益的信心。第六,**异构 ASIC 尚未落地**:作者在结论里承认 ASIC 可以进一步消除通信开销并提升能效,但本文只给出“未来方向“而未实现。
独立分析的弱点
独立分析可识别出以下可改进的弱点。**弱点一:HBM 触发掉速后的策略偏保守**。当前 DSA/LServe 在 >1M token 时直接 fallback 到 GPU-only,浪费了 FPGA 的可用算力;改进方向是实现 HBM-aware 分层存储调度,例如把“热 key“动态迁回 BRAM,或在 FPGA 上实现 HBM 带宽感知的 sliding-window 索引,把超过 L2 容量的部分用 ring buffer 处理。**弱点二:MemAgent 在 batch>2 时反而劣化**。Table 4 显示 MemAgent 在 BS=4 时已跌至 0.93×、BS=8 跌至 0.49×、BS=32 仅 0.13×,作者只给出“动态切换到 GPU-centric“的被动方案;改进方向是引入 FPGA 上的 batched decode kernel,利用 weight reuse 提升 dense 算子吞吐,或在 batch 大时把 prepare memory 中可向量化的部分(如 embedding 查表)也用 HLS SIMD 指令展开。**弱点三:缺少统一 benchmark**。所有方法在自家原始工作负载上评估,跨方法的延迟归一化与可比较性较弱;改进方向是构建一个统一的 LLM 推理优化 benchmark(类似 LongBench、InfiniteBench)并把本系统的 kernel library 接入到 OpenCompass、lm-evaluation-harness 等评测框架中。**弱点四:能耗建模偏粗**。Table 3 的能耗数据是 geomean,缺少按工作负载(短 vs 长上下文)和 batch size 的细分;改进方向是提供 per-workload、per-step 的能耗分解,结合 FPGA 内部 power rail 测量(而非整机功耗估算)以支持碳足迹核算。**弱点五:缺少对 KV cache 压缩率与精度的消融**。论文假设压缩索引向量后质量不退化,但缺少端到端 perplexity / 任务精度的对比表;改进方向是引入下游 benchmark(如 LongBench、SCROLLS、HotpotQA)证明压缩索引在异构系统上仍能保持模型质量。**弱点六:与现有 GPU-only 推理栈的兼容性差**。本系统要求用户自己组合 kernel 并通过 PCIe API 通信,难以嵌入 vLLM/SGLang;改进方向是把 kernel library 包装成兼容 vLLM Attention backend 的 plugin,从而无缝替换现有的 FlashAttention、PagedAttention 后端。
未来方向
作者明确提出了三条未来方向。**方向一:异构 ASIC 设计**。当前 FPGA 在 16nm 工艺下性能受限,作者建议把 GPU-FPGA 协同的设计原则固化到 ASIC 中(例如把 FPGA 上的 Compute Relevancy+Retrieval 模块与 GPU 通过片上 NoC 直连),从而消除 PCIe 通信开销并进一步降低能耗。**方向二:设计自动化**。当前需要用户手动组合 kernel,作者计划通过 design automation 让用户只需声明四阶段算子即可自动生成 FPGA bitstream 与 GPU kernel;这与 AMD 的 MLIR-AIE、Intel 的 oneAPI 等高层次综合工具结合是可行路径。**方向三:新算法的适配**。本文 Pipeline 抽象应能容纳未来新算法(例如 hierarchical memory、neural cache、speculative memory),未来工作应把这些方法纳入同一分析骨架。基于结果的延伸方向还包括:(a)把异构思想扩展到 multi-GPU + multi-FPGA 集群,研究跨节点 memory processing 切分;(b)把能耗指标纳入优化目标,研究“每焦耳有效 token“的服务成本模型;(c)探索稀疏注意力与 RAG 的融合范式(如 SLM-based RAG、retrieval-augmented sparse attention),并在同一异构系统上评估;(d)把本文 kernel library 接入 RLHF/在线学习场景,使 TTT 这类当前被放弃的方法也能在异构系统上获益。
复现评估
复现性整体良好,但仍有门槛。**代码与数据**:作者在脚注 2 提供了 GitHub 仓库链接 https://github.com/OswaldHe/HeteroLLM,承诺开源 kernel library 与部署脚本,但仓库未在正文给出具体提交时间或 commit hash。**算法实现**:所对比的基线方法均来自原始作者的代码(vLLM、TileLang、LServe 官方仓库、DRAGIN、MemAgent、LaCT 等),并对部分代码做了 AMD ROCm 适配(如 LServe 通过 HIPIFY 移植 CUDA kernel)。**硬件要求**:主实验使用 AMD MI210 GPU + Alveo U55C FPGA 节点(AMD EPYC 7v13 host CPU),这是商用可购硬件但 FPGA 板卡价格昂贵且需要 PCIe P2P 与 DMA 驱动支持;NVIDIA A100 数据基于 profiling 估算,需复现者自行搭建异构节点。**工具链**:FPGA kernel 使用 Xilinx Vitis HLS 2024.2 + Vivado 2024.2,GPU 侧使用 ROCm 6.3;这两套工具链的学习曲线陡峭,且对操作系统、内核版本有依赖。**实验数据量**:RAG 实验使用 Wikipedia dump 评估,最大到 20M 文档,对存储和检索索引构建有一定要求。**复现难度评估**:在已有 AMD GPU + Alveo FPGA 的实验室中可较完整复现(难度中等);在仅有 NVIDIA GPU 的实验室中只能复现 GPU-only baseline 与 AMD MI210 profiling 部分(难度较低);FPGA 综合与 bitstream 生成的工程量较大。**与现有系统的兼容性**:作者明确指出部分方法(如新版 LaCT、MemAgent v2)可能未及时适配,且本系统与 vLLM 等推理栈的集成尚未完成,建议复现者留意仓库更新。
论文图表
Figure 1 是一张 GPU-FPGA 异构系统的概念图:左边是 AMD MI210 GPU(承担 compute-bound 与 dense 步骤),右边是 Alveo U55C FPGA(承担 memory-bound 与 irregular 步骤),中间通过 PCIe 互连;右侧文字给出在 SA-R 与 DSA 上的总体加速(1.2–1.8×)和能耗改善(1.3–4.7×)总结。
这是全文的视觉提纲,让读者一眼理解本文硬件方案的整体形态和收益概览;放在 motivation 部分最能展示作者为什么要走 GPU-FPGA 路线。
Figure 3 折线图,横轴是序列长度(4K、16K、64K、256K、1M),纵轴是 memory processing 占 decoding time 的比例;多条曲线分别对应 DSA、SA-R、LServe 等稀疏注意力方法;显示从 4K 的 1%–11% 增长到 1M 的 22%–81%。
提供 sparse attention 的 profiling 实测数据,支撑 Claim 1(内存处理是 critical path);放在 motivation 部分用于量化瓶颈。
Figure 4 折线图,横轴是文档数量(500K 到 20M),纵轴是 memory processing 占总延迟的比例;对比单阶段 RAG(DRAGIN/FLARE/FS-RAG,40%–61%)与两阶段 RAG(含 reranker,曲线在 500K 起点更高、增速更慢)。
支撑 RAG 场景下内存处理也是主要瓶颈,且揭示两阶段 RAG 因 reranker 存在不同特性;放在 motivation 部分。
Figure 5 是双子图:左图给出 Titans/HMT/LaCT 在不同上下文长度下的 memory processing 占比(即使是短上下文也占大头,因为有多层线性投影);右图给出 MemAgent 的 Prepare Memory 高达 97% 的极端情况。
补充说明合成记忆与参数化记忆的瓶颈分布,特别是 MemAgent 的 Prepare Memory 极端值,是 motivation 与 method 之间的桥梁。