arXiv'26 | PLENA：Agent 时代的 LLM 推理，被两堵 Memory Wall 卡住了

arXiv'26 | PLENA：Agent 时代的 LLM 推理，被两堵 Memory Wall 卡住了

原文：Combating the Memory Walls: Optimization Pathways for Long-Context Agentic LLM Inference

---

1. 前言

你有没有想过，同样一个 70B 模型，跑 chatbot 和跑 agent，硬件的"体感"完全是两回事？

Chatbot 场景：用户问一句，模型答几百 token，皆大欢喜。Agent 场景：模型要吞下一整个网页 DOM、一长串 tool-call 轨迹，动辄 90K token 的 context，再慢慢往外吐几千 token。这篇来自剑桥、Imperial College London 和爱丁堡的工作先甩了一个数据：agentic workload 的单次推理 token 消耗平均是 chatbot 的 100 倍，极端情况 1000 倍。

这张图值得逐个子图拆开看。

（a）token 用量：注意纵轴是对数刻度。左边 Chatbot（IFeval、GSM8K）单次推理 1-2K token；中间 Coding 任务上到 3K-20K；右边 CUA（Computer Use Agent，操作电脑的 agent）和 WUA（Web Use Agent）直接干到 100K 量级——OSWorld 里操作 LibreOffice、Chrome 浏览网页这类任务，要把整个屏幕状态/DOM 树喂给模型。深浅两色的堆叠还告诉你一件事：agent 任务的 token 大头在 prefill（深色），因为环境观测（网页、屏幕截图的文本表示）远比模型自己的输出长。

（b）计算热点随生成 token 数迁移：横轴是生成的 token 数，纵轴是 FFN 和 Attention 各占总 FLOPs 的百分比，两条线在 19K 附近交叉。为什么会迁移？关键在两者的计算量随 context 长度 L 的增长速率不同：

• FFN 的计算量是"每个 token 过一遍固定大小的矩阵"，处理一个新 token 的开销不随 L 变——线性于 L
• Attention 里每个新 token 要和前面所有 L 个 token 的 K/V 做点积——单 token 开销线性于 L，累计就是平方于 L

L 小的时候 FFN 的大常数占优（70B 模型的 FFN 矩阵很宽），L 大了 Attention 的平方项必然反超。chatbot 生成几百 token 根本走不到交叉点，agent 生成几万 token（LongWriter 负载 prefill 5K 后 decode 到 85K）就结结实实跨过去了——这意味着只优化 FFN（大多数加速器的默认重点）或只优化 attention 都不够，两个都得管。

（c）显存构成随 context 切换：横轴还是 token 数，虚线是模型权重（固定 ~60 GB），实线是 KV Cache。KV Cache 的体积 = 层数 × 2（K 和 V）× token 数 × hidden 维度，严格线性于 L，所以那条线一路爬升，在 64K-128K 之间越过权重线——128K 时单 batch 的 FP16 KV 约 39 GB。这就是"weight-dominant → KV-dominant"的切换点：短 context 时显存优化该压权重（量化/压缩），长 context 时该压 KV，而 agent 负载天然活在 KV-dominant 区。

于是就撞上了论文标题里的两堵墙：

• 带宽墙（bandwidth wall）：KV 和权重要以极高带宽反复读，HBM 带宽喂不饱计算单元
• 容量墙（capacity wall）：KV Cache 随 context 疯长，很快吃光 HBM，batch size 被压到个位数

这两堵墙一起把 GPU/TPU 的算力利用率摁在地上摩擦。作者的回应是 PLENA——一套软硬件协同设计的推理加速器系统，从三条 pathway 同时拆墙。

2. 为什么现有硬件在 agent 负载上躺平

先说一个容易被忽略的架构错配。TPU 这类加速器用的是方形 systolic array（比如 128×128），它假设 GEMM 的 M、K、N 三个维度尺寸相近。但 agent 推理是什么情况？容量墙压着 batch size 起不来，decode 阶段的 GEMM 是典型的 fat GEMM：M（batch 相关维度）只有个位数到几十，K、N 却是几千上万。方形阵列里大片 PE 全程围观，利用率惨不忍睹。

上图（a）很直观：同样数量的乘法器，方形阵列（比如 4×4 里只有上面 2 行在干活）一半 PE 围观，扁平阵列（2×8）全部激活——因为短边贴着 batch 维度 M 的实际大小。

（b）是整篇论文的"作战地图"，信息密度很高，值得逐条线读（场景：LLaMA-3 70B，128K context；横轴 batch size，纵轴可达 GFLOPs/s，均为对数刻度）：

