LithOS: An Operating System for Efficient Machine Learning on GPUs

Patrick H. Coppock, Brian Zhang, Eliot H. Solomon, Vasilis Kypriotis, Leon Yang†, Bikash Sharma†, Dan Schatzberg†, Todd C. Mowry, Dimitrios Skarlatos | Carnegie Mellon University / †Meta | SOSP '25 (31st ACM Symposium on Operating Systems Principles), Seoul, Oct 13–16 2025 | DOI: 10.1145/3731569.3764818

Problem

GPUs are indispensable for datacenter ML yet run at low utilization — 52% at Microsoft, 10% at Alibaba, below 30% for many Meta inference services, and ~40% for Llama 3 training — while per-GPU power now exceeds 1,000 W, making the status quo economically and energetically unsustainable. Dedicating a GPU to a single workload still leaves cores idle on communication stalls, low batch sizes, and overprovisioning for dynamic loads. Existing GPU-sharing solutions are either non-transparent (require framework/app changes), coarse-grained (whole inference requests, training batches, or operators), or both — causing head-of-line (HoL) blocking, interference, and idle capacity. MPS multiplexes contexts but interferes; MIG gives strong but coarse GPC-level isolation with

5 s reconfiguration; temporal schedulers (time slicing, TGS, Clockwork) run one job at a time. The authors argue that high utilization with isolation, energy efficiency, and transparency requires an operating-system approach.

Core Insight

By transparently interposing at the CUDA Driver API and decoupling kernel submission from execution, an OS can schedule ML work spatially and temporally at the granularity of individual Texture Processing Clusters (TPCs) and sub-kernel "atoms" — enabling TPC stealing, hardware right-sizing, and fine-grained DVFS without any changes to models, frameworks, or compilers.

Method

LithOS (~5000 lines of Rust, built atop MPS) interposes at the CUDA Driver API via a dynamically linked library, LibLithOS, that mimics native CUDA, and maintains a system-wide view of GPU state. Four core components sit over the GPU device driver:

Apps (TensorRT HP-A, JAX HP-B, PyTorch BE)
        |
     LibLithOS  (mimics native CUDA)
        |
  +----------- LithOS core -----------------------------+
  | TPC Scheduler | Kernel Atomizer | Right-sizing | DVFS|
  |  + Stealing   |  (atoms/Prelude)|  l=m/t+b     |f_final|
  |        \____ Online Latency Prediction ____/         |
  +-----------------------------------------------------+
        |
   GPU Device Driver  ->  GPU Hardware [TPC0..TPCn] + Memory

TPC Scheduler: dynamic per-TPC spatial allocation (no MIG-style reconfiguration overhead) with TPC Stealing (idle TPCs lent to other apps), guarded by per-TPC timers, a latency predictor, capped outstanding atoms, and lower hardware stream priority on stolen TPCs. Work flows through per-stream launch queues, TPC quotas, and sync queues with a Tracker thread throttling at a tunable 100 µs outstanding-work limit.
Kernel Atomizer: transparently splits kernels into atoms (subsets of thread blocks) without source/PTX access via a Prelude kernel (Algorithm 1) that range-checks block_idx using an AtomMetadata struct; atom count = predicted duration / tunable atom_duration (250–500 µs effective). Reduces HoL blocking and allows mid-kernel TPC reconfiguration.
Hardware Right-sizing: fits an Amdahl-form curve l = m/t + b (latency vs TPC count) from two points, plus an occupancy-based filtering heuristic, to find the minimum TPCs that raise latency by at most a latency slip parameter k.
Power Management (DVFS): a sequence-based model weights each kernel by its runtime fraction (w), estimates sensitivity via a first-order Taylor approximation, and assigns f_final = f_max / (1 + k/S), where S = Σ w·s.
Online Latency Prediction: learns per-operator latencies (indexed by ordinal position in the data-flow graph, per launch queue) with no offline profiling; feeds all four components.

The prototype supports Ampere and Hopper (extends libsmctrl, reverse-engineers the QMD struct and Hopper Thread Block Clusters), supports CUDA Graphs, and disables stealing/atomization for cross-block-sync or persistent kernels.

Experimental Setup

