论文阅读：TightLLM

TightLLM 是一个面向高吞吐的单卡 LLM 推理系统，发表在 IEEE Transactions on Computers 2025 年 7 月刊。它要解决的场景是：模型太大、单张 GPU 装不下，只能把权重和 KV cache 卸载（offload）到 CPU 内存，靠 GPU 和 CPU 之间不断搬数据来完成推理。此时真正卡住吞吐的不是算力，而是这条 GPU↔CPU 的数据传输。TightLLM 的核心观察是：现有系统用一套固定不变的传输-计算重叠策略，而推理过程中的传输量其实一直在变，固定策略盖不住；应该让 offloading 策略随传输量的变化而自适应调整。

这篇工作在做什么#

一句话：把 offloading 从静态策略变成自适应策略，并用「重算换传输」的思路让计算和传输尽量刚好重叠。

具体有两个组件：

KV distributor（KV 分配器）：不再把整份 KV cache 都从 CPU 搬回 GPU，而是搬一部分、在 GPU 上重新计算另一部分。搬多少、算多少由一个整数线性规划（ILP）动态决定，目标是让计算时间刚好等于传输时间，两者完全重叠、GPU 不空转。这一招被作者称为 trade-compute-for-transfer（用计算换传输）。
Weight loader（权重加载器）：把一个 transformer block 的权重切成若干片，分摊到该 block 要处理的多个 batch 上流水线加载，从而把「一次性加载整块权重」这个大开销藏进多批计算里。

在 OPT 和 Llama 系列模型、消费级（RTX 3090）和服务器级（A800）两种机器上评测，解码阶段吞吐提升 1.3–23 倍，预填充阶段 1.2–22 倍；大规模离线任务的完成时间缩短 59.6%–94.9%。

前置：LLM 推理的显存与 offloading#

要理解这篇工作，先要清楚推理为什么这么吃显存，以及 offloading 到底在搬什么。

推理的两个阶段#

decoder-only 的 LLM 由多个 transformer block 堆叠而成，每个 block 内是一层注意力（attention）加一层前馈网络（FFN），前一个 block 的输出作为后一个 block 的输入。推理分两个阶段：

prefill（预填充）：把输入的 prompt 一次性喂进去，在一次前向里处理完所有输入 token，产出回复的第一个 token。
decoding（解码）：自回归地一次生成一个 token，每生成一个就把它接到序列末尾再生成下一个。

两个阶段都要反复用到前面所有 token 的键（ $K$ ）和值（ $V$ ）矩阵。为了避免每步都重算，系统会把这些矩阵缓存下来，就是 KV cache。

三块显存#

推理时 GPU 显存主要被三部分占用：

模型权重（model weights）：占大头，对某个确定的模型是固定值。
KV cache：随序列长度和batch size 线性增长，要等整个输入处理完才释放。
中间激活（intermediate activation）：block 之间传递的中间量，相对可以忽略。

关键在于第 2 项会涨。论文用 OPT-30B 举例：序列长度从 0.5K 涨到 8K，显存需求从 61 GB 涨到 140 GB；batch size 从 4 涨到 64，显存从 66 GB 涨到 225 GB。而 RTX 3090 只有 24 GB、A800 只有 80 GB——OPT-30B 在任何序列长度或 batch 下都超出 3090，序列超过 2K 或 batch 超过 8 就超出 A800。单卡显存和实际需求之间有巨大的缺口。

offloading 及其代价#

弥补这个缺口有三条路：多卡模型并行（需要多 GPU）、模型压缩/量化（牺牲精度），以及把权重和 KV cache 卸载到 CPU 内存——后者不需要额外硬件、也不损失精度，最适合个人用户，是近年比较流行的做法。

offloading 的做法是：GPU 显存里只保留当前这一步计算需要的那部分权重和 KV cache，其余都放在 CPU 内存；算完当前 block，就把下一 block 需要的权重和 KV cache 从 CPU 搬进来，同时把新产生的 KV cache 搬回 CPU。

offloading 的显存布局与数据流

代价显而易见：每一步都要在 GPU 和 CPU 之间搬数据，这条 PCIe 链路带宽有限，传输开销（transfer overhead）成了新瓶颈。offloading 会拉高单次请求的延迟，但对非交互式任务（批量翻译整个文库、文档处理、离线评测等，对单条请求的延迟不敏感、允许长达 24 小时的宽松时限）来说完全可以接受，这类任务真正在乎的是吞吐。OpenAI、Anthropic 都为这类场景提供了专门的高吞吐 API。所以 offloading 系统的首要目标就是：在有限显存下把吞吐做高。

