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 倍。

agentic 负载的 token 用量远超 chatbot

上图（a）里 OSWorld、Chrome 这些 Computer Use Agent 任务的 token 用量直接飙到 100K 量级。更有意思的是（b）和（c）：随着 context 变长，计算热点从 FFN 迁移到 Attention（LLaMA-3-70B 上交叉点约在 19K 生成 token），而显存占用则从 weight-dominant 切换到 KV-dominant——128K context 下单个 batch 的 FP16 KV Cache 就要 39 GB，比很多模型权重本身还大。

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

带宽墙（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 全程围观，利用率惨不忍睹。

方形阵列 vs 扁平阵列的可达 FLOPs 对比

上图（a）很直观：同样数量的乘法器，方形阵列一半 PE 在闲着，扁平（flattened）阵列全部激活。（b）则展示了两条优化路径叠加的效果——扁平化抬高带宽墙下的利用率，量化把容量墙往右推（batch 能开更大）。

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 统一调度：

PLENA 加速器整体架构

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）：

三种配置下 prefill/decode 的时序分解

从 Config 1（朴素 self-attention、无 prefetch、128×128 方阵）到 Config 3（FlashAttention + prefetch + 1024×16 扁平阵列）：

Prefill：2.2 倍加速，平均计算利用率从 30.5% 提到 68.1%，FFN 阶段的计算利用率干到 99.4%
Decode：10.2 倍加速——方阵配置下 decode 的平均计算利用率只有可怜的 0.7%，扁平化后 attention 部分的 SA 利用率从 0.5% 提到 6.1%（decode 本质 memory-bound，这已经是把带宽用满后的结果）

系统级对比（相同 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 量化压缩——读完你会发现”搜网络架构”和”搜芯片架构”在方法上惊人地一致。