NEO: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference

framework 2411.01142
cpu-offloadingkv-cachegpu-memoryonline-inferenceasymmetric-pipelining
← Synthesis ← framework Survey Source ↗

§1 TL;DR #

Neo 将部分 decode attention 计算和 KVCache 从 GPU 卸载到本机 CPU，通过 asymmetric pipelining（非对称子批次重叠）和 load-aware scheduling（贪心回退到 GPU-only）提升在线推理吞吐。在内存受限的 T4 上达 7.5× 吞吐提升，A10G 上 26%，H100 上 14%，且不牺牲延迟。

§2 痛点 / 方法 / 结果 #

Q1 痛点 #

在线 LLM 推理中 GPU 吞吐高度依赖 batch size，但 KVCache 占用大量 GPU 内存导致实际 batch size 受限，GPU 计算单元严重闲置——这是 "GPU memory crisis"。先前 CPU offloading 方案（FlexGen、PowerInfer、HeteGen）用延迟换吞吐，不适合在线场景；FastDecode 完全卸载到 CPU 但需要 8 台远程 CPU 服务器处理单 GPU 的 attention，成本过高且 CPU 成为瓶颈。

核心 insight：decode attention 是 memory-bandwidth-bounded，GPU-CPU 在带宽上的差距（600 vs 200 GB/s）远小于算力差距（125 TFLOPS vs 1.2 TFLOPS），CPU 是 decode attention 的合理卸载目标。

Q2 方法 #