现有系统压传输开销的两招#

为了尽量把传输开销藏起来，现有系统主要用两个技巧：

多批次推理（multi-batch inference）：一个 block 的权重加载进来后，连续服务多个 batch 再切换到下一 block，把「加载权重」这件事分摊到多个 batch 上，减少权重反复重载的次数。
传输-计算重叠（trans-comp overlap）：在计算当前 batch 的同时，预取（preload）下一 batch 要用的权重或 KV cache，让传输时间被计算时间盖住。

问题在于，这两招用的都是一套固定的流水线安排：不管当前这一步要搬多少数据，重叠的方式都一样。这就是所谓的静态 offloading（static offloading），也是 TightLLM 要动的地方。

问题：静态 offloading 为什么打不满吞吐#

作者做了一个实验来说明静态策略的局限：在 A800 上用 Accelerate 和 FlexGen 跑 OPT-175B，序列长度 0.5K–8K、batch size 8。这个配置需要 343 GB（0.5K）到 614 GB（8K）显存，远超 A800 的 80 GB。结果 FlexGen 和 Accelerate 的吞吐比 TightLLM 低 35.9%–95.4%（预填充）和 60.4%–95.6%（解码）；序列到 8K 时 Accelerate 直接 OOM（因为它只把权重卸载到 CPU，KV cache 全留在 GPU）。

根因在于：权重和 KV cache 的传输量在推理过程中是动态变化的，而静态策略对此视而不见。当某一步的传输时间超过了计算时间，GPU 就只能干等数据搬完，出现空转，吞吐随之下降。这个动态变化在两个阶段各有表现。

静态流水线为什么盖不住传输

解码阶段的传输瓶颈来自权重加载和 KV cache 加载：

一个 block 里的第一个 batch：既要加载模型权重，又要加载 KV cache。而一个 block 权重的加载时间通常远大于它的计算时间，计算根本盖不住权重传输，形成瓶颈。
非第一个 batch：只需要加载 KV cache，而 KV cache 的加载量随已生成 token 数线性增长。解码早期，KV cache 加载量小于计算量，能被盖住；但随着解码进行，KV cache 越来越大，加载时间逐渐超过计算时间，从某一步（第 $M$ 步）到后面（第 $N$ 步）迅速膨胀成瓶颈。

预填充阶段的传输瓶颈同样来自权重和 KV cache，但方向相反：预填充不需要从 CPU 加载 KV cache，而是要把一次性生成的大量 KV cache 卸载（offload）回 CPU。加上「当前层的 KV cache 没卸载完，下一层的 KV cache 没地方放、只能等」这个依赖，卸载会阻塞后续计算。这里的瓶颈到底是权重加载还是 KV 卸载，取决于 prompt 长度：prompt 长，KV 卸载是主瓶颈；prompt 短，权重加载是主瓶颈。

于是问题归结为两个：(1) 怎么应对解码时越来越长的 KV cache 加载时间？(2) 怎么把巨大的权重加载时间藏起来？ 静态策略回答不了，因为它没法随传输量的变化调整。

核心 idea：用计算换传输#

TightLLM 的关键洞见是：随着序列变长，计算时间的增速远慢于传输时间的增速，两者的差距越拉越大，GPU 上会空出越来越多的算力。既然 GPU 闲着，不如让它重新计算一部分 KV cache，这样就能少从 CPU 搬回同样一部分——用本来浪费掉的计算，换掉一段传输。

这一步和量化、稀疏加载 KV 之类的做法有本质区别：重算得到的 KV cache 和从 CPU 搬回来的完全一致，不损失任何精度。它换来的是传输量下降，代价只是多花一些 GPU 计算。

围绕这个思路，TightLLM 设计了两个组件——KV distributor 负责「搬多少、算多少」的动态分配，weight loader 负责把权重加载切碎藏进计算。系统在 CPU 内存里保存一份完整的权重（weight pool）和 KV cache（KV cache pool），GPU 显存只放当前需要的部分，两个组件在此之上做调度。

方法一：KV distributor#

KV distributor 的任务是把 KV cache 里的 token 动态地分成两拨：一拨从 CPU/磁盘加载，另一拨在 GPU 上重算。分配比例通过一个整数线性规划（ILP）求解（论文用 Gurobi），目标是让计算时间和传输时间尽量相等，从而完全重叠。

解码阶段：加载前段 + 重算后段#