Component	Value
GPU	1× NVIDIA A100 SXM4 (108 SMs, 40 GB)
Host	30 CPU cores, 216 GB RAM (Lambda Labs)
Software	Ubuntu 22.04, CUDA 12.6, Rust 1.83.0-nightly, Python 3.10, PyTorch 2.3, TensorRT 10.1, TensorRT-LLM 0.11.0, Triton 24.07
Implementation	~5000 LOC Rust, built on top of MPS
NVIDIA baselines	Time slicing, MPS, stream Priority, MIG
Research baselines	TGS, REEF (app mods), Orion (app mods + offline profiling)
HP inference	RetinaNet, YOLOv5, ResNet-50 v1.5, Llama 3 8B, GPT-J 6B, BERT-Large (Triton dynamic batching; RetinaNet on ONNX Runtime)
BE training	ResNet-50, MobileNetV2, VGG-19, DLRM, BERT-Large
LLM trace	Microsoft Azure trace
Latency constraints	MLPerf datacenter inference benchmark (2.3×–7.4× of baseline latency)
Metrics	SLO attainment, goodput/throughput, P₉₅/P₉₉ tail latency, capacity savings, energy savings

Headline Quantitative Results

Inference-only multitenancy (2 HP + 1 BE):

SLO attainment and throughput (Fig. 14): LithOS reaches 100% SLO + 1.0 throughput; MPS 45% SLO at 1.11 throughput; MIG 0.71; thread limits 0.58; TGS leads priority-only at 84% SLO.
HP B goodput (Fig. 15): LithOS 0.50 vs MIG 0.37; only LithOS gives best HP throughput while sustaining high BE.
HP A tail latency (Fig. 16): LithOS 13× better than MPS, 4× better than Orion, 1.2× better than TGS.

Hybrid inference-training multitenancy (Fig. 17):

LithOS holds tail latency within 20% of ideal; 2.34× better than REEF and 1.18× better than TGS on average; up to 13.54× / avg 4.7× better than MPS. (MPS 5.83× ideal latency / 60% throughput; stream priority 2.89× / 68%; TGS avg 1.41×; REEF avg 2.89×, tails to 8.93×.)
Service throughput within 1% of load worst-case; training throughput 34× vs TGS; aggregate throughput 1.35× vs TGS, overall 1.23×–1.57× (avg 1.38×).

Right-sizing (Fig. 18): up to 51%, mean 26% capacity savings; with slip 1.1, P₉₉ +4% and throughput −4%; fitted-curve R² 0.92–0.99.

DVFS (Fig. 19): up to 46%, mean 26% energy savings (baseline mostly at 1410 MHz); with slip 1.1, mean P₉₉ +7%.

Ablation (Fig. 20/21): TPC scheduler → 1.38× ideal tail; +atomization → 1.19× avg (up to 1.55×) at 10% throughput overhead. Atomization case study: LithOS beats REEF 6.5× (vary batch) / 3.9× (vary sequence length), with atomization contributing 2× / 1.3×; HP tail within 14% (1 ms) / 7% (0.45 ms) of ideal. Latency-prediction misprediction: HP 0.9% / 0.38% (P₉₉ errors 49 / 31 µs), BE 14% / 11%. Overhead: LithOS +4% (atomization <1%) vs TGS/REEF ~2%, Orion 6%.

Limitations

Built on top of MPS and on reverse-engineering (QMD, libsmctrl, TPC masking, Hopper Thread Block Clusters) because current APIs do not expose TPC/SM control; native driver support is argued feasible but absent.
Frequency switching is slow (~50 ms), so DVFS is impractical for very short kernels; sub-ms DVFS and per-SM power control remain future hardware needs.
No hardware preemption exists; atomization is a software workaround. Special kernels (cross-block sync, persistent) require disabling stealing/atomization.
BE latency prediction is less accurate (14% / 11% misprediction).
Bandwidth/memory isolation is not implemented (estimated 4–13% gains); compute isolation is prioritized.

Open Problems Called Out

Native driver/hardware support: kernel-to-SM assignment, hardware preemption, cache/memory partitioning, NUMA-style placement, richer context management.
Fine-grained (sub-ms) DVFS and per-SM, spatially-applied power control for diverse, input-dependent future workloads.
Extending OS-level management to other GPU resources — memory, bandwidth, PCIe, SSDs, networking — and to intra-SM heterogeneity (e.g., tensor cores).
More involved modeling to improve best-effort latency prediction and cover even more diverse GPU workloads.