• 两条灰虚线（Theoretical）：方形阵列和 PLENA 的理论峰值算力。注意 PLENA 的理论线略低——扁平阵列多了 adder tree 等开销，纸面峰值其实吃亏
• 绿实线（Reachable Square）：方形阵列实际可达性能。它远低于自己的理论线，且随 batch 增大缓慢爬坡——小 batch 时 fat GEMM 的形状错配让大片 PE 空转，这段差距就是"带宽墙 + 形状错配"偷走的算力
• 蓝实线（Reachable PLENA W/O Quant）：不加量化的 PLENA。同样 batch 下明显高于绿线——这段垂直差距就是 Pathway 1（扁平化）的收益：形状对了，PE 不空转
• 第一道竖灰墙（batch≈4）：Memory Bandwidth Wall 之前的区域，性能随 batch 线性涨（每个新请求都在分摊权重读取）
• 第二道竖灰墙（batch≈16）：Memory Capacity Wall——KV Cache 把 HBM 塞满，batch physically 开不上去，蓝线到此为止
• 青色实线（Reachable PLENA W Quant）：加上 Pathway 2（非对称量化）后，权重和 KV 变小，同样 HBM 能塞下更大 batch——容量墙被向右推（图中红色箭头），青线得以继续沿着更高的 batch 爬升，最终逼近理论峰值

一句话总结这张图：扁平化是"垂直方向"抬高利用率，量化是"水平方向"推开容量墙，两条 pathway 正交叠加，缺一条都到不了右上角。

3. PLENA 的三条 pathway

3.1 Pathway 1：扁平 systolic array

把 64×64 的方阵拍扁成 4×1024 这类形状（BLEN×MLEN，BLEN 对齐 batch 维度 M）。fat GEMM 的 M 小，那就让阵列的短边等于 M，长边吃 K/N 维度的并行。每个 cycle 从 Matrix SRAM 和 Vector SRAM 各取一条 MLEN 宽的输入，切成子向量喂给各个 sub-array，输出走 adder tree 归约。

对 FlashAttention 还有个额外彩蛋：扁平阵列可以多个 attention head 并行处理，长 context + 小 batch 的场景下利用率反而更高。

3.2 Pathway 2：非对称量化

这里"非对称"不是指 zero-point，而是指权重（W）/激活（A）/KV Cache 三者用不同精度的策略：

• 激活对量化误差最敏感 → 片上用高精度 FP 存在 Vector SRAM
• 权重和 KV 不那么敏感 → 用低精度 MX 格式（MXFP/MXINT）压进 Matrix SRAM 和 HBM
• 对 KV 追加 Hadamard 旋转压异常值，K/V 只被 attention GEMM 消费，所以 inverse transform 可以推迟到进 Matrix SRAM 时再做

算法侧配了两个 trick：output-norm 引导的分块 clipping 搜索和选择性旋转（不像 QuaRot 全局把旋转融进权重，而是逐层搜索要不要在线旋转）。效果：W4A4KV4 下 LLaMA-3-8B 六个零样本任务平均 70.39 分，QuaRot 只有 65.18（FP16 是 73.22）；70B 上差距更大（76.20 vs 69.21）。微缩放算术和旋转并不天然兼容，论文的消融实验证明盲套 QuaRot 反而伤精度，这个观察挺有价值。

3.3 Pathway 3：原生 FlashAttention 支持

FlashAttention 在 GPU 上是软件功劳，在专用加速器上得硬件配合：PLENA 设计了 IO-aware 的融合 attention 数据通路，中间激活（S、P 矩阵）全程留在片上，配合 transposable matrix SRAM（转置/非转置访问都不加开销）和 memory prefetch，把 off-chip 流量砍到最低。

整体架构如下，Matrix Unit（扁平阵列）+ Vector Unit（归约/逐元素）+ Scalar Unit 三件套，由 decoder 的 system-pipeline controller 统一调度：

4. 不止是芯片：全栈 + 自动设计空间探索

我个人认为这篇工作最"论文之外"的价值在工具链：PLENA 配了自定义 ISA、编译器、事务级模拟器和自动化 DSE（设计空间探索）流程，且承诺全部开源。

DSE 用 BoTorch 做多目标贝叶斯优化，在 perplexity / latency / 面积三个目标上搜 Pareto 前沿：

BoTorch 明显优于随机采样和 TPE。搜出来的设计点很有意思：追极致低延迟的配置用 MXINT8 + 1024 宽 MLEN（面积 203 mm²），追小面积的配置只要 23.6 mm² 还能保住 6.60 的 perplexity——架构参数和量化格式是联动搜索的，这其实就是把 NAS 那套方法论搬到了加速器设计上。

5. 效果

三条 pathway 的消融看下面这张时序分解图（LLaMA-3.3-70B，batch 16，prefill 90K token / decode 98K context）：

