🤖 AI Summary
This work addresses the performance bottleneck in modern deep learning caused by the overhead of frequent GPU kernel launches due to numerous small-scale tensor operations. To mitigate this, the authors propose a persistent GPU kernel runtime system that continuously schedules computational tasks via a host-side task queue and dynamically loads operators through NVRTC-based just-in-time compilation and function pointer injection, thereby eliminating redundant kernel launches. The design introduces an innovative dual-slot aliasing mechanism for concurrent updates and a generic tensor abstraction, enabling transparent integration with PyTorch via TorchDispatch. Evaluated on representative workloads dominated by small operations, the system achieves up to 15.3× speedup over standard PyTorch, substantially improves GPU utilization, and maintains strong ecosystem compatibility.
📝 Abstract
Modern deep learning workloads often consist of many small tensor operations, especially in inference, attention, and micro-batched training. In these settings, kernel launch overhead can become a major bottleneck, sometimes exceeding the actual computation time.
We present GPUOS, a GPU runtime JIT system that reduces launch overhead using a persistent kernel architecture with runtime operator injection. GPUOS runs a single long-lived GPU kernel that continuously processes tasks from a host-managed work queue, eliminating repeated kernel launches. To support diverse operations, GPUOS uses NVIDIA NVRTC to just-in-time compile operators at runtime and inject them into the running kernel through device function pointer tables. This design enables operator updates without restarting the kernel or recompiling the system.
GPUOS introduces four key ideas: (1) a persistent worker kernel with atomic task queues, (2) a runtime operator injection mechanism based on NVRTC and relocatable device code, (3) a dual-slot aliasing scheme for safe concurrent operator updates, and (4) transparent PyTorch integration through TorchDispatch that batches micro-operations into unified submissions. The system supports arbitrary tensor shapes, strides, data types, and broadcasting through a generic tensor abstraction.
Experiments show that GPUOS achieves up to 15.3x speedup over standard PyTorch on workloads dominated by small operations, including micro-batched inference and attention patterns. GPUOS improves utilization while remaining compatible with the PyTorch ecosystem.