LithOS (SOSP'25)
GPU multi-tenancy
- Vendor-provided GPU-sharing mechanisms are transparent but have an isolation vs. utilization trade-off: MPS achieves high utilization with little performance isolation, whereas MIG provides strong isolation via spatially disjoint partitions but can leave resources underutilized.
- LithOS motivates that an OS-level SW layer can enable more flexible and fine-grained spatio-temporal GPU resource manangement while preserving transparency.
Key techniques
- Intercept kernal launch driver APIs.
- User submits kernel by annotating (a) number of requested TPCs; and (b) performance factor $k$ that accounts for tolerable slowdown when adjusting TPC quota.
- The kernel is not immediately submitted to GPU, but to LithOS’s per-stream launch queues.
- Decide the number of allocated TPCs (“TPC quota”).
- For each kernel, estimate the number of “useful” TPCs by dividing number of thread blocks with occupancy per TPC.
cuOccupancyMaxActiveBlocksPerMultiprocessor()seems to be the driver API to query the denominator.
- If it is less than the number of requested TPCs, use the estimation as the TCP quota.
- Otherwise, use a learned linear interpolation, which projects the kernel latency when using one-TPC to all-TPCs, in order to find the minimal TPC count that satisfies $k$.
- For each kernel, estimate the number of “useful” TPCs by dividing number of thread blocks with occupancy per TPC.
- Atomize kernels into standard atom units
- Partition the kernel into threadblocks, called “atoms”, which is a scheduling unit of LithOS. An atom runs for roughly 250-500us when given the allocated TPC count.
- How to align atom’s execution latency is under-specified in the paper. Offline atom latency prediction can be hard; probably profiling it would make more sense.
- Schedule the kernels onto TPCs and perform TPC stealing.
- LithOS specifies which TPC to schedule an atom.
- Texture Processing Unit (TPC): SM < TPC < GPC
- Leverages recent CU mask, Prelude, and QMD struct (driver’s translation of kernel, containing metadata such as kernel entrypoint addr, grid/block configs, and masks).
- Idle TPCs of a kernel can be stolen to schedule other kernels, effectively their atoms, but only during original kernel’s idle scheduling rounds. As a scheduling round matches one atom duration (few 100 us), this bounds the priority inversion period.
- LithOS specifies which TPC to schedule an atom.
- Track outstanding atoms with sync queues.
- LithOS throttles dispatch into GPU from launch queues until outstanding work drops below a threshold (100us).
Comments
- LithOS convincigly demonstrates the effectiveness of SW/OS-controlled GPU scheduling, but it heavily depends on reverse-engineered, vendor-/architecture-specific mechanisms. It is strong as a research demonstration of a future GPU OS interface, but faces barriers to production deployment.