先说怎么读这张图：每个配置两条横向泳道，Comp 是计算单元忙碌区段、Mem 是访存忙碌区段；绿色是 Attention 阶段、紫色是 FFN 阶段；横轴是时间，条越短越快。左列 prefill（90K token）、右列 decode（98K context），三行是三个递进配置。

• Config 1（朴素 self-attention、无 prefetch、128×128 方阵）：prefill 泳道上 Comp 稀稀拉拉、Mem 密密麻麻（平均计算利用率 30.5%、访存 51.6%）——S、P 这些中间矩阵在片上放不下，反复往 HBM 搬，计算在等数据。decode 侧更惨，Comp 利用率 0.7%：batch 16 的 decode 每次只算 16 个 token 的新 Q，却要扫全部 98K 的 KV，纯搬运
• Config 2（换上 FlashAttention，仍是方阵、无 prefetch）：attention 的访存利用率从 52.0% 掉到 1.6%——中间激活留在片上了，HBM 流量确实被砍掉。但注意 Comp 利用率反而略降（30.5%→28.8%）：流量省了，计算该等还是等，因为没有 prefetch，每段计算开始前都要同步等数据到位；方阵形状的错配也原样保留
• Config 3（FlashAttention + prefetch + 1024×16 扁平阵列）：Comp 泳道几乎连成实心条（prefill 平均 68.1%，FFN 段 99.4%/100%）——prefetch 把"等数据"藏进了上一段计算，扁平阵列把"PE 空转"抹掉。最终 prefill 2.2 倍于 Config 1；decode 侧 10.2 倍，attention 的 SA 利用率从 0.5% 提到 6.1%——别嫌 6.1% 低，decode 本质 memory-bound，这已经是带宽用满后计算单元能达到的物理上限了

这张图的教训比数字本身有价值：Config 1→2 说明只砍流量不做重叠，收益兑现不了；2→3 说明重叠和形状匹配才是把纸面节省变成真实加速的临门一脚。

系统级对比（相同 HBM 容量预算、batch 开到各自能塞下的最大值）：

• LLaMA-3.1-8B agentic 负载（90K 输入/8K 输出）：吞吐 1.45 倍于 A100，能效 2.65 倍
• LLaMA-3.3-70B 同负载：吞吐 2.21 倍于 A100（也高于 H100 的 2.04 倍），能效 4.04 倍
• 对 TPU v6e 的优势更夸张，agentic 场景下 TPU 只有 A100 的 0.47-0.53 倍吞吐，PLENA 是它的 4.7 倍

注意一个诚实的细节：PLENA 的 TTFT（首 token 延迟）普遍比 GPU 长（70B 上 43s vs A100 的 29s），它赢在大 batch 吞吐和能效。这是给云端 agent 批量服务设计的芯片，不是给交互式对话设计的。

6. 一点个人 take

这篇文章我最喜欢的其实是第一章的 workload 分析——它把"agent 推理和 chatbot 推理是两种硬件负载"这件事用数据钉死了：token 量差两个数量级、计算热点从 FFN 挪到 attention、显存从 weight-dominant 变成 KV-dominant。做推理优化的同行不妨把这三条当 checklist，很多为 chatbot 调的系统配置（比如 batch 策略、KV Cache 配额）直接搬到 agent 场景就是负优化。

第二个感想："batch size 由容量决定、利用率由形状决定" 这个链条值得反复咀嚼。我们在 GPU 上做推理优化时习惯把 batch 当自由变量，但 128K context 下 KV Cache 一口气吃掉 39 GB，batch 根本不是你想开就能开的。PLENA 的解法是量化推容量墙 + 扁平阵列适配小 M，两手抓。GPU 上对应的思路其实就是 KV 量化 + 针对小 batch 的 kernel 特化，方向是相通的。

最后泼点冷水：这是模拟器 + RTL 层面的结果，没有流片；而且 fat GEMM 的形状随模型架构变化（MoE、GQA 比例都会影响），固定 4×1024 的阵列形状能不能吃遍未来模型，要打个问号——好在他们的 DSE 框架本身就是拿来回答这个问题的。

欢迎评论区交流。

---

顺带一提，PLENA 用 BoTorch 在 perplexity/延迟/面积三个目标上自动搜索加速器配置，这套多目标贝叶斯优化的方法论和 NAS/AutoML 是同一套内功。我们把这个方向的积累写成了一本书《动手学 AutoML：从 NAS 到大语言模型优化实战》，进阶篇专门讲搜索策略（进化/贝叶斯/强化学习/DARTS），应用篇也覆盖了 LLM 量化压缩——读完你会发现"搜网络架构"和"搜芯片架构"在方法上惊人地一致。