解码时每步只新生成一个 token 的 KV，主要瓶颈是加载已有的 KV cache。KV distributor 的做法是：先把序列前段 token 的 KV cache 从 CPU 搬回来，再基于这些数据在 GPU 上重算后段 token 的 KV cache。重算的 token 数记作 $N_r$ ，序列总长记作 $N_{ctx}$ ，于是有 $N_r$ 个 token 走重算、 $N_{ctx}-N_r$ 个 token 走加载。

KV distributor：用重算换传输

这样做把「加载 batch $i$ 的 KV cache」这段传输，换成了「在计算 batch $i-1$ 时顺带重算」的计算。 $N_r$ 调大，计算变多、要加载的 KV 变少； $N_r$ 调小则反之。ILP 就是要找到那个让计算与传输刚好重叠的 $N_r$ ，消除 GPU 空转。

优化建模。吞吐定义为「生成的 token 数 / 总延迟」，在 $m$ 步迭代、 $n$ 层、batch size $b_s$ 下：

\max\; G_{token} = \max \frac{b_s\, m}{\sum_{i=1}^{m}\sum_{j=1}^{n} T_{batch}^{(i,j)}}

最大化吞吐等价于最小化每个 batch 的延迟。而一个 batch 的延迟取决于计算和传输里更慢的那个（两者重叠进行）：

\min\; T_{batch} = \min\bigl\{\max\{T_{compute}^{Attn}+T_{compute}^{FFN},\; T_{cache}\}\bigr\}

计算时间由重算 $N_r$ 个 token 的注意力和 FFN 决定。注意力层的 FLOPs 由 QKV 投影、掩码打分、输出投影三部分组成，合计 $b_s N_r\, 2 d_{attn}(4 d_{model}+(n_{ctx}+N_r))$ ；FFN 的 FLOPs 为 $b_s N_r\, 4 d_{model} d_{FFN}$ 。除以峰值算力 $FLOPS$ 与模型算力利用率 $MFU$ （由离线 profiling 得到），就是计算时间：

T_{compute}^{Attn} = \frac{b_s N_r\, 2 d_{attn}\bigl(4 d_{model}+(n_{ctx}+N_r)\bigr)}{FLOPS \times MFU^{Attn}}

T_{compute}^{FFN} = \frac{b_s N_r\, 4 d_{model} d_{FFN}}{FLOPS \times MFU^{FFN}}

传输时间则是加载那 $N_{ctx}-N_r$ 个未重算 token 的 KV cache 的时间，分为从 CPU 内存读和从磁盘读两段（ $C_c$ 、 $C_d$ 分别是 KV cache 存在 CPU 内存和磁盘上的比例， $Bdw_{cg}$ 、 $Bdw_{dc}$ 是 CPU→GPU 和磁盘→CPU 的带宽）：

T_{cache} = \frac{4(C_d+C_c) b_s (N_{ctx}-N_r) d_{model}}{Bdw_{cg}} + \frac{4 C_d\, b_s (N_{ctx}-N_r) d_{model}}{Bdw_{dc}}

作者补充测量了 FLOPS 随序列长度的变化：FFN 的 FLOPS 先陡升后趋于稳定，注意力的 FLOPS 在短序列时快速上升、长序列时因内存访问开销变大而下降。正因为实际算力随序列变化、不是常数，最优 $N_r$ 只能当作整数规划来解，而不是套一个闭式公式。

重算 KV cache：到底在算什么，GPU 上要多存什么#

“重算 KV cache” 里藏着一个关键问题：要算出某个 token 在第 $\ell$ 层的 K、V，得先有它在这一层的输入（hidden state）。这个输入从哪来，决定了两种完全不同的实现，代价也天差地别。

方案一：把每一层的输入都留在 GPU 上。 若保留了后段 token 在每层的 hidden state $h^{(\ell)}$ ，重算 KV 就只是一次投影 $K^{(\ell)}=W_K^{(\ell)}h^{(\ell)}$ 、 $V^{(\ell)}=W_V^{(\ell)}h^{(\ell)}$ ，几乎不花算力。但代价在显存：一个 token 每层的 hidden state 是 $d_{model}$ 个数，而它每层的 KV 是 $2 d_{model}$ 个数（K、V 各一份）。存下所有层的输入，体积差不多就是 KV cache 的一半——想省的显存和传输只省掉一半，而且这些 hidden state 同样随序列变长而膨胀，迟早又得往 CPU 搬。瓶颈没有真正解决。

