论文阅读：Improving Throughput-oriented LLM Inference with CPU Computations

一层每次要搬什么#

现在的 LLM 基本都是 decoder-only 架构，每个 decoder 层由一个多头注意力（MHA）和一个前馈网络（FFN）组成。推理时，为了算完一层，需要从慢存储搬三样东西到 GPU：

• MHA 权重：算注意力用；
• FFN 权重：算前馈网络用；
• KV cache：注意力要读的历史 key/value。

其中 MHA 的计算依赖 MHA 权重和 KV cache 都到齐，而 FFN 只要它自己的权重到了就能开算。KV cache 的大小由下式给出：

其中 $b$ 是 batch size，$t$ 是总序列长度（prompt + 已生成），$hd$ 是每个注意力头的隐藏维度，$p$ 是精度字节数。关键在于它随 batch size 和序列长度线性增长——吞吐导向场景把 batch 开得很大，KV cache 就从一个可忽略的中间量变成了必须认真搬运的大块数据。

两个阶段：prefill 与 generation#

一次生成分两个阶段。Prefill 把整段 prompt 跑一遍、建立 KV cache，它一次处理所有输入 token，是重计算的（compute-bound）：512 个 token 的 prefill 计算量大约是 generation 单步的 512 倍。Generation 阶段则每次只生成一个新 token，每步都要把整层权重和 KV cache 读一遍，却只做很少的计算，是重访存的（memory-bound）。offloading 推理的瓶颈主要出在 generation 阶段。

三种 workload 排布#

FlexGen 提供了两种把 batch 铺进有限显存的方式，加上压缩，一共三种典型排布：

• 单流水线（single-pipelined）：把尽量多的 batch 塞进 GPU，逐层遍历权重。计算能被传输完全覆盖，但因为每个 batch 每步只产出一个 token，逐层搬权重的传输成了瓶颈。
• 多流水线（multi-pipelined，即 FlexGen 的 zigzag 调度）：把 batch 分成若干组，一层权重加载后先服务完所有组，再进入下一层，从而摊薄权重的重复加载。代价是每组都要搬自己的 KV cache，多次 KV cache 传输会拖住 MHA。
• 多流水线 + 压缩：把权重或 KV cache 压缩以减少传输量，但 GPU 侧要额外做一次解压缩。

三种排布里，作者对 FlexGen OPT-30B 的实测都指向同一个结论：数据传输是瓶颈，GPU 存在明显空闲。下图是单流水线情形的直观示意——PCIe 上的传输块（宽）远长于 GPU 上的计算块（窄），GPU 大部分时间在等数据。

3.1 Prefill：不 offload，只把 KV 异步回传#

Prefill 是重计算的，把它丢给慢得多的 CPU 会严重拖慢整体，所以这篇工作不把 prefill 的计算 offload 到 CPU。它在 prefill 阶段做的是另一件事：一旦某层的 KV 矩阵算出来，就立即异步地把它从 GPU 传回主机（device-to-host），用一个独立进程处理这个回传，让它和 GPU 的计算重叠。

这一点正好补上了 FlexGen 的一个短板：FlexGen 在 prefill 阶段是同步回传 KV cache 的，batch 一大就会卡住计算。用独立进程做异步回传后，随 batch 增大的这段传输阻塞被缓解，prefill 阶段也拿到了延迟收益。

3.2 让 attention 在 CPU 上跑得快：去掉多余的 reshape#

把注意力交给 CPU，还有一个现实障碍：现有 decoder 层的实现是为 GPU 高度优化的，直接搬到 CPU 上跑并不高效。作者以 OPT 的注意力实现为例，它对 KV cache 做了两次 reshape：第一次是算 KV 时的常规变形，第二次是为了让后续访问在 GPU 上更高效而做的重排。这第二次 reshape 在 GPU 上代价很小，但在 CPU 上会触发一次很大的内存拷贝。

作者的做法是换一种 KV cache 的存放形状，把它直接排成

这样只需要第一次 reshape，第二次就完全省掉了。实测这一改动让 CPU 上（BF16）的注意力层延迟降低约 30%~35%。注意力在 CPU 上算得越快，动态划分模块就能把越多的工作分给 CPU，反过来又进一步压低了 GPU 的计算和传输。

3.3 切分点：一层怎么在 CPU 和 GPU 之间切#

