LithOS On GPU
LithOS on GPU
lithos_analysis
lithos＿analysis
LithOS 论文细节分析
基于论文 LithOS：An Operating System for Efficient Machine Learning on GPUs、CMU 作者博文、NVIDIA Hopper 架构资料与 NVIDIA MPS 文档整理而成。 实现层的关键路径讲清楚： Driver API interposition－＞virtual streams－＞TPC Scheduler－＞stealing－＞atomization－ ＞right－sizing－＞MPS－＞GPU execution
目录
1．一句话结论 2．LithOS 解决了什么问题 3．LithOS 在软件栈中的位置 4．为什么是 TPC，不是 SM 5．GPC／TPC／SM／block／warp 关系 6．LithOS 的整体架构 7．CUDA Driver API interposition 是如何工作的 8．MPS 的作用以及它和 LithOS 的关系 9．TPC Scheduler／stealing／atomization 10．kernel 如何切成 atom，如何保证结果正确 11．right－sizing：kernel 需要几个 TPC 是怎么估出来的 12．cache／shared memory／block independence 对 LithOS 的意义 13．评测图怎么看：为什么 LithOS 同时兼顾吞吐和 SLO 14．代码是否开源，以及能还原到什么程度 15．实现边界与未公开细节 16．参考文献
一句话结论
LithOS 不是替代 NVIDIA GPU 硬件调度器，而是在 CUDA Driver API 之上插入一层 软件 OS层，把 kernel 的提交 和 kernel 真正进入 GPU 执行 解耦，然后在软件里做： －TPC 粒度的空间调度；
- kernel atomization，把长 kernel 按 thread－block 范围切成多个 atom；
- 在线 latency predictor＋right－sizing，估算＂够用的最小 TPC 数＂；
- MPS 之上的多进程并发，而不是单纯 context time－slice；
- stealing＋power management，提高 GPU 利用率与能效。
换句话说，LithOS 更像一个 GPU OS layer，而不只是一个 scheduler patch。
LithOS 解决了什么问题
论文的出发点很明确：数据中心 GPU 很贵，但很多场景下利用率并不高。问题不只是＂任务少＂，而是现有共享方式太粗：
- 整卡独占：资源浪费；
- 纯时间片：高优任务尾延迟差；
- MIG 之类静态切分：隔离好，但不够灵活；
- MPS：吞吐通常高，但容易出现性能干扰。
LithOS 的核心判断是： GPU 需要一种更像操作系统的资源管理方式，而不是完全依赖驱动和硬件内部的默认调度。论文明确提出了四个核心机制：
1．TPC Scheduler 2．Kernel Atomization 3．Hardware Right－Sizing 4．Transparent Power Management
LithOS 在软件栈中的位置
LithOS 不是在 CUDA kernel 内部工作，也不是修改模型图优化器。它工作在 host 侧 CUDA Driver API interposition 这一层。
典型路径可以理解为：
Application / Framework
-> CUDA Driver API
-> LibLithOS (interposition)
-> launch queues / scheduler / atomizer
-> NVIDIA driver / MPS
-> GPU
这意味着：
- 应用、框架、runtime 仍然按原样调用 CUDA；
- LithOS 先接住这些调用；
- 再决定是否立刻提交、是否推迟、是否切 atom、是否换 stream、给多少 TPC。这个位置非常关键，因为一旦 kernel 已经提交给 GPU，很多资源属性就改不了了。 LithOS 的价值就在于：把＂何时提交＂和＂提交成什么形态＂这件事，提前到 host 侧软件层统一处理。
为什么是 TPC，不是 SM
这几乎是理解 LithOS 的第一关键点。
核心结论
SM 是性能建模单位，TPC 是 Lith OS 可控的分配单位。 也就是说：
- block 最终还是跑在 SM 上；
- 但 LithOS 施加约束、做隔离、做 stealing 时，粒度是 TPC；
- 在 H100 这类架构上，可近似看成 $1 \mathrm{TPC}=2 \mathrm{SM}$ 。
为什么不能直接按 SM
这不是因为 LithOS＂不想＂按 SM，而是因为在它面向的 NVIDIA 硬件上，可稳定控制的最细边界就在 TPC。 公开材料里，CMU 作者博文直接写到： Hardware limitations require that LithOS schedule pairs of SMs，or，TPCs． 因此：
- 执行单位仍是 $S M$ ；
- 分配单位 是 TPC；
- 对 LithOS 来说，最小可控粒度通常不是 1 个 SM ，而是 1 个 TPC 。
工程上怎么理解
如果一个 kernel 经过 occupancy／latency 估算后＂理论上需要 3 个 SM ＂：
- 这只是 SM 级性能模型 的结果；
- 最终分配时，LithOS 需要把它 量化到 TPC 桶 里；
- 在 H100 这类机器上，它更像会给 2 个 SM 或 4 个 SM 对应的 1 或 2 个 TPC。
因此，你可以把它理解为：
- 算账 用 SM；
- 发钱 按 TPC。
GPC／TPC／SM／block／warp 关系
下面这张图是结合 Hopper 文档和论文讨论重绘的硬件层级图。
Redrawn GPU hierarchy for H100：GPC－＞TPC－＞SM Based on NVIDIA Hopper documentation and LithOS paper discussions
SM：the execution unit that runs thread blocks and warps． TPC：the allocation unit LithOS can control on H100－class GPUs；typically 2 SMs／TPC． GPC：a higher－level cluster that groups multiple TPCs． 图 1．H100 类 GPU 的层级关系重绘图。GPC 是更高一级簇，TPC 是由多个 SM 组成的分配边界， SM 是实际执行 block／warp 的单元。重绘依据：Hopper 资料与 LithOS 论文对 TPC 调度的描述。见文献［3］［4］。
1．GPC／TPC／SM
硬件层级通常可理解为：
GPC > TPC > SM
- SM：真正执行 block／warp 的计算单元；
- TPC：比 SM 更上的硬件簇；在 H100 上通常是 $2 \mathrm{SM} / 1 \mathrm{TPC}$ ；
- GPC：更大的分区，包含多个 TPC。
2．block 和 warp
在 CUDA 中：
- 一个 kernel 有很多 thread blocks（CTA）
- 一个 block 里有很多线程
- 在 NVIDIA GPU 上通常 32 个线程组成 1 个 warp
因此关系可以写成：
Kernel -> Grid -> Blocks -> Warps -> Threads
3．block 和 SM
关键规则：
- 一个 block 只会在一个 SM 上执行；
- block 不会拆到多个 SM；
- block 内线程可共享 shared memory；
- warp 是 SM 内真正的执行／调度粒度。
这也是为什么 LithOS 的 atomization 可以按 block range 切，而不是去切 warp 或切 block 内部状态。
LithOS 的整体架构
下面这张图是根据论文 Figure 7 和 CMU 博文重绘的实现视图。
LithOS architecture（redrawn from paper Figure 7 and CMU blog）
Host－side interposition＋software scheduling＋GPU execution via MPS
Submission is decoupled from execution：app launches are first queued in LibLithOS． Dispatcher threads choose when to submit work，how many TPCs to allocate，and whether to atomize long kernels． GPU execution remains on NVIDIA hardware；LithOS does not replace low－level CTA／warp scheduling．
图 2．LithOS 架构重绘图。上层应用通过 LibLithOS 进入软件调度层，软件层包含 virtual streams、TPC scheduler、atomizer、latency predictor、tracker 等模块，底部仍然通过 NVIDIA driver／MPS 执行。重绘依据：论文 Figure 7 与博文实现说明。见文献［1］［2］。
这个架构最关键的实现思想
1．应用继续像平常一样用 CUDA 2．LibLithOS 拦截 Driver API 3．Launch 先进入 virtual streams／launch queues 4．Dispatcher thread 决定真正提交时机 5．调度器决定给多少 TPC、是否 stealing、是否 atomization 6．必要时用 Prelude／wrapper 形式发射 7．底层仍通过 NVIDIA driver／MPS／GPU 真正执行 因此 LithOS 的关键不是＂把 GPU 内核重写掉＂，而是： 把原本隐藏在 driver／hardware 内部的多租户调度权，尽可能上移到软件层。
CUDA Driver API interposition 是如何工作的
1．拦截的是 host API，不是改 kernel 逻辑
LithOS 所谓 interposition，拦截的是 host 侧的 Driver API，例如： －cuInit －cuStreamCreate －cuModuleLoad －cuModuleGetFunction －cuLaunchKernel 它接住的是 函数调用链，而不是去修改已经编译好的 PTX／SASS 数学逻辑。
2．Linux 上最常见的实现思路
工程上常见做法是：
- 先放一个自己的动态库；
- 对外提供与 CUDA Driver API 同名的函数；
- 内部通过 dlsym（RTLD＿NEXT，．．．）或 driver entry point 机制找到真正的 NVIDIA 实现；
- 先做自己的逻辑，再决定是否调用真实函数。
可以把它想成下面这种伪代码：
CUresult cuLaunchKernel(...) {
record_launch_metadata(...);
enqueue_to_virtual_stream(...);
// 对应用先表现为"提交成功"
return CUDA_SUCCESS;
}
之后由 LithOS 自己的 dispatcher thread 再决定：
- 何时真正调用真实 cuLaunchKernel
- 是直接发原 kernel
- 还是发 Prelude kernel
- 用哪个 stream
- 给多少 TPC
- 是否切 atom
3．为什么这很重要
因为如果应用直接把 kernel 发给 GPU： －优先级难改；
- TPC 分配难改；
- 无法切 atom；
- 无法根据系统状态重排提交顺序。
所以 interposition 的目标不是＂改 kernel 算法＂，而是抢回 launch control。
MPS 的作用以及它和 LithOS 的关系
下面这张图是根据 NVIDIA MPS 文档重绘的 client－server 路径图。
MPS client－server path（redrawn from NVIDIA MPS architecture docs）
Without MPS，contexts are often time－sliced；with MPS，multiple clients can overlap on the GPU．
Without MPS
GPU schedule： A ｜ B ｜ A ｜ B
Different contexts can serialize via time slicing．
GPU：overlapping kernel execution
Example ove
| A1 | B1 | C1 | A2 | B2 |
|---|
NVIDIA docs：MPS is a binary－compatible client－server implementation of the CUDA API． Pre－Volta clients funnel work through the MPS server；on Volta＋clients submit more directly while the server mediates shared resources．
图 3．MPS client－server 路径重绘图。没有 MPS 时，不同 context 更容易在 GPU 级别 time－ slice；有 MPS 时，多个 client 的工作可以通过共享资源路径并发执行。重绘依据：NVIDIA MPS文档的 client－server architecture 说明。见文献［5］［6］。
1．MPS 是什么
MPS（Multi－Process Service）是 NVIDIA 的多进程 CUDA 共享机制。 官方文档把它描述为：
核心意义是：
- 多个进程各自有自己的 CUDA context；
- 如果没有 MPS，这些 context 在 GPU 上往往更接近 time－slicing；
- 有 MPS 后，多进程工作可以更好地并发。
2．MPS 为什么重要
如果没有 MPS：
- 进程 A 有 context A
- 进程 B 有 context B
- 它们在 GPU 上的执行容易出现串行化／时间片化
如果有 MPS：
- 多个 client 的工作可以更好地重叠；
- 更容易把 GPU 资源填满；
- 减少 context switching 带来的浪费。
3．LithOS 为什么建在 MPS 之上
LithOS 自己做的是更高一层的策略调度，但它仍然需要底层能够让不同 context 的 work 真正并发。 如果底层还是只能粗粒度 time－slice，那么上层的软件调度再聪明，空间共享能力也会大打折扣。
所以可以把它理解成：
- MPS：提供多进程 GPU 并发底座
- LithOS：在这个底座上继续做更细的 TPC 级调度、atomization、right－sizing
TPC Scheduler／stealing／atomization
1．TPC Scheduler 不是硬件 CTA scheduler
原生 GPU 的 CTA／block 落位到 SM，是底层执行系统负责的。 LithOS 的 TPC Scheduler 并不直接决定：
- 这个 CTA 一定去 SM ＃17
- 下一个 warp 一定去 SM ＃32
它做的是更上层的事： －哪个 kernel／atom 先进入 GPU
- 给它几个 TPC
- 哪些 TPC 是保底 quota
- 哪些可以被 steal
- 什么时候延迟 best－effort work
2．它怎么知道什么＂空闲＂
LithOS 并不是直接盯着某个硬件寄存器看＂某个 SM 是否刚释放＂。 它更像维护了一本软件账本： 1．先知道自己发了什么 work； 2．通过 tracker／sync queue 知道哪些 work 完成； 3．通过在线 latency predictor 和 per－TPC timer 估计哪些资源快空闲。 所以它对＂空闲＂的认知是：
- 确定空闲：完成事件已到
- 预测快空闲：timer 推断快结束
- 暂时别借：长 kernel 或关键高优任务正在占用
3．stealing 是什么
stealing 的本质是：
某个应用虽然＂配到了＂一些TPC，但当前并没有把它们都吃满，于是 LithOS 暂时把空着的 TPC 借给别的 workload。
注意，这不是抢一个正在忙的 TPC，而是借暂时空着或可快速归还的那部分资源。
kernel 如何切成 atom，如何保证结果正确
下面这张图是根据论文对 atomization 的算法描述重绘的 block－range 切分图。
Kernel atomization by block range（redrawn from paper algorithm discussion）
Example：a grid of 64 thread blocks is split into 4 atoms．Prelude launches the same kernel shape but executes only a selected block ri
图 4．atomization 的 block－range 视图重绘图。一个 grid 的 thread blocks 被分成多个不重叠区间，每个 atom 只负责一个区间。Prelude 根据 block＿idx 判断当前 block 是否属于该 atom。重绘依据：论文对 Kernel Atomizer 和 Prelude kernel 的说明。见文献［1］。
1．atomization 的基本思路
LithOS 不是把一个 block 的内部状态切成两半，也不是切 warp。它切的是：
一个 kernel 的 thread－block 集合
例如一个 kernel 有 64 个 block： －Atom 0：［0，16） －Atom 1：［16，32） －Atom 2：［32，48） －Atom 3：［48，64） 每个 atom 各自 launch 一次，但只让自己负责的那段 block 真正执行。
2．Prelude kernel 是怎么用的
LithOS 不是直接修改原 kernel 本体，而是让一次 launch 改成：
- 发一个 Prelude／wrapper kernel
- Prelude 把 blockIdx 线性化成全局 block＿idx
- 如果 block＿idx 落在当前 atom 的范围里，就调用原 kernel 入口
- 不在范围里就直接退出
因此它改变的是 launch 包装方式，而不是原算子的数学逻辑。
3．为什么 atom 之间还能保证正确
因为普通 CUDA 编程模型要求：
- thread blocks 应能独立执行
- blocks 之间不能依赖一个特定执行顺序
- block 内可以同步；block 间通常不能直接依赖 shared memory
所以 LithOS 只要保证：
- 各 atom 的 block range 不重叠
- 所有 block 最终 都执行一次
那么语义上就仍然等价于原来的整个 kernel 执行了一次。
4．中间结果保存在哪里
这里要分清两类状态：
A．block 内部临时状态
例如： －寄存器 －shared memory －thread－local 临时变量
这些状态的生命周期本来就只到该 block 结束。 LithOS 不会 把它们跨 atom 保存。
B．跨 block 可见结果
如果 kernel 本来会把结果写入 output tensor／global memory：
- Atom 0 写自己负责的部分
- Atom 1 再写下一段
- 最终输出仍然累积到同一份 global memory／tensor buffer 中
所以 LithOS 不需要引入一套额外＂原子级上下文保存＂，因为它切的是 block 集合，不是 block 内部执行过程。
right－sizing：kernel 需要几个 TPC 是怎么估出来的
right－sizing 的目标不是问： ｜这个kernel 最多能占多少TPC？ 而是问： 在性能损失可接受的前提下，这个kernel 最少 需要多少 TPC？
1．先求一个＂有用上界＂
论文给出的 heuristic 是：
- 先看 kernel 一共有多少 thread blocks；
- 再看每个 TPC 大概能同时驻留多少 block（occupancy per TPC）；
- 用这两者估一个 useful TPC upper bound
直觉上讲：
- 如果一个 kernel 的 block 本来就不多；
- 或每个 TPC 已经能驻留很多 block；
- 再继续加 TPC，收益可能很小。
2．再用在线模型找＂够用的最小值＂
论文还提到：
- 使用 kernel 在 1 个 TPC 和 全部 TPC 下的时延点；
- 拟合出一个近似 Amdahl 风格的缩放曲线；
- 再根据 latency slip 约束，选出＂最少但够用＂的 TPC 数。
因此：
- SM 级 occupancy 信息 提供建模输入；
- TPC 级配额 是最终输出。
这就是为什么前面会说： SM 是模型单位，TPC 是分配单位。
1．block 和 block 之间共享什么
普通 CUDA 下：
- shared memory 是 block 级资源，别的 block 不能直接用；
- L1 更接近 SM 本地 cache；
- L2 更接近全设备共享 cache。
也就是说：
- Block A 不能直接读 Block B 的 shared memory；
- 但如果两者访问同一份 global memory，可能在 L2 上间接受益。
2．这为什么对 LithOS 很关键 LithOS 的 atomization 之所以通常成立，就是因为它依赖了普通 CUDA 的一个重要假设： 不同blocks 不应依赖彼此的即时执行顺序，也不该直接依赖彼此的 shared memory 状态。 因此把一个 kernel 的 blocks 分多批次发射，语义通常仍是正确的。 当然，涉及 cooperative kernels／persistent kernels／特殊全局同步语义时，支持边界就会变复杂，这也是论文专门强调的限制点之一。
评测图怎么看：为什么 LithOS 同时兼顾吞吐和 SLO
下面这张图是论文 Figure 13 的摘录。
Figure 13．SLO attainment and throughput by system．
图 5．论文 Figure 13 摘录：不同系统的 SLO attainment 与 Throughput。从图上看，LithOS 处于右上角区域，说明它在该实验里同时兼顾了更高的吞吐与更高的高优任务 SLO 达成率。见文献 ［1］。
1．纵轴：SLO attainment 是什么 这里的 SLO 是 Service Level Objective，也就是服务目标。 在论文该实验里，两个高优应用各有自己的目标：
- 一个更偏 latency SLO
- 一个更偏 throughput SLO
图里的 SLO attainment（％）可以理解为： 高优应用对各自服务目标的综合达成程度。 $100 %$ 的含义是：两个高优任务都完全满足了各自目标。 2．为什么 MPS 吞吐高但 SLO 低 MPS 的特点是：
- 多进程并发能力强；
- 资源利用率通常好； －但它主要解决的是共享执行，不是高优任务 QoS 保护。 所以 MPS 往往会表现成：
- 吞吐高；
- 但高优请求容易被长 kernel、干扰负载拖慢；
- 因而 SLO attainment 不高。
3．为什么 MIG／Limits SLO 高但吞吐一般
MIG 或严格 limits 这类方案的特点是：
- 把资源硬隔离得比较明确；
- 高优任务比较容易稳定满足 SLO；
- 但如果某个分区闲着，资源难以灵活借给别的任务。
所以它们常表现成：
- SLO 高；
- 吞吐未必最高。
4．为什么 LithOS 能落在右上角
LithOS 同时结合了：
- TPC 级空间隔离（不像纯 MPS 那么容易互相踩）
- stealing（不像纯硬隔离那样浪费空闲）
- atomization（高优任务到来时不必长期等 BE 长 kernel 自然结束）
- right－sizing（不让 kernel 无脑吃满资源）
因此它不是单纯偏向＂吞吐＂或单纯偏向＂SLO＂，而是尽量同时把两者做高。
代码是否开源，以及能还原到什么程度
我没有查到 LithOS 的公开官方源码仓库。 公开可见的资料主要是：
- arXiv／论文 PDF
- SOSP 论文条目
- CMU作者博文
- 相关 slides
从公开材料能确认的实现点包括：
- 原型 implemented in Rust
- 工作在 CUDA Driver API interposition 层
- 依赖 MPS 支持多 context 并发
- 关键机制包括 virtual streams，TPC scheduling，atomization，right－sizing，power management
但没有公开 repo 可逐文件审阅。
因此，本分析文档里对实现模块的解释分成两类：
1．论文／博文明确说明的部分 2．基于分开信息的工程合理推断
实现边界与未公开细节
1．已经明确公开的
论文与博文已经清楚说明了： －TPC 粒度调度； －kernel atomization； －LibLithOS／interposition； －launch queues／tracker／predictor；
- right－sizing 与 power management；
- 建立在 MPS 之上。
2．没有完全公开写透的
公开论文没有把下面这些低层细节完整写透：
- 具体通过什么底层机制实现 per－launch TPC masking
- 和不同版本 NVIDIA driver／私有接口的细节耦合程度
- 对各种特殊 kernel（如 cooperative／persistent）的完整支持矩阵
- 具体源码模块边界与内部数据结构
因此对这些问题，最稳妥的表述应是：
1．Patrick H．Coppock，Brian Zhang，Eliot H．Solomon，et al．LithOS：An Operating System for Efficient Machine Learning on GPUs．arXiv：2504．15465， 2025. https：／／arxiv．org／abs／2504．15465 2．Patrick H．Coppock．LithOS：An Operating System for Efficient Machine Learning on GPUs． CMU CSD PhD Blog， 2025. https：／／www．cs．cmu．edu／～csd－phd－blog／2025／lithos／ 3．NVIDIA．NVIDIA Hopper Architecture In－Depth． https：／／developer．nvidia．com／blog／nvidia－hopper－architecture－in－depth／ 4．NVIDIA．NVIDIA H100 Tensor Core GPU Architecture Whitepaper． https：／／www．advancedclustering．com／wp－content／uploads／2022／03／gtc22－whitepaper－ hopper．pdf 5．NVIDIA．Multi－Process Service（MPS）documentation． https：／／docs．nvidia．com／deploy／mps／index．html 6．NVIDIA．CUDA Multi－Process Service Overview（PDF）． https：／／docs．nvidia．com／deploy／pdf／CUDA＿Multi＿Process＿Service＿Overview．pdf 7．NVIDIA．CUDA Driver API－Driver Entry Point Access． https：／／docs．nvidia．com／cuda／cuda－driver－api／group＿＿CUDA＿＿DRIVER＿＿ENTRY＿＿POINT．html 8．NVIDIA．CUDA C Programming Guide． https：／／docs．nvidia．com／cuda／cuda－c－programming－guide／
图片说明
- images／01＿h100＿hierarchy＿redrawn．png：基于 NVIDIA Hopper 资料与论文讨论重绘
- images／02＿lithos＿architecture＿redrawn．png：基于论文 Figure 7 与 CMU 博文重绘
- images／03＿mps＿architecture＿redrawn．png：基于 NVIDIA MPS 文档重绘
- images／04＿tpc＿stealing＿atomization＿redrawn．png：基于论文 Figure 9 重绘
- images／05＿fig13＿paper＿excerpt．png：论文 Figure 13 的摘录截图
- images／06＿atomization＿block＿ranges＿redrawn．png：基于论文对 atomization／Prelude 的说明重绘
评论
加载评论中...