方案二（TightLLM 采用）：GPU 上什么中间结果都不留，从头把这段 token 重算一遍。 只需要知道后段是哪些 token（它们的 id 本来就在，embedding 查表即得）。重算时，把这 $N_r$ 个 token 的 embedding 送进第一层，像一次小型 prefill 一样让它们逐层走完整个前向：每一层都做完整的注意力（对前段已载入的 KV 加上后段自己刚重算出的 KV 做 attention），再走 FFN；本层的 K、V 是这趟前向的副产品，算完立即用于本层注意力，用完不落盘、不搬运。

前面的计算量公式正好印证了这一点：重算的开销里既有完整注意力的三部分（QKV、mask、project），又有 FFN，而不只是 KV 投影——这说明它走的是完整前向，也就是方案二，而非方案一。

重算 KV cache 的实现：把后段 token 逐层重新前向一遍

为什么宁可多算、也不多存？ 因为解码阶段是传输 / 显存受限的：GPU 算力大量闲置，真正卡住的是 PCIe 带宽和显存容量（这也是”随序列变长，计算时间远低于传输时间”的由来）。方案二把几乎不占空间的 token id 留在 GPU，用富余的算力换掉稀缺的传输与显存；方案一恰好反过来，省了本就富余的算力，却继续占着稀缺的显存。这与训练里的 activation checkpointing（用重算换显存）是同一种思想：不存中间结果，需要时重新算。

两种实现的存储代价对比

所以”要在 GPU 上多存什么”？ 答案是几乎不用多存——只有那 $N_r$ 个待重算 token 的 id / embedding（ $N_r \times d_{model}$ ，且只有一份），相比想省下搬运的 KV cache（ $2 \times N_r \times d_{model} \times L$ ， $L$ 为层数）完全可以忽略。

这也解释了预填充阶段的那个配合（见下一节）：既然这批 token 反正要在解码时重算，预填充时算出它们的 KV 后就直接丢掉、不卸载到 CPU——它们的 KV 从一开始就没占用过搬运预算。

预填充阶段：干脆丢弃要重算的 KV#

预填充阶段没有 KV 加载，只有 KV 卸载。这里 KV distributor 用了一个巧妙的配合：预填充时，那些计划在解码阶段被重算的 token，其 KV cache 根本不保存、不卸载——反正解码时会重算出来，现在存了也是白存。少卸载这部分 KV，就直接砍掉了预填充阶段的一段传输开销，也缓解了「当前层 KV 没卸载完、下一层没法算」的阻塞。

方法二：weight loader#

weight loader 针对的是另一个瓶颈：一个 block 的权重加载时间远超单个 batch 的计算时间，靠一个 batch 的计算根本盖不住。它的思路是：既然一个 block 的权重加载进来后要服务多个 batch，那就把权重切成若干片，分摊到这些 batch 上流水线加载——在计算当前 batch 时，顺带加载下一 block 的一小片权重，每个 batch 只承担一小段传输。

解码阶段#

按要分摊的 batch 数把权重切片。每算一个 batch，就加载一片权重加上下一 batch 要用的 KV cache；等所有权重片都加载完，再同步、进入下一个 block。

短序列：切换 block 时如果不流水线，整块权重的加载会让第一个 batch 严重停顿。把权重切成（比如）4 片、每个 batch 加载 1 片，就能把这段延迟摊平藏掉。
长序列：KV cache 加载量变大，光靠切权重也盖不住全部传输。此时 weight loader 会和 KV distributor 联动——对前面的 batch 多重算一些 KV（ $N_r$ 调大）、对后面的 batch 少重算，在每个 batch 内部重新平衡传输与计算，把 GPU 空转压到最小。

更新后的建模。长序列下要把权重加载也纳入延迟，于是传输侧变成 KV cache 加载加上注意力、FFN 两部分权重加载：

\min\bigl\{\max\{T_{compute}^{Attn}+T_{compute}^{FFN},\; T_{cache}+T_{weight}^{FFN}+T_{weight}^{Attn}\}\bigr\}

其中权重加载时间被 $b_n$ （分摊的 batch 数）除掉，体现「分摊」的效果（ $W_c$ 、 $W_d$ 是权重存在 CPU 内存和磁盘上的比例， $d_{mlp}$ 是 MLP 维度）：

T_{weight}^{Attn} = \frac{4 d_{model} d_{attn}}{b_n} \times \left(\frac{W_d+W_c}{Bdw_{cg}} + \frac{W_d}{Bdw_{dc}}\right)