Asymmetric GPU-CPU pipelining（§3.1）：

  1. KVCache 分为 GPU-cache 和 CPU-cache；每个请求的 KVCache 完整存放在一处。优先使用 GPU-cache 最大化 GPU 内存利用。
  2. 将请求分为两个 非对称子批次：
  3. Batch-0：所有 prefill 请求 + GPU-decode 请求 + 少量 CPU-decode 请求（long linear stage on GPU, short CPU attention）
  4. Batch-1：大部分 CPU-decode 请求（short linear stage, long CPU attention）
  5. 两个子批次交替执行：batch-0 的 GPU linear stage 与 batch-1 的 CPU attention 重叠，反之亦然。非对称设计使 GPU 始终繁忙，避免 symmetric pipelining 中 GPU 等待 CPU 的 bubble。
  6. Layer-wise KV swapping：每层计算后立即传输 KV，与下一层计算重叠。

    7. Load-aware scheduling（§3.2）：

    迭代时间模型：

    $$T \approx T_{tr} = L \times \left( \max\{T_{po_0} + T_{pr_0},\; T_{ca_1}\} + \max\{T_{po_1} + T_{pr_1} + T_{ga_0},\; T_{ca_0}\} \right)$$

    约束条件（GPU 不空等 CPU）：

    $$T_{l_0} \geq T_{ca_1} \quad \text{and} \quad T_{l_1} + T_{ga_0} \geq T_{ca_0}$$

    调度器每次迭代执行 6 步决策：初始化 → 调度 GPU-decode → 调度 prefill → 调度 CPU-decode → 裁减 prefill（避免 swap-out 开销使 CPU 空等）→ 贪心选择（两批次方案 vs GPU-only 方案取 $T_{tr}/x$ 更大者）。

    Greedy 原则保证 Neo 永远不差于 GPU-only baseline。

    核心技术壁垒：asymmetric batch division 的 insight——将 prefill（GPU 密集）和 CPU-decode 放在互补子批次中，使得两个 max 项中 GPU 端的 linear stage 总是足够长以覆盖 CPU attention，从而消除 pipeline bubble。这要求对每次迭代的 GPU/CPU 时间精确建模并动态调整批次构成。

    Q3 结果 #

    • T4 + LLaMA-2-7B：在线吞吐 7.5×（T4 仅 16GB HBM，baseline batch size 极小）
    • A10G + LLaMA-3.1-8B：6.4%–26% throughput 提升，最高 CPU 配置（g5.16xlarge）达 79.3%
    • 2×H100 + LLaMA-3.1-70B：14.3% throughput 提升
    • 延迟：所有 percentile 与 GPU-only 基线持平
    • vs FastDecode+：Neo 始终优于 baseline，FastDecode+ 在长输出时降至 baseline 的 <60%

    §3 架构 / 方法图 #

    sequenceDiagram participant Sched as Request Scheduler (CPU) participant PQ as Prefill WaitQueue participant GQ as GPU Decode RunQueue participant CQ as CPU Decode RunQueue participant GPU as GPU (Linear + GPU Attn) participant CPU as CPU (PACPU Attention) Sched->>Sched: Per-iteration scheduling (6-step) Sched->>GPU: Batch-0: prefill + GPU-decode + few CPU-decode Sched->>CPU: Batch-1: most CPU-decode requests par Batch-0 GPU linear ↔ Batch-1 CPU attention GPU->>GPU: po₀ + pr₀ (linear stage, long) CPU->>CPU: ca₁ (CPU attention, fits under GPU time) end par Batch-1 GPU linear + GPU attn ↔ Batch-0 CPU attention GPU->>GPU: po₁ + pr₁ + ga₀ (linear + GPU attn) CPU->>CPU: ca₀ (CPU attention for batch-0 CPU requests) end GPU->>Sched: Iteration complete, return to scheduling

    关键组件：

    • GPU-cache / CPU-cache：分页 KVCache，每个请求的 KV 完整存在于 GPU 或 CPU 一侧。GPU-cache 优先分配。
    • PACPU (Paged-Attention-for-CPU)：C++ torch extension → ISPC 编译的向量化 CPU attention kernel，支持 AVX2/AVX512/Neon，paged attention 算法减少碎片。
    • Scheduling engine：每次迭代构建两种方案（asymmetric vs GPU-only），通过 $T_{tr}/x$ 比较选择更优方案（Greedy 原则）。

    §4 作者证明 #

    记号表 #

    符号含义
    $T$单次迭代总时间
    $T_{tr}$Transformer 层时间（占 >95%）
    $L$Transformer 层数
    $T_{po_x}$Batch-x 的 post-projection 时间
    $T_{pr_x}$Batch-x 的 pre-projection 时间
    $T_{ga_x}$Batch-x 的 GPU attention 时间
    $T_{ca_x}$Batch-x 的 CPU attention 时间
    $T_{l_x}$$T_{po_x} + T_{pr_x}$（batch-x linear 时间）
    $x$总 batch size

    方程物理意义 #

    主方程 $T \approx L \times (\max\{T_{l_0}, T_{ca_1}\} + \max\{T_{l_1} + T_{ga_0}, T_{ca_0}\})$：

    每层执行两个 pipeline stage。第一个 stage 中 GPU 执行 batch-0 的 linear ops，与 CPU 执行 batch-1 的 attention 重叠——取 max。第二个 stage 中 GPU 执行 batch-1 的 linear + batch-0 的 GPU attention，与 CPU 执行 batch-0 的 CPU attention 重叠。非对称性体现在 batch-0 的 linear stage 天然更长（包含 prefill + GPU-decode），能覆盖 batch-1 的 CPU attention。

    约束 $T_{l_0} \geq T_{ca_1}$ 和 $T_{l_1} + T_{ga_0} \geq T_{ca_0}$ 保证 Hiding CPU 原则：GPU 永不空等 CPU。

    6 项检查 #

    1. 迭代时间模型可验证：离线 profiling 估计各组件时间，线性插值。论文承认精度有限（短输出时 Neo 可能略差于 baseline）。
    2. Greedy 保证：每次迭代比较两种方案取最优——数学上保证 throughput ≥ GPU-only。实际中由于 profiling 误差和 scheduling 开销，存在轻微违反。
    3. CPU 带宽是绑定约束：sensitivity study 证实 throughput gain 正相关于 CPU 内存带宽而非核数（Fig 10a）。
    4. SwiftLLM vs vLLM 基线验证：单 GPU 性能持平；2-GPU 差 8.8%——多 GPU 结果存在基线偏弱问题。
    5. FastDecode+ 对比公平：作者自行实现 FastDecode+ 并赋予 Neo 的调度优势，仍然因完全卸载导致 CPU 过载。
    6. 层级 KV swap 开销未深度分析：论文提及 layer-wise swapping 但未定量剥离 PCIe 争用影响。

      7. §5 实验与数据 #

      在线延迟-负载曲线（Fig 6） #

      SettingNeo Max Throughput GainLatency Profile
      2×H100 + 70B + AC14.3% (at 2s latency)低 RPS 时持平，高 RPS 时优势明显
      A10G + 8B + AC6.4% (at 2s latency)类似
      T4 + 7B + OSC563% (at 1s latency)T4 HBM 极小，vLLM batch size 受限

      不同 input/output 长度（Fig 9） #

      • 短 output：Neo 可能略差（尝试 offload 的开销 > 收益）
      • 中等 output：gain 达峰值（GPU/CPU 时间恰好平衡）
      • 长 output：gain 逐渐下降，回落到接近 baseline
      • 极长 output：偶尔略差于 baseline（profiling 不精确）

      CPU 容量 sensitivity（Fig 10a） #

      EC2 InstanceCPU CoresPeak BW ApproxMax Throughput Gain
      g5.2xlarge4~100 GB/s12.2%
      g5.4xlarge8~100 GB/s13.3%
      g5.8xlarge16~200 GB/s29.7%
      g5.16xlarge32~400 GB/s79.3%

      带宽翻倍则 gain 近似翻倍——验证了 CPU 内存带宽是绑定约束而非算力。

      §6 论证链 #

      StepClaimEvidenceStrength
      1GPU 吞吐受限于 KVCache 导致的 batch size 瓶颈Fig 1 (NanoFlow), §2.1 引用 vLLM 数据Strong — well-established
      2Decode attention 是 memory-bandwidth-bounded，CPU 带宽差距远小于算力差距§2.2 定量比较 (600 vs 200 GB/s, 125 vs 1.2 TFLOPS)Strong — 硬件规格事实
      3Symmetric pipelining (FastDecode) 导致 GPU 内存浪费 + GPU 空等 CPU + 分批困难§3.1 三点分析 + Fig 3, 4Medium — 分析清晰但缺少 micro-benchmark 验证
      4Asymmetric pipelining 消除上述三个问题§3.1 Fig 5 pipeline 设计Medium — 设计合理但验证主要通过端到端结果
      5Load-aware scheduling 保证 never-worse-than-baseline§3.2 Greedy 原则 + 方程推导Strong — 数学保证（受 profiling 精度限制）
      6端到端验证：T4 上 7.5×，A10G 26%，H100 14%Fig 6, 9, 10Strong — 多硬件、多模型、多 workload

      §7 实现 cross-reference #

      • 基于 SwiftLLM 实现（~2K LoC 简化 vLLM），Neo 增加 4K Python + 1.5K C++
      • CPU kernel：PACPU (C++ torch extension → ISPC)，支持 AVX2/AVX512/Neon
      • GPU kernel：替换 Triton-JIT 为 CUDA C++ 以减少 kernel launch overhead（Python GIL 阻塞问题）
      • 多 GPU：Ray actors + NCCL，每个 actor 持有独立 CPU KVCache 分区
      • 开源承诺但未给出 URL

      关键实现细节：

      1. PACPU 的 Flash Decoding 并行策略：沿 request 维度分 partition，每个 task 访问连续内存 block，dispatch 到不同 CPU 线程，最后 reduce partial output。消除跨线程内存争用。
      2. Kernel launch overhead 是隐藏瓶颈：Python GIL 使 CPU data plane kernel 不能与 GPU control plane kernel launching 并行，被迫串行化。替换 Triton → CUDA C++ 是实测后的必要优化。

        3. 框架特定分析 #

        系统范围 #

        • 阶段覆盖：prefill + decode（不分离，共用 GPU 实例）
        • Serving 模式：selective batching（继承 Orca），连续批处理
        • 并行维度：TP（多 GPU），CPU offloading 作为额外并行维度
        • 部署模式：单节点，仅使用本机 CPU（不使用远程 CPU）

        调度与资源管理 #

        • 调度粒度：iteration-level（每次迭代重新决策）
        • 抢占策略：支持 GPU↔CPU 间 swap in/out（iteration 间）
        • 过载策略：无显式过载处理（非 disaggregated，无 early rejection）
        • KVCache 管理：paged attention 算法（继承 vLLM），GPU-cache + CPU-cache 分开管理
        • Swap 触发：调度器根据 GPU 内存余量决定 swap 方向

        胜负场景 #

        Workload RegimeNeoGPU-only BaselineWhy
        小内存 GPU (T4 16GB)7.5× 吞吐batch size 极小卸载大量 KVCache 到 CPU，释放 GPU batch capacity
        中等 GPU (A10G 24GB)6%–26%基线GPU 内存仍然受限但不极端
        大内存 GPU (H100 80GB)14%基线GPU 已能容纳较大 batch，offloading 边际收益小
        短 output持平或略差基线offloading overhead > 收益
        极长 output接近基线基线大部分迭代回退到 GPU-only 模式
        强 CPU 配置最高 79.3%-CPU BW 翻倍 → gain 翻倍

        部署上下文 #

        • Serving stage：prefill + decode 一体（非 disaggregated）
        • 硬件亲和性：内存受限 GPU 受益最大（T4 >> A10G >> H100）；CPU BW 越高越好
        • 生态集成：基于 SwiftLLM（简化 vLLM），论文计划集成到 vLLM
        • 迁移路径：替换 inference engine → SwiftLLM+Neo；未来可能为 vLLM 插件