SSD 三级卸载先驱

元数据

论文标题:FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU
出处:ICML 2023
作者方向:Stanford / CMU / UC Berkeley 联合
关键词:single-GPU LLM、offloading、throughput-oriented inference、policy search

一句话精髓

第一个把”模型权重 + KV + activations”全部当成可在 HBM-DRAM-SSD 之间搬运的张量,通过线性规划自动找最优放置策略,实现单 GPU 跑超大模型。

解决的问题

2023 年初的痛点:

175B 参数模型至少要几百 GB 显存,单 GPU 完全装不下
已有方案(DeepSpeed-Inference、HuggingFace Accelerate)可以卸载,但靠人手调且策略保守
学术界缺乏一份”如何把张量在多级存储间搬”的系统性最优方案

FlexGen 把这个问题形式化为优化问题。

关键 idea

设计	内容
三类张量统一管理	weights、KV cache、activations 各有自己的 placement 决策
块级流水	每个 transformer block 独立决定加载源,前向时流水
计算-加载-卸载三阶段重叠	经典的双缓冲思路放大版
线性规划求解放置	给定 GPU/CPU/SSD 容量与带宽,LP 求最优块大小 + 放置
量化复合	把 4-bit 量化 + 卸载 + 流水组合起来加倍收益

🧠 关键洞察:LLM 推理的张量搬运是高度规则的(每 layer 拉权重 → 算 → 写 KV),非常适合 LP 这种”全局已知”的优化方法,不需要在线决策。

关键架构图

  ┌─────────────────────────────────────┐
  │   Policy Solver(LP / 启发式)       │
  │   输入:GPU 内存、CPU 内存、SSD、模型│
  │   输出:每个块在哪一层、何时移动     │
  └────────────┬────────────────────────┘
               │ 静态调度
               ▼
  HBM:      [Wt] [KV] [Act]
   ↑↓ async load/store(双缓冲)
  DRAM:     [Wt cache] [KV spill]
   ↑↓ async
  SSD:      [Wt overflow] [KV cold]

关键数据(论文报告)

单卡(T4 / 1×A100)上跑 OPT-175B,throughput 比此前最强 offload 方案提升一个数量级
在 batch 推理场景下,一台单卡服务器能服务的总 token/s 达到此前认为不可能的水平
量化 + 卸载组合相比单一手段,额外提升约 2-3×

局限

批量吞吐导向,延迟不友好——专为离线大批量任务设计,在线低延迟用不上
假设 workload 静态(prompt 长度分布固定),多用户混合时 LP 失效
未考虑多卡协同(纯单卡)
量化策略偏保守,后续工作(KIVI / AWQ)在 KV 量化上更激进

对本项目的启示

🌟 方法论标杆——FlexGen 证明了”放置决策可以形式化为优化问题”。这正是本项目”token 单位成本建模”的方法论原型:

从静态 LP 到在线优化:FlexGen 假设 workload 静态,生产环境必须做在线学习或模型预测控制
从单卡到分离式:把 FlexGen 的”GPU/CPU/SSD”扩展到”HBM/DRAM/远端 RDMA 池/SSD”,约束条件多一档
从权重/KV/Act 到多类型:把变量集合扩到 KV + 向量索引 + 多模态 blob——每类有自己的访问模式约束
延迟约束加进 LP:FlexGen 主要约束容量,我们要把 SLO(TTFT、TPOT)写进目标函数
复用 FlexGen 的策略搜索框架:作为 baseline 实现,后续把它升级成在线版本

横向对比

系统	决策方式	在线 / 离线	多卡	多类型数据
DeepSpeed-Inference	启发式人调	在线	✅	weights+act
HuggingFace Accelerate	简易策略	在线	✅	weights
FlexGen	LP 静态最优	离线	❌	weights+KV+act
PowerInfer / PowerInfer-2	hot/cold 神经元统计	在线	❌	weights
本项目目标	在线优化 + 模型驱动	在线	✅ + 分离式	KV+向量+blob

待精读

LP 求解的具体形式(变量、约束、目标函数)——这是 token 成本建模的直接参考
各档量化下的精度损失实测
FlexGen 后续是否有”在线版本”工作

08. FlexGen — 单卡跑大模型推理的离线优化

搜索