T_{weight}^{FFN} = \frac{2 d_{model} d_{mlp}}{b_n} \times \left(\frac{W_d+W_c}{Bdw_{cg}} + \frac{W_d}{Bdw_{dc}}\right)

预填充阶段#

预填充里权重加载和 KV cache 卸载可以并行进行，于是出现两种可能的瓶颈：KV 卸载开销大于权重加载时，瓶颈在 KV 卸载；反之在权重加载。权重加载量是固定的，所以瓶颈由 prompt 长度决定——prompt 越长，要卸载的 KV 越多，KV 卸载越可能成为主瓶颈。weight loader 对这两种情况都做了处理。

合起来：自适应 offloading 算法#

把两个组件串起来就是论文的 Algorithm 1。对每一批请求，先估算平均序列长度，据此算出 KV cache 传输时间 $T_{cache}$ 、权重加载时间 $T_{weight}$ 、计算时间 $T_{compute}$ ，再解出最优重算 token 数 $N_r$ 。然后逐层、逐 batch 地推进：

每层的第一个 batch 负责加载下一层的权重片；
解码阶段，加载下一 batch 的 KV cache，并按 $N_r$ 决定重算多少；
预填充阶段，计算当前 batch，并把上一 batch（扣除要重算的那部分后）的 KV cache 卸载出去。

整套流程的关键就在于 $N_r$ 、权重切片数、加载/卸载比例都是根据当前状态动态算出来的，而不是写死的——这正是「自适应」相对「静态」的差别。

效果#

评测在两台机器上进行：消费级 $C_1$ （RTX 3090 + 24 核 CPU + 128 GB 内存）、服务器级 $C_2$ （A800 + 26 核 CPU + 256 GB 内存）。模型用 OPT（13B/30B/66B/175B）和 Llama-2（13B/70B），prompt 长度 0.5K–8K，每条生成 32 token。基线是 Hugging Face Accelerate 和 FlexGen。整个系统约 5000 行代码，基于 PyTorch，用多个 CUDA stream 并发地跑计算和传输。

解码阶段：吞吐提升 1.3–23 倍。 $C_1$ 上平均最高 10.4 倍， $C_2$ 上最高 19.9 倍（ $C_2$ 跑更大的模型，传输瓶颈更重，TightLLM 收益更大）。序列越长收益越明显：0.5K 时 2.4–12.4 倍，4K 时 3.3–22.8 倍。就 OPT-175B 而言，单 token 延迟降低 59.7%–95.7%——重算让计算时间增大到原来的 3.1–5.4 倍，但传输时间下降了 68.9%–97.5%，净效果是延迟大降。
预填充阶段：吞吐提升 1.2–22 倍。这里的吞吐不随序列长度单调变化，因为瓶颈在「KV 卸载」和「权重加载」之间来回切换（模型大/序列短时是 KV 卸载，模型小/序列长时是权重加载）。单 token 延迟降低 35.9%–95.5%；预填充没有重算、计算量不变，传输时间下降 64.6%–99.1%。
大规模任务：在翻译、文本分类、特征抽取、摘要四个数据集上，端到端完成时间缩短 59.6%–94.9%。
消融实验：解码阶段去掉 KV distributor 吞吐降 12.6%–43.1%，去掉 weight loader 降 4.4%–74.8%；预填充阶段去掉 KV distributor 降最多 46.6%，去掉 weight loader 降最多 38.2%。两个组件各自都有明显贡献，且在不同阶段主次不同。

局限与讨论#

KV cache 加载的最小粒度是「整层」。注意力需要把该层所有相关 token 的 KV 一次性放进 GPU 才能算，没法只加载一个更小的 token 窗口。序列并行可以把注意力按 token 切分到多台机器，但在单卡上按 token 切会降低计算效率。更细粒度的传输-计算平衡是作者留的未来工作。
运行时动态调度有额外开销。TightLLM 走的是基于离线建模（profiling + ILP）的路线，而不是完全在运行时动态决策。作者也指出，纯运行时的细粒度调度虽然控制力更强，但引入的信息采集、参数同步等开销不一定划算，未必总能超过基于模型的方法。
和压缩/量化正交，可叠加。TightLLM 的重算是无损的，和量化、稀疏、KV cache 压缩这些有损或改变计算量的方法是正交的，可以结合起来进一步降低显存、提高吞吐。相比之下，InfiniGen（稀疏加载 KV）、PowerInfer（按 FFN 稀疏卸载权重）这类工作要么损失精度、要么没充分利用 GPU 算力，而 TightLLM 在不损失精度的前提下动态调整算-传比例，是它区别于同类工作的地方。