把一层里按顺序排列的权重看成一个序列，两个切分点 $fp$（first split point）和 $sp$（second split point）把它切成三段：$[{:}fp]$ 和 $[sp{:}]$ 交给 CPU 算，中间的 $[fp{:}sp]$ 交给 GPU 算（对应权重要搬过 PCIe）。放到上一节的时间线上，$[{:}fp]$ 就是整块 MHA，$[fp{:}sp]$ 是 FFN 主体，$[sp{:}]$ 是分回给 CPU 的 FFN 尾段。

对应的两侧耗时可以写出来。设 $C[i]$、$G[i]$、$H[i]$ 分别是第 $i$ 块权重在 CPU 上的计算、在 GPU 上的计算、以及搬到 GPU 的传输延迟，$b$ 是 batch size，$H[KV]$ 是 KV cache 的传输延迟：

CPU 侧每个 batch 都要重算，所以乘 $b$；GPU 侧的传输只做一次，但计算要对每个 batch 重复。由于 CPU 和 GPU 是并行工作的，一层的实际耗时是两者的较大值，最优切分就是让这个较大值尽量小。压缩情形只需在 GPU 侧再加上解压缩延迟 $D[KV]$，Direct Storage Access 情形则要减去那部分绕开主机、直接从 NVMe 进 GPU 的传输——形式几乎一样，只是把对应的项加进来或减出去。

3.4 多流水线、压缩与 Direct Storage Access#

上面是单流水线的基本形态。当 DRAM 足够、要把 batch 开到超过显存容量时，就切换到多流水线，这里有两个实现细节：

• pinned memory 的取舍。把可分页内存拷进锁页（pinned）内存能让传输异步、更快，但这次拷贝本身很贵。所以作者让常驻权重留在可分页内存，只有推理期间真正要传给 GPU 的数据才放进锁页内存。
• 别每个 batch 都卡住 GPU。如果 KV cache 需要传给 GPU，就不能每组都同步等它传完，而要把工作划分成「让 CPU 一直忙到最后一块 KV cache 传完为止」，把 GPU 侧的传输时间全程盖住。

压缩的情形下，GPU 要多做一次解压缩。解压缩可以和 host-to-device 传输重叠：第一块数据一到就开始解压，不必等所有块传完。因为解压缩推迟了 GPU 上其它计算的开始，CPU 就能相应地多分到一份计算。

Direct Storage Access（DSA，即 NVIDIA DirectStorage / GPU Direct Storage）允许把 NVMe 上的数据直接读进 GPU，绕开主机内存，也就省掉了 host-to-device 这一跳和解压缩（权重以未压缩形式存在 NVMe 上）。作者为每个 NVMe 设备开一个独立进程并行加载，而不是一个进程串行地把所有权重读进 DRAM。

3.5 动态划分：自动找切分点#

前面所有排布都假设切分点是已知的。但对一个给定的模型 / 配置 / 硬件，这些切分点要手动测、而且每换一个用例就得重来。这篇工作把它自动化，做成一个动态工作划分模块。

流程分两期。预处理先做硬件 profiling：把单层里各块权重的 CPU 计算、GPU 计算、host-to-device 传输、解压缩、以及 NVMe→GPU 传输延迟都实测出来。这里有个容易忽略的点——多个传输、以及传输与 CPU 计算同时进行时会相互抢占资源、彼此拖慢，所以延迟必须在「并发」条件下测，而不能拿单独测的值去线性叠加。NVMe 这一侧尤其明显：不同 NVMe 挂在不同 PCIe 通道上、带宽各异，简单按单独测得的带宽线性分配传输量会得到次优解，作者用一个迭代过程逐步微调分给每个 NVMe 的传输量来逼近平衡。

拿到延迟表后，用动态规划搜切分点。状态就是前面的两侧耗时取较大者：

遍历所有可能的 $(fp, sp)$，挑出让一层端到端最快的那一对。之后由 Device Map Creator 逐个权重决定去向：常驻 DRAM、传给 GPU、还是放 NVMe。

运行时用三个进程：一个负责计算、一个负责 host-to-device、一个负责 device-to-host（开了 DSA 时每个 NVMe 设备再各配一个进程）。用多进程而非多线程有两个好处：一是不同 CUDA context 里的传输可以处理未锁页的数据、省掉锁页拷贝；二是能在同一进程内灵活管理动态切分点、把 host-to-device 传输和 CPU 计算重叠。进程之间用 barrier 同步：等权重传完才开始算、算完 KV 才开始回传。到了 FFN 层，因为 GPU 侧不产生要回传主机的数据，device-to-host 的等待